-- | 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
  , upgradeAs
  , 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
  ( ConnectionEndpoint (..)
  , NativeMuxer (..)
  , 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
  , streamReadChunk :: Int -> IO ByteString
streamReadChunk = IORef NoiseSession
-> IORef ByteString -> StreamIO -> Int -> IO ByteString
decryptAndReadChunk 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 Noise frames from the raw stream until one decrypts to a
-- non-empty plaintext, and return that plaintext. 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.
nextPlaintext :: IORef NoiseSession -> StreamIO -> IO ByteString
nextPlaintext :: IORef NoiseSession -> StreamIO -> IO ByteString
nextPlaintext IORef NoiseSession
recvRef StreamIO
rawIO = do
  ct <- StreamIO -> IO ByteString
readFramedMessage StreamIO
rawIO
  if BS.null ct
    then nextPlaintext recvRef rawIO -- zero-length frame: no message to decrypt
    else do
      sess <- readIORef recvRef
      case decryptMessage sess ct of
        Left String
err -> String -> IO ByteString
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO ByteString) -> String -> IO ByteString
forall a b. (a -> b) -> a -> b
$ String
"nextPlaintext: decrypt failed: " 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 IORef NoiseSession -> StreamIO -> IO ByteString
nextPlaintext IORef NoiseSession
recvRef StreamIO
rawIO -- empty transport message (keepalive)
            else ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
pt

-- | Read and decrypt a byte from the Noise channel: pop the buffer if
-- it has bytes, otherwise decrypt the next frame and buffer the rest.
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
  bs <- if BS.null buf then nextPlaintext recvRef rawIO else pure buf
  writeIORef bufRef (BS.tail bs)
  pure (BS.head bs)

-- | Chunk-level read from the Noise channel: hand back up to @n@ bytes
-- of the buffered plaintext (a decrypted frame is already a chunk),
-- decrypting the next frame only when the buffer is empty. Bytes
-- beyond @n@ stay buffered for the next read.
decryptAndReadChunk :: IORef NoiseSession -> IORef ByteString -> StreamIO -> Int -> IO ByteString
decryptAndReadChunk :: IORef NoiseSession
-> IORef ByteString -> StreamIO -> Int -> IO ByteString
decryptAndReadChunk IORef NoiseSession
recvRef IORef ByteString
bufRef StreamIO
rawIO Int
n = do
  buf <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
bufRef
  bs <- if BS.null buf then nextPlaintext recvRef rawIO else pure buf
  let (front, rest) = BS.splitAt n bs
  writeIORef bufRef rest
  pure front

-- | 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; an IORef buffer holds
-- the bytes a byte- or chunk-level read did not consume.
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
  let -- Buffered bytes if any, otherwise the next yamux chunk.
      nextChunk = do
        buf <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
readBuf
        if BS.null buf
          then do
            result <- streamRead yamuxStream
            case result of
              Left YamuxError
err -> String -> IO ByteString
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO ByteString) -> String -> IO ByteString
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 ByteString
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"yamuxStreamRead: empty chunk"
                | Bool
otherwise -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
chunk
          else pure buf
  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
        chunk <- nextChunk
        writeIORef readBuf (BS.tail chunk)
        pure (BS.head chunk)
    , streamReadChunk = \Int
n -> do
        chunk <- IO ByteString
nextChunk
        let (front, rest) = BS.splitAt n chunk
        writeIORef readBuf rest
        pure front
    , streamClose = do
        _ <- YS.streamClose yamuxStream  -- Sends FIN flag
        pure ()
    }

-- | Upgrade a raw connection, taking every role from the direction.
--
-- Pipeline: mss(/noise) -> Noise XX -> mss(/yamux/1.0.0) -> Yamux.
-- 'Outbound' runs the initiator/client side of all three, 'Inbound' the
-- responder/server side. go-libp2p derives the same way
-- (@isServer := dir == network.DirInbound@ in its upgrader), which is
-- what lets a TCP simultaneous connect flip roles: the peer that must
-- act as the server passes 'Inbound' even though it called connect().
upgradeAs :: Direction -> KeyPair -> RawConnection -> IO Connection
upgradeAs :: Direction -> KeyPair -> RawConnection -> IO Connection
upgradeAs Direction
dir KeyPair
identityKP RawConnection
rawConn = case RawConnection -> ConnectionEndpoint
rcEndpoint RawConnection
rawConn of
  ByteStreamEndpoint StreamIO
rawIO -> Direction -> KeyPair -> RawConnection -> StreamIO -> IO Connection
upgradeByteStream Direction
dir KeyPair
identityKP RawConnection
rawConn StreamIO
rawIO
  NativeMuxerEndpoint NativeMuxer
native -> Direction -> RawConnection -> NativeMuxer -> IO Connection
nativeToConnection Direction
dir RawConnection
rawConn NativeMuxer
native

upgradeByteStream :: Direction -> KeyPair -> RawConnection -> StreamIO -> IO Connection
upgradeByteStream :: Direction -> KeyPair -> RawConnection -> StreamIO -> IO Connection
upgradeByteStream Direction
dir KeyPair
identityKP RawConnection
rawConn StreamIO
rawIO = do
  let isServer :: Bool
isServer = Direction
dir Direction -> Direction -> Bool
forall a. Eq a => a -> a -> Bool
== Direction
Inbound
      negotiate :: StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiate = if Bool
isServer then StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateResponder else StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateInitiator
      role :: String
role = if Bool
isServer then String
"upgradeInbound" else String
"upgradeOutbound"

  -- Step 1: multistream-select -> "/noise"
  secResult <- StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiate 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
role String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": /noise negotiation failed")

  -- Step 2: Noise XX handshake
  (noiseSess, HandshakeResult remotePeerId _remotePK) <-
    performStreamHandshake identityKP dir 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 <- negotiate 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
role String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": /yamux/1.0.0 negotiation failed")

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

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

nativeToConnection :: Direction -> RawConnection -> NativeMuxer -> IO Connection
nativeToConnection :: Direction -> RawConnection -> NativeMuxer -> IO Connection
nativeToConnection Direction
dir RawConnection
rawConn NativeMuxer
native = do
  stateVar <- ConnState -> IO (TVar ConnState)
forall a. a -> IO (TVar a)
newTVarIO ConnState
ConnOpen
  let muxer = MuxerSession
        { muxOpenStream :: IO StreamIO
muxOpenStream = NativeMuxer -> IO StreamIO
nativeOpenStream NativeMuxer
native
        , muxAcceptStream :: IO StreamIO
muxAcceptStream = NativeMuxer -> IO StreamIO
nativeAcceptStream NativeMuxer
native
        , muxClose :: IO ()
muxClose = NativeMuxer -> IO ()
nativeClose NativeMuxer
native
        }
  pure Connection
    { connPeerId = nativePeerId native
    , connDirection = dir
    , connLocalAddr = rcLocalAddr rawConn
    , connRemoteAddr = rcRemoteAddr rawConn
    , connSecurity = nativeSecurity native
    , connMuxer = nativeMuxerProtocol native
    , connSession = muxer
    , connState = stateVar
    }

-- | Upgrade an outbound (dialer) raw connection.
upgradeOutbound :: KeyPair -> RawConnection -> IO Connection
upgradeOutbound :: KeyPair -> RawConnection -> IO Connection
upgradeOutbound = Direction -> KeyPair -> RawConnection -> IO Connection
upgradeAs Direction
Outbound

-- | Upgrade an inbound (listener) raw connection.
upgradeInbound :: KeyPair -> RawConnection -> IO Connection
upgradeInbound :: KeyPair -> RawConnection -> IO Connection
upgradeInbound = Direction -> KeyPair -> RawConnection -> IO Connection
upgradeAs Direction
Inbound