-- | multistream-select protocol negotiation.
--
-- Implements Initiator and Responder roles for negotiating
-- which protocol to use over a connection or stream.
module LibP2P.MultistreamSelect.Negotiation
  ( NegotiationResult (..)
  , ProtocolId
  , StreamIO (..)
  , negotiateInitiator
  , negotiateResponder
  , mkMemoryStreamPair
  , readExactBounded
  ) where

import Control.Concurrent.STM
import Control.Exception (IOException, catch)
import Control.Monad (replicateM)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Text (Text)
import Data.Word (Word64, Word8)
import LibP2P.Core.Varint (decodeUvarint, maxVarintBytes)
import LibP2P.MultistreamSelect.Wire

-- | A protocol identifier (e.g. "/noise", "/yamux/1.0.0").
type ProtocolId = Text

-- | Maximum accepted multistream-select message length in bytes.
-- Negotiation runs on raw, unauthenticated connections before any
-- handshake, so the declared length must be capped before allocating
-- or reading the payload. go-multistream rejects messages over 1024
-- bytes ("incoming message was too large"); protocol ids are far shorter.
maxMessageLength :: Word64
maxMessageLength :: Word64
maxMessageLength = Word64
1024

-- | Result of a negotiation attempt.
data NegotiationResult
  = Accepted !ProtocolId
  | NoProtocol
  deriving (Int -> NegotiationResult -> ShowS
[NegotiationResult] -> ShowS
NegotiationResult -> String
(Int -> NegotiationResult -> ShowS)
-> (NegotiationResult -> String)
-> ([NegotiationResult] -> ShowS)
-> Show NegotiationResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NegotiationResult -> ShowS
showsPrec :: Int -> NegotiationResult -> ShowS
$cshow :: NegotiationResult -> String
show :: NegotiationResult -> String
$cshowList :: [NegotiationResult] -> ShowS
showList :: [NegotiationResult] -> ShowS
Show, NegotiationResult -> NegotiationResult -> Bool
(NegotiationResult -> NegotiationResult -> Bool)
-> (NegotiationResult -> NegotiationResult -> Bool)
-> Eq NegotiationResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NegotiationResult -> NegotiationResult -> Bool
== :: NegotiationResult -> NegotiationResult -> Bool
$c/= :: NegotiationResult -> NegotiationResult -> Bool
/= :: NegotiationResult -> NegotiationResult -> Bool
Eq)

-- | Abstraction for stream I/O to enable testing with in-memory buffers.
data StreamIO = StreamIO
  { StreamIO -> ByteString -> IO ()
streamWrite    :: ByteString -> IO ()
  , StreamIO -> IO Word8
streamReadByte :: IO Word8   -- ^ Read exactly one byte (blocks until available)
  , StreamIO -> IO ()
streamClose    :: IO ()      -- ^ Close/half-close the stream (signals EOF to remote)
  }

-- | Create an in-memory stream pair for testing using STM TQueue.
-- Writes to stream A appear as reads on stream B and vice versa.
mkMemoryStreamPair :: IO (StreamIO, StreamIO)
mkMemoryStreamPair :: IO (StreamIO, StreamIO)
mkMemoryStreamPair = do
  queueAtoB <- IO (TQueue Word8)
forall a. IO (TQueue a)
newTQueueIO :: IO (TQueue Word8)
  queueBtoA <- newTQueueIO :: IO (TQueue Word8)
  let writeToQueue TQueue Word8
q ByteString
bs = (Word8 -> IO ()) -> [Word8] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> (Word8 -> STM ()) -> Word8 -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TQueue Word8 -> Word8 -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Word8
q) (ByteString -> [Word8]
BS.unpack ByteString
bs)
      readFromQueue TQueue a
q = STM a -> IO a
forall a. STM a -> IO a
atomically (TQueue a -> STM a
forall a. TQueue a -> STM a
readTQueue TQueue a
q)
  pure
    ( StreamIO (writeToQueue queueAtoB) (readFromQueue queueBtoA) (pure ())
    , StreamIO (writeToQueue queueBtoA) (readFromQueue queueAtoB) (pure ())
    )

-- | Chunk size for 'readExactBounded'. Bounds the transient boxed-list
-- allocation per read step regardless of the requested length.
readChunkSize :: Int
readChunkSize :: Int
readChunkSize = Int
32768

-- | Read exactly @n@ bytes from a stream, bounded by @maxLen@.
--
-- Shared by every length-delimited protocol in the stack (see issue
-- #169): the declared length is validated against the caller's
-- protocol-defined cap before a single byte is read or allocated, so a
-- hostile length prefix cannot trigger an unbounded allocation. Bytes
-- are accumulated in chunks of at most 'readChunkSize', keeping
-- transient memory use proportional to the chunk size, not to @n@.
--
-- I/O failures during the read (stream reset, EOF) are returned as
-- 'Left' instead of propagating as 'IOException's.
readExactBounded
  :: StreamIO
  -> Int  -- ^ Maximum acceptable length (protocol-defined cap)
  -> Int  -- ^ Number of bytes to read
  -> IO (Either String ByteString)
readExactBounded :: StreamIO -> Int -> Int -> IO (Either String ByteString)
readExactBounded StreamIO
stream Int
maxLen Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
      Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readExactBounded: negative length: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n))
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxLen =
      Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readExactBounded: requested " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n
                  String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" bytes exceeds maximum " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
maxLen))
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either String ByteString
forall a b. b -> Either a b
Right ByteString
BS.empty)
  | Bool
otherwise =
      (ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (ByteString -> Either String ByteString)
-> ([ByteString] -> ByteString)
-> [ByteString]
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> ByteString
BS.concat ([ByteString] -> Either String ByteString)
-> IO [ByteString] -> IO (Either String ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO [ByteString]
go Int
n) IO (Either String ByteString)
-> (IOException -> IO (Either String ByteString))
-> IO (Either String ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
e :: IOException) ->
        Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readExactBounded: read failed: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> IOException -> String
forall a. Show a => a -> String
show IOException
e))
  where
    go :: Int -> IO [ByteString]
    go :: Int -> IO [ByteString]
go Int
0 = [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go Int
remaining = do
      let m :: Int
m = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
readChunkSize Int
remaining
      chunk <- [Word8] -> ByteString
BS.pack ([Word8] -> ByteString) -> IO [Word8] -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO Word8 -> IO [Word8]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
m (StreamIO -> IO Word8
streamReadByte StreamIO
stream)
      (chunk :) <$> go (remaining - m)

-- | Read a complete multistream-select message from a stream.
-- Reads varint length byte-by-byte, then reads the full payload.
-- The declared length is validated against 'maxMessageLength' before
-- any payload byte is read.
readMessage :: StreamIO -> IO (Either String Text)
readMessage :: StreamIO -> IO (Either String Text)
readMessage StreamIO
stream = do
  varintResult <- StreamIO -> IO (Either String ByteString)
readVarint StreamIO
stream
  case varintResult of
    Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
    Right ByteString
varintBytes ->
      case ByteString -> Either String (Word64, ByteString)
decodeUvarint ByteString
varintBytes of
        Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
        Right (Word64
len, ByteString
_)
          | Word64
len Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
maxMessageLength ->
              Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
"readMessage: incoming message too large (max 1024 bytes)")
          | Bool
otherwise -> do
              payloadOrErr <-
                StreamIO -> Int -> Int -> IO (Either String ByteString)
readExactBounded StreamIO
stream (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
maxMessageLength) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
len)
              case payloadOrErr of
                Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
                Right ByteString
payload ->
                  case ByteString -> Either String (Text, ByteString)
decodeMessage (ByteString
varintBytes ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
payload) of
                    Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
                    Right (Text
msg, ByteString
_) -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Either String Text
forall a b. b -> Either a b
Right Text
msg)

-- | Read a varint one byte at a time from the stream.
-- The read loop is bounded at 'maxVarintBytes' (9 bytes per the
-- unsigned-varint spec) so a peer streaming continuation bytes (0x80)
-- cannot keep us reading and accumulating forever.
--
-- Like 'readExactBounded', I/O failures (stream reset, EOF from a peer
-- that disconnected mid-negotiation) are returned as 'Left', so the
-- negotiation functions report them as 'NoProtocol' instead of leaking
-- an exception.
readVarint :: StreamIO -> IO (Either String ByteString)
readVarint :: StreamIO -> IO (Either String ByteString)
readVarint StreamIO
stream =
  Int -> [Word8] -> IO (Either String ByteString)
go Int
0 [] IO (Either String ByteString)
-> (IOException -> IO (Either String ByteString))
-> IO (Either String ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
e :: IOException) ->
    Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readVarint: read failed: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> IOException -> String
forall a. Show a => a -> String
show IOException
e))
  where
    go :: Int -> [Word8] -> IO (Either String ByteString)
    go :: Int -> [Word8] -> IO (Either String ByteString)
go Int
n [Word8]
acc
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxVarintBytes =
          Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left String
"readVarint: varint too long (exceeds 9 bytes)")
      | Bool
otherwise = do
          b <- StreamIO -> IO Word8
streamReadByte StreamIO
stream
          if b < 0x80
            then pure (Right (BS.pack (reverse (b : acc))))
            else go (n + 1) (b : acc)

-- | Write a multistream-select message to a stream.
writeMessage :: StreamIO -> Text -> IO ()
writeMessage :: StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
msg = StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream (Text -> ByteString
encodeMessage Text
msg)

-- | Negotiate as the Initiator.
--
-- Pipelines the multistream header and the first protocol proposal in a
-- single write before reading anything, as the multistream-select spec
-- recommends ("the initiator SHOULD pipeline the multistream protocol
-- id and the desired protocol id in the same packet"): this saves one
-- round trip per negotiation. It then reads the header echo and the
-- reply to the optimistic proposal; on @na@ it falls back to proposing
-- the remaining protocols sequentially.
negotiateInitiator :: StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateInitiator :: StreamIO -> [Text] -> IO NegotiationResult
negotiateInitiator StreamIO
stream [] = do
  -- Nothing to propose: announce the header only (as before pipelining)
  -- and fail the negotiation after checking the peer's echo.
  StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
multistreamHeader
  result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
  case result of
    Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
    Right Text
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
negotiateInitiator StreamIO
stream (Text
firstProto : [Text]
rest) = do
  StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream (Text -> ByteString
encodeMessage Text
multistreamHeader ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeMessage Text
firstProto)
  headerReply <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
  case headerReply of
    Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
    Right Text
header
      | Text
header Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
multistreamHeader -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
      | Bool
otherwise -> Text -> IO NegotiationResult -> IO NegotiationResult
awaitReply Text
firstProto ([Text] -> IO NegotiationResult
tryProtocols [Text]
rest)
  where
    tryProtocols :: [Text] -> IO NegotiationResult
tryProtocols [] = NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
    tryProtocols (Text
proto : [Text]
remaining) = do
      StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
proto
      Text -> IO NegotiationResult -> IO NegotiationResult
awaitReply Text
proto ([Text] -> IO NegotiationResult
tryProtocols [Text]
remaining)

    -- Read the responder's answer to an already-sent proposal: an echo
    -- accepts it, @na@ runs the fallback, anything else is a protocol
    -- violation.
    awaitReply :: Text -> IO NegotiationResult -> IO NegotiationResult
awaitReply Text
proto IO NegotiationResult
onNa = do
      result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
      case result of
        Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
        Right Text
response
          | Text
response Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
proto -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> NegotiationResult
Accepted Text
proto)
          | Text
response Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
naMessage -> IO NegotiationResult
onNa
          | Bool
otherwise -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol

-- | Negotiate as the Responder.
-- Receives header, then responds to the initiator's proposal.
negotiateResponder :: StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateResponder :: StreamIO -> [Text] -> IO NegotiationResult
negotiateResponder StreamIO
stream [Text]
supported = do
  result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
  case result of
    Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
    Right Text
header
      | Text
header Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
multistreamHeader -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
      | Bool
otherwise -> do
          StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
multistreamHeader
          IO NegotiationResult
handleProposals
  where
    handleProposals :: IO NegotiationResult
handleProposals = do
      result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
      case result of
        Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
        Right Text
proposal
          | Text
proposal Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
supported -> do
              StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
proposal
              NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> NegotiationResult
Accepted Text
proposal)
          | Bool
otherwise -> do
              StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
naMessage
              IO NegotiationResult
handleProposals