-- | DHT RPC message encoding/decoding (protobuf).
--
-- Wire format from docs/09-dht.md:
--   Message framing: [uvarint length][protobuf message]
--   Message fields: type(1), key(2), record(3), closerPeers(8), providerPeers(9)
--   Record fields: key(1), value(2), timeReceived(5)
--   Peer fields: id(1), addrs(2), connection(3)
--
-- Uses proto3-wire for protobuf encoding/decoding, same pattern as Identify.Message.
module Network.LibP2P.DHT.Message
  ( -- * Types
    MessageType (..)
  , DHTRecord (..)
  , DHTPeer (..)
  , DHTMessage (..)
    -- * Protobuf encode/decode (no framing)
  , encodeDHTMessage
  , decodeDHTMessage
    -- * Wire framing (uvarint length prefix)
  , encodeFramed
  , decodeFramed
    -- * Stream I/O helpers
  , writeFramedMessage
  , readFramedMessage
    -- * Constants
  , maxDHTMessageSize
    -- * Defaults
  , emptyDHTMessage
  ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import Data.Text (Text)
import qualified Data.Text.Lazy as TL
import Data.Word (Word32)
import Proto3.Wire.Decode (Parser, RawMessage, ParseError, at, one, repeated, embedded, embedded', parse)
import qualified Proto3.Wire.Decode as Decode
import Proto3.Wire.Encode (MessageBuilder)
import qualified Proto3.Wire.Encode as Encode
import Proto3.Wire.Types (FieldNumber (..))
import Network.LibP2P.Core.Varint (encodeUvarint, decodeUvarint)
import Network.LibP2P.DHT.Types (ConnectionType (..))
import Network.LibP2P.MultistreamSelect.Negotiation (StreamIO (..))

-- | Maximum DHT message size: 64 KiB.
maxDHTMessageSize :: Int
maxDHTMessageSize :: Int
maxDHTMessageSize = Int
64 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024

-- | DHT RPC message type (protobuf enum).
data MessageType
  = PutValue      -- ^ 0
  | GetValue      -- ^ 1
  | AddProvider   -- ^ 2
  | GetProviders  -- ^ 3
  | FindNode      -- ^ 4
  deriving (Int -> MessageType -> ShowS
[MessageType] -> ShowS
MessageType -> String
(Int -> MessageType -> ShowS)
-> (MessageType -> String)
-> ([MessageType] -> ShowS)
-> Show MessageType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MessageType -> ShowS
showsPrec :: Int -> MessageType -> ShowS
$cshow :: MessageType -> String
show :: MessageType -> String
$cshowList :: [MessageType] -> ShowS
showList :: [MessageType] -> ShowS
Show, MessageType -> MessageType -> Bool
(MessageType -> MessageType -> Bool)
-> (MessageType -> MessageType -> Bool) -> Eq MessageType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MessageType -> MessageType -> Bool
== :: MessageType -> MessageType -> Bool
$c/= :: MessageType -> MessageType -> Bool
/= :: MessageType -> MessageType -> Bool
Eq, Int -> MessageType
MessageType -> Int
MessageType -> [MessageType]
MessageType -> MessageType
MessageType -> MessageType -> [MessageType]
MessageType -> MessageType -> MessageType -> [MessageType]
(MessageType -> MessageType)
-> (MessageType -> MessageType)
-> (Int -> MessageType)
-> (MessageType -> Int)
-> (MessageType -> [MessageType])
-> (MessageType -> MessageType -> [MessageType])
-> (MessageType -> MessageType -> [MessageType])
-> (MessageType -> MessageType -> MessageType -> [MessageType])
-> Enum MessageType
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: MessageType -> MessageType
succ :: MessageType -> MessageType
$cpred :: MessageType -> MessageType
pred :: MessageType -> MessageType
$ctoEnum :: Int -> MessageType
toEnum :: Int -> MessageType
$cfromEnum :: MessageType -> Int
fromEnum :: MessageType -> Int
$cenumFrom :: MessageType -> [MessageType]
enumFrom :: MessageType -> [MessageType]
$cenumFromThen :: MessageType -> MessageType -> [MessageType]
enumFromThen :: MessageType -> MessageType -> [MessageType]
$cenumFromTo :: MessageType -> MessageType -> [MessageType]
enumFromTo :: MessageType -> MessageType -> [MessageType]
$cenumFromThenTo :: MessageType -> MessageType -> MessageType -> [MessageType]
enumFromThenTo :: MessageType -> MessageType -> MessageType -> [MessageType]
Enum, MessageType
MessageType -> MessageType -> Bounded MessageType
forall a. a -> a -> Bounded a
$cminBound :: MessageType
minBound :: MessageType
$cmaxBound :: MessageType
maxBound :: MessageType
Bounded)

-- | DHT record (protobuf Message.Record).
data DHTRecord = DHTRecord
  { DHTRecord -> ByteString
recKey          :: !ByteString   -- ^ field 1
  , DHTRecord -> ByteString
recValue        :: !ByteString   -- ^ field 2
  , DHTRecord -> Text
recTimeReceived :: !Text         -- ^ field 5, RFC 3339
  } deriving (Int -> DHTRecord -> ShowS
[DHTRecord] -> ShowS
DHTRecord -> String
(Int -> DHTRecord -> ShowS)
-> (DHTRecord -> String)
-> ([DHTRecord] -> ShowS)
-> Show DHTRecord
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DHTRecord -> ShowS
showsPrec :: Int -> DHTRecord -> ShowS
$cshow :: DHTRecord -> String
show :: DHTRecord -> String
$cshowList :: [DHTRecord] -> ShowS
showList :: [DHTRecord] -> ShowS
Show, DHTRecord -> DHTRecord -> Bool
(DHTRecord -> DHTRecord -> Bool)
-> (DHTRecord -> DHTRecord -> Bool) -> Eq DHTRecord
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DHTRecord -> DHTRecord -> Bool
== :: DHTRecord -> DHTRecord -> Bool
$c/= :: DHTRecord -> DHTRecord -> Bool
/= :: DHTRecord -> DHTRecord -> Bool
Eq)

-- | DHT peer info (protobuf Message.Peer).
data DHTPeer = DHTPeer
  { DHTPeer -> ByteString
dhtPeerId       :: !ByteString       -- ^ field 1: raw Peer ID bytes
  , DHTPeer -> [ByteString]
dhtPeerAddrs    :: ![ByteString]     -- ^ field 2: raw multiaddr bytes
  , DHTPeer -> ConnectionType
dhtPeerConnType :: !ConnectionType   -- ^ field 3: connection capability
  } deriving (Int -> DHTPeer -> ShowS
[DHTPeer] -> ShowS
DHTPeer -> String
(Int -> DHTPeer -> ShowS)
-> (DHTPeer -> String) -> ([DHTPeer] -> ShowS) -> Show DHTPeer
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DHTPeer -> ShowS
showsPrec :: Int -> DHTPeer -> ShowS
$cshow :: DHTPeer -> String
show :: DHTPeer -> String
$cshowList :: [DHTPeer] -> ShowS
showList :: [DHTPeer] -> ShowS
Show, DHTPeer -> DHTPeer -> Bool
(DHTPeer -> DHTPeer -> Bool)
-> (DHTPeer -> DHTPeer -> Bool) -> Eq DHTPeer
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DHTPeer -> DHTPeer -> Bool
== :: DHTPeer -> DHTPeer -> Bool
$c/= :: DHTPeer -> DHTPeer -> Bool
/= :: DHTPeer -> DHTPeer -> Bool
Eq)

-- | DHT RPC message (protobuf Message).
data DHTMessage = DHTMessage
  { DHTMessage -> MessageType
msgType          :: !MessageType       -- ^ field 1
  , DHTMessage -> ByteString
msgKey           :: !ByteString        -- ^ field 2
  , DHTMessage -> Maybe DHTRecord
msgRecord        :: !(Maybe DHTRecord) -- ^ field 3
  , DHTMessage -> [DHTPeer]
msgCloserPeers   :: ![DHTPeer]         -- ^ field 8
  , DHTMessage -> [DHTPeer]
msgProviderPeers :: ![DHTPeer]         -- ^ field 9
  } deriving (Int -> DHTMessage -> ShowS
[DHTMessage] -> ShowS
DHTMessage -> String
(Int -> DHTMessage -> ShowS)
-> (DHTMessage -> String)
-> ([DHTMessage] -> ShowS)
-> Show DHTMessage
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DHTMessage -> ShowS
showsPrec :: Int -> DHTMessage -> ShowS
$cshow :: DHTMessage -> String
show :: DHTMessage -> String
$cshowList :: [DHTMessage] -> ShowS
showList :: [DHTMessage] -> ShowS
Show, DHTMessage -> DHTMessage -> Bool
(DHTMessage -> DHTMessage -> Bool)
-> (DHTMessage -> DHTMessage -> Bool) -> Eq DHTMessage
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DHTMessage -> DHTMessage -> Bool
== :: DHTMessage -> DHTMessage -> Bool
$c/= :: DHTMessage -> DHTMessage -> Bool
/= :: DHTMessage -> DHTMessage -> Bool
Eq)

-- | Default empty message (FIND_NODE with empty key, no record, no peers).
emptyDHTMessage :: DHTMessage
emptyDHTMessage :: DHTMessage
emptyDHTMessage = DHTMessage
  { msgType :: MessageType
msgType          = MessageType
PutValue
  , msgKey :: ByteString
msgKey           = ByteString
BS.empty
  , msgRecord :: Maybe DHTRecord
msgRecord        = Maybe DHTRecord
forall a. Maybe a
Nothing
  , msgCloserPeers :: [DHTPeer]
msgCloserPeers   = []
  , msgProviderPeers :: [DHTPeer]
msgProviderPeers = []
  }

-- Encoding

-- | Encode a DHTMessage to protobuf wire format (no length prefix).
encodeDHTMessage :: DHTMessage -> ByteString
encodeDHTMessage :: DHTMessage -> ByteString
encodeDHTMessage DHTMessage
msg = LazyByteString -> ByteString
BL.toStrict (LazyByteString -> ByteString) -> LazyByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ MessageBuilder -> LazyByteString
Encode.toLazyByteString (MessageBuilder -> LazyByteString)
-> MessageBuilder -> LazyByteString
forall a b. (a -> b) -> a -> b
$
     FieldNumber -> Word32 -> MessageBuilder
Encode.uint32 (Word64 -> FieldNumber
FieldNumber Word64
1) (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (MessageType -> Int
forall a. Enum a => a -> Int
fromEnum (DHTMessage -> MessageType
msgType DHTMessage
msg)))
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> Word -> Maybe ByteString -> MessageBuilder
optBytes Word
2 (ByteString -> Maybe ByteString
nonEmpty (DHTMessage -> ByteString
msgKey DHTMessage
msg))
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> Maybe DHTRecord -> MessageBuilder
optRecord (DHTMessage -> Maybe DHTRecord
msgRecord DHTMessage
msg)
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> Word -> (DHTPeer -> MessageBuilder) -> [DHTPeer] -> MessageBuilder
forall a. Word -> (a -> MessageBuilder) -> [a] -> MessageBuilder
repEmbedded Word
8 DHTPeer -> MessageBuilder
encodeDHTPeer (DHTMessage -> [DHTPeer]
msgCloserPeers DHTMessage
msg)
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> Word -> (DHTPeer -> MessageBuilder) -> [DHTPeer] -> MessageBuilder
forall a. Word -> (a -> MessageBuilder) -> [a] -> MessageBuilder
repEmbedded Word
9 DHTPeer -> MessageBuilder
encodeDHTPeer (DHTMessage -> [DHTPeer]
msgProviderPeers DHTMessage
msg)
  where
    optBytes :: Word -> Maybe ByteString -> MessageBuilder
    optBytes :: Word -> Maybe ByteString -> MessageBuilder
optBytes Word
_ Maybe ByteString
Nothing  = MessageBuilder
forall a. Monoid a => a
mempty
    optBytes Word
n (Just ByteString
v) = FieldNumber -> ByteString -> MessageBuilder
Encode.byteString (Word64 -> FieldNumber
FieldNumber (Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)) ByteString
v

    nonEmpty :: ByteString -> Maybe ByteString
    nonEmpty :: ByteString -> Maybe ByteString
nonEmpty ByteString
bs
      | ByteString -> Bool
BS.null ByteString
bs = Maybe ByteString
forall a. Maybe a
Nothing
      | Bool
otherwise  = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs

    optRecord :: Maybe DHTRecord -> MessageBuilder
    optRecord :: Maybe DHTRecord -> MessageBuilder
optRecord Maybe DHTRecord
Nothing  = MessageBuilder
forall a. Monoid a => a
mempty
    optRecord (Just DHTRecord
r) = FieldNumber -> MessageBuilder -> MessageBuilder
Encode.embedded (Word64 -> FieldNumber
FieldNumber Word64
3) (DHTRecord -> MessageBuilder
encodeDHTRecord DHTRecord
r)

    repEmbedded :: Word -> (a -> MessageBuilder) -> [a] -> MessageBuilder
    repEmbedded :: forall a. Word -> (a -> MessageBuilder) -> [a] -> MessageBuilder
repEmbedded Word
n a -> MessageBuilder
enc = (a -> MessageBuilder) -> [a] -> MessageBuilder
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
x -> FieldNumber -> MessageBuilder -> MessageBuilder
Encode.embedded (Word64 -> FieldNumber
FieldNumber (Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)) (a -> MessageBuilder
enc a
x))

-- | Encode a DHTRecord sub-message.
encodeDHTRecord :: DHTRecord -> MessageBuilder
encodeDHTRecord :: DHTRecord -> MessageBuilder
encodeDHTRecord DHTRecord
rec =
     Word -> Maybe ByteString -> MessageBuilder
optBytes Word
1 (ByteString -> Maybe ByteString
nonEmpty (DHTRecord -> ByteString
recKey DHTRecord
rec))
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> Word -> Maybe ByteString -> MessageBuilder
optBytes Word
2 (ByteString -> Maybe ByteString
nonEmpty (DHTRecord -> ByteString
recValue DHTRecord
rec))
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> Word -> Maybe Text -> MessageBuilder
optText Word
5 (Text -> Maybe Text
nonEmpty' (DHTRecord -> Text
recTimeReceived DHTRecord
rec))
  where
    optBytes :: Word -> Maybe ByteString -> MessageBuilder
    optBytes :: Word -> Maybe ByteString -> MessageBuilder
optBytes Word
_ Maybe ByteString
Nothing  = MessageBuilder
forall a. Monoid a => a
mempty
    optBytes Word
n (Just ByteString
v) = FieldNumber -> ByteString -> MessageBuilder
Encode.byteString (Word64 -> FieldNumber
FieldNumber (Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)) ByteString
v

    optText :: Word -> Maybe Text -> MessageBuilder
    optText :: Word -> Maybe Text -> MessageBuilder
optText Word
_ Maybe Text
Nothing  = MessageBuilder
forall a. Monoid a => a
mempty
    optText Word
n (Just Text
v) = FieldNumber -> Text -> MessageBuilder
Encode.text (Word64 -> FieldNumber
FieldNumber (Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)) (Text -> Text
TL.fromStrict Text
v)

    nonEmpty :: ByteString -> Maybe ByteString
    nonEmpty :: ByteString -> Maybe ByteString
nonEmpty ByteString
bs
      | ByteString -> Bool
BS.null ByteString
bs = Maybe ByteString
forall a. Maybe a
Nothing
      | Bool
otherwise  = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs

    nonEmpty' :: Text -> Maybe Text
    nonEmpty' :: Text -> Maybe Text
nonEmpty' Text
t
      | Text
t Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
""   = Maybe Text
forall a. Maybe a
Nothing
      | Bool
otherwise = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
t

-- | Encode a DHTPeer sub-message.
encodeDHTPeer :: DHTPeer -> MessageBuilder
encodeDHTPeer :: DHTPeer -> MessageBuilder
encodeDHTPeer DHTPeer
peer =
     Word -> Maybe ByteString -> MessageBuilder
optBytes Word
1 (ByteString -> Maybe ByteString
nonEmpty (DHTPeer -> ByteString
dhtPeerId DHTPeer
peer))
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> (ByteString -> MessageBuilder) -> [ByteString] -> MessageBuilder
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\ByteString
a -> FieldNumber -> ByteString -> MessageBuilder
Encode.byteString (Word64 -> FieldNumber
FieldNumber Word64
2) ByteString
a) (DHTPeer -> [ByteString]
dhtPeerAddrs DHTPeer
peer)
  MessageBuilder -> MessageBuilder -> MessageBuilder
forall a. Semigroup a => a -> a -> a
<> FieldNumber -> Word32 -> MessageBuilder
Encode.uint32 (Word64 -> FieldNumber
FieldNumber Word64
3) (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ConnectionType -> Int
forall a. Enum a => a -> Int
fromEnum (DHTPeer -> ConnectionType
dhtPeerConnType DHTPeer
peer)))
  where
    optBytes :: Word -> Maybe ByteString -> MessageBuilder
    optBytes :: Word -> Maybe ByteString -> MessageBuilder
optBytes Word
_ Maybe ByteString
Nothing  = MessageBuilder
forall a. Monoid a => a
mempty
    optBytes Word
n (Just ByteString
v) = FieldNumber -> ByteString -> MessageBuilder
Encode.byteString (Word64 -> FieldNumber
FieldNumber (Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)) ByteString
v

    nonEmpty :: ByteString -> Maybe ByteString
    nonEmpty :: ByteString -> Maybe ByteString
nonEmpty ByteString
bs
      | ByteString -> Bool
BS.null ByteString
bs = Maybe ByteString
forall a. Maybe a
Nothing
      | Bool
otherwise  = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs

-- Decoding

-- | Decode a DHTMessage from protobuf wire format.
decodeDHTMessage :: ByteString -> Either ParseError DHTMessage
decodeDHTMessage :: ByteString -> Either ParseError DHTMessage
decodeDHTMessage = Parser RawMessage DHTMessage
-> ByteString -> Either ParseError DHTMessage
forall a. Parser RawMessage a -> ByteString -> Either ParseError a
parse Parser RawMessage DHTMessage
dhtMessageParser

dhtMessageParser :: Parser RawMessage DHTMessage
dhtMessageParser :: Parser RawMessage DHTMessage
dhtMessageParser = MessageType
-> ByteString
-> Maybe DHTRecord
-> [DHTPeer]
-> [DHTPeer]
-> DHTMessage
DHTMessage
  (MessageType
 -> ByteString
 -> Maybe DHTRecord
 -> [DHTPeer]
 -> [DHTPeer]
 -> DHTMessage)
-> Parser RawMessage MessageType
-> Parser
     RawMessage
     (ByteString
      -> Maybe DHTRecord -> [DHTPeer] -> [DHTPeer] -> DHTMessage)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word32 -> MessageType
toMessageType (Word32 -> MessageType)
-> Parser RawMessage Word32 -> Parser RawMessage MessageType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser RawField Word32 -> FieldNumber -> Parser RawMessage Word32
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive Word32 -> Word32 -> Parser RawField Word32
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one Parser RawPrimitive Word32
Decode.uint32 Word32
0) (Word64 -> FieldNumber
FieldNumber Word64
1))
  Parser
  RawMessage
  (ByteString
   -> Maybe DHTRecord -> [DHTPeer] -> [DHTPeer] -> DHTMessage)
-> Parser RawMessage ByteString
-> Parser
     RawMessage
     (Maybe DHTRecord -> [DHTPeer] -> [DHTPeer] -> DHTMessage)
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField ByteString
-> FieldNumber -> Parser RawMessage ByteString
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive ByteString
-> ByteString -> Parser RawField ByteString
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one Parser RawPrimitive ByteString
Decode.byteString ByteString
BS.empty) (Word64 -> FieldNumber
FieldNumber Word64
2)
  Parser
  RawMessage
  (Maybe DHTRecord -> [DHTPeer] -> [DHTPeer] -> DHTMessage)
-> Parser RawMessage (Maybe DHTRecord)
-> Parser RawMessage ([DHTPeer] -> [DHTPeer] -> DHTMessage)
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField (Maybe DHTRecord)
-> FieldNumber -> Parser RawMessage (Maybe DHTRecord)
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawMessage DHTRecord -> Parser RawField (Maybe DHTRecord)
forall a. Parser RawMessage a -> Parser RawField (Maybe a)
embedded Parser RawMessage DHTRecord
dhtRecordParser) (Word64 -> FieldNumber
FieldNumber Word64
3)
  Parser RawMessage ([DHTPeer] -> [DHTPeer] -> DHTMessage)
-> Parser RawMessage [DHTPeer]
-> Parser RawMessage ([DHTPeer] -> DHTMessage)
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField [DHTPeer]
-> FieldNumber -> Parser RawMessage [DHTPeer]
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive DHTPeer -> Parser RawField [DHTPeer]
forall a. Parser RawPrimitive a -> Parser RawField [a]
repeated (Parser RawMessage DHTPeer -> Parser RawPrimitive DHTPeer
forall a. Parser RawMessage a -> Parser RawPrimitive a
embedded' Parser RawMessage DHTPeer
dhtRecordedPeerParser)) (Word64 -> FieldNumber
FieldNumber Word64
8)
  Parser RawMessage ([DHTPeer] -> DHTMessage)
-> Parser RawMessage [DHTPeer] -> Parser RawMessage DHTMessage
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField [DHTPeer]
-> FieldNumber -> Parser RawMessage [DHTPeer]
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive DHTPeer -> Parser RawField [DHTPeer]
forall a. Parser RawPrimitive a -> Parser RawField [a]
repeated (Parser RawMessage DHTPeer -> Parser RawPrimitive DHTPeer
forall a. Parser RawMessage a -> Parser RawPrimitive a
embedded' Parser RawMessage DHTPeer
dhtRecordedPeerParser)) (Word64 -> FieldNumber
FieldNumber Word64
9)

-- | Parse a DHTRecord sub-message.
dhtRecordParser :: Parser RawMessage DHTRecord
dhtRecordParser :: Parser RawMessage DHTRecord
dhtRecordParser = ByteString -> ByteString -> Text -> DHTRecord
DHTRecord
  (ByteString -> ByteString -> Text -> DHTRecord)
-> Parser RawMessage ByteString
-> Parser RawMessage (ByteString -> Text -> DHTRecord)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser RawField ByteString
-> FieldNumber -> Parser RawMessage ByteString
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive ByteString
-> ByteString -> Parser RawField ByteString
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one Parser RawPrimitive ByteString
Decode.byteString ByteString
BS.empty) (Word64 -> FieldNumber
FieldNumber Word64
1)
  Parser RawMessage (ByteString -> Text -> DHTRecord)
-> Parser RawMessage ByteString
-> Parser RawMessage (Text -> DHTRecord)
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField ByteString
-> FieldNumber -> Parser RawMessage ByteString
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive ByteString
-> ByteString -> Parser RawField ByteString
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one Parser RawPrimitive ByteString
Decode.byteString ByteString
BS.empty) (Word64 -> FieldNumber
FieldNumber Word64
2)
  Parser RawMessage (Text -> DHTRecord)
-> Parser RawMessage Text -> Parser RawMessage DHTRecord
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField Text -> FieldNumber -> Parser RawMessage Text
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive Text -> Text -> Parser RawField Text
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one (Text -> Text
TL.toStrict (Text -> Text)
-> Parser RawPrimitive Text -> Parser RawPrimitive Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser RawPrimitive Text
Decode.text) Text
"") (Word64 -> FieldNumber
FieldNumber Word64
5)

-- | Parse a DHTPeer sub-message.
dhtRecordedPeerParser :: Parser RawMessage DHTPeer
dhtRecordedPeerParser :: Parser RawMessage DHTPeer
dhtRecordedPeerParser = ByteString -> [ByteString] -> ConnectionType -> DHTPeer
DHTPeer
  (ByteString -> [ByteString] -> ConnectionType -> DHTPeer)
-> Parser RawMessage ByteString
-> Parser RawMessage ([ByteString] -> ConnectionType -> DHTPeer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser RawField ByteString
-> FieldNumber -> Parser RawMessage ByteString
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive ByteString
-> ByteString -> Parser RawField ByteString
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one Parser RawPrimitive ByteString
Decode.byteString ByteString
BS.empty) (Word64 -> FieldNumber
FieldNumber Word64
1)
  Parser RawMessage ([ByteString] -> ConnectionType -> DHTPeer)
-> Parser RawMessage [ByteString]
-> Parser RawMessage (ConnectionType -> DHTPeer)
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser RawField [ByteString]
-> FieldNumber -> Parser RawMessage [ByteString]
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive ByteString -> Parser RawField [ByteString]
forall a. Parser RawPrimitive a -> Parser RawField [a]
repeated Parser RawPrimitive ByteString
Decode.byteString) (Word64 -> FieldNumber
FieldNumber Word64
2)
  Parser RawMessage (ConnectionType -> DHTPeer)
-> Parser RawMessage ConnectionType -> Parser RawMessage DHTPeer
forall a b.
Parser RawMessage (a -> b)
-> Parser RawMessage a -> Parser RawMessage b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Word32 -> ConnectionType
toConnectionType (Word32 -> ConnectionType)
-> Parser RawMessage Word32 -> Parser RawMessage ConnectionType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser RawField Word32 -> FieldNumber -> Parser RawMessage Word32
forall a. Parser RawField a -> FieldNumber -> Parser RawMessage a
at (Parser RawPrimitive Word32 -> Word32 -> Parser RawField Word32
forall a. Parser RawPrimitive a -> a -> Parser RawField a
one Parser RawPrimitive Word32
Decode.uint32 Word32
0) (Word64 -> FieldNumber
FieldNumber Word64
3))

-- | Convert Word32 to MessageType (clamped to valid range).
toMessageType :: Word32 -> MessageType
toMessageType :: Word32 -> MessageType
toMessageType Word32
n
  | Word32
n Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (MessageType -> Int
forall a. Enum a => a -> Int
fromEnum (MessageType
forall a. Bounded a => a
maxBound :: MessageType)) = Int -> MessageType
forall a. Enum a => Int -> a
toEnum (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
n)
  | Bool
otherwise = MessageType
PutValue  -- default for unknown

-- | Convert Word32 to ConnectionType (clamped to valid range).
toConnectionType :: Word32 -> ConnectionType
toConnectionType :: Word32 -> ConnectionType
toConnectionType Word32
n
  | Word32
n Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ConnectionType -> Int
forall a. Enum a => a -> Int
fromEnum (ConnectionType
forall a. Bounded a => a
maxBound :: ConnectionType)) = Int -> ConnectionType
forall a. Enum a => Int -> a
toEnum (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
n)
  | Bool
otherwise = ConnectionType
NotConnected  -- default for unknown

-- Wire framing

-- | Encode a DHTMessage with uvarint length prefix.
encodeFramed :: DHTMessage -> ByteString
encodeFramed :: DHTMessage -> ByteString
encodeFramed DHTMessage
msg =
  let payload :: ByteString
payload = DHTMessage -> ByteString
encodeDHTMessage DHTMessage
msg
      lenPrefix :: ByteString
lenPrefix = Word64 -> ByteString
encodeUvarint (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
payload))
  in ByteString
lenPrefix ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
payload

-- | Decode a DHTMessage from uvarint-length-prefixed bytes.
decodeFramed :: Int -> ByteString -> Either String DHTMessage
decodeFramed :: Int -> ByteString -> Either String DHTMessage
decodeFramed Int
maxSize ByteString
bs = do
  (len, rest) <- ByteString -> Either String (Word64, ByteString)
decodeUvarint ByteString
bs
  let msgLen = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
len :: Int
  if msgLen > maxSize
    then Left $ "DHT message too large: " ++ show msgLen ++ " > " ++ show maxSize
    else if BS.length rest < msgLen
      then Left $ "DHT message truncated: expected " ++ show msgLen ++ " bytes, got " ++ show (BS.length rest)
      else case decodeDHTMessage (BS.take msgLen rest) of
        Left ParseError
err -> String -> Either String DHTMessage
forall a b. a -> Either a b
Left (String -> Either String DHTMessage)
-> String -> Either String DHTMessage
forall a b. (a -> b) -> a -> b
$ String
"DHT protobuf decode error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ParseError -> String
forall a. Show a => a -> String
show ParseError
err
        Right DHTMessage
msg -> DHTMessage -> Either String DHTMessage
forall a b. b -> Either a b
Right DHTMessage
msg

-- Stream I/O helpers

-- | Write a framed DHT message to a stream.
writeFramedMessage :: StreamIO -> DHTMessage -> IO ()
writeFramedMessage :: StreamIO -> DHTMessage -> IO ()
writeFramedMessage StreamIO
stream DHTMessage
msg = StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream (DHTMessage -> ByteString
encodeFramed DHTMessage
msg)

-- | Read a framed DHT message from a stream.
-- Reads the uvarint length prefix, then the protobuf payload.
readFramedMessage :: StreamIO -> Int -> IO (Either String DHTMessage)
readFramedMessage :: StreamIO -> Int -> IO (Either String DHTMessage)
readFramedMessage StreamIO
stream Int
maxSize = do
  -- Read varint bytes one at a time (up to 10 bytes)
  varintBytes <- StreamIO -> IO ByteString
readVarintBytes StreamIO
stream
  case decodeUvarint varintBytes of
    Left String
err -> Either String DHTMessage -> IO (Either String DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String DHTMessage
forall a b. a -> Either a b
Left (String -> Either String DHTMessage)
-> String -> Either String DHTMessage
forall a b. (a -> b) -> a -> b
$ String
"DHT varint decode error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
err)
    Right (Word64
len, ByteString
_) -> do
      let msgLen :: Int
msgLen = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
len :: Int
      if Int
msgLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxSize
        then Either String DHTMessage -> IO (Either String DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String DHTMessage
forall a b. a -> Either a b
Left (String -> Either String DHTMessage)
-> String -> Either String DHTMessage
forall a b. (a -> b) -> a -> b
$ String
"DHT message too large: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
msgLen String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" > " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
maxSize)
        else do
          payload <- StreamIO -> Int -> IO ByteString
readExact StreamIO
stream Int
msgLen
          case decodeDHTMessage payload of
            Left ParseError
err -> Either String DHTMessage -> IO (Either String DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String DHTMessage
forall a b. a -> Either a b
Left (String -> Either String DHTMessage)
-> String -> Either String DHTMessage
forall a b. (a -> b) -> a -> b
$ String
"DHT protobuf decode error: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ParseError -> String
forall a. Show a => a -> String
show ParseError
err)
            Right DHTMessage
msg -> Either String DHTMessage -> IO (Either String DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DHTMessage -> Either String DHTMessage
forall a b. b -> Either a b
Right DHTMessage
msg)

-- | Read exactly n bytes from a stream.
readExact :: StreamIO -> Int -> IO ByteString
readExact :: StreamIO -> Int -> IO ByteString
readExact StreamIO
stream Int
n = [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) -> [Int] -> IO [Word8]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (IO Word8 -> Int -> IO Word8
forall a b. a -> b -> a
const (StreamIO -> IO Word8
streamReadByte StreamIO
stream)) [Int
1 .. Int
n]

-- | Read unsigned varint bytes from a stream (up to 10 bytes).
readVarintBytes :: StreamIO -> IO ByteString
readVarintBytes :: StreamIO -> IO ByteString
readVarintBytes StreamIO
stream = [Word8] -> Int -> IO ByteString
forall {t}. (Ord t, Num t) => [Word8] -> t -> IO ByteString
go [] (Int
0 :: Int)
  where
    go :: [Word8] -> t -> IO ByteString
go [Word8]
acc t
n
      | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
>= t
10 = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Word8] -> ByteString
BS.pack ([Word8] -> [Word8]
forall a. [a] -> [a]
reverse [Word8]
acc))  -- max varint length
      | Bool
otherwise = do
          b <- StreamIO -> IO Word8
streamReadByte StreamIO
stream
          if b < 0x80
            then pure (BS.pack (reverse (b : acc)))
            else go (b : acc) (n + 1)