{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Test.HTTP2.Client.ManagerSpec where

import Control.Concurrent (threadDelay)
import Control.Concurrent.Async
import Control.Concurrent.MVar
import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as Builder
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Lazy.Char8 as LBS
import Data.IORef
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, isJust)
import Data.Streaming.Network (bindPortTCP, bindRandomPortTCP)
import Data.Unique
import Foreign.Marshal.Alloc (mallocBytes)
import GHC.IO.Exception
import HTTP2.Client.Manager
import HTTP2.Client.Manager.Internal
import Network.HTTP.Types
import qualified Network.HTTP2.Client as Client
import qualified Network.HTTP2.Client as HTTP2
import Network.HTTP2.Server (defaultServerConfig)
import qualified Network.HTTP2.Server as Server
import qualified Network.HTTP2.Server.Internal as Server
import Network.Socket
import qualified Network.Socket as NS
import qualified OpenSSL.Session as SSL
import qualified StmContainers.Map as StmMap
import System.Random (randomRIO)
import qualified System.TimeManager
import Test.Hspec

echoTest :: Http2Manager -> TLSEnabled -> Int -> Expectation
echoTest = echoTest' "localhost" "/echo" "some body"

echoTest' :: ByteString -> ByteString -> Builder.Builder -> Http2Manager -> TLSEnabled -> Int -> Expectation
echoTest' hostname path msg mgr tlsEnabled serverPort =
  withHTTP2Request mgr (tlsEnabled, hostname, serverPort) (Client.requestBuilder "GET" path [] msg) $ \res -> do
    Client.responseStatus res `shouldBe` Just status200
    readResponseBody res `shouldReturn` Builder.toLazyByteString msg

-- | server delays by 0..100ms between every two lines, so this is good for randomized load testing.
--
-- TODO: implement the delay below!
multiLineEchoTest :: Http2Manager -> TLSEnabled -> Int -> Expectation
multiLineEchoTest = echoTest' "localhost" "/multiline" "1\n2\n3\n4\n5\n6\n"

spec :: Spec
spec = do
  describe "HTTP2.Client.Manager [without TLS]" $ do
    specTemplate Nothing
  -- TODO: Perhaps there are some TLS specific test cases we might want to
  -- write. Like ensure that the hostname validation works (it currently
  -- does't!).
  describe "HTTP2.Client.Manager [with TLS]" $ do
    localhostCtx <- runIO $ loadServerSSLContext
    specTemplate (Just localhostCtx)

    it "should error appropriately when the server presents a wrong certificate" $ do
      ctx <- loadWrongServerSSLContext
      mgr <- mkTestManager
      withTestServer (Just ctx) $ \TestServer {..} ->
        echoTest mgr True serverPort `shouldThrow` (\(_ :: SSL.SomeSSLException) -> True)

#if DENABLE-TRAILING-DOT-TEST
    it "should allow accepting a certificate without a trailing dot" $ do
      mgr <- setSSLRemoveTrailingDot True <$> mkTestManager
      withTestServer (Just localhostCtx) $ \TestServer {..} ->
        echoTest' "localhost." "/echo" "some body" mgr True serverPort
#endif

specTemplate :: Maybe SSL.SSLContext -> Spec
specTemplate mCtx = do
  it "should be able to make an HTTP2 request" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      mgr <- mkTestManager

      echoTest mgr (isJust mCtx) serverPort

      readIORef acceptedConns `shouldReturn` 1

  it "should fail appropriately when a server is not available" $ do
    mgr <- mkTestManager

    -- Assumes that nothing is running on this port
    echoTest mgr (isJust mCtx) 42420 `shouldThrow` (\IOError {..} -> ioe_type == NoSuchThing)

  it "should be able to re-use an HTTP2 connection for multiple requests" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      mgr <- mkTestManager

      echoTest mgr (isJust mCtx) serverPort
      echoTest mgr (isJust mCtx) serverPort

      readIORef acceptedConns `shouldReturn` 1

  it "should process many concurrent requests correctly under different timing situations" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      let numserial = 8
          -- don't incease `numparallel` beyond max streams until
          -- https://github.com/kazu-yamamoto/http2/issues/71 is resolved!
          numparallel = 99
      mgr <- mkTestManager
      let onethread _ = replicateM_ numserial (multiLineEchoTest mgr (isJust mCtx) serverPort)
      mapConcurrently_ onethread [(0 :: Int) .. numparallel]

  it "shouldn't try to re-use a disconnected connection" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      mgr <- mkTestManager

      echoTest mgr (isJust mCtx) serverPort
      disconnectTargetWithTimeout mgr (isJust mCtx, "localhost", serverPort) 1_000_000
      echoTest mgr (isJust mCtx) serverPort

      readIORef acceptedConns `shouldReturn` 2

  it "should re-use connections even for concurrent requests" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      mgr <- mkTestManager

      -- Do 1 request before the concurrent ones, otherwise most of them would
      -- end up making multiple connections.
      echoTest mgr (isJust mCtx) serverPort
      mapConcurrently_ id $ replicate 10 (echoTest mgr (isJust mCtx) serverPort)

      readIORef acceptedConns `shouldReturn` 1

  it "should re-use the connection even if one of the requests is stuck forever" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      mgr <- mkTestManager
      infiniteRespHeaderRecieved <- newEmptyMVar
      chunkOfInfiniteTest <- newEmptyMVar

      infiniteRespThread <- async $ withHTTP2Request mgr (isJust mCtx, "localhost", serverPort) (Client.requestNoBody "GET" "/inifite" []) $ \res -> do
        Client.responseStatus res `shouldBe` Just status200
        putMVar infiniteRespHeaderRecieved ()

        -- The test server writes "foo\n" 1000 times in each chunk. We do this
        -- to ensure that this request is still alive even after an echo test
        -- which started after this and finished before this.
        takeMVar chunkOfInfiniteTest
        Client.getResponseBodyChunk res `shouldReturn` BS.concat (replicate 1000 "foo\n")

      takeMVar infiniteRespHeaderRecieved
      echoTest mgr (isJust mCtx) serverPort

      putMVar chunkOfInfiniteTest ()
      -- We wait on the thread here to ensure that any expectation failures in
      -- the thread are caught.
      wait infiniteRespThread

      readIORef acceptedConns `shouldReturn` 1

  it "should re-use the connection even if an exception is thrown while handling a response" $ do
    withTestServer mCtx $ \TestServer {..} -> do
      mgr <- mkTestManager

      let exceptionThrower =
            withHTTP2Request mgr (isJust mCtx, "localhost", serverPort) (Client.requestBuilder "GET" "/echo" [] "some body") $ \_ ->
              throw TestException

      -- Also test if the exception is propagated correctly and doesn't cause
      -- other errors
      exceptionThrower `shouldThrow` (== TestException)
      echoTest mgr (isJust mCtx) serverPort

      readIORef acceptedConns `shouldReturn` 1

  it "should fail with appropriate error when a dead connection is used" $ do
    mgr <- mkTestManager
    Just deadConn <- withTestServer mCtx $ \TestServer {..} -> do
      echoTest mgr (isJust mCtx) serverPort
      readIORef acceptedConns `shouldReturn` 1
      atomically $ StmMap.lookup (isJust mCtx, "localhost", serverPort) (connections mgr)

    let brokenRequest = sendRequestWithConnection deadConn (Client.requestBuilder "GET" "/echo" [] "some body") $ \_ -> do
          expectationFailure "Expected no response when request is made to a dead server"
        -- It is hard to control which one of these will be thrown
        expectedException (fromException -> Just ConnectionAlreadyClosed) = True
        expectedException (fromException -> Just Client.ConnectionIsClosed) = True
        expectedException (fromException -> Just (ioe_type -> ResourceVanished)) = True
        expectedException _ = False
    brokenRequest `shouldThrow` expectedException

  it "should create a new connection when the server restarts" $ do
    mgr <- mkTestManager
    port <- withTestServer mCtx $ \TestServer {..} -> do
      echoTest mgr (isJust mCtx) serverPort
      readIORef acceptedConns `shouldReturn` 1
      pure serverPort

    -- See "should fail with appropriate error when a dead connection is used"
    -- to know what happens when we don't wait for the background thread to go
    -- away.
    Just deadConn <-
      atomically $ StmMap.lookup (isJust mCtx, "localhost", port) (connections mgr)
    void $ waitCatch $ backgroundThread deadConn

    withTestServerOnPort mCtx port $ \TestServer {..} -> do
      echoTest mgr (isJust mCtx) port
      -- this is still 1 because we have a new 'TestServer'
      readIORef acceptedConns `shouldReturn` 1

data TestException = TestException
  deriving (Show, Eq)

instance Exception TestException

data TestServer = TestServer
  { serverThread :: Async (),
    acceptedConns :: IORef Int,
    liveConns :: IORef (Map Unique (Async ())),
    serverPort :: Int
  }

loadServerSSLContext :: IO SSL.SSLContext
loadServerSSLContext = do
  ctx <- SSL.context
  SSL.contextSetCertificateFile ctx "test/resources/localhost.pem"
  SSL.contextSetPrivateKeyFile ctx "test/resources/localhost-key.pem"
  SSL.contextSetALPNProtos ctx ["h2"]
  SSL.contextSetCiphers ctx "HIGH"
  sslCheck <- SSL.contextCheckPrivateKey ctx
  sslCheck `shouldBe` True
  pure ctx

loadWrongServerSSLContext :: IO SSL.SSLContext
loadWrongServerSSLContext = do
  ctx <- SSL.context
  SSL.contextSetCertificateFile ctx "test/resources/localhost.example.com.pem"
  SSL.contextSetPrivateKeyFile ctx "test/resources/localhost.example.com-key.pem"
  SSL.contextSetALPNProtos ctx ["h2"]
  SSL.contextSetCiphers ctx "HIGH"
  sslCheck <- SSL.contextCheckPrivateKey ctx
  sslCheck `shouldBe` True
  pure ctx

mkTestManager :: IO Http2Manager
mkTestManager = do
  mgr <- defaultHttp2Manager
  SSL.contextSetCAFile (sslContext mgr) "test/resources/unit-ca.pem"
  pure mgr

withTestServer :: Maybe SSL.SSLContext -> (TestServer -> IO a) -> IO a
withTestServer mCtx action = do
  bracket (bindRandomPortTCP "*") (close . snd) $ \(port, sock) ->
    withTestServerOnSocket mCtx action (port, sock)

withTestServerOnPort :: Maybe SSL.SSLContext -> Int -> (TestServer -> IO a) -> IO a
withTestServerOnPort mCtx serverPort action = do
  bracket (bindPortTCP serverPort "*") close $ \listenSock ->
    withTestServerOnSocket mCtx action (serverPort, listenSock)

withTestServerOnSocket :: Maybe SSL.SSLContext -> (TestServer -> IO a) -> (Int, Socket) -> IO a
withTestServerOnSocket mCtx action (serverPort, listenSock) = do
  acceptedConns <- newIORef 0
  liveConns <- newIORef mempty
  let cleanupServer serverThread = do
        cancel serverThread
        mapM_ cancel =<< readIORef liveConns
  bracket (async $ testServerOnSocket mCtx listenSock acceptedConns liveConns) cleanupServer $ \serverThread ->
    action TestServer {..}

allocServerConfig :: Either Socket SSL.SSL -> IO Server.Config
allocServerConfig (Left sock) = HTTP2.allocSimpleConfig sock 4096
allocServerConfig (Right ssl) = do
  buf <- mallocBytes bufsize
  timmgr <- System.TimeManager.initialize $ 30 * 1000000
  -- Sometimes the frame header says that the payload length is 0. Reading 0
  -- bytes multiple times seems to be causing errors in openssl. I cannot figure
  -- out why. The previous implementation didn't try to read from the socket
  -- when trying to read 0 bytes, so special handling for 0 maintains that
  -- behaviour.
  let readData prevChunk 0 = pure prevChunk
      readData prevChunk n = do
        -- Handling SSL.ConnectionAbruptlyTerminated as a stream end
        -- (some sites terminate SSL connection right after returning the data).
        chunk <- SSL.read ssl n `catch` \(_ :: SSL.ConnectionAbruptlyTerminated) -> pure mempty
        let chunkLen = BS.length chunk
        if
          | chunkLen == 0 || chunkLen == n ->
              pure (prevChunk <> chunk)
          | chunkLen > n ->
              error "openssl: SSL.read returned more bytes than asked for, this is probably a bug"
          | otherwise ->
              readData (prevChunk <> chunk) (n - chunkLen)

  let s = fromMaybe (error "http2-manager: SSL without socket") $ SSL.sslSocket ssl
  mysa <- NS.getSocketName s
  peersa <- NS.getPeerName s
  pure
    Server.Config
      { Server.confWriteBuffer = buf,
        Server.confBufferSize = bufsize,
        Server.confSendAll = SSL.write ssl,
        Server.confReadN = readData mempty,
        Server.confPositionReadMaker = Server.defaultPositionReadMaker,
        Server.confTimeoutManager = timmgr,
        Server.confMySockAddr = mysa,
        Server.confPeerSockAddr = peersa,
        HTTP2.confReadNTimeout = False
      }

testServerOnSocket :: Maybe SSL.SSLContext -> Socket -> IORef Int -> IORef (Map Unique (Async ())) -> IO ()
testServerOnSocket mCtx listenSock connsCounter conns = do
  listen listenSock 1024
  forever $ do
    (sock, _) <- accept listenSock
    serverCfgParam <- case mCtx of
      Nothing -> pure $ Left sock
      Just ctx -> do
        ssl <- SSL.connection ctx sock
        SSL.accept ssl
        pure (Right ssl)
    connKey <- newUnique
    modifyIORef connsCounter (+ 1)
    let shutdownSSL = case serverCfgParam of
          Left _ -> pure ()
          Right ssl -> SSL.shutdown ssl SSL.Bidirectional
        cleanup cfg = do
          Server.freeSimpleConfig cfg `finally` (shutdownSSL `finally` close sock)
    thread <- async $ bracket (allocServerConfig serverCfgParam) cleanup $ \cfg -> do
      Server.run defaultServerConfig cfg testServer `finally` modifyIORef conns (Map.delete connKey)
    modifyIORef conns $ Map.insert connKey thread

testServer :: Server.Request -> Server.Aux -> (Server.Response -> [Server.PushPromise] -> IO ()) -> IO ()
testServer req _ respWriter = do
  case Server.requestPath req of
    Just "/echo" -> do
      reqBody <- readRequestBody req
      respWriter (Server.responseBuilder status200 [] (Builder.lazyByteString reqBody)) []
    Just "/inifite" -> do
      let infiniteBSWriter :: (Builder.Builder -> IO ()) -> IO () -> IO ()
          infiniteBSWriter bsWriter flush = do
            bsWriter $ Builder.lazyByteString $ LBS.concat $ replicate 1000 "foo\n"
            flush
            infiniteBSWriter bsWriter flush
          infiniteResponse = Server.responseStreaming status200 [] infiniteBSWriter
      respWriter infiniteResponse []
    Just "/multiline" -> do
      reqBodyLines <- LBS.lines <$> readRequestBody req
      let delayedBSWriter :: [LBS.ByteString] -> (Builder.Builder -> IO ()) -> IO () -> IO ()
          delayedBSWriter [] _ _ = pure ()
          delayedBSWriter (line : moreLines) bsWriter flush = do
            bsWriter $ Builder.lazyByteString (line <> "\n")
            flush
            threadDelay =<< randomRIO (0, 100_000)
            delayedBSWriter moreLines bsWriter flush
          delayedResponse = Server.responseStreaming status200 [] (delayedBSWriter reqBodyLines)
      respWriter delayedResponse []
    _ -> do
      respWriter (Server.responseNoBody status404 []) []

readResponseBody :: Client.Response -> IO LBS.ByteString
readResponseBody res = readChunks (Client.getResponseBodyChunk res)

readRequestBody :: Server.Request -> IO LBS.ByteString
readRequestBody req = readChunks (Server.getRequestBodyChunk req)

readChunks :: IO ByteString -> IO LBS.ByteString
readChunks action = LBS.fromChunks <$> go []
  where
    go chunks = do
      action >>= \case
        "" -> pure chunks
        c -> go (chunks <> [c])
