-- | Connection upgrade pipeline for the Switch.
--
-- Transforms a raw transport connection into a fully upgraded
-- (secure + multiplexed) Connection by executing a 4-step pipeline:
--   1. multistream-select: negotiate security protocol ("/noise")
--   2. Noise XX handshake: encrypted channel + remote PeerId
--   3. multistream-select: negotiate muxer ("/yamux/1.0.0")
--   4. Yamux session init: multiplexed streams
--
-- Implements the connection upgrading pipeline.
module LibP2P.Switch.Upgrade
  ( -- * Streaming handshake
    performStreamHandshake
    -- * Encrypted StreamIO
  , noiseSessionToStreamIO
    -- * Yamux → MuxerSession adapter
  , yamuxToMuxerSession
    -- * Full upgrade pipeline
  , upgradeOutbound
  , upgradeInbound
    -- * Helpers (exported for testing)
  , readExact
  , readFramedMessage
  , writeFramedMessage
  ) where

import Control.Concurrent.Async (async, cancel, race, waitCatch)
import Control.Concurrent.STM (atomically, isEmptyTQueue, newTVarIO, retry)
import Control.Exception (SomeException, catch)
import Control.Monad (unless)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Word (Word8)
import LibP2P.Core.Binary (readWord16BE)
import LibP2P.Crypto.Key (KeyPair (..))
import LibP2P.Crypto.PeerId (fromPublicKey)
import LibP2P.Yamux.Frame (maxStreamWindowSize)
import LibP2P.Yamux.Session (closeSession, newSession, recvLoop, sendLoop)
import qualified LibP2P.Yamux.Session as Yamux
import LibP2P.Yamux.Stream (streamRead)
import qualified LibP2P.Yamux.Stream as YS
import LibP2P.Yamux.Types (SessionRole (..), YamuxSession (ysessSendCh), YamuxStream)
import LibP2P.MultistreamSelect.Negotiation
  ( NegotiationResult (..)
  , StreamIO (..)
  , negotiateInitiator
  , negotiateResponder
  , readExactBounded
  )
import LibP2P.Noise.Framing (chunkPlaintext, encodeFrame)
import LibP2P.Noise.Handshake
  ( HandshakeResult (..)
  , buildHandshakePayload
  , decodeNoisePayload
  , encodeNoisePayload
  , getRemoteNoiseStaticKey
  , initHandshakeInitiator
  , initHandshakeResponder
  , readHandshakeMsg
  , verifyStaticKey
  , writeHandshakeMsg
  )
import LibP2P.Noise.Session
  ( NoiseSession
  , decryptMessage
  , encryptMessage
  , mkNoiseSession
  )
import LibP2P.Switch.Types
  ( ConnState (..)
  , Connection (..)
  , Direction (..)
  , MuxerSession (..)
  )
import LibP2P.Transport (RawConnection (..))
import System.Timeout (timeout)
import qualified LibP2P.Crypto.Protobuf as Proto
import qualified LibP2P.Noise.Handshake as HS

-- | Read exactly n bytes from a StreamIO.
--
-- Exception-style wrapper over 'readExactBounded' for the Noise frame
-- reader and the yamux read callback, which expect failures as
-- exceptions. Defence in depth against remote memory exhaustion: the
-- request size is bounded by maxStreamWindowSize (callers must validate
-- lengths against flow control before reading).
readExact :: StreamIO -> Int -> IO ByteString
readExact :: StreamIO -> Int -> IO ByteString
readExact StreamIO
stream Int
n =
  (String -> IO ByteString)
-> (ByteString -> IO ByteString)
-> Either String ByteString
-> IO ByteString
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> IO ByteString
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (Either String ByteString -> IO ByteString)
-> IO (Either String ByteString) -> IO ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< StreamIO -> Int -> Int -> IO (Either String ByteString)
readExactBounded StreamIO
stream (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
maxStreamWindowSize) Int
n

-- | Read a 2-byte-BE-length-prefixed Noise frame from a StreamIO.
readFramedMessage :: StreamIO -> IO ByteString
readFramedMessage :: StreamIO -> IO ByteString
readFramedMessage StreamIO
stream = do
  lenBytes <- StreamIO -> Int -> IO ByteString
readExact StreamIO
stream Int
2
  let len = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Word16
readWord16BE ByteString
lenBytes) :: Int
  if len == 0
    then pure BS.empty
    else readExact stream len

-- | Write a 2-byte-BE-length-prefixed Noise frame to a StreamIO.
-- Fails (instead of truncating the length prefix) if the message exceeds
-- the 65535-byte Noise message cap.
writeFramedMessage :: StreamIO -> ByteString -> IO ()
writeFramedMessage :: StreamIO -> ByteString -> IO ()
writeFramedMessage StreamIO
stream ByteString
msg =
  (String -> IO ())
-> (ByteString -> IO ()) -> Either String ByteString -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream) (ByteString -> Either String ByteString
encodeFrame ByteString
msg)

-- | Perform a Noise XX handshake over a StreamIO using framed messages.
-- Returns (NoiseSession, HandshakeResult) with the remote PeerId.
performStreamHandshake
  :: KeyPair -> Direction -> StreamIO -> IO (NoiseSession, HandshakeResult)
performStreamHandshake :: KeyPair
-> Direction -> StreamIO -> IO (NoiseSession, HandshakeResult)
performStreamHandshake KeyPair
identityKP Direction
dir StreamIO
stream = case Direction
dir of
  Direction
Outbound -> KeyPair -> StreamIO -> IO (NoiseSession, HandshakeResult)
performInitiatorHandshake KeyPair
identityKP StreamIO
stream
  Direction
Inbound  -> KeyPair -> StreamIO -> IO (NoiseSession, HandshakeResult)
performResponderHandshake KeyPair
identityKP StreamIO
stream

-- | Initiator (dialer) side of the Noise XX handshake.
--
-- Message flow:
--   1. Initiator → Responder: e (empty payload)
--   2. Responder → Initiator: e, ee, s, es (responder identity payload)
--   3. Initiator → Responder: s, se (initiator identity payload)
performInitiatorHandshake :: KeyPair -> StreamIO -> IO (NoiseSession, HandshakeResult)
performInitiatorHandshake :: KeyPair -> StreamIO -> IO (NoiseSession, HandshakeResult)
performInitiatorHandshake KeyPair
identityKP StreamIO
stream = do
  (hsState0, noiseStaticPub) <- KeyPair -> IO (HandshakeState, ByteString)
initHandshakeInitiator KeyPair
identityKP

  -- Message 1: → (empty payload)
  (msg1, hsState1) <- either (fail . ("initiator msg1 write: " <>)) pure $
    writeHandshakeMsg hsState0 BS.empty
  writeFramedMessage stream msg1

  -- Message 2: ← (responder's identity payload)
  msg2 <- readFramedMessage stream
  (payload2, hsState2) <- either (fail . ("initiator msg2 read: " <>)) pure $
    readHandshakeMsg hsState1 msg2

  -- Decode responder's identity
  remoteNP <- either (fail . ("initiator decode payload: " <>)) pure $
    decodeNoisePayload payload2
  remotePubKey <- either (fail . ("initiator decode pubkey: " <>)) pure $
    Proto.decodePublicKey (HS.npIdentityKey remoteNP)
  let remotePeerId = PublicKey -> PeerId
fromPublicKey PublicKey
remotePubKey

  -- Verify identity_sig: binds identity key to Noise static key
  case getRemoteNoiseStaticKey hsState2 of
    Maybe ByteString
Nothing -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"performInitiatorHandshake: remote Noise static key unavailable after msg2"
    Just ByteString
remoteNoisePub ->
      if Bool -> Bool
not (PublicKey -> ByteString -> ByteString -> Bool
verifyStaticKey PublicKey
remotePubKey ByteString
remoteNoisePub (NoisePayload -> ByteString
HS.npIdentitySig NoisePayload
remoteNP))
        then String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"performInitiatorHandshake: identity signature verification failed"
        else () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  -- Message 3: → (initiator's identity payload)
  identPayload <- either (fail . ("initiator payload build: " <>)) pure $
    encodeNoisePayload <$> buildHandshakePayload identityKP noiseStaticPub
  (msg3, hsStateFinal) <- either (fail . ("initiator msg3 write: " <>)) pure $
    writeHandshakeMsg hsState2 identPayload
  writeFramedMessage stream msg3

  let noiseSession = CacophonyState -> NoiseSession
mkNoiseSession (HandshakeState -> CacophonyState
HS.hsNoiseState HandshakeState
hsStateFinal)
  pure (noiseSession, HandshakeResult remotePeerId remotePubKey)

-- | Responder (listener) side of the Noise XX handshake.
performResponderHandshake :: KeyPair -> StreamIO -> IO (NoiseSession, HandshakeResult)
performResponderHandshake :: KeyPair -> StreamIO -> IO (NoiseSession, HandshakeResult)
performResponderHandshake KeyPair
identityKP StreamIO
stream = do
  (hsState0, noiseStaticPub) <- KeyPair -> IO (HandshakeState, ByteString)
initHandshakeResponder KeyPair
identityKP

  -- Message 1: ← (empty payload)
  msg1 <- readFramedMessage stream
  (_payload1, hsState1) <- either (fail . ("responder msg1 read: " <>)) pure $
    readHandshakeMsg hsState0 msg1

  -- Message 2: → (responder's identity payload)
  identPayload <- either (fail . ("responder payload build: " <>)) pure $
    encodeNoisePayload <$> buildHandshakePayload identityKP noiseStaticPub
  (msg2, hsState2) <- either (fail . ("responder msg2 write: " <>)) pure $
    writeHandshakeMsg hsState1 identPayload
  writeFramedMessage stream msg2

  -- Message 3: ← (initiator's identity payload)
  msg3 <- readFramedMessage stream
  (payload3, hsStateFinal) <- either (fail . ("responder msg3 read: " <>)) pure $
    readHandshakeMsg hsState2 msg3

  -- Decode initiator's identity
  remoteNP <- either (fail . ("responder decode payload: " <>)) pure $
    decodeNoisePayload payload3
  remotePubKey <- either (fail . ("responder decode pubkey: " <>)) pure $
    Proto.decodePublicKey (HS.npIdentityKey remoteNP)
  let remotePeerId = PublicKey -> PeerId
fromPublicKey PublicKey
remotePubKey

  -- Verify identity_sig: binds identity key to Noise static key
  case getRemoteNoiseStaticKey hsStateFinal of
    Maybe ByteString
Nothing -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"performResponderHandshake: remote Noise static key unavailable after msg3"
    Just ByteString
remoteNoisePub ->
      if Bool -> Bool
not (PublicKey -> ByteString -> ByteString -> Bool
verifyStaticKey PublicKey
remotePubKey ByteString
remoteNoisePub (NoisePayload -> ByteString
HS.npIdentitySig NoisePayload
remoteNP))
        then String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"performResponderHandshake: identity signature verification failed"
        else () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  let noiseSession = CacophonyState -> NoiseSession
mkNoiseSession (HandshakeState -> CacophonyState
HS.hsNoiseState HandshakeState
hsStateFinal)
  pure (noiseSession, HandshakeResult remotePeerId remotePubKey)

-- | Create an encrypted StreamIO from a NoiseSession and raw StreamIO.
--
-- Uses separate IORefs for send/recv session state (each direction's
-- CipherState is independent in Noise). A read buffer (IORef ByteString)
-- bridges Noise's message-boundary decryption with StreamIO's byte-level reads.
noiseSessionToStreamIO
  :: IORef NoiseSession    -- ^ Send session state
  -> IORef NoiseSession    -- ^ Recv session state
  -> IORef ByteString      -- ^ Read buffer (decrypted but unconsumed bytes)
  -> StreamIO              -- ^ Raw (unencrypted) StreamIO
  -> StreamIO
noiseSessionToStreamIO :: IORef NoiseSession
-> IORef NoiseSession -> IORef ByteString -> StreamIO -> StreamIO
noiseSessionToStreamIO IORef NoiseSession
sendRef IORef NoiseSession
recvRef IORef ByteString
bufRef StreamIO
rawIO = StreamIO
  { streamWrite :: ByteString -> IO ()
streamWrite = IORef NoiseSession -> StreamIO -> ByteString -> IO ()
encryptAndWrite IORef NoiseSession
sendRef StreamIO
rawIO
  , streamReadByte :: IO Word8
streamReadByte = IORef NoiseSession -> IORef ByteString -> StreamIO -> IO Word8
decryptAndReadByte IORef NoiseSession
recvRef IORef ByteString
bufRef StreamIO
rawIO
  , streamClose :: IO ()
streamClose = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()  -- Encryption layer does not own the connection
  }

-- | Encrypt plaintext and write as framed Noise messages.
-- A Noise message is capped at 65535 bytes, so plaintext is split into
-- chunks of at most 65519 bytes (65535 minus the 16-byte AEAD tag) and
-- each chunk is encrypted and framed as its own Noise transport message.
encryptAndWrite :: IORef NoiseSession -> StreamIO -> ByteString -> IO ()
encryptAndWrite :: IORef NoiseSession -> StreamIO -> ByteString -> IO ()
encryptAndWrite IORef NoiseSession
sendRef StreamIO
rawIO ByteString
plaintext =
  (ByteString -> IO ()) -> [ByteString] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ByteString -> IO ()
encryptChunk (ByteString -> [ByteString]
chunkPlaintext ByteString
plaintext)
  where
    encryptChunk :: ByteString -> IO ()
encryptChunk ByteString
chunk = do
      sess <- IORef NoiseSession -> IO NoiseSession
forall a. IORef a -> IO a
readIORef IORef NoiseSession
sendRef
      case encryptMessage sess chunk of
        Left String
err -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"encryptAndWrite: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
err
        Right (ByteString
ct, NoiseSession
sess') -> do
          IORef NoiseSession -> NoiseSession -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef NoiseSession
sendRef NoiseSession
sess'
          StreamIO -> ByteString -> IO ()
writeFramedMessage StreamIO
rawIO ByteString
ct

-- | Read and decrypt a byte from the Noise channel.
-- If the buffer has bytes, return the first. Otherwise, read Noise
-- frames from the raw stream until one decrypts to a non-empty
-- plaintext, and buffer the result. A transport message with an empty
-- plaintext (a frame carrying only the AEAD tag) is legal — some
-- implementations send it as a keepalive — and a zero-length frame
-- carries no Noise message at all; both yield zero application bytes,
-- so reading continues at the next frame.
decryptAndReadByte :: IORef NoiseSession -> IORef ByteString -> StreamIO -> IO Word8
decryptAndReadByte :: IORef NoiseSession -> IORef ByteString -> StreamIO -> IO Word8
decryptAndReadByte IORef NoiseSession
recvRef IORef ByteString
bufRef StreamIO
rawIO = do
  buf <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
bufRef
  if BS.null buf
    then fillFromNextFrame
    else popByte buf
  where
    popByte :: ByteString -> IO Word8
popByte ByteString
bs = do
      IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
bufRef (HasCallStack => ByteString -> ByteString
ByteString -> ByteString
BS.tail ByteString
bs)
      Word8 -> IO Word8
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HasCallStack => ByteString -> Word8
ByteString -> Word8
BS.head ByteString
bs)
    fillFromNextFrame :: IO Word8
fillFromNextFrame = do
      ct <- StreamIO -> IO ByteString
readFramedMessage StreamIO
rawIO
      if BS.null ct
        then fillFromNextFrame -- zero-length frame: no message to decrypt
        else do
          sess <- readIORef recvRef
          case decryptMessage sess ct of
            Left String
err -> String -> IO Word8
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO Word8) -> String -> IO Word8
forall a b. (a -> b) -> a -> b
$ String
"decryptAndReadByte: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
err
            Right (ByteString
pt, NoiseSession
sess') -> do
              IORef NoiseSession -> NoiseSession -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef NoiseSession
recvRef NoiseSession
sess'
              if ByteString -> Bool
BS.null ByteString
pt
                then IO Word8
fillFromNextFrame -- empty transport message (keepalive)
                else ByteString -> IO Word8
popByte ByteString
pt

-- | Bounded window given to the send loop to flush the GoAway frame
-- before the transport is closed underneath it.
goAwayFlushTimeoutUs :: Int
goAwayFlushTimeoutUs :: Int
goAwayFlushTimeoutUs = Int
200000

-- | Wrap a YamuxSession as a MuxerSession.
-- Starts sendLoop and recvLoop as background threads.
-- The MuxerSession provides open/accept stream operations that
-- produce StreamIO-compatible streams.
--
-- The supplied close action closes the underlying transport; muxClose
-- runs it after sending GoAway and stopping the session loops, so
-- closing a connection actually releases the socket.
yamuxToMuxerSession :: YamuxSession -> IO () -> IO MuxerSession
yamuxToMuxerSession :: YamuxSession -> IO () -> IO MuxerSession
yamuxToMuxerSession YamuxSession
yamuxSess IO ()
closeTransport = do
  -- Start background loops
  sendLoopA <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (YamuxSession -> IO ()
sendLoop YamuxSession
yamuxSess)
  recvLoopA <- async (recvLoop yamuxSess)
  pure MuxerSession
    { muxOpenStream = do
        result <- Yamux.openStream yamuxSess
        case result of
          Right YamuxStream
stream -> YamuxStream -> IO StreamIO
yamuxStreamToStreamIO YamuxStream
stream
          Left YamuxError
err -> String -> IO StreamIO
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO StreamIO) -> String -> IO StreamIO
forall a b. (a -> b) -> a -> b
$ String
"muxOpenStream: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> YamuxError -> String
forall a. Show a => a -> String
show YamuxError
err
    , muxAcceptStream = do
        -- Fail the accept as soon as the receive loop dies (remote
        -- GoAway followed by EOF, transport error, or local close), so
        -- the Switch's stream accept loop exits and tears down the
        -- connection instead of blocking forever.
        result <- race (waitCatch recvLoopA) (Yamux.acceptStream yamuxSess)
        case result of
          Left Either SomeException ()
_ -> String -> IO StreamIO
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"muxAcceptStream: session terminated"
          Right (Right YamuxStream
stream) -> YamuxStream -> IO StreamIO
yamuxStreamToStreamIO YamuxStream
stream
          Right (Left YamuxError
err) -> String -> IO StreamIO
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO StreamIO) -> String -> IO StreamIO
forall a b. (a -> b) -> a -> b
$ String
"muxAcceptStream: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> YamuxError -> String
forall a. Show a => a -> String
show YamuxError
err
    , muxClose = do
        -- Queue GoAway, give the send loop a bounded window to flush
        -- it, then stop the loops and close the underlying transport.
        closeSession yamuxSess
        _ <- timeout goAwayFlushTimeoutUs $ atomically $ do
          empty <- isEmptyTQueue (ysessSendCh yamuxSess)
          unless empty retry
        cancel sendLoopA
        cancel recvLoopA
        closeTransport `catch` \(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    }

-- | Convert a YamuxStream to StreamIO with a read buffer.
-- Yamux delivers data in chunks via streamRead, but StreamIO requires
-- byte-by-byte reads. An IORef buffer bridges this gap.
yamuxStreamToStreamIO :: YamuxStream -> IO StreamIO
yamuxStreamToStreamIO :: YamuxStream -> IO StreamIO
yamuxStreamToStreamIO YamuxStream
yamuxStream = do
  readBuf <- ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
BS.empty
  pure StreamIO
    { streamWrite = \ByteString
bs -> do
        result <- YamuxStream -> ByteString -> IO (Either YamuxError ())
YS.streamWrite YamuxStream
yamuxStream ByteString
bs
        case result of
          Right () -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Left YamuxError
err -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"yamuxStreamWrite: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> YamuxError -> String
forall a. Show a => a -> String
show YamuxError
err
    , streamReadByte = do
        buf <- readIORef readBuf
        if BS.null buf
          then do
            result <- streamRead yamuxStream
            case result of
              Left YamuxError
err -> String -> IO Word8
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO Word8) -> String -> IO Word8
forall a b. (a -> b) -> a -> b
$ String
"yamuxStreamRead: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> YamuxError -> String
forall a. Show a => a -> String
show YamuxError
err
              Right ByteString
chunk
                | ByteString -> Bool
BS.null ByteString
chunk -> String -> IO Word8
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"yamuxStreamRead: empty chunk"
                | ByteString -> Int
BS.length ByteString
chunk Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 -> Word8 -> IO Word8
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HasCallStack => ByteString -> Word8
ByteString -> Word8
BS.head ByteString
chunk)
                | Bool
otherwise -> do
                    IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
readBuf (HasCallStack => ByteString -> ByteString
ByteString -> ByteString
BS.tail ByteString
chunk)
                    Word8 -> IO Word8
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HasCallStack => ByteString -> Word8
ByteString -> Word8
BS.head ByteString
chunk)
          else do
            writeIORef readBuf (BS.tail buf)
            pure (BS.head buf)
    , streamClose = do
        _ <- YS.streamClose yamuxStream  -- Sends FIN flag
        pure ()
    }

-- | Upgrade an outbound (dialer) raw connection.
-- Pipeline: mss(/noise) → Noise XX → mss(/yamux/1.0.0) → Yamux client
upgradeOutbound :: KeyPair -> RawConnection -> IO Connection
upgradeOutbound :: KeyPair -> RawConnection -> IO Connection
upgradeOutbound KeyPair
identityKP RawConnection
rawConn = do
  let rawIO :: StreamIO
rawIO = RawConnection -> StreamIO
rcStreamIO RawConnection
rawConn

  -- Step 1: multistream-select → "/noise"
  secResult <- StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateInitiator StreamIO
rawIO [ProtocolId
"/noise"]
  case secResult of
    Accepted ProtocolId
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    NegotiationResult
NoProtocol -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"upgradeOutbound: /noise negotiation failed"

  -- Step 2: Noise XX handshake (initiator)
  (noiseSess, HandshakeResult remotePeerId _remotePK) <-
    performStreamHandshake identityKP Outbound rawIO

  -- Step 3: Create encrypted StreamIO
  sendRef <- newIORef noiseSess
  recvRef <- newIORef noiseSess
  bufRef  <- newIORef BS.empty
  let encryptedIO = IORef NoiseSession
-> IORef NoiseSession -> IORef ByteString -> StreamIO -> StreamIO
noiseSessionToStreamIO IORef NoiseSession
sendRef IORef NoiseSession
recvRef IORef ByteString
bufRef StreamIO
rawIO

  -- Step 4: multistream-select → "/yamux/1.0.0" (over encrypted channel)
  muxResult <- negotiateInitiator encryptedIO ["/yamux/1.0.0"]
  case muxResult of
    Accepted ProtocolId
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    NegotiationResult
NoProtocol -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"upgradeOutbound: /yamux/1.0.0 negotiation failed"

  -- Step 5: Initialize Yamux session (client = odd IDs)
  let yamuxWrite = StreamIO -> ByteString -> IO ()
streamWrite StreamIO
encryptedIO
      yamuxRead  = \Int
n -> StreamIO -> Int -> IO ByteString
readExact StreamIO
encryptedIO Int
n
  yamuxSess <- newSession RoleClient yamuxWrite yamuxRead
  muxer <- yamuxToMuxerSession yamuxSess (rcClose rawConn)

  -- Build Connection
  stateVar <- newTVarIO ConnOpen
  pure Connection
    { connPeerId     = remotePeerId
    , connDirection  = Outbound
    , connLocalAddr  = rcLocalAddr rawConn
    , connRemoteAddr = rcRemoteAddr rawConn
    , connSecurity   = "/noise"
    , connMuxer      = "/yamux/1.0.0"
    , connSession    = muxer
    , connState      = stateVar
    }

-- | Upgrade an inbound (listener) raw connection.
-- Pipeline: mss(/noise) → Noise XX → mss(/yamux/1.0.0) → Yamux server
upgradeInbound :: KeyPair -> RawConnection -> IO Connection
upgradeInbound :: KeyPair -> RawConnection -> IO Connection
upgradeInbound KeyPair
identityKP RawConnection
rawConn = do
  let rawIO :: StreamIO
rawIO = RawConnection -> StreamIO
rcStreamIO RawConnection
rawConn

  -- Step 1: multistream-select → "/noise"
  secResult <- StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateResponder StreamIO
rawIO [ProtocolId
"/noise"]
  case secResult of
    Accepted ProtocolId
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    NegotiationResult
NoProtocol -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"upgradeInbound: /noise negotiation failed"

  -- Step 2: Noise XX handshake (responder)
  (noiseSess, HandshakeResult remotePeerId _remotePK) <-
    performStreamHandshake identityKP Inbound rawIO

  -- Step 3: Create encrypted StreamIO
  sendRef <- newIORef noiseSess
  recvRef <- newIORef noiseSess
  bufRef  <- newIORef BS.empty
  let encryptedIO = IORef NoiseSession
-> IORef NoiseSession -> IORef ByteString -> StreamIO -> StreamIO
noiseSessionToStreamIO IORef NoiseSession
sendRef IORef NoiseSession
recvRef IORef ByteString
bufRef StreamIO
rawIO

  -- Step 4: multistream-select → "/yamux/1.0.0" (over encrypted channel)
  muxResult <- negotiateResponder encryptedIO ["/yamux/1.0.0"]
  case muxResult of
    Accepted ProtocolId
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    NegotiationResult
NoProtocol -> String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"upgradeInbound: /yamux/1.0.0 negotiation failed"

  -- Step 5: Initialize Yamux session (server = even IDs)
  let yamuxWrite = StreamIO -> ByteString -> IO ()
streamWrite StreamIO
encryptedIO
      yamuxRead  = \Int
n -> StreamIO -> Int -> IO ByteString
readExact StreamIO
encryptedIO Int
n
  yamuxSess <- newSession RoleServer yamuxWrite yamuxRead
  muxer <- yamuxToMuxerSession yamuxSess (rcClose rawConn)

  -- Build Connection
  stateVar <- newTVarIO ConnOpen
  pure Connection
    { connPeerId     = remotePeerId
    , connDirection  = Inbound
    , connLocalAddr  = rcLocalAddr rawConn
    , connRemoteAddr = rcRemoteAddr rawConn
    , connSecurity   = "/noise"
    , connMuxer      = "/yamux/1.0.0"
    , connSession    = muxer
    , connState      = stateVar
    }