-- | Noise message framing: 2-byte big-endian length prefix.
--
-- All Noise messages (handshake and transport) are framed as:
-- [2 bytes BE length][noise_message]
--
-- Per the libp2p Noise spec, a Noise message has a maximum length of
-- 65535 bytes. Plaintext larger than one message allows must be split
-- across multiple Noise transport messages before encryption (see
-- 'chunkPlaintext'); 'encodeFrame' rejects oversized messages instead
-- of silently truncating the length prefix.
module LibP2P.Noise.Framing
  ( encodeFrame
  , decodeFrame
  , chunkPlaintext
  , maxNoiseMessageSize
  , maxNoisePlaintextSize
  ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import LibP2P.Core.Binary (readWord16BE, word16BE)

-- | Maximum Noise message size (limited by 2-byte length prefix).
maxNoiseMessageSize :: Int
maxNoiseMessageSize :: Int
maxNoiseMessageSize = Int
65535

-- | Maximum plaintext per Noise transport message: the 65535-byte Noise
-- message cap minus the 16-byte ChaChaPoly1305 authentication tag.
-- Matches the chunking threshold used by go-libp2p.
maxNoisePlaintextSize :: Int
maxNoisePlaintextSize :: Int
maxNoisePlaintextSize = Int
maxNoiseMessageSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
16

-- | Split plaintext into chunks of at most 'maxNoisePlaintextSize' bytes,
-- so each chunk plus its AEAD tag fits in a single Noise message.
-- Empty input yields a single empty chunk so callers still emit one frame.
chunkPlaintext :: ByteString -> [ByteString]
chunkPlaintext :: ByteString -> [ByteString]
chunkPlaintext ByteString
bs
  | ByteString -> Int
BS.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxNoisePlaintextSize = [ByteString
bs]
  | Bool
otherwise =
      let (ByteString
chunk, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
maxNoisePlaintextSize ByteString
bs
       in ByteString
chunk ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: ByteString -> [ByteString]
chunkPlaintext ByteString
rest

-- | Encode a Noise message with a 2-byte big-endian length prefix.
-- Rejects messages larger than 'maxNoiseMessageSize' — the 2-byte prefix
-- cannot represent them, and truncating the length would corrupt every
-- subsequent frame boundary on the connection.
encodeFrame :: ByteString -> Either String ByteString
encodeFrame :: ByteString -> Either String ByteString
encodeFrame ByteString
msg
  | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxNoiseMessageSize =
      String -> Either String ByteString
forall a b. a -> Either a b
Left (String -> Either String ByteString)
-> String -> Either String ByteString
forall a b. (a -> b) -> a -> b
$
        String
"encodeFrame: message length " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
len
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" exceeds maximum Noise message size " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
maxNoiseMessageSize
  | Bool
otherwise = ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (ByteString -> Either String ByteString)
-> ByteString -> Either String ByteString
forall a b. (a -> b) -> a -> b
$ Word16 -> ByteString
word16BE (Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
msg
  where
    len :: Int
len = ByteString -> Int
BS.length ByteString
msg

-- | Decode a framed Noise message. Returns the message and remaining bytes.
decodeFrame :: ByteString -> Either String (ByteString, ByteString)
decodeFrame :: ByteString -> Either String (ByteString, ByteString)
decodeFrame ByteString
bs
  | ByteString -> Int
BS.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = String -> Either String (ByteString, ByteString)
forall a b. a -> Either a b
Left String
"decodeFrame: need at least 2 bytes for length"
  | Bool
otherwise =
      let len :: Int
len = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Word16
readWord16BE ByteString
bs) :: Int
          rest :: ByteString
rest = Int -> ByteString -> ByteString
BS.drop Int
2 ByteString
bs
       in if ByteString -> Int
BS.length ByteString
rest Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
            then String -> Either String (ByteString, ByteString)
forall a b. a -> Either a b
Left (String -> Either String (ByteString, ByteString))
-> String -> Either String (ByteString, ByteString)
forall a b. (a -> b) -> a -> b
$ String
"decodeFrame: expected " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
len String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" bytes but got " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (ByteString -> Int
BS.length ByteString
rest)
            else (ByteString, ByteString) -> Either String (ByteString, ByteString)
forall a b. b -> Either a b
Right (Int -> ByteString -> ByteString
BS.take Int
len ByteString
rest, Int -> ByteString -> ByteString
BS.drop Int
len ByteString
rest)