Remove serialization layer for CustomPayload (spec change)

status-im · Jan 27, 2025 · 938407b · 938407b
1 parent 60a72ed
commit 938407b
Show file tree

Hide file tree

Showing 5 changed files with 118 additions and 124 deletions.
diff --git a/fluffy/network/wire/messages.nim b/fluffy/network/wire/messages.nim
@@ -46,11 +46,13 @@ type
 
   PingMessage* = object
     enrSeq*: uint64
-    customPayload*: ByteList[2048]
+    payload_type*: uint16
+    payload*: ByteList[1100]
 
   PongMessage* = object
     enrSeq*: uint64
-    customPayload*: ByteList[2048]
+    payload_type*: uint16
+    payload*: ByteList[1100]
 
   FindNodesMessage* = object
     distances*: List[uint16, 256]

diff --git a/fluffy/network/wire/ping_extensions.nim b/fluffy/network/wire/ping_extensions.nim
@@ -9,10 +9,6 @@
 
 import ssz_serialization
 
-type CustomPayloadExtensionsFormat* = object
-  `type`*: uint16
-  payload*: ByteList[1100]
-
 const
   # Extension types
   CapabilitiesType* = 0'u16
@@ -43,47 +39,25 @@ type
     error_code*: uint16
     message*: ByteList[MAX_ERROR_BYTE_LENGTH]
 
+  CustomPayload* =
+    CapabilitiesPayload | BasicRadiusPayload | HistoryRadiusPayload | ErrorPayload
+
   ErrorCode* = enum
     ExtensionNotSupported = 0
     RequestedDataNotFound = 1
     FailedToDecodePayload = 2
     SystemError = 3
 
-func encodeCustomPayload*(payload: CapabilitiesPayload): ByteList[2048] =
-  let
-    encodedPayload = SSZ.encode(payload)
-    customPayload = CustomPayloadExtensionsFormat(
-      `type`: CapabilitiesType, payload: ByteList[1100](encodedPayload)
-    )
-  ByteList[2048](SSZ.encode(customPayload))
-
-func encodeCustomPayload*(payload: BasicRadiusPayload): ByteList[2048] =
-  let
-    encodedPayload = SSZ.encode(payload)
-    customPayload = CustomPayloadExtensionsFormat(
-      `type`: BasicRadiusType, payload: ByteList[1100](encodedPayload)
-    )
-  return ByteList[2048](SSZ.encode(customPayload))
-
-func encodeCustomPayload*(payload: HistoryRadiusPayload): ByteList[2048] =
-  let
-    encodedPayload = SSZ.encode(payload)
-    customPayload = CustomPayloadExtensionsFormat(
-      `type`: HistoryRadiusType, payload: ByteList[1100](encodedPayload)
-    )
-  return ByteList[2048](SSZ.encode(customPayload))
-
-func encodeCustomPayload*(payload: ErrorPayload): ByteList[2048] =
-  let
-    encodedPayload = SSZ.encode(payload)
-    customPayload = CustomPayloadExtensionsFormat(
-      `type`: ErrorType, payload: ByteList[1100](encodedPayload)
-    )
-  return ByteList[2048](SSZ.encode(customPayload))
-
-func encodeErrorPayload*(code: ErrorCode): ByteList[2048] =
-  encodeCustomPayload(
-    ErrorPayload(
-      error_code: uint16(ord(code)), message: ByteList[MAX_ERROR_BYTE_LENGTH].init(@[])
-    )
+func encodePayload*(payload: CustomPayload): ByteList[1100] =
+  ByteList[1100].init(SSZ.encode(payload))
+
+func encodeErrorPayload*(code: ErrorCode): (uint16, ByteList[1100]) =
+  (
+    ErrorType,
+    encodePayload(
+      ErrorPayload(
+        error_code: uint16(ord(code)),
+        message: ByteList[MAX_ERROR_BYTE_LENGTH].init(@[]),
+      )
+    ),
   )
diff --git a/fluffy/network/wire/portal_protocol.nim b/fluffy/network/wire/portal_protocol.nim
@@ -336,59 +336,64 @@ func truncateEnrs(
   enrs
 
 proc handlePingExtension(
-    p: PortalProtocol, encodedCustomPayload: ByteList[2048], srcId: NodeId
-): ByteList[2048] =
-  let customPayload = decodeSsz(
-    encodedCustomPayload.asSeq(), CustomPayloadExtensionsFormat
-  ).valueOr:
-    # invalid custom payload format, send back FailedToDecodePayload
-    return encodeErrorPayload(ErrorCode.FailedToDecodePayload)
-
-  if customPayload.`type` notin p.pingExtensionCapabilities:
+    p: PortalProtocol,
+    payloadType: uint16,
+    encodedPayload: ByteList[1100],
+    srcId: NodeId,
+): (uint16, ByteList[1100]) =
+  if payloadType notin p.pingExtensionCapabilities:
     return encodeErrorPayload(ErrorCode.ExtensionNotSupported)
 
-  case customPayload.`type`
+  case payloadType
   of CapabilitiesType:
-    let payload = decodeSsz(customPayload.payload.asSeq(), CapabilitiesPayload).valueOr:
+    let payload = decodeSsz(encodedPayload.asSeq(), CapabilitiesPayload).valueOr:
       return encodeErrorPayload(ErrorCode.FailedToDecodePayload)
 
     p.radiusCache.put(srcId, payload.data_radius)
 
-    encodeCustomPayload(
-      CapabilitiesPayload(
-        client_info: ByteList[MAX_CLIENT_INFO_BYTE_LENGTH].init(@[]),
-        data_radius: p.dataRadius(),
-        capabilities: List[uint16, MAX_CAPABILITIES_LENGTH].init(
-          p.pingExtensionCapabilities.toSeq()
-        ),
-      )
+    (
+      payloadType,
+      encodePayload(
+        CapabilitiesPayload(
+          client_info: ByteList[MAX_CLIENT_INFO_BYTE_LENGTH].init(@[]),
+          data_radius: p.dataRadius(),
+          capabilities: List[uint16, MAX_CAPABILITIES_LENGTH].init(
+            p.pingExtensionCapabilities.toSeq()
+          ),
+        )
+      ),
     )
   of BasicRadiusType:
-    let payload = decodeSsz(customPayload.payload.asSeq(), BasicRadiusPayload).valueOr:
+    let payload = decodeSsz(encodedPayload.asSeq(), BasicRadiusPayload).valueOr:
       return encodeErrorPayload(ErrorCode.FailedToDecodePayload)
 
     p.radiusCache.put(srcId, payload.data_radius)
 
-    encodeCustomPayload(HistoryRadiusPayload(data_radius: p.dataRadius()))
+    (payloadType, encodePayload(HistoryRadiusPayload(data_radius: p.dataRadius())))
   of HistoryRadiusType:
-    let payload = decodeSsz(customPayload.payload.asSeq(), HistoryRadiusPayload).valueOr:
+    let payload = decodeSsz(encodedPayload.asSeq(), HistoryRadiusPayload).valueOr:
       return encodeErrorPayload(ErrorCode.FailedToDecodePayload)
 
     p.radiusCache.put(srcId, payload.data_radius)
 
-    encodeCustomPayload(
-      HistoryRadiusPayload(data_radius: p.dataRadius(), ephemeral_header_count: 0)
+    (
+      payloadType,
+      encodePayload(
+        HistoryRadiusPayload(data_radius: p.dataRadius(), ephemeral_header_count: 0)
+      ),
     )
   else:
     encodeErrorPayload(ErrorCode.ExtensionNotSupported)
 
 proc handlePing(p: PortalProtocol, ping: PingMessage, srcId: NodeId): seq[byte] =
   # TODO: Need to think about the effect of malicious actor sending lots of
   # pings from different nodes to clear the LRU.
+  let (payloadType, payload) =
+    handlePingExtension(p, ping.payload_type, ping.payload, srcId)
+
   encodeMessage(
     PongMessage(
-      enrSeq: p.localNode.record.seqNum,
-      customPayload: handlePingExtension(p, ping.customPayload, srcId),
+      enrSeq: p.localNode.record.seqNum, payload_type: payloadType, payload: payload
     )
   )
 
@@ -697,7 +702,7 @@ proc reqResponse[Request: SomeMessage, Response: SomeMessage](
 proc pingImpl*(
     p: PortalProtocol, dst: Node
 ): Future[PortalResult[PongMessage]] {.async: (raises: [CancelledError]).} =
-  let pingCustomPayload = encodeCustomPayload(
+  let pingPayload = encodePayload(
     CapabilitiesPayload(
       client_info: ByteList[MAX_CLIENT_INFO_BYTE_LENGTH].init(@[]),
       data_radius: p.dataRadius(),
@@ -706,8 +711,11 @@ proc pingImpl*(
     )
   )
 
-  let ping =
-    PingMessage(enrSeq: p.localNode.record.seqNum, customPayload: pingCustomPayload)
+  let ping = PingMessage(
+    enrSeq: p.localNode.record.seqNum,
+    payload_type: CapabilitiesType,
+    payload: pingPayload,
+  )
 
   return await reqResponse[PingMessage, PongMessage](p, dst, ping)
 
@@ -758,15 +766,12 @@ proc ping*(
 
     let pong = pongResponse.get()
 
-    let customPayload = decodeSsz(
-      pong.customPayload.asSeq(), CustomPayloadExtensionsFormat
-    ).valueOr:
-      return err("Pong message contains invalid custom payload")
-
-    if customPayload.`type` != CapabilitiesType:
-      return err("Pong message contains invalid custom payload")
+    # Note: currently only decoding as capabilities payload as this is the only
+    # one that we support sending.
+    if pong.payload_type != CapabilitiesType:
+      return err("Pong message contains invalid or error payload")
 
-    let payload = decodeSsz(customPayload.payload.asSeq(), CapabilitiesPayload).valueOr:
+    let payload = decodeSsz(pong.payload.asSeq(), CapabilitiesPayload).valueOr:
       return err("Pong message contains invalid CapabilitiesPayload")
 
     p.radiusCache.put(dst.id, payload.data_radius)