QUICTransport.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}

module Network.Transport.QUIC.Internal.QUICTransport
  ( -- * QUICTransport
    QUICTransport,
    newQUICTransport,
    foldOpenEndPoints,
    transportConfig,
    transportInputSocket,
    transportState,

    -- ** Configuration
    QUICTransportConfig (..),
    defaultQUICTransportConfig,

    -- * TransportState
    TransportState (..),
    localEndPoints,
    nextEndPointId,

    -- * LocalEndPoint
    LocalEndPoint,
    localAddress,
    localEndPointId,
    localEndPointState,
    localQueue,
    nextConnInId,
    nextSelfConnOutId,
    newLocalEndPoint,
    closeLocalEndpoint,

    -- * LocalEndPointState
    LocalEndPointState (..),
    ValidLocalEndPointState,
    incomingConnections,
    outgoingConnections,
    nextConnectionCounter,

    -- ** ConnectionCounter
    ConnectionCounter,

    -- * RemoteEndPoint
    RemoteEndPoint (..),
    remoteEndPointAddress,
    remoteEndPointId,
    remoteServerConnId,
    remoteEndPointState,
    closeRemoteEndPoint,
    createRemoteEndPoint,
    createConnectionTo,

    -- ** Remote endpoint state
    RemoteEndPointState (..),
    ValidRemoteEndPointState (..),
    remoteStream,
    remoteStreamIsClosed,
    Direction (..),

    -- * Re-exports
    (^.),
  )
where

import Control.Concurrent.Async (forConcurrently_)
import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar, readMVar, tryPutMVar)
import Control.Concurrent.STM.TQueue (TQueue, writeTQueue)
import Control.Exception (bracketOnError)
import Control.Monad (forM_)
import Control.Monad.STM (atomically)
import Data.Function ((&))
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Word (Word32)
import Lens.Micro.Platform (makeLenses, (%~), (+~), (^.))
import Network.QUIC (Stream)
import Network.Socket (HostName, ServiceName, Socket)
import Network.Socket qualified as N
import Network.TLS (Credential)
import Network.Transport (ConnectErrorCode (ConnectFailed), EndPointAddress, Event (EndPointClosed, ErrorEvent), EventErrorCode (EventConnectionLost), NewEndPointErrorCode (NewEndPointFailed), TransportError (TransportError))
import Network.Transport.QUIC.Internal.Client (streamToEndpoint)
import Network.Transport.QUIC.Internal.Messaging
  ( ClientConnId,
    ServerConnId,
    firstNonReservedServerConnId,
    sendCloseConnection,
    sendCloseEndPoint,
  )
import Network.Transport.QUIC.Internal.QUICAddr (EndPointId, QUICAddr (..), encodeQUICAddr)

{- The QUIC transport has three levels of statefullness:

1. The transport itself

The transport contains state required to create new endpoints, and close them. This includes,
for example, a container of existing endpoints.

2. Endpoints

An endpoint has some state regarding the connections it has. An endpoint may have zero or more
connection, and must have state to be able to create new connections, and close existing ones.

3. Connections

Finally, each connection between endpoint has some state, needed to receive data.
-}

-- | Represents the configuration used by the entire transport.
data QUICTransportConfig = QUICTransportConfig
  { -- | Host name, which can be an IP address or a domain name.
    hostName :: HostName,
    -- | Port or service name. The default is port 443.
    serviceName :: ServiceName,
    -- | At least one set of credentials is required.
    credentials :: NonEmpty Credential,
    -- | Note that if your credentials is self-signed, you will have
    -- to turn off 'validateCredentials'. This should only be set to 'False'
    -- in tests, or in a private network.
    validateCredentials :: Bool
  }
  deriving (Eq, Show)

defaultQUICTransportConfig :: HostName -> NonEmpty Credential -> QUICTransportConfig
defaultQUICTransportConfig host creds =
  QUICTransportConfig
    { hostName = host,
      serviceName = "443",
      credentials = creds,
      validateCredentials = True
    }

data QUICTransport = QUICTransport
  { _transportConfig :: QUICTransportConfig,
    _transportInputSocket :: Socket,
    _transportState :: MVar TransportState
  }

data TransportState
  = TransportStateValid ValidTransportState
  | TransportStateClosed

data ValidTransportState = ValidTransportState
  { _localEndPoints :: !(Map EndPointId LocalEndPoint),
    _nextEndPointId :: !EndPointId
  }

-- | Create a new QUICTransport
newQUICTransport :: QUICTransportConfig -> IO QUICTransport
newQUICTransport config = do
  addr <- NE.head <$> N.getAddrInfo (Just N.defaultHints) (Just (hostName config)) (Just (serviceName config))
  bracketOnError
    ( N.socket
        (N.addrFamily addr)
        N.Datagram -- QUIC is based on UDP
        N.defaultProtocol
    )
    N.close
    $ \socket -> do
      N.setSocketOption socket N.ReuseAddr 1
      N.withFdSocket socket N.setCloseOnExecIfNeeded
      N.bind socket (N.addrAddress addr)

      port <- N.socketPort socket
      QUICTransport
        config{serviceName=show port}
        socket
        <$> newMVar (TransportStateValid $ ValidTransportState mempty 1)

data LocalEndPoint = OpenLocalEndPoint
  { _localAddress :: !EndPointAddress,
    _localEndPointId :: !EndPointId,
    _localEndPointState :: !(MVar LocalEndPointState),
    -- | Queue used to receive events
    _localQueue :: !(TQueue Event)
  }

-- | A 'ConnectionCounter' uniquely identifies a connections within the context of an endpoint.
-- This allows to hold multiple separate connections between two endpoint addresses.
--
-- NOTE: I tried to use the `StreamId` type from the `quic` library, but it was
-- clearly not unique per stream. I don't understand if this was intentional or not.
newtype ConnectionCounter = ConnectionCounter Word32
  deriving newtype (Eq, Show, Ord, Bounded, Enum, Real, Integral, Num)

data LocalEndPointState
  = LocalEndPointStateValid ValidLocalEndPointState
  | LocalEndPointStateClosed
  deriving (Show)

data ValidLocalEndPointState = ValidLocalEndPointState
  { _incomingConnections :: Map (EndPointAddress, ConnectionCounter) RemoteEndPoint,
    _outgoingConnections :: Map (EndPointAddress, ConnectionCounter) RemoteEndPoint,
    _nextSelfConnOutId :: !ClientConnId,
    -- | We identify connections by remote endpoint address, AND ConnectionCounter,
    --    to support multiple connections between the same two endpoint addresses
    _nextConnInId :: !ServerConnId,
    _nextConnectionCounter :: ConnectionCounter
  }
  deriving (Show)

data RemoteEndPoint = RemoteEndPoint
  { _remoteEndPointAddress :: !EndPointAddress,
    _remoteEndPointId :: !EndPointId,
    _remoteEndPointState :: !(MVar RemoteEndPointState)
  }

remoteServerConnId :: RemoteEndPoint -> ServerConnId
remoteServerConnId = fromIntegral . _remoteEndPointId

instance Show RemoteEndPoint where
  show (RemoteEndPoint address _ _) = "<RemoteEndPoint @ " <> show address <> ">"

data RemoteEndPointState
  = -- | In the short window between a connection
    --      being initiated and the handshake completing
    RemoteEndPointInit
  | RemoteEndPointValid ValidRemoteEndPointState
  | RemoteEndPointClosed

data ValidRemoteEndPointState = ValidRemoteEndPointState
  { _remoteStream :: Stream,
    _remoteStreamIsClosed :: MVar ()
  }

makeLenses ''QUICTransport
makeLenses ''TransportState
makeLenses ''ValidTransportState
makeLenses ''LocalEndPoint
makeLenses ''LocalEndPointState
makeLenses ''ValidLocalEndPointState
makeLenses ''RemoteEndPoint
makeLenses ''ValidRemoteEndPointState

-- | Fold over all open local endpoitns of a transport
foldOpenEndPoints :: QUICTransport -> (LocalEndPoint -> IO a) -> IO [a]
foldOpenEndPoints quicTransport f =
  readMVar (quicTransport ^. transportState) >>= \case
    TransportStateClosed -> pure []
    TransportStateValid st ->
      mapM f (Map.elems $ st ^. localEndPoints)

newLocalEndPoint :: QUICTransport -> TQueue Event -> IO (Either (TransportError NewEndPointErrorCode) LocalEndPoint)
newLocalEndPoint quicTransport newLocalQueue = do
  modifyMVar (quicTransport ^. transportState) $ \case
    TransportStateClosed -> pure (TransportStateClosed, Left $ TransportError NewEndPointFailed "Transport closed")
    TransportStateValid validState -> do
      let newEndPointId = validState ^. nextEndPointId

      newLocalState <-
        newMVar
          ( LocalEndPointStateValid $
              ValidLocalEndPointState
                { _incomingConnections = mempty,
                  _outgoingConnections = mempty,
                  _nextConnInId = firstNonReservedServerConnId,
                  _nextSelfConnOutId = 0,
                  _nextConnectionCounter = 0
                }
          )
      let openEndpoint =
            OpenLocalEndPoint
              { _localAddress =
                  encodeQUICAddr
                    ( QUICAddr
                        (hostName $ quicTransport ^. transportConfig)
                        (serviceName $ quicTransport ^. transportConfig)
                        newEndPointId
                    ),
                _localEndPointId = newEndPointId,
                _localEndPointState = newLocalState,
                _localQueue = newLocalQueue
              }

      pure
        ( TransportStateValid
            ( validState
                & localEndPoints %~ Map.insert newEndPointId openEndpoint
                & nextEndPointId +~ 1
            ),
          Right openEndpoint
        )

closeLocalEndpoint ::
  QUICTransport ->
  LocalEndPoint ->
  IO ()
closeLocalEndpoint quicTransport localEndPoint = do
  modifyMVar_ (quicTransport ^. transportState) $ \case
    TransportStateClosed -> pure TransportStateClosed
    TransportStateValid vst ->
      pure . TransportStateValid $
        vst
          & localEndPoints
            %~ Map.delete (localEndPoint ^. localEndPointId)

  mPreviousState <- modifyMVar (localEndPoint ^. localEndPointState) $ \case
    LocalEndPointStateClosed -> pure (LocalEndPointStateClosed, Nothing)
    LocalEndPointStateValid st -> pure (LocalEndPointStateClosed, Just st)

  -- Close outgoing remote endpoints before incoming. The peer's handleIncomingMessages
  -- reader writes ConnectionClosed in response to our outgoing close; its listenForClose
  -- writes ErrorEvent in response to our incoming close. Processing outgoing first gives
  -- the peer's event queue the expected ConnectionClosed-before-ErrorEvent ordering.
  forM_ mPreviousState $ \vst -> do
    forConcurrently_ (vst ^. outgoingConnections) tryCloseRemoteStream
    forConcurrently_ (vst ^. incomingConnections) tryCloseRemoteStream
  atomically $ writeTQueue (localEndPoint ^. localQueue) EndPointClosed
  where
    tryCloseRemoteStream :: RemoteEndPoint -> IO ()
    tryCloseRemoteStream remoteEndPoint = do
      mCleanup <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
        RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing)
        RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing)
        RemoteEndPointValid vst ->
          pure
            ( RemoteEndPointClosed,
              Just $ do
                _ <- sendCloseEndPoint (vst ^. remoteStream)
                _ <- tryPutMVar (vst ^. remoteStreamIsClosed) ()
                pure ()
            )

      case mCleanup of
        Nothing -> pure ()
        Just cleanup -> cleanup

-- | Attempt to close a remote endpoint. If the remote endpoint is in
-- any non-valid state (e.g. already closed), then nothing happens.
--
-- Otherwise, a control message is sent to the remote end to nicely ask to
-- close this connection.
closeRemoteEndPoint :: Direction -> RemoteEndPoint -> IO ()
closeRemoteEndPoint direction remoteEndPoint = do
  mAct <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
    RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing)
    RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing)
    RemoteEndPointValid (ValidRemoteEndPointState stream isClosed) ->
      let cleanup = do
            _ <- case direction of
              Outgoing -> sendCloseConnection stream
              Incoming -> sendCloseEndPoint stream
            _ <- tryPutMVar isClosed ()
            pure ()
       in pure (RemoteEndPointClosed, Just cleanup)

  case mAct of
    Nothing -> pure ()
    Just act -> act

data Direction
  = Outgoing
  | Incoming
  deriving (Eq, Show, Ord, Enum, Bounded)

-- | Create a remote end point in the 'init' state.
--
-- The resulting remote end point is NOT set up, such that
-- it could be set up separately to /receive/ messages, or /send/ them.
createRemoteEndPoint ::
  LocalEndPoint ->
  EndPointAddress ->
  Direction ->
  IO (Either (TransportError ConnectErrorCode) (RemoteEndPoint, ConnectionCounter))
createRemoteEndPoint localEndPoint remoteAddress direction = do
  modifyMVar (localEndPoint ^. localEndPointState) $ \case
    LocalEndPointStateClosed -> pure (LocalEndPointStateClosed, Left $ TransportError ConnectFailed "endpoint is closed")
    LocalEndPointStateValid st -> do
      remoteEndPoint <-
        RemoteEndPoint
          remoteAddress
          -- The design of using the next Server connection ID
          -- as the RemoteId comes from the TCP transport

          (fromIntegral $ st ^. nextConnInId)
          <$> newMVar RemoteEndPointInit
      pure
        ( LocalEndPointStateValid $
            st
              & (if direction == Incoming then incomingConnections else outgoingConnections) %~ Map.insert (remoteAddress, st ^. nextConnectionCounter) remoteEndPoint
              & nextConnectionCounter +~ 1
              & nextConnInId +~ 1,
          Right (remoteEndPoint, st ^. nextConnectionCounter)
        )

-- | Create a remote end point, set up as a client that connects
-- to the remote 'EndPointAddress'.
createConnectionTo ::
  NonEmpty Credential ->
  -- | Validate credentials
  Bool ->
  LocalEndPoint ->
  EndPointAddress ->
  IO (Either (TransportError ConnectErrorCode) RemoteEndPoint)
createConnectionTo creds validateCreds localEndPoint remoteAddress = do
  createRemoteEndPoint localEndPoint remoteAddress Outgoing >>= \case
    Left err -> pure $ Left err
    Right (remoteEndPoint, _) -> do
      -- TODO: each call to @connect@ currently opens a dedicated QUIC connection
      -- and carries a single logical connection on its stream. Preferred
      -- architecture: one QUIC connection per (local endpoint, peer endpoint)
      -- pair, with each logical connection carried on its own stream. Streams
      -- already give us independent flow control and avoid head-of-line blocking.
      streamToEndpoint
        creds
        validateCreds
        (localEndPoint ^. localAddress)
        remoteAddress
        (surfaceConnectionLost remoteEndPoint)
        >>= \case
          Left exc -> pure $ Left exc
          Right (closeStream, stream) -> do
            let validState =
                  RemoteEndPointValid $
                    ValidRemoteEndPointState
                      { _remoteStream = stream,
                        _remoteStreamIsClosed = closeStream
                      }
            modifyMVar_
              (remoteEndPoint ^. remoteEndPointState)
              (\_ -> pure validState)
            pure $ Right remoteEndPoint
  where
    -- Idempotent: surfaces EventConnectionLost exactly once, only if the remote
    -- endpoint was still Valid when invoked. Called from multiple termination
    -- sites (peer-initiated close, QUIC exception, forked-thread finally) so that
    -- no close path can leave us silent — the state-transition gate dedupes them.
    surfaceConnectionLost remoteEndPoint = do
      mAct <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
        RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing)
        RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing)
        RemoteEndPointValid (ValidRemoteEndPointState stream isClosed) ->
          let cleanup = do
                _ <- sendCloseConnection stream
                _ <- tryPutMVar isClosed ()
                onConnectionLost
           in pure (RemoteEndPointClosed, Just cleanup)
      case mAct of
        Nothing -> pure ()
        Just act -> act
    onConnectionLost =
      atomically
        . writeTQueue (localEndPoint ^. localQueue)
        . ErrorEvent
        $ TransportError
          (EventConnectionLost remoteAddress)
          "Connection reset"