From aa1e738a97d1995c1ddb821123e4842409218b0c Mon Sep 17 00:00:00 2001
From: Kim De Mey <kim.demey@gmail.com>
Date: Thu, 19 Sep 2024 17:53:11 +0200
Subject: [PATCH] Add discv5 constants to know allowed max talkresp message
 size (#732)

---
 eth/p2p/discoveryv5/encoding.nim        | 16 +++++++++
 eth/p2p/discoveryv5/protocol.nim        |  2 +-
 tests/p2p/test_discoveryv5_encoding.nim | 45 +++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/eth/p2p/discoveryv5/encoding.nim b/eth/p2p/discoveryv5/encoding.nim
index 75e55ff..8c0a5a1 100644
--- a/eth/p2p/discoveryv5/encoding.nim
+++ b/eth/p2p/discoveryv5/encoding.nim
@@ -55,6 +55,22 @@ const
   # the UDP payload and the UDP header is not taken into account.
   # https://github.com/ethereum/devp2p/blob/26e380b1f3a57db16fbdd4528dde82104c77fa38/discv5/discv5-wire.md#udp-communication
   maxDiscv5PacketSize* = 1280
+  # Following constants can be used to calculate the overhead of a packet and
+  # thus the maximum size of a payload that can be sent over talkresp.
+  discv5OrdinaryPacketOverhead* = # total 87 bytes
+    16 + # IV size
+    55 + # header size
+    16 # HMAC
+  # talkResp message = msgId + rlp: [request-id, response]
+  discv5TalkRespOverhead* = # total 16 bytes
+    1 + # talkResp msg id
+    3 + # rlp encoding outer list, max length will be encoded in 2 bytes
+    9 + # request id (max = 8) + 1 byte from rlp encoding byte string
+    3 # rlp encoding response byte string, max length in 2 bytes
+  # TalkResp message is a response message so the session is established and a
+  # ordinary discv5 packet is used for size calculation.
+  maxDiscv5TalkRespPayload* = maxDiscv5PacketSize - discv5OrdinaryPacketOverhead -
+    discv5TalkRespOverhead
 
 type
   AESGCMNonce* = array[gcmNonceSize, byte]
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index bbd30f3..5945853 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -89,7 +89,7 @@ import
     ip_vote, nodes_verification]
 
 export
-  results, node, enr, encoding.maxDiscv5PacketSize
+  results, node, enr, encoding.maxDiscv5PacketSize, encoding.maxDiscv5TalkRespPayload
 
 declareCounter discovery_message_requests_outgoing,
   "Discovery protocol outgoing message requests", labels = ["response"]
diff --git a/tests/p2p/test_discoveryv5_encoding.nim b/tests/p2p/test_discoveryv5_encoding.nim
index 0406560..3fe66ae 100644
--- a/tests/p2p/test_discoveryv5_encoding.nim
+++ b/tests/p2p/test_discoveryv5_encoding.nim
@@ -158,6 +158,24 @@ suite "Discovery v5.1 Protocol Message Encodings":
       message.kind == talkResp
       message.talkResp.response == "hi".toBytes()
 
+  test "Talk Response Payload limit":
+    let
+      payload = repeat(5.byte, maxDiscv5TalkRespPayload)
+      tr = TalkRespMessage(response: payload)
+      reqId = RequestId(id: @[1.byte, 2, 3, 4, 5, 6, 7, 8]) # max requestId = 8 bytes
+
+    let encoded = encodeMessage(tr, reqId)
+    check encoded.len() == maxDiscv5TalkRespPayload + discv5TalkRespOverhead
+
+    let decoded = decodeMessage(encoded)
+    check decoded.isOk()
+
+    let message = decoded.get()
+    check:
+      message.reqId == reqId
+      message.kind == talkResp
+      message.talkResp.response == payload
+
   test "Ping with too large RequestId":
     let
       enrSeq = 1'u64
@@ -626,3 +644,30 @@ suite "Discovery v5.1 Additional Encode/Decode":
       decoded.get().messageOpt.get().kind == ping
       decoded.get().messageOpt.get().ping.enrSeq == 0
       decoded[].requestNonce == nonce
+
+  test "Encode / Decode Ordinary Message Packet - TalkResp Payload limit":
+    let
+      payload = repeat(5.byte, maxDiscv5TalkRespPayload)
+      m = TalkRespMessage(response: payload)
+      reqId = RequestId.init(rng[])
+      message = encodeMessage(m, reqId)
+
+    # Need to manually add the secrets that normally get negotiated in the
+    # handshake packet.
+    var secrets: HandshakeSecrets
+    codecA.sessions.store(nodeB.id, nodeB.address.get(), secrets.recipientKey,
+      secrets.initiatorKey)
+    codecB.sessions.store(nodeA.id, nodeA.address.get(), secrets.initiatorKey,
+      secrets.recipientKey)
+
+    let (data, nonce) = encodeMessagePacket(rng[], codecA, nodeB.id,
+      nodeB.address.get(), message)
+
+    check data.len() == maxDiscv5PacketSize
+
+    let decoded = codecB.decodePacket(nodeA.address.get(), data)
+    check:
+      decoded.isOk()
+      decoded[].flag == OrdinaryMessage
+      decoded[].messageOpt.isSome()
+      decoded[].requestNonce == nonce