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

-- 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 HTTP2.Client.Manager.Internal where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.ByteString
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C8
import Data.IORef
import Data.Map
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import Data.Streaming.Network
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import Data.Unique
import Foreign.Marshal.Alloc (mallocBytes)
import GHC.IO.Exception
import qualified Network.HPACK as HPACK
import qualified Network.HTTP2.Client as HTTP2
import qualified Network.HTTP2.Client.Internal as HTTP2
import qualified Network.Socket as NS
import qualified OpenSSL.Session as SSL
import qualified StmContainers.Map as StmMap
import System.IO.Error
import qualified System.TimeManager
import System.Timeout
import Prelude

data HTTP2Conn = HTTP2Conn
  { backgroundThread :: Async (),
    disconnect :: IO (),
    connectionActionMVar :: MVar ConnectionAction
  }

type TLSEnabled = Bool

type HostName = ByteString

type Port = Int

type Target = (TLSEnabled, HostName, Port)

data ConnectionAction
  = SendRequest Request
  | CloseConnection

data Request = Request
  { -- | The request to be sent.
    request :: HTTP2.Request,
    -- | Consumer for the response, must not exit until the response body is
    -- completely consumed.
    --
    -- The response consumer has to return 'IO ()' because we want to processes
    -- different requests on one connection and 'HTTP2.run' ties the return type
    -- of response consumer to the return type of itself. Even if the response
    -- consumer returned something else we would need another empty MVar to
    -- write the result, this is being dealt with in
    -- 'sendRequestWithConnection'.
    responseConsumer :: HTTP2.Response -> IO (),
    -- | There are exceptions which cannot be communicated in a continuation
    -- becasue they can be raised even before the continuation starts. We also
    -- need a way to communicate any exceptions raised by the continuation
    -- itself. This 'MVar' will be written to in any of those cases.
    exceptionMVar :: MVar SomeException
  }

-- | FUTUREWORK: Support HTTPS, perhaps ALPN negotiation can also be used to
-- HTTP1. I think HTTP1 vs HTTP2 can not be negotated without TLS, so perhaps
-- this manager will default to HTTP2.
data Http2Manager = Http2Manager
  { connections :: StmMap.Map Target HTTP2Conn,
    cacheLimit :: Int,
    -- | In microseconds, defaults to 30s
    tcpConnectionTimeout :: Int,
    sslContext :: SSL.SSLContext,
    sslRemoveTrailingDot :: Bool
  }

defaultHttp2Manager :: IO Http2Manager
defaultHttp2Manager = do
  ctx <- SSL.context
  SSL.contextSetVerificationMode ctx $
    SSL.VerifyPeer
      { vpFailIfNoPeerCert = True,
        -- Only relvant when running as server
        vpClientOnce = False,
        vpCallback = Nothing
      }
  SSL.contextSetALPNProtos ctx ["h2"]
  http2ManagerWithSSLCtx ctx

http2ManagerWithSSLCtx :: SSL.SSLContext -> IO Http2Manager
http2ManagerWithSSLCtx sslContext = do
  connections <- StmMap.newIO
  let cacheLimit = 20
      tcpConnectionTimeout = 30_000_000
      sslRemoveTrailingDot = False
  pure $ Http2Manager {..}

-- | Warning: This won't affect already established connections
setCacheLimit :: Int -> Http2Manager -> Http2Manager
setCacheLimit cl mgr = mgr {cacheLimit = cl}

-- | Warning: This won't affect already established connections
setSSLContext :: SSL.SSLContext -> Http2Manager -> Http2Manager
setSSLContext ctx mgr = mgr {sslContext = ctx}

-- | Remove traling dots in hostname while verifying hostname in the certificate
-- presented by the server. For instance, when connecting with
-- 'foo.example.com.' (Note the trailing dot) by default most SSL libraries fail
-- hostname verification if the server has a certificate for 'foo.example.com'
-- (Note the lack of a trailing dot). Setting this flag makes the hostname
-- verification succeed for these hosts. However, this will make the hostname
-- verification fail if the host presents a certificate which does have a
-- trailing dot.
--
-- Discussion about why this is not implemented as a flag on 'SSLContext':
-- https://github.com/openssl/openssl/issues/11560
--
-- Warning: This won't affect already established connections
setSSLRemoveTrailingDot :: Bool -> Http2Manager -> Http2Manager
setSSLRemoveTrailingDot b mgr = mgr {sslRemoveTrailingDot = b}

-- | In microseconds
setTCPConnectionTimeout :: Int -> Http2Manager -> Http2Manager
setTCPConnectionTimeout n mgr = mgr {tcpConnectionTimeout = n}

-- | Does not check whether connection is actually running. Users should use
-- 'withHTTP2Request'. This function is good for testing.
sendRequestWithConnection :: HTTP2Conn -> HTTP2.Request -> (HTTP2.Response -> IO r) -> IO r
sendRequestWithConnection conn req k = do
  result :: MVar r <- newEmptyMVar
  threadKilled :: MVar SomeException <- newEmptyMVar
  putMVar (connectionActionMVar conn) (SendRequest (Request req (putMVar result <=< k) threadKilled))

  let waitResult = takeMVar result
      waitError = takeMVar threadKilled
      waitDeath = do
        res <- waitCatch (backgroundThread conn)
        pure $
          case res of
            Left e -> e
            Right _ -> SomeException ConnectionAlreadyClosed

  race waitResult (race waitError waitDeath) >>= \case
    Left r -> pure r
    Right (Left e) -> throwIO e
    Right (Right e) -> throwIO e

-- | Make an HTTP2 request, if it is the first time the 'Http2Manager' sees this
-- target, it creates the connection and keeps it around for
-- any subsequent requests. Subsequest requests try to use this connection, in
-- case the connection is already dead (e.g. the background thread has
-- finished), a new connection is created.
--
-- It is important that the continuation provided by the caller of this function
-- consumes the response body completely before it returns.
--
-- NOTE: If many concurrent requests are made to the same server using a single
-- instance of 'Http2Manager', it could cause the manager to make multiple
-- connections to the server. Eventually only one connection will be kept open.
-- This, in theory, would cause some contention over 'STM' based 'Map' that the
-- 'Http2Manager' keeps and so could decrease throughput. In cases where many
-- concurrent requests are to be made, it might be best to ensure that a
-- connection exists using 'connectIfNotAlreadyConnected' before making all the
-- requests.
withHTTP2Request :: Http2Manager -> Target -> HTTP2.Request -> (HTTP2.Response -> IO a) -> IO a
withHTTP2Request mgr target req k = do
  conn <- getOrMakeConnection mgr target
  sendRequestWithConnection conn req k

-- | Temporary workaround for https://github.com/kazu-yamamoto/http2/issues/102
withHTTP2RequestOnSingleUseConn :: Http2Manager -> Target -> HTTP2.Request -> (HTTP2.Response -> IO a) -> IO a
withHTTP2RequestOnSingleUseConn mgr target req k =
  withHTTP2RequestOnSingleUseConnWithHook mgr target req (\_ -> pure ()) k

withHTTP2RequestOnSingleUseConnWithHook ::
  Http2Manager ->
  Target ->
  HTTP2.Request ->
  (NS.SockAddr -> IO ()) ->
  (HTTP2.Response -> IO a) ->
  IO a
withHTTP2RequestOnSingleUseConnWithHook Http2Manager {..} target req onConnect k = do
  sendReqMVar <- newEmptyMVar
  thread <- liftIO . async $ startPersistentHTTP2ConnectionWithHook sslContext target cacheLimit sslRemoveTrailingDot tcpConnectionTimeout onConnect sendReqMVar
  let newConn = HTTP2Conn thread (putMVar sendReqMVar CloseConnection) sendReqMVar
  sendRequestWithConnection newConn req $ \resp -> do
    k resp <* disconnect newConn

-- | Connects to a server if it is not already connected, useful when making
-- many concurrent requests. This way the first few requests don't have to fight
-- for making a connection This way the first few requests don't have to fight
-- for making a connection.
connectIfNotAlreadyConnected :: Http2Manager -> Target -> IO ()
connectIfNotAlreadyConnected mgr target = void $ getOrMakeConnection mgr target

-- | Gets a connection if it exists and is alive, otherwise connects to the
-- given 'Target'.
getOrMakeConnection :: Http2Manager -> Target -> IO HTTP2Conn
getOrMakeConnection mgr@Http2Manager {..} target = do
  mConn <- atomically $ getConnection mgr target
  maybe connect pure mConn
  where
    -- Ensures that any old connection is preserved. This is required to ensure
    -- that concurrent calls to this function don't cause the connections to
    -- leak.
    insertNewConn :: HTTP2Conn -> STM (Bool, HTTP2Conn)
    insertNewConn newConn = do
      StmMap.lookup target connections >>= \case
        Just existing -> pure (False, existing)
        Nothing -> do
          StmMap.insert newConn target connections
          pure (True, newConn)

    connect :: IO HTTP2Conn
    connect = do
      sendReqMVar <- newEmptyMVar
      thread <- liftIO . async $ startPersistentHTTP2Connection sslContext target cacheLimit sslRemoveTrailingDot tcpConnectionTimeout sendReqMVar
      let newConn = HTTP2Conn thread (putMVar sendReqMVar CloseConnection) sendReqMVar
      (inserted, finalConn) <- atomically $ insertNewConn newConn
      unless inserted $ do
        -- It is possible that the connection won't leak because it is waiting
        -- on an MVar and as soon as it gets removed from the map and GC collects
        -- the 'HTTP2Conn', the connection thread _should_ in theory get
        -- 'BlockedIndefinitelyOnMVar' exception. So perhaps this is useless?
        disconnect newConn
      pure finalConn

-- | Removes connection from map if it is not alive anymore
getConnection :: Http2Manager -> Target -> STM (Maybe HTTP2Conn)
getConnection mgr target = do
  StmMap.lookup target (connections mgr) >>= \case
    Nothing -> pure Nothing
    Just conn ->
      -- If there is a connection for the target, ensure that it is alive
      -- before using it.
      pollSTM (backgroundThread conn) >>= \case
        Nothing -> pure (Just conn)
        Just _ -> do
          -- Maybe there is value in logging any exceptions we
          -- recieve here. But logging in STM will be tricky, and the threads
          -- running requests on the connection which got an exception would've
          -- anyway recieved the exception, so maybe it is not as valueable.
          StmMap.delete target (connections mgr)
          pure Nothing

-- | Disconnects HTTP2 connection if there exists one. Will hang around until
-- all the ongoing requests complete. This would throw an error if the
-- background thread maintaining the connection throws an error, e.g. there was
-- a TLS error or the connection was already disconnected with error.
disconnectTarget :: Http2Manager -> Target -> IO ()
disconnectTarget mgr target = do
  mConn <- atomically $ getConnection mgr target
  case mConn of
    Nothing -> pure ()
    Just conn -> do
      disconnect conn
      wait (backgroundThread conn)
        `finally` atomically (StmMap.delete target (connections mgr))

-- | Disconnects HTTP2 connection if there exists one. If the background thread
-- running the connection does not finish within 1 second, it is canceled.
-- Errors from the background thread running the connection are not propagated.
--
-- NOTE: Any requests in progress might not finish correctly.
disconnectTargetWithTimeout :: Http2Manager -> Target -> Int -> IO ()
disconnectTargetWithTimeout mgr target microSeconds = do
  mConn <- atomically $ getConnection mgr target
  case mConn of
    Nothing -> pure ()
    Just conn -> do
      disconnect conn

      -- Wait with timeout using two threads:
      -- 1. background thread which _should_ be exiting soon
      -- 2. sleep for given number of microseconds.
      --
      -- whenever one of them finishes, the other is canceled. Errors are
      -- ignored.
      --
      -- All of this to say wait max 1 second for the background thread to
      -- finish.
      let waitWithTimeout = do
            waitOneSec <- async $ threadDelay microSeconds
            void $ waitAnyCatchCancel [waitOneSec, backgroundThread conn]

      waitWithTimeout
        `finally` atomically (StmMap.delete target (connections mgr))

startPersistentHTTP2Connection ::
  SSL.SSLContext ->
  Target ->
  Int ->
  Bool ->
  Int ->
  MVar ConnectionAction ->
  IO ()
startPersistentHTTP2Connection ctx target cl removeTrailingDot tcpConnectTimeout sendReqMVar =
  startPersistentHTTP2ConnectionWithHook ctx target cl removeTrailingDot tcpConnectTimeout (\_ -> pure ()) sendReqMVar

startPersistentHTTP2ConnectionWithHook ::
  SSL.SSLContext ->
  Target ->
  -- cacheLimit
  Int ->
  -- sslRemoveTrailingDot
  Bool ->
  -- | TCP connect timeout in microseconds
  Int ->
  (NS.SockAddr -> IO ()) ->
  -- MVar used to communicate requests or the need to close the connection.  (We could use a
  -- queue here to queue several requests, but since the requestor has to wait for the
  -- response, it might as well block before sending off the request.)
  MVar ConnectionAction ->
  IO ()
startPersistentHTTP2ConnectionWithHook ctx (tlsEnabled, hostname, port) cl removeTrailingDot tcpConnectTimeout onConnect sendReqMVar = do
  liveReqs <- newIORef mempty
  let clientConfig =
        HTTP2.defaultClientConfig
          { HTTP2.scheme = if tlsEnabled then "https" else "http",
            HTTP2.authority = C8.unpack hostname,
            HTTP2.cacheLimit = cl
          }
      -- Sends error to requests which show up too late, i.e. after the
      -- connection is already closed
      tooLateNotifier e = do
        forever $ do
          takeMVar sendReqMVar >>= \case
            SendRequest Request {..} -> do
              -- No need to get stuck here
              void $ tryPutMVar exceptionMVar (SomeException e)
            CloseConnection -> pure ()

      -- Sends errors to the request threads when an error occurs
      cleanupThreadsWith (SomeException e) = do
        -- Is it really OK to cancel the remaining threads even when throwing
        -- 'ConnectionAlreadyClosed', there is a chance that they could finish,
        -- but how would we know here?
        mapM_ (\(thread, _) -> cancelWith thread e) =<< readIORef liveReqs
        -- Spawns a thread that will hang around for 1 second to deal with
        -- the race betwen main thread sending a request and this thread
        -- already having stoped waiting for new requests. Sending requests
        -- after 'handleRequests' has finsihed just causes the main thread
        -- to hang until recieving 'BlockedIndefinitelyOnMVar'.
        --
        -- 1 second is hopefully enough to ensure that this thread is seen
        -- as finished.
        void $ async $ race_ (tooLateNotifier e) (threadDelay 1_000_000)
        throwIO e

      hostnameForTLS =
        if removeTrailingDot
          then fromMaybe hostname (stripSuffix "." hostname)
          else hostname
      transportConfig
        | tlsEnabled = Just $ TLSParams ctx hostnameForTLS
        | otherwise = Nothing

  handle cleanupThreadsWith $
    bracket connectTCPWithTimeout NS.close $ \sock -> do
      onConnect =<< NS.getPeerName sock
      bracket (mkTransport sock transportConfig) cleanupTransport $ \transport ->
        bracket (allocHTTP2Config transport) HTTP2.freeSimpleConfig $ \http2Cfg -> do
          let runAction = HTTP2.run clientConfig http2Cfg $ \sendReq _aux -> do
                handleRequests liveReqs sendReq
          -- Any request threads still hanging about after 'runAction' finishes
          -- are canceled with 'ConnectionAlreadyClosed'.
          flip finally (cleanupThreadsWith (SomeException ConnectionAlreadyClosed)) $
            -- Any exceptions thrown will get re-thrown to any running requests,
            -- handle at the top level is not good as 'finally' wrapping this
            -- function would kill all threads with some other exception.
            handle cleanupThreadsWith $
              runAction
  where
    handleRequests :: IORef LiveReqs -> SendReqFn -> IO ()
    handleRequests liveReqs sendReq = do
      let waitAndFork = do
            reqOrStop <- takeMVar sendReqMVar
            case reqOrStop of
              SendRequest r@(Request {..}) -> do
                processRequest liveReqs sendReq r
                  `catches` exceptionHandlers exceptionMVar
                waitAndFork
              CloseConnection -> do
                mapM_ (wait . fst) =<< readIORef liveReqs
      waitAndFork

    processRequest :: IORef LiveReqs -> SendReqFn -> Request -> IO ()
    processRequest liveReqs sendReq Request {..} = do
      unique <- newUnique
      thread <- async $ do
        let actionWithHandlers =
              sendReq request responseConsumer
                `catches` exceptionHandlers exceptionMVar
            cleanup = do
              atomicModifyIORef' liveReqs (\m -> (Map.delete unique m, ()))
        actionWithHandlers `finally` cleanup
      atomicModifyIORef' liveReqs (\m -> (Map.insert unique (thread, exceptionMVar) m, ()))

    -- Specially handle 'ConnectionAlreadyClosed' otherwise it shows up as 'SomeAsyncException'
    tooLateHandler threadKilled e@ConnectionAlreadyClosed =
      putMVar threadKilled (SomeException e)
    generalHandler threadKilled e = putMVar threadKilled e
    exceptionHandlers threadKilled = [Handler $ tooLateHandler threadKilled, Handler $ generalHandler threadKilled]

    connectTCPWithTimeout :: IO NS.Socket
    connectTCPWithTimeout = do
      mSock <- timeout tcpConnectTimeout $ fst <$> getSocketTCP hostname port
      case mSock of
        Just sock -> pure sock
        Nothing -> do
          let errStr =
                "TCP connection with "
                  <> Text.unpack (Text.decodeUtf8 hostname)
                  <> ":"
                  <> show port
                  <> " took longer than "
                  <> show tcpConnectTimeout
                  <> " microseconds"
          throwIO $ mkIOError TimeExpired errStr Nothing Nothing

type LiveReqs = Map Unique (Async (), MVar SomeException)

type SendReqFn = HTTP2.Request -> (HTTP2.Response -> IO ()) -> IO ()

data Transport
  = InsecureTransport NS.Socket
  | SecureTransport SSL.SSL

data TLSParams = TLSParams
  { context :: SSL.SSLContext,
    hostname :: HostName
  }

mkTransport :: NS.Socket -> Maybe TLSParams -> IO Transport
mkTransport sock Nothing = pure $ InsecureTransport sock
mkTransport sock (Just TLSParams {..}) = do
  ssl <- SSL.connection context sock
  let hostnameStr = Text.unpack $ Text.decodeUtf8 hostname
  -- Perhaps a hook at enable/disable or customize this would be nice.
  -- OpenSSL also supports a callback.
  SSL.setTlsextHostName ssl hostnameStr
  SSL.enableHostnameValidation ssl hostnameStr
  SSL.connect ssl
  pure $ SecureTransport ssl

cleanupTransport :: Transport -> IO ()
cleanupTransport (InsecureTransport _) = pure ()
cleanupTransport (SecureTransport ssl) = SSL.shutdown ssl SSL.Unidirectional

data ConnectionAlreadyClosed = ConnectionAlreadyClosed
  deriving (Show)

instance Exception ConnectionAlreadyClosed

bufsize :: HPACK.BufferSize
bufsize = 4096

allocHTTP2Config :: Transport -> IO HTTP2.Config
allocHTTP2Config (InsecureTransport sock) = do
  res <- try $ HTTP2.allocSimpleConfig sock bufsize
  case res of
    Left (e :: SomeException) -> do
      throwIO e
    Right conf -> do
      pure conf
allocHTTP2Config (SecureTransport 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 acc 0 = pure acc
      readData acc 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 (acc <> chunk)
          | chunkLen > n ->
              error "openssl: SSL.read returned more bytes than asked for, this is probably a bug"
          | otherwise ->
              readData (acc <> chunk) (n - chunkLen)
  let s = fromMaybe (error "http2-manager: SSL without socket") $ SSL.sslSocket ssl
  mysa <- NS.getSocketName s
  peersa <- NS.getPeerName s

  pure
    HTTP2.Config
      { HTTP2.confWriteBuffer = buf,
        HTTP2.confBufferSize = bufsize,
        HTTP2.confSendAll = SSL.write ssl,
        HTTP2.confReadN = readData mempty,
        HTTP2.confPositionReadMaker = HTTP2.defaultPositionReadMaker,
        HTTP2.confTimeoutManager = timmgr,
        HTTP2.confMySockAddr = mysa,
        HTTP2.confPeerSockAddr = peersa,
        HTTP2.confReadNTimeout = False
      }
