From 222d61c52ea7508a3800222065cf78cf4c24f270 Mon Sep 17 00:00:00 2001
From: ScottyPoi <scott.simpson@ethereum.org>
Date: Tue, 14 Jan 2025 10:15:17 -0700
Subject: [PATCH] include radius in type 0 payload

---
 .../portalnetwork/src/networks/network.ts     | 58 ++++++++++---------
 .../src/wire/payloadExtensions.ts             |  3 +
 2 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/packages/portalnetwork/src/networks/network.ts b/packages/portalnetwork/src/networks/network.ts
index ded3ddfca..4f01d0a80 100644
--- a/packages/portalnetwork/src/networks/network.ts
+++ b/packages/portalnetwork/src/networks/network.ts
@@ -253,6 +253,7 @@ export abstract class BaseNetwork extends EventEmitter {
       case PingPongPayloadExtensions.CLIENT_INFO_RADIUS_AND_CAPABILITIES: {
         payload = ClientInfoAndCapabilities.serialize({
           ClientInfo: encodeClientInfo(this.portal.clientInfo),
+          DataRadius: this.nodeRadius,
           Capabilities: this.capabilities,
         })
         break
@@ -323,10 +324,11 @@ export abstract class BaseNetwork extends EventEmitter {
         )
         switch (type) {
           case PingPongPayloadExtensions.CLIENT_INFO_RADIUS_AND_CAPABILITIES: {
-            const { ClientInfo, Capabilities } = ClientInfoAndCapabilities.deserialize(payload)
+            const { ClientInfo, Capabilities, DataRadius } = ClientInfoAndCapabilities.deserialize(payload)
             this.logger.extend('PONG')(
               `Client ${shortId(enr.nodeId)} is ${decodeClientInfo(ClientInfo).clientName} node with capabilities: ${Capabilities}`,
             )
+            this.routingTable.updateRadius(enr.nodeId, DataRadius)
             break
           }
           case PingPongPayloadExtensions.BASIC_RADIUS_PAYLOAD: {
@@ -374,38 +376,42 @@ export abstract class BaseNetwork extends EventEmitter {
     }
     const { type, payload } = CustomPayloadExtensionsFormat.deserialize(pingMessage.customPayload)
     let pongPayload: Uint8Array
-    switch (type) {
-      case PingPongPayloadExtensions.CLIENT_INFO_RADIUS_AND_CAPABILITIES: {
-        pongPayload = this.pingPongPayload(type)
-        break
-      }
-      case PingPongPayloadExtensions.BASIC_RADIUS_PAYLOAD: {
-        const { dataRadius } = BasicRadius.deserialize(payload)
-        this.routingTable.updateRadius(src.nodeId, dataRadius)
-        pongPayload = this.pingPongPayload(type)
-        break
-      }
-      case PingPongPayloadExtensions.HISTORY_RADIUS_PAYLOAD: {
-        if (this.networkId !== NetworkId.HistoryNetwork) {
+    if (this.capabilities.includes(type)) {
+      switch (type) {
+        case PingPongPayloadExtensions.CLIENT_INFO_RADIUS_AND_CAPABILITIES: {
+          const { DataRadius } = ClientInfoAndCapabilities.deserialize(payload)
+          this.routingTable.updateRadius(src.nodeId, DataRadius)
+          pongPayload = this.pingPongPayload(type)
+          break
+        }
+        case PingPongPayloadExtensions.BASIC_RADIUS_PAYLOAD: {
+          const { dataRadius } = BasicRadius.deserialize(payload)
+          this.routingTable.updateRadius(src.nodeId, dataRadius)
+          pongPayload = this.pingPongPayload(type)
+          break
+        }
+        case PingPongPayloadExtensions.HISTORY_RADIUS_PAYLOAD: {
+          const { dataRadius } = HistoryRadius.deserialize(payload)
+          this.routingTable.updateRadius(src.nodeId, dataRadius)
+          pongPayload = this.pingPongPayload(type)
+          break
+        }
+        default: {
           pongPayload = ErrorPayload.serialize({
             errorCode: 0,
             message: hexToBytes(
-              fromAscii('HISTORY_RADIUS extension not supported on this network'),
+              fromAscii(`${this.constructor.name} does not support PING extension type: ${type}`),
             ),
           })
-        } else {
-          const { dataRadius } = HistoryRadius.deserialize(payload)
-          this.routingTable.updateRadius(src.nodeId, dataRadius)
-          pongPayload = this.pingPongPayload(type)
         }
-        break
-      }
-      default: {
-        pongPayload = ErrorPayload.serialize({
-          errorCode: 0,
-          message: hexToBytes(fromAscii(`Unsupported PING extension type: ${type}`)),
-        })
       }
+    } else {
+      pongPayload = ErrorPayload.serialize({
+        errorCode: 0,
+        message: hexToBytes(
+          fromAscii(`${this.constructor.name} does not support PING extension type: ${type}`),
+        ),
+      })
     }
     const customPayload = CustomPayloadExtensionsFormat.serialize({
       type,
diff --git a/packages/portalnetwork/src/wire/payloadExtensions.ts b/packages/portalnetwork/src/wire/payloadExtensions.ts
index 94503ab87..efd8d662f 100644
--- a/packages/portalnetwork/src/wire/payloadExtensions.ts
+++ b/packages/portalnetwork/src/wire/payloadExtensions.ts
@@ -91,8 +91,11 @@ export const MAX_CAPABILITIES_LENGTH = 400
 
 export const Capabilities = new ListBasicType(PingPongPayloadType, MAX_CAPABILITIES_LENGTH)
 
+export const DataRadius = new UintBigintType(32)
+
 export const ClientInfoAndCapabilities = new ContainerType({
     ClientInfo,
+    DataRadius,
     Capabilities
 })