-- | Yamux session management: create, openStream, acceptStream, ping, goaway.
--
-- Implements the session-level Yamux protocol per HashiCorp yamux spec.md.
-- The session manages a collection of multiplexed streams over a single
-- underlying transport connection.
--
-- Two background loops run per session:
--   recvLoop: reads 12-byte headers from transport, dispatches to streams
--   sendLoop: dequeues from ysessSendCh, writes to transport
module LibP2P.Yamux.Session
  ( newSession
  , closeSession
  , openStream
  , acceptStream
  , ping
  , sendGoAway
  , recvLoop
  , sendLoop
  , acceptBacklog
  ) where

import Control.Concurrent.STM
import Control.Exception (finally)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map
import Data.Word (Word32)
import LibP2P.Yamux.Frame
import LibP2P.Yamux.Types
import Numeric.Natural (Natural)

-- | Maximum number of inbound streams buffered while the application
-- is not accepting (go-yamux AcceptBacklog). The spec requires this
-- buffer to be bounded to mitigate DoS; excess SYNs are reset.
acceptBacklog :: Natural
acceptBacklog :: Natural
acceptBacklog = Natural
256

-- | Create a new Yamux session over a transport connection.
-- Client uses odd stream IDs starting at 1, server uses even starting at 2.
newSession :: SessionRole -> (ByteString -> IO ()) -> (Int -> IO ByteString) -> IO YamuxSession
newSession :: SessionRole
-> (ByteString -> IO ())
-> (Int -> IO ByteString)
-> IO YamuxSession
newSession SessionRole
role ByteString -> IO ()
writeFn Int -> IO ByteString
readFn = do
  let startId :: Word32
startId = case SessionRole
role of
        SessionRole
RoleClient -> Word32
1
        SessionRole
RoleServer -> Word32
2
  nextId <- Word32 -> IO (TVar Word32)
forall a. a -> IO (TVar a)
newTVarIO Word32
startId
  streams <- newTVarIO Map.empty
  acceptCh <- newTBQueueIO acceptBacklog
  sendCh <- newTQueueIO
  shutdown <- newTVarIO False
  remoteGoAway <- newTVarIO Nothing
  pings <- newTVarIO Map.empty
  nextPingId <- newTVarIO 1
  pure
    YamuxSession
      { ysessRole = role
      , ysessNextStreamId = nextId
      , ysessStreams = streams
      , ysessAcceptCh = acceptCh
      , ysessSendCh = sendCh
      , ysessShutdown = shutdown
      , ysessRemoteGoAway = remoteGoAway
      , ysessPings = pings
      , ysessNextPingId = nextPingId
      , ysessWrite = writeFn
      , ysessRead = readFn
      }

-- | Gracefully close the session by sending GoAway Normal.
closeSession :: YamuxSession -> IO ()
closeSession :: YamuxSession -> IO ()
closeSession YamuxSession
sess = YamuxSession -> GoAwayCode -> IO ()
sendGoAway YamuxSession
sess GoAwayCode
GoAwayNormal

-- | Open a new outbound stream. Allocates the next stream ID and sends SYN.
-- Returns YamuxSessionShutdown after a local GoAway, or YamuxGoAway with
-- the received code after a remote GoAway.
openStream :: YamuxSession -> IO (Either YamuxError YamuxStream)
openStream :: YamuxSession -> IO (Either YamuxError YamuxStream)
openStream YamuxSession
sess = do
  -- Check shutdown state
  status <- STM (Either YamuxError ()) -> IO (Either YamuxError ())
forall a. STM a -> IO a
atomically (STM (Either YamuxError ()) -> IO (Either YamuxError ()))
-> STM (Either YamuxError ()) -> IO (Either YamuxError ())
forall a b. (a -> b) -> a -> b
$ do
    shut <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar Bool
ysessShutdown YamuxSession
sess)
    remote <- readTVar (ysessRemoteGoAway sess)
    pure $ case (shut, remote) of
      (Bool
True, Maybe GoAwayCode
_) -> YamuxError -> Either YamuxError ()
forall a b. a -> Either a b
Left YamuxError
YamuxSessionShutdown
      (Bool
False, Just GoAwayCode
code) -> YamuxError -> Either YamuxError ()
forall a b. a -> Either a b
Left (GoAwayCode -> YamuxError
YamuxGoAway GoAwayCode
code)
      (Bool
False, Maybe GoAwayCode
Nothing) -> () -> Either YamuxError ()
forall a b. b -> Either a b
Right ()
  case status of
    Left YamuxError
err -> Either YamuxError YamuxStream -> IO (Either YamuxError YamuxStream)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (YamuxError -> Either YamuxError YamuxStream
forall a b. a -> Either a b
Left YamuxError
err)
    Right () -> do
      -- Allocate stream ID (atomically increment by 2)
      sid <- STM Word32 -> IO Word32
forall a. STM a -> IO a
atomically (STM Word32 -> IO Word32) -> STM Word32 -> IO Word32
forall a b. (a -> b) -> a -> b
$ do
        nextId <- TVar Word32 -> STM Word32
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar Word32
ysessNextStreamId YamuxSession
sess)
        writeTVar (ysessNextStreamId sess) (nextId + 2)
        pure nextId
      -- Create stream in SYNSent state
      stream <- newStream sess sid StreamSYNSent
      -- Register stream
      atomically $ modifyTVar' (ysessStreams sess) (Map.insert sid stream)
      -- Send SYN frame (Data frame with SYN flag, no payload)
      let hdr =
            YamuxHeader
              { yhVersion :: Word8
yhVersion = Word8
0
              , yhType :: FrameType
yhType = FrameType
FrameData
              , yhFlags :: Flags
yhFlags = Flags
defaultFlags {flagSYN = True}
              , yhStreamId :: Word32
yhStreamId = Word32
sid
              , yhLength :: Word32
yhLength = Word32
0
              }
      atomically $ writeTQueue (ysessSendCh sess) (hdr, BS.empty)
      pure (Right stream)

-- | Accept an inbound stream. Blocks until a remote SYN arrives.
-- Returns YamuxSessionShutdown if the session is shut down.
acceptStream :: YamuxSession -> IO (Either YamuxError YamuxStream)
acceptStream :: YamuxSession -> IO (Either YamuxError YamuxStream)
acceptStream YamuxSession
sess = do
  stream <- STM YamuxStream -> IO YamuxStream
forall a. STM a -> IO a
atomically (STM YamuxStream -> IO YamuxStream)
-> STM YamuxStream -> IO YamuxStream
forall a b. (a -> b) -> a -> b
$ TBQueue YamuxStream -> STM YamuxStream
forall a. TBQueue a -> STM a
readTBQueue (YamuxSession -> TBQueue YamuxStream
ysessAcceptCh YamuxSession
sess)
  -- Send ACK (WindowUpdate frame with ACK flag)
  let hdr =
        YamuxHeader
          { yhVersion :: Word8
yhVersion = Word8
0
          , yhType :: FrameType
yhType = FrameType
FrameWindowUpdate
          , yhFlags :: Flags
yhFlags = Flags
defaultFlags {flagACK = True}
          , yhStreamId :: Word32
yhStreamId = YamuxStream -> Word32
ysStreamId YamuxStream
stream
          , yhLength :: Word32
yhLength = Word32
0
          }
  atomically $ writeTQueue (ysessSendCh sess) (hdr, BS.empty)
  -- Transition to Established only from SYNReceived. The remote may have
  -- already half-closed (FIN) before we accepted; that state must survive.
  atomically $ do
    st <- readTVar (ysState stream)
    case st of
      StreamState
StreamSYNReceived -> TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamEstablished
      StreamState
_ -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  pure (Right stream)

-- | Send a Ping and wait for the ACK response.
-- Ping uses StreamID 0 and the Length field carries an opaque value.
ping :: YamuxSession -> IO (Either YamuxError ())
ping :: YamuxSession -> IO (Either YamuxError ())
ping YamuxSession
sess = do
  (pingId, waiter) <- STM (Word32, PingWaiter) -> IO (Word32, PingWaiter)
forall a. STM a -> IO a
atomically (STM (Word32, PingWaiter) -> IO (Word32, PingWaiter))
-> STM (Word32, PingWaiter) -> IO (Word32, PingWaiter)
forall a b. (a -> b) -> a -> b
$ do
    pid <- TVar Word32 -> STM Word32
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar Word32
ysessNextPingId YamuxSession
sess)
    writeTVar (ysessNextPingId sess) (pid + 1)
    w <- newEmptyTMVar
    modifyTVar' (ysessPings sess) (Map.insert pid w)
    pure (pid, w)
  -- Send Ping SYN frame
  let hdr =
        YamuxHeader
          { yhVersion :: Word8
yhVersion = Word8
0
          , yhType :: FrameType
yhType = FrameType
FramePing
          , yhFlags :: Flags
yhFlags = Flags
defaultFlags {flagSYN = True}
          , yhStreamId :: Word32
yhStreamId = Word32
0
          , yhLength :: Word32
yhLength = Word32
pingId
          }
  atomically $ writeTQueue (ysessSendCh sess) (hdr, BS.empty)
  -- Wait for ACK (or a session-failure notification)
  result <- atomically $ takeTMVar waiter
  -- Cleanup
  atomically $ modifyTVar' (ysessPings sess) (Map.delete pingId)
  pure result

-- | Send a GoAway frame with the specified error code.
-- Sets ysessShutdown to True so no new streams can be opened.
sendGoAway :: YamuxSession -> GoAwayCode -> IO ()
sendGoAway :: YamuxSession -> GoAwayCode -> IO ()
sendGoAway YamuxSession
sess GoAwayCode
code = do
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxSession -> TVar Bool
ysessShutdown YamuxSession
sess) Bool
True
  let errCode :: Word32
errCode = GoAwayCode -> Word32
goAwayCodeToWord32 GoAwayCode
code
  let hdr :: YamuxHeader
hdr =
        YamuxHeader
          { yhVersion :: Word8
yhVersion = Word8
0
          , yhType :: FrameType
yhType = FrameType
FrameGoAway
          , yhFlags :: Flags
yhFlags = Flags
defaultFlags
          , yhStreamId :: Word32
yhStreamId = Word32
0
          , yhLength :: Word32
yhLength = Word32
errCode
          }
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TQueue (YamuxHeader, ByteString)
-> (YamuxHeader, ByteString) -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue (YamuxSession -> TQueue (YamuxHeader, ByteString)
ysessSendCh YamuxSession
sess) (YamuxHeader
hdr, ByteString
BS.empty)

-- | Receive loop: reads 12-byte headers from transport and dispatches frames.
-- This loop runs until the transport connection is closed or an error occurs.
--
-- Whenever the loop terminates -- transport EOF or error (ysessRead
-- throws), a fatal protocol error, or cancellation -- the session is
-- torn down via failSession so that no reader, writer or ping waiter
-- is left blocked forever on a session that can no longer make
-- progress.
recvLoop :: YamuxSession -> IO ()
recvLoop :: YamuxSession -> IO ()
recvLoop YamuxSession
sess = IO ()
go IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` YamuxSession -> IO ()
failSession YamuxSession
sess
  where
    go :: IO ()
go = do
      -- Read 12-byte header
      headerBytes <- YamuxSession -> Int -> IO ByteString
ysessRead YamuxSession
sess Int
headerSize
      case decodeHeader headerBytes of
        -- Malformed header (unknown frame type): tell the peer why we
        -- are leaving before terminating, as go-yamux does
        Left String
_err -> YamuxSession -> GoAwayCode -> IO ()
sendGoAway YamuxSession
sess GoAwayCode
GoAwayProtocol
        Right YamuxHeader
hdr ->
          -- Verify version
          if YamuxHeader -> Word8
yhVersion YamuxHeader
hdr Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0
            then YamuxSession -> GoAwayCode -> IO ()
sendGoAway YamuxSession
sess GoAwayCode
GoAwayProtocol
            else do
              continue <- YamuxSession -> YamuxHeader -> IO Bool
dispatchFrame YamuxSession
sess YamuxHeader
hdr
              when continue go

-- | Tear down the session once the receive loop can no longer make
-- progress (transport EOF mid-frame, transport error, or a fatal
-- protocol error). Mirrors go-yamux, which closes every stream with
-- the session error on exit: new streams are refused, every registered
-- stream is reset so blocked readers and writers observe an error
-- instead of hanging, and every pending ping fails with
-- YamuxSessionShutdown.
failSession :: YamuxSession -> IO ()
failSession :: YamuxSession -> IO ()
failSession YamuxSession
sess = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxSession -> TVar Bool
ysessShutdown YamuxSession
sess) Bool
True
  streams <- TVar (Map Word32 YamuxStream) -> STM (Map Word32 YamuxStream)
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar (Map Word32 YamuxStream)
ysessStreams YamuxSession
sess)
  mapM_ (\YamuxStream
s -> TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
s) StreamState
StreamReset) (Map.elems streams)
  writeTVar (ysessStreams sess) Map.empty
  pings <- readTVar (ysessPings sess)
  mapM_ (\PingWaiter
w -> PingWaiter -> Either YamuxError () -> STM Bool
forall a. TMVar a -> a -> STM Bool
tryPutTMVar PingWaiter
w (YamuxError -> Either YamuxError ()
forall a b. a -> Either a b
Left YamuxError
YamuxSessionShutdown)) (Map.elems pings)
  writeTVar (ysessPings sess) Map.empty

-- | Dispatch a decoded frame to the appropriate handler.
-- Returns False when a fatal protocol error occurred and the receive
-- loop must terminate (go-yamux treats these as session-fatal).
dispatchFrame :: YamuxSession -> YamuxHeader -> IO Bool
dispatchFrame :: YamuxSession -> YamuxHeader -> IO Bool
dispatchFrame YamuxSession
sess YamuxHeader
hdr = case YamuxHeader -> FrameType
yhType YamuxHeader
hdr of
  FrameType
FrameData -> YamuxSession -> YamuxHeader -> IO Bool
handleDataFrame YamuxSession
sess YamuxHeader
hdr
  FrameType
FrameWindowUpdate -> YamuxSession -> YamuxHeader -> IO Bool
handleWindowUpdate YamuxSession
sess YamuxHeader
hdr
  FrameType
FramePing -> YamuxSession -> YamuxHeader -> IO ()
handlePing YamuxSession
sess YamuxHeader
hdr IO () -> IO Bool -> IO Bool
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  FrameType
FrameGoAway -> YamuxSession -> YamuxHeader -> IO ()
handleGoAway YamuxSession
sess YamuxHeader
hdr IO () -> IO Bool -> IO Bool
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

-- | Handle a Data frame: validate declared length, read payload, manage
-- stream state, deliver data. Returns False on fatal protocol error.
handleDataFrame :: YamuxSession -> YamuxHeader -> IO Bool
handleDataFrame :: YamuxSession -> YamuxHeader -> IO Bool
handleDataFrame YamuxSession
sess YamuxHeader
hdr = do
  let sid :: Word32
sid = YamuxHeader -> Word32
yhStreamId YamuxHeader
hdr
      flags :: Flags
flags = YamuxHeader -> Flags
yhFlags YamuxHeader
hdr
      declaredLen :: Word32
declaredLen = YamuxHeader -> Word32
yhLength YamuxHeader
hdr
  -- Handle SYN first so the flow-control check below sees the new stream
  synOk <-
    if Flags -> Bool
flagSYN Flags
flags
      then YamuxSession -> Word32 -> IO Bool
acceptInboundSYN YamuxSession
sess Word32
sid
      else Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  if not synOk
    then do
      sendGoAway sess GoAwayProtocol
      pure False
    else do
      -- Flow control: validate the declared length against the receive
      -- window BEFORE reading the payload off the transport. A frame that
      -- overruns the window is a protocol error and must not cause the
      -- session to buffer attacker-controlled amounts of memory.
      reserved <- reserveRecvWindow sess sid declaredLen
      if not reserved
        then do
          sendGoAway sess GoAwayProtocol
          pure False
        else do
          payload <-
            if declaredLen > 0
              then ysessRead sess (fromIntegral declaredLen)
              else pure BS.empty
          -- Handle ACK flag: transition SYNSent -> Established
          when (flagACK flags) $ handleAckFlag sess sid
          -- Deliver payload to stream buffer (window already reserved)
          when (BS.length payload > 0) $ do
            mStream <- lookupStream sess sid
            case mStream of
              Just YamuxStream
stream -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TQueue ByteString -> ByteString -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue (YamuxStream -> TQueue ByteString
ysRecvBuf YamuxStream
stream) ByteString
payload
              Maybe YamuxStream
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure () -- unknown stream: discard
          -- Handle FIN flag
          when (flagFIN flags) $ applyRemoteFin sess sid
          -- Handle RST flag
          when (flagRST flags) $ applyRemoteRst sess sid
          pure True

-- | Handle a WindowUpdate frame: update send window, manage stream
-- lifecycle. Returns False on fatal protocol error.
handleWindowUpdate :: YamuxSession -> YamuxHeader -> IO Bool
handleWindowUpdate :: YamuxSession -> YamuxHeader -> IO Bool
handleWindowUpdate YamuxSession
sess YamuxHeader
hdr = do
  let sid :: Word32
sid = YamuxHeader -> Word32
yhStreamId YamuxHeader
hdr
      flags :: Flags
flags = YamuxHeader -> Flags
yhFlags YamuxHeader
hdr
      delta :: Word32
delta = YamuxHeader -> Word32
yhLength YamuxHeader
hdr
  -- Handle SYN flag: create new inbound stream (with parity + duplicate validation)
  synOk <-
    if Flags -> Bool
flagSYN Flags
flags
      then YamuxSession -> Word32 -> IO Bool
acceptInboundSYN YamuxSession
sess Word32
sid
      else Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  if not synOk
    then do
      sendGoAway sess GoAwayProtocol
      pure False
    else do
      -- Handle ACK flag
      when (flagACK flags) $ handleAckFlag sess sid
      -- Update send window
      when (delta > 0) $ do
        mStream <- lookupStream sess sid
        case mStream of
          Just YamuxStream
stream -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            w <- TVar Word32 -> STM Word32
forall a. TVar a -> STM a
readTVar (YamuxStream -> TVar Word32
ysSendWindow YamuxStream
stream)
            writeTVar (ysSendWindow stream) (w + delta)
          Maybe YamuxStream
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      -- Handle FIN flag
      when (flagFIN flags) $ applyRemoteFin sess sid
      -- Handle RST flag
      when (flagRST flags) $ applyRemoteRst sess sid
      pure True

-- | Look up a stream by ID.
lookupStream :: YamuxSession -> Word32 -> IO (Maybe YamuxStream)
lookupStream :: YamuxSession -> Word32 -> IO (Maybe YamuxStream)
lookupStream YamuxSession
sess Word32
sid = Word32 -> Map Word32 YamuxStream -> Maybe YamuxStream
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Word32
sid (Map Word32 YamuxStream -> Maybe YamuxStream)
-> IO (Map Word32 YamuxStream) -> IO (Maybe YamuxStream)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map Word32 YamuxStream) -> IO (Map Word32 YamuxStream)
forall a. TVar a -> IO a
readTVarIO (YamuxSession -> TVar (Map Word32 YamuxStream)
ysessStreams YamuxSession
sess)

-- | Validate and register an inbound SYN (parity + duplicate check).
-- Returns False on protocol error; the caller sends GoAway and stops.
-- When the accept backlog is full the SYN is rejected with RST instead
-- of being buffered (spec.md, ACK backlog: the buffer MUST be bounded);
-- that is not a protocol error, so the session stays alive.
acceptInboundSYN :: YamuxSession -> Word32 -> IO Bool
acceptInboundSYN :: YamuxSession -> Word32 -> IO Bool
acceptInboundSYN YamuxSession
sess Word32
sid = do
  valid <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ YamuxSession -> Word32 -> STM Bool
validateInboundSYN YamuxSession
sess Word32
sid
  if not valid
    then pure False
    else do
      stream <- newStream sess sid StreamSYNReceived
      atomically $ do
        full <- isFullTBQueue (ysessAcceptCh sess)
        if full
          then writeTQueue (ysessSendCh sess) (rstHeader sid, BS.empty)
          else do
            modifyTVar' (ysessStreams sess) (Map.insert sid stream)
            writeTBQueue (ysessAcceptCh sess) stream
      pure True

-- | Data frame carrying only the RST flag for the given stream.
rstHeader :: Word32 -> YamuxHeader
rstHeader :: Word32 -> YamuxHeader
rstHeader Word32
sid =
  YamuxHeader
    { yhVersion :: Word8
yhVersion = Word8
0
    , yhType :: FrameType
yhType = FrameType
FrameData
    , yhFlags :: Flags
yhFlags = Flags
defaultFlags {flagRST = True}
    , yhStreamId :: Word32
yhStreamId = Word32
sid
    , yhLength :: Word32
yhLength = Word32
0
    }

-- | Reserve receive window for a declared Data-frame length before the
-- payload is read off the transport. Returns False on a flow-control
-- violation (declared length exceeds the stream's receive window or the
-- absolute maxStreamWindowSize bound); the session must then terminate.
reserveRecvWindow :: YamuxSession -> Word32 -> Word32 -> IO Bool
reserveRecvWindow :: YamuxSession -> Word32 -> Word32 -> IO Bool
reserveRecvWindow YamuxSession
sess Word32
sid Word32
len
  | Word32
len Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0 = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
> Word32
maxStreamWindowSize = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
  | Bool
otherwise = STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
      streams <- TVar (Map Word32 YamuxStream) -> STM (Map Word32 YamuxStream)
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar (Map Word32 YamuxStream)
ysessStreams YamuxSession
sess)
      case Map.lookup sid streams of
        -- Unknown stream: payload is read and discarded, bounded by the
        -- maxStreamWindowSize check above (defence in depth)
        Maybe YamuxStream
Nothing -> Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
        Just YamuxStream
stream -> do
          w <- TVar Word32 -> STM Word32
forall a. TVar a -> STM a
readTVar (YamuxStream -> TVar Word32
ysRecvWindow YamuxStream
stream)
          if len > w
            then pure False
            else do
              writeTVar (ysRecvWindow stream) (w - len)
              pure True

-- | ACK flag: transition SYNSent -> Established.
handleAckFlag :: YamuxSession -> Word32 -> IO ()
handleAckFlag :: YamuxSession -> Word32 -> IO ()
handleAckFlag YamuxSession
sess Word32
sid = do
  mStream <- YamuxSession -> Word32 -> IO (Maybe YamuxStream)
lookupStream YamuxSession
sess Word32
sid
  case mStream of
    Just YamuxStream
stream -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      st <- TVar StreamState -> STM StreamState
forall a. TVar a -> STM a
readTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream)
      case st of
        StreamState
StreamSYNSent -> TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamEstablished
        StreamState
_ -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Maybe YamuxStream
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Shared FIN transition (spec.md, Closing a stream). The remote may
-- half-close from any pre-close state: nothing in the spec ties FIN to
-- the local ACK state, and go-libp2p pipelines SYN, data and FIN in one
-- burst, so SYNSent/SYNReceived must transition like Established.
applyRemoteFin :: YamuxSession -> Word32 -> IO ()
applyRemoteFin :: YamuxSession -> Word32 -> IO ()
applyRemoteFin YamuxSession
sess Word32
sid = do
  mStream <- YamuxSession -> Word32 -> IO (Maybe YamuxStream)
lookupStream YamuxSession
sess Word32
sid
  case mStream of
    Just YamuxStream
stream -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      st <- TVar StreamState -> STM StreamState
forall a. TVar a -> STM a
readTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream)
      case st of
        StreamState
StreamSYNSent -> TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamRemoteClose
        StreamState
StreamSYNReceived -> TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamRemoteClose
        StreamState
StreamEstablished -> TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamRemoteClose
        StreamState
StreamLocalClose -> do
          -- Both sides FIN'd: the stream is dead, reclaim its map slot
          TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamClosed
          TVar (Map Word32 YamuxStream)
-> (Map Word32 YamuxStream -> Map Word32 YamuxStream) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (YamuxSession -> TVar (Map Word32 YamuxStream)
ysessStreams YamuxSession
sess) (Word32 -> Map Word32 YamuxStream -> Map Word32 YamuxStream
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Word32
sid)
        StreamState
_ -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Maybe YamuxStream
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Shared RST transition: any state -> Reset. A reset stream is dead,
-- so its map slot is reclaimed in the same transaction; the application
-- still holds the handle and observes StreamReset through it.
applyRemoteRst :: YamuxSession -> Word32 -> IO ()
applyRemoteRst :: YamuxSession -> Word32 -> IO ()
applyRemoteRst YamuxSession
sess Word32
sid = do
  mStream <- YamuxSession -> Word32 -> IO (Maybe YamuxStream)
lookupStream YamuxSession
sess Word32
sid
  case mStream of
    Just YamuxStream
stream -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      TVar StreamState -> StreamState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxStream -> TVar StreamState
ysState YamuxStream
stream) StreamState
StreamReset
      TVar (Map Word32 YamuxStream)
-> (Map Word32 YamuxStream -> Map Word32 YamuxStream) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (YamuxSession -> TVar (Map Word32 YamuxStream)
ysessStreams YamuxSession
sess) (Word32 -> Map Word32 YamuxStream -> Map Word32 YamuxStream
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Word32
sid)
    Maybe YamuxStream
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Handle a Ping frame (StreamID must be 0).
-- SYN: echo back with ACK flag and same opaque value.
-- ACK: resolve the matching pending ping.
handlePing :: YamuxSession -> YamuxHeader -> IO ()
handlePing :: YamuxSession -> YamuxHeader -> IO ()
handlePing YamuxSession
sess YamuxHeader
hdr
  | Flags -> Bool
flagSYN (YamuxHeader -> Flags
yhFlags YamuxHeader
hdr) = do
      -- Echo back Ping with ACK
      let respHdr :: YamuxHeader
respHdr =
            YamuxHeader
              { yhVersion :: Word8
yhVersion = Word8
0
              , yhType :: FrameType
yhType = FrameType
FramePing
              , yhFlags :: Flags
yhFlags = Flags
defaultFlags {flagACK = True}
              , yhStreamId :: Word32
yhStreamId = Word32
0
              , yhLength :: Word32
yhLength = YamuxHeader -> Word32
yhLength YamuxHeader
hdr -- echo opaque value
              }
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TQueue (YamuxHeader, ByteString)
-> (YamuxHeader, ByteString) -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue (YamuxSession -> TQueue (YamuxHeader, ByteString)
ysessSendCh YamuxSession
sess) (YamuxHeader
respHdr, ByteString
BS.empty)
  | Flags -> Bool
flagACK (YamuxHeader -> Flags
yhFlags YamuxHeader
hdr) = do
      -- Resolve pending ping
      let pingId :: Word32
pingId = YamuxHeader -> Word32
yhLength YamuxHeader
hdr
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        pMap <- TVar (Map Word32 PingWaiter) -> STM (Map Word32 PingWaiter)
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar (Map Word32 PingWaiter)
ysessPings YamuxSession
sess)
        case Map.lookup pingId pMap of
          Just PingWaiter
waiter -> PingWaiter -> Either YamuxError () -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar PingWaiter
waiter (() -> Either YamuxError ()
forall a b. b -> Either a b
Right ())
          Maybe PingWaiter
Nothing -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Handle a GoAway frame (StreamID must be 0).
-- Preserve the received error code so callers can distinguish a clean
-- shutdown (0x00) from a protocol (0x01) or internal (0x02) error.
-- A code outside the spec-defined range is itself a protocol violation
-- and is recorded as GoAwayProtocol.
handleGoAway :: YamuxSession -> YamuxHeader -> IO ()
handleGoAway :: YamuxSession -> YamuxHeader -> IO ()
handleGoAway YamuxSession
sess YamuxHeader
hdr = do
  let code :: GoAwayCode
code = case Word32 -> Maybe GoAwayCode
word32ToGoAwayCode (YamuxHeader -> Word32
yhLength YamuxHeader
hdr) of
        Just GoAwayCode
c -> GoAwayCode
c
        Maybe GoAwayCode
Nothing -> GoAwayCode
GoAwayProtocol
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Maybe GoAwayCode) -> Maybe GoAwayCode -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (YamuxSession -> TVar (Maybe GoAwayCode)
ysessRemoteGoAway YamuxSession
sess) (GoAwayCode -> Maybe GoAwayCode
forall a. a -> Maybe a
Just GoAwayCode
code)

-- | Send loop: dequeues frames from ysessSendCh and writes to transport.
sendLoop :: YamuxSession -> IO ()
sendLoop :: YamuxSession -> IO ()
sendLoop YamuxSession
sess = IO ()
forall {b}. IO b
go
  where
    go :: IO b
go = do
      (hdr, payload) <- STM (YamuxHeader, ByteString) -> IO (YamuxHeader, ByteString)
forall a. STM a -> IO a
atomically (STM (YamuxHeader, ByteString) -> IO (YamuxHeader, ByteString))
-> STM (YamuxHeader, ByteString) -> IO (YamuxHeader, ByteString)
forall a b. (a -> b) -> a -> b
$ TQueue (YamuxHeader, ByteString) -> STM (YamuxHeader, ByteString)
forall a. TQueue a -> STM a
readTQueue (YamuxSession -> TQueue (YamuxHeader, ByteString)
ysessSendCh YamuxSession
sess)
      ysessWrite sess (encodeHeader hdr)
      when (BS.length payload > 0) $ ysessWrite sess payload
      go

-- | Create a new YamuxStream with the given initial state.
newStream :: YamuxSession -> Word32 -> StreamState -> IO YamuxStream
newStream :: YamuxSession -> Word32 -> StreamState -> IO YamuxStream
newStream YamuxSession
sess Word32
sid StreamState
initState = do
  stateVar <- StreamState -> IO (TVar StreamState)
forall a. a -> IO (TVar a)
newTVarIO StreamState
initState
  sendWin <- newTVarIO initialWindowSize
  recvWin <- newTVarIO initialWindowSize
  recvBuf <- newTQueueIO
  sendNotify <- newEmptyTMVarIO
  pure
    YamuxStream
      { ysStreamId = sid
      , ysState = stateVar
      , ysSendWindow = sendWin
      , ysRecvWindow = recvWin
      , ysRecvBuf = recvBuf
      , ysSendNotify = sendNotify
      , ysSession = sess
      }

-- | Validate an inbound SYN stream ID for parity and uniqueness.
-- Returns True if valid, False if protocol error (caller must send GoAway).
-- Remote peers must use the opposite parity: client expects even, server expects odd.
validateInboundSYN :: YamuxSession -> Word32 -> STM Bool
validateInboundSYN :: YamuxSession -> Word32 -> STM Bool
validateInboundSYN YamuxSession
sess Word32
sid = do
  let validParity :: Bool
validParity = case YamuxSession -> SessionRole
ysessRole YamuxSession
sess of
        -- Server expects odd IDs (from client)
        SessionRole
RoleServer -> Word32 -> Bool
forall a. Integral a => a -> Bool
odd Word32
sid
        -- Client expects even IDs (from server)
        SessionRole
RoleClient -> Word32 -> Bool
forall a. Integral a => a -> Bool
even Word32
sid
  if Word32
sid Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0 Bool -> Bool -> Bool
|| Bool -> Bool
not Bool
validParity
    then Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    else do
      streams <- TVar (Map Word32 YamuxStream) -> STM (Map Word32 YamuxStream)
forall a. TVar a -> STM a
readTVar (YamuxSession -> TVar (Map Word32 YamuxStream)
ysessStreams YamuxSession
sess)
      pure (not (Map.member sid streams))

-- | Helper: execute action when condition is True.
when :: Bool -> IO () -> IO ()
when :: Bool -> IO () -> IO ()
when Bool
True IO ()
action = IO ()
action
when Bool
False IO ()
_ = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()