feat: chunkify handleFindNode

src/service/service.ts

@@ -782,8 +782,8 @@ export class Discv5 extends (EventEmitter as { new (): Discv5EventEmitter }) {
    */
   private handleFindNode(nodeAddr: INodeAddress, message: IFindNodeMessage): void {
     const { id, distances } = message;
-    let nodes: ENR[] = [];
-    for (const distance of new Set(distances)) {
+    const nodes: ENR[] = [];
+    topLoop: for (const distance of new Set(distances)) {
       // filter out invalid distances
       if (distance < 0 || distance > 256) {
         continue;
@@ -795,33 +795,30 @@ export class Discv5 extends (EventEmitter as { new (): Discv5EventEmitter }) {
           this.enr.encodeToValues(this.keypair.privateKey);
         }
         nodes.push(this.enr);
+        if (nodes.length >= 16) {
+          break topLoop;
+        }
       } else {
-        nodes.push(...this.kbuckets.valuesOfDistance(distance));
-      }
-    }
-    // limit response to 16 nodes
-    nodes = nodes.slice(0, 16);
-    if (nodes.length === 0) {
-      log("Sending empty NODES response to %o", nodeAddr);
-      try {
-        this.sessionService.sendResponse(nodeAddr, createNodesMessage(id, 1, nodes));
-        this.metrics?.sentMessageCount.inc({ type: MessageType[MessageType.NODES] });
-      } catch (e) {
-        log("Failed to send a NODES response. Error: %s", (e as Error).message);
+        for (const node of this.kbuckets.valuesOfDistance(distance)) {
+          nodes.push(node);
+          if (nodes.length >= 16) {
+            break topLoop;
+          }
+        }
       }
-      return;
     }
+
     // Responses assume that a session is established.
     // Thus, on top of the encoded ENRs the packet should be a regular message.
     // A regular message has a tag (32 bytes), an authTag (12 bytes)
     // and the NODES response has an ID (8 bytes) and a total (8 bytes).
     // The encryption adds the HMAC (16 bytes) and can be at most 16 bytes larger
     // So, the total empty packet size can be at most 92
     const nodesPerPacket = Math.floor((MAX_PACKET_SIZE - 92) / MAX_RECORD_SIZE);
-    const total = Math.ceil(nodes.length / nodesPerPacket);
+    const nodesChunks = chunkify(nodes, nodesPerPacket);
+    const total = nodesChunks.length;
     log("Sending %d NODES responses to %o", total, nodeAddr);
-    for (let i = 0; i < nodes.length; i += nodesPerPacket) {
-      const _nodes = nodes.slice(i, i + nodesPerPacket);
+    for (const _nodes of nodesChunks) {
       try {
         this.sessionService.sendResponse(nodeAddr, createNodesMessage(id, total, _nodes));
         this.metrics?.sentMessageCount.inc({ type: MessageType[MessageType.NODES] });
@@ -1032,3 +1029,18 @@ export class Discv5 extends (EventEmitter as { new (): Discv5EventEmitter }) {
     this.connectionUpdated(nodeId, { type: ConnectionStatusType.Disconnected });
   };
 }
+
+export function chunkify<T>(arr: T[], itemsPerChunk: number): T[][] {
+  const chunkCount = Math.ceil(arr.length / itemsPerChunk);
+  if (chunkCount <= 1) {
+    return [arr];
+  }
+
+  const result: T[][] = [];
+
+  for (let i = 0; i < arr.length; i += itemsPerChunk) {
+    result.push(arr.slice(i, Math.min(arr.length, i + itemsPerChunk)));
+  }
+
+  return result;
+}

test/unit/service/service.test.ts

@@ -2,7 +2,7 @@
 import { expect } from "chai";
 import { multiaddr } from "@multiformats/multiaddr";
 
-import { Discv5 } from "../../../src/service/service.js";
+import { chunkify, Discv5 } from "../../../src/service/service.js";
 import { ENR } from "../../../src/enr/index.js";
 import { generateKeypair, KeypairType, createPeerIdFromKeypair } from "../../../src/keypair/index.js";
 
@@ -66,3 +66,108 @@ describe("Discv5", async () => {
     await service1.stop();
   });
 });
+
+describe("chunkify", function () {
+  const itemsPerChunk = 3;
+  const testCases: { arrLength: number; expected: number[][] }[] = [
+    { arrLength: 0, expected: [[]] },
+    { arrLength: 1, expected: [[0]] },
+    { arrLength: 2, expected: [[0, 1]] },
+    { arrLength: 3, expected: [[0, 1, 2]] },
+    { arrLength: 4, expected: [[0, 1, 2], [3]] },
+    {
+      arrLength: 5,
+      expected: [
+        [0, 1, 2],
+        [3, 4],
+      ],
+    },
+    {
+      arrLength: 6,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+      ],
+    },
+    {
+      arrLength: 7,
+      expected: [[0, 1, 2], [3, 4, 5], [6]],
+    },
+    {
+      arrLength: 8,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+        [6, 7],
+      ],
+    },
+    {
+      arrLength: 9,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+        [6, 7, 8],
+      ],
+    },
+    {
+      arrLength: 10,
+      expected: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]],
+    },
+    {
+      arrLength: 11,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+        [6, 7, 8],
+        [9, 10],
+      ],
+    },
+    {
+      arrLength: 12,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+        [6, 7, 8],
+        [9, 10, 11],
+      ],
+    },
+    {
+      arrLength: 13,
+      expected: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], [12]],
+    },
+    {
+      arrLength: 14,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+        [6, 7, 8],
+        [9, 10, 11],
+        [12, 13],
+      ],
+    },
+    {
+      arrLength: 15,
+      expected: [
+        [0, 1, 2],
+        [3, 4, 5],
+        [6, 7, 8],
+        [9, 10, 11],
+        [12, 13, 14],
+      ],
+    },
+    {
+      arrLength: 16,
+      expected: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], [12, 13, 14], [15]],
+    },
+  ];
+  for (const { arrLength, expected } of testCases) {
+    it(`array ${arrLength} length`, () => {
+      expect(
+        chunkify(
+          Array.from({ length: arrLength }, (_, i) => i),
+          itemsPerChunk
+        )
+      ).to.be.deep.equal(expected);
+    });
+  }
+});