diff --git a/src/service/service.ts b/src/service/service.ts index 4d4b0cad..2cc2ea73 100644 --- a/src/service/service.ts +++ b/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,22 +795,19 @@ 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) @@ -818,10 +815,10 @@ export class Discv5 extends (EventEmitter as { new (): Discv5EventEmitter }) { // 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(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; +} diff --git a/test/unit/service/service.test.ts b/test/unit/service/service.test.ts index 7666884d..b5fa037e 100644 --- a/test/unit/service/service.test.ts +++ b/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); + }); + } +});