Remove modeling of client connection ID in the QUIC scheme

Author: LaurentRDC

packages/network-transport-quic/src/Network/Transport/QUIC/Internal.hs

@@ -34,13 +34,11 @@ import Control.Monad (unless, when)
 import Data.Bifunctor (Bifunctor (first))
 import Data.Binary qualified as Binary (decodeOrFail)
 import Data.ByteString (ByteString, fromStrict)
-import Data.Foldable (forM_)
 import Data.Function ((&))
 import Data.Functor ((<&>))
 import Data.IORef (newIORef, readIORef, writeIORef)
 import Data.List.NonEmpty (NonEmpty)
 import Data.Map.Strict qualified as Map
-import Data.Maybe (isNothing)
 import Lens.Micro.Platform ((+~))
 import Network.QUIC qualified as QUIC
 import Network.TLS (Credential)
@@ -62,8 +60,7 @@ import Network.Transport
   )
 import Network.Transport.QUIC.Internal.Configuration (credentialLoadX509)
 import Network.Transport.QUIC.Internal.Messaging
-  ( ClientConnId,
-    MessageReceived (..),
+  ( MessageReceived (..),
     createConnectionId,
     decodeMessage,
     encodeMessage,
@@ -101,7 +98,6 @@ import Network.Transport.QUIC.Internal.QUICTransport
     nextSelfConnOutId,
     remoteEndPointAddress,
     remoteEndPointState,
-    remoteIncoming,
     remoteServerConnId,
     remoteStream,
     transportConfig,
@@ -182,17 +178,15 @@ handleNewStream quicTransport stream = do
                     (remoteEndPoint, _) <- either throwIO pure =<< createRemoteEndPoint ourEndPoint remoteAddress Incoming
                     doneMVar <- newEmptyMVar
 
-                    clientConnId <- either (throwIO . userError) (pure . fromIntegral) =<< recvWord32 stream
                     let serverConnId = remoteServerConnId remoteEndPoint
-                        connectionId = createConnectionId serverConnId clientConnId
+                        -- One logical connection per stream; clientConnId is always 0.
+                        connectionId = createConnectionId serverConnId 0
 
                     let st =
                           RemoteEndPointValid $
                             ValidRemoteEndPointState
                               { _remoteStream = stream,
-                                _remoteStreamIsClosed = doneMVar,
-                                _remoteIncoming = Just clientConnId,
-                                _remoteNextConnOutId = 0
+                                _remoteStreamIsClosed = doneMVar
                               }
                     modifyMVar_
                       (remoteEndPoint ^. remoteEndPointState)
@@ -253,12 +247,8 @@ handleIncomingMessages ourEndPoint remoteEndPoint =
     release (Left err) = closeRemoteEndPoint Incoming remoteEndPoint >> prematureExit err
     release (Right _) = closeRemoteEndPoint Incoming remoteEndPoint
 
-    connectionId = createConnectionId serverConnId
-
-    writeConnectionClosedSTM connId =
-      writeTQueue
-        ourQueue
-        (ConnectionClosed (connectionId connId))
+    -- One logical connection per stream; clientConnId is always 0.
+    connectionId = createConnectionId serverConnId 0
 
     go = either prematureExit loop
 
@@ -268,34 +258,31 @@ handleIncomingMessages ourEndPoint remoteEndPoint =
           Left errmsg -> do
             -- Throwing will trigger 'prematureExit'
             throwIO $ userError $ "(handleIncomingMessages) Failed with: " <> errmsg
-          Right (Message connId bytes) -> handleMessage connId bytes >> loop stream
+          Right (Message bytes) -> handleMessage bytes >> loop stream
           Right StreamClosed -> throwIO $ userError "(handleIncomingMessages) Stream closed"
-          Right (CloseConnection connId) -> do
-            atomically (writeConnectionClosedSTM connId)
+          Right CloseConnection -> do
+            atomically (writeTQueue ourQueue (ConnectionClosed connectionId))
             mAct <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
               RemoteEndPointInit -> pure (RemoteEndPointClosed, Nothing)
               RemoteEndPointClosed -> pure (RemoteEndPointClosed, Nothing)
-              RemoteEndPointValid (ValidRemoteEndPointState _ isClosed _ _) -> do
+              RemoteEndPointValid (ValidRemoteEndPointState _ isClosed) -> do
                 pure (RemoteEndPointClosed, Just $ putMVar isClosed ())
             case mAct of
               Nothing -> pure ()
               Just cleanup -> cleanup
           Right CloseEndPoint -> do
-            connIds <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
-              RemoteEndPointValid vst -> do
-                pure (RemoteEndPointClosed, vst ^. remoteIncoming)
-              other -> pure (other, Nothing)
-            unless
-              (isNothing connIds)
-              ( atomically $
-                  forM_
-                    connIds
-                    (writeTQueue ourQueue . ConnectionClosed . connectionId)
-              )
+            -- handleIncomingMessages only runs on incoming remote endpoints, so if
+            -- the state was still Valid there is exactly one logical connection to
+            -- surface as closed.
+            wasValid <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
+              RemoteEndPointValid _ -> pure (RemoteEndPointClosed, True)
+              other -> pure (other, False)
+            when wasValid $
+              atomically $ writeTQueue ourQueue (ConnectionClosed connectionId)
 
-    handleMessage :: ClientConnId -> [ByteString] -> IO ()
-    handleMessage clientConnId payload =
-      atomically (writeTQueue ourQueue (Received (connectionId clientConnId) payload))
+    handleMessage :: [ByteString] -> IO ()
+    handleMessage payload =
+      atomically (writeTQueue ourQueue (Received connectionId payload))
 
     prematureExit :: IOException -> IO ()
     prematureExit exc = do
@@ -363,24 +350,24 @@ newConnection ourEndPoint creds validateCreds remoteAddress _reliability _connec
     else
       createConnectionTo creds validateCreds ourEndPoint remoteAddress >>= \case
         Left err -> pure $ Left err
-        Right (remoteEndPoint, connId) -> do
+        Right remoteEndPoint -> do
           connAlive <- newIORef True
           pure
             . Right
             $ Connection
-              { send = sendConn remoteEndPoint connAlive connId,
-                close = closeConn remoteEndPoint connAlive connId
+              { send = sendConn remoteEndPoint connAlive,
+                close = closeConn remoteEndPoint connAlive
               }
   where
     ourAddress = ourEndPoint ^. localAddress
-    sendConn remoteEndPoint connAlive connId packets =
+    sendConn remoteEndPoint connAlive packets =
       readMVar (remoteEndPoint ^. remoteEndPointState) >>= \case
         RemoteEndPointInit -> undefined
         RemoteEndPointValid vst ->
           readIORef connAlive >>= \case
             False -> pure . Left $ TransportError SendClosed "Connection closed"
             True ->
-              sendMessage (vst ^. remoteStream) connId packets
+              sendMessage (vst ^. remoteStream) packets
                 <&> first (TransportError SendFailed . show)
         RemoteEndPointClosed -> do
           readIORef connAlive >>= \case
@@ -390,9 +377,9 @@ newConnection ourEndPoint creds validateCreds remoteAddress _reliability _connec
             -- 'connAlive' IORefs.
             False -> pure . Left $ TransportError SendClosed "Connection closed"
             True -> pure . Left $ TransportError SendFailed "Remote endpoint closed"
-    closeConn remoteEndPoint connAlive connId = do
+    closeConn remoteEndPoint connAlive = do
       mCleanup <- modifyMVar (remoteEndPoint ^. remoteEndPointState) $ \case
-        RemoteEndPointValid vst@(ValidRemoteEndPointState stream isClosed _ _) -> do
+        RemoteEndPointValid vst@(ValidRemoteEndPointState stream isClosed) -> do
           readIORef connAlive >>= \case
             False -> pure (RemoteEndPointValid vst, Nothing)
             True -> do
@@ -401,7 +388,7 @@ newConnection ourEndPoint creds validateCreds remoteAddress _reliability _connec
               -- safe against races with the finally in streamToEndpoint that can
               -- also signal isClosed on QUIC.Client.run exit.
               let cleanup = do
-                    _ <- sendCloseConnection connId stream
+                    _ <- sendCloseConnection stream
                     _ <- tryPutMVar isClosed ()
                     pure ()
               pure (RemoteEndPointClosed, Just cleanup)

packages/network-transport-quic/src/Network/Transport/QUIC/Internal/Client.hs

@@ -17,7 +17,7 @@ import Network.QUIC qualified as QUIC
 import Network.QUIC.Client qualified as QUIC.Client
 import Network.Transport (ConnectErrorCode (ConnectNotFound), EndPointAddress, TransportError (..))
 import Network.Transport.QUIC.Internal.Configuration (Credential, mkClientConfig)
-import Network.Transport.QUIC.Internal.Messaging (ClientConnId, MessageReceived (..), handshake, receiveMessage)
+import Network.Transport.QUIC.Internal.Messaging (MessageReceived (..), handshake, receiveMessage)
 import Network.Transport.QUIC.Internal.QUICAddr (QUICAddr (QUICAddr), decodeQUICAddr)
 
 streamToEndpoint ::
@@ -28,8 +28,6 @@ streamToEndpoint ::
   EndPointAddress ->
   -- | Their address
   EndPointAddress ->
-  -- | Client-allocated connection id to send as part of the handshake
-  ClientConnId ->
   -- | Called when the QUIC connection or stream ends without us having initiated the
   -- close. Must be idempotent (the caller typically gates on remote endpoint state so
   -- that repeated invocations are safe) — this handler is invoked from multiple sites
@@ -44,7 +42,7 @@ streamToEndpoint ::
           QUIC.Stream
         )
     )
-streamToEndpoint creds validateCreds ourAddress theirAddress clientConnId onConnLoss =
+streamToEndpoint creds validateCreds ourAddress theirAddress onConnLoss =
   case decodeQUICAddr theirAddress of
     Left errmsg -> pure $ Left (TransportError ConnectNotFound errmsg)
     Right (QUICAddr hostname servicename _) -> do
@@ -58,7 +56,7 @@ streamToEndpoint creds validateCreds ourAddress theirAddress clientConnId onConn
             QUIC.waitEstablished conn
             restore $
               bracket (QUIC.stream conn) QUIC.closeStream $ \stream -> do
-                handshake (ourAddress, theirAddress) clientConnId stream
+                handshake (ourAddress, theirAddress) stream
                   >>= either
                     (\_ -> putMVar streamMVar (Left $ TransportError ConnectNotFound "handshake failed"))
                     (\_ -> putMVar streamMVar (Right stream))
@@ -104,7 +102,7 @@ streamToEndpoint creds validateCreds ourAddress theirAddress clientConnId onConn
         Right StreamClosed -> mask_ $ do
           _ <- tryPutMVar doneMVar ()
           onConnLoss
-        Right (CloseConnection _) ->
+        Right CloseConnection ->
           -- Peer closed the logical connection cleanly; no ErrorEvent.
           () <$ tryPutMVar doneMVar ()
         Right CloseEndPoint -> mask_ $ do

packages/network-transport-quic/src/Network/Transport/QUIC/Internal/Messaging.hs

@@ -50,20 +50,19 @@ import Network.Transport (ConnectionId, EndPointAddress)
 import Network.Transport.Internal (decodeWord32, encodeWord32)
 import Network.Transport.QUIC.Internal.QUICAddr (QUICAddr (QUICAddr), decodeQUICAddr)
 
--- | Send a message to a remote endpoint ID
+-- | Send a message on the stream.
 --
 -- This function is thread-safe; while the data is sending, asynchronous
 -- exceptions are masked, to be rethrown after the data is sent.
 sendMessage ::
   Stream ->
-  ClientConnId ->
   [ByteString] ->
   IO (Either QUIC.QUICException ())
-sendMessage stream connId messages =
+sendMessage stream messages =
   try
     ( QUIC.sendStreamMany
         stream
-        (encodeMessage connId messages)
+        (encodeMessage messages)
     )
 
 -- | Receive a message, including its local destination endpoint ID
@@ -87,18 +86,15 @@ receiveMessage stream = mask $ \restore ->
 -- The encoding is composed of a header, and the payloads.
 -- The message header is composed of:
 -- 1. A control byte, to determine how the message should be parsed.
--- 2. A 32-bit word that encodes the endpoint ID of the destination endpoint;
--- 3. A 32-bit word that encodes the number of frames in the message
+-- 2. A 32-bit word that encodes the number of frames in the message
 --
 -- The payload frames are each prepended with the length of the frame.
 encodeMessage ::
-  ClientConnId ->
   [ByteString] ->
   [ByteString]
-encodeMessage connId messages =
+encodeMessage messages =
   BS.concat
     [ BS.singleton messageControlByte,
-      encodeWord32 (fromIntegral connId),
       encodeWord32 (fromIntegral $ length messages)
     ]
     : [encodeWord32 (fromIntegral $ BS.length message) <> message | message <- messages]
@@ -116,13 +112,12 @@ decodeMessage get =
   where
     go ctrl
       | ctrl == closeEndPointControlByte = pure $ Right CloseEndPoint
-      | ctrl == closeConnectionControlByte = Right . CloseConnection . fromIntegral <$> getWord32
+      | ctrl == closeConnectionControlByte = pure $ Right CloseConnection
       | ctrl == messageControlByte = do
-          connId <- getWord32
           numMessages <- getWord32
           messages <- replicateM (fromIntegral numMessages) $ do
             getWord32 >>= get . fromIntegral
-          pure . Right $ Message (fromIntegral connId) messages
+          pure . Right $ Message messages
       | otherwise = pure $ Left $ "Unsupported control byte: " <> show ctrl
     getWord32 = get 4 <&> decodeWord32
 
@@ -146,10 +141,8 @@ getAllBytes get n = go n mempty
           else go (m - BS.length bytes) (acc <> [bytes])
 
 data MessageReceived
-  = Message
-      {-# UNPACK #-} !ClientConnId
-      {-# UNPACK #-} ![ByteString]
-  | CloseConnection !ClientConnId
+  = Message {-# UNPACK #-} ![ByteString]
+  | CloseConnection
   | CloseEndPoint
   | StreamClosed
   deriving (Show, Eq)
@@ -220,13 +213,12 @@ closeConnectionControlByte :: ControlByte
 closeConnectionControlByte = 255
 
 -- | Send a message to close the connection.
-sendCloseConnection :: ClientConnId -> Stream -> IO (Either QUIC.QUICException ())
-sendCloseConnection connId stream =
+sendCloseConnection :: Stream -> IO (Either QUIC.QUICException ())
+sendCloseConnection stream =
   try
     ( QUIC.sendStream
         stream
-        ( BS.concat [BS.singleton closeConnectionControlByte, encodeWord32 (fromIntegral connId)]
-        )
+        (BS.singleton closeConnectionControlByte)
     )
 
 -- | Send a message to close the connection.
@@ -240,23 +232,21 @@ sendCloseEndPoint stream =
     )
 
 -- | Handshake protocol that a client, connecting to a remote endpoint,
--- has to perform. Two round-trips:
+-- has to perform:
 --
 -- 1. client -> server: address payload
 -- 2. server -> client: ack1 (handshake payload accepted)
--- 3. client -> server: clientConnId
--- 4. server -> client: ack2 (ConnectionOpened has been enqueued on server's endpoint)
+-- 3. server -> client: ack2 (ConnectionOpened has been enqueued on server's endpoint)
 --
 -- The ack2 step is load-bearing: when this function returns, the server has
 -- already written ConnectionOpened to its local queue. Without it, the server's
 -- @connect@ would return before the peer's queue has the ConnectionOpened event,
 -- which races with subsequent sends on other connections.
 handshake ::
   (EndPointAddress, EndPointAddress) ->
-  ClientConnId ->
   Stream ->
   IO (Either () ())
-handshake (ourAddress, theirAddress) clientConnId stream =
+handshake (ourAddress, theirAddress) stream =
   case decodeQUICAddr theirAddress of
     Left errmsg -> throwIO $ userError ("Could not decode QUIC address: " <> errmsg)
     Right (QUICAddr _ _ serverEndPointId) -> do
@@ -273,12 +263,7 @@ handshake (ourAddress, theirAddress) clientConnId stream =
           Right _ ->
             recvAck stream >>= \case
               Left () -> pure $ Left ()
-              Right () ->
-                try @SomeException
-                  (QUIC.sendStream stream (BS.toStrict (Binary.encode clientConnId)))
-                  >>= \case
-                    Left _ -> pure $ Left ()
-                    Right () -> recvAck stream
+              Right () -> recvAck stream
 
 -- | Part of the connection ID that is client-allocated.
 newtype ClientConnId = ClientConnId Word32