feat(static-sharding): filter peer connections per shards (#1626)

* add interface for `ShardInfo`

* enr: add deserialization logic & setup getters

* add sharding related utils

* utils: add shard<-> bytes conversion helpers

* pass `pubSubTopics` to `Waku`

* add `rs`/`rsv` details during discovery

* connection-manager: discard irrelevant peers

* add tests for static sharding - peer exchange

* update `ConnectionManager` tests to account for topic validity

* add js suffix to import

* address some comments

* move shardInfo encoding to ENR

* test: update for new API

* enr: add tests for serialisation & deserialisation

* address comment

* update test

* move getPeershardInfo to ConnectionManager and return ShardInfo instead of bytes

* update encoding and decoding relay shards to also factor for shards>64

* relay shard encoding decoding: use DataView and verbose spec tests

* improve tests for relay shard encoding decoding

* rm: only

* improve log message for unconfigured pubsub topic

* minor improvement

* fix: buffer <> Uint8array problems with shard decoding

* fix: test

* rm: only

packages/core/package.json

@@ -73,6 +73,7 @@
   },
   "dependencies": {
     "@noble/hashes": "^1.3.2",
+    "@waku/enr": "^0.0.17",
     "@waku/interfaces": "0.0.18",
     "@waku/proto": "0.0.5",
     "@waku/utils": "0.0.11",

packages/core/src/lib/connection_manager.ts

@@ -1,17 +1,22 @@
 import type { PeerId } from "@libp2p/interface/peer-id";
 import type { PeerInfo } from "@libp2p/interface/peer-info";
 import type { Peer } from "@libp2p/interface/peer-store";
+import type { PeerStore } from "@libp2p/interface/peer-store";
 import { CustomEvent, EventEmitter } from "@libp2p/interfaces/events";
+import { decodeRelayShard } from "@waku/enr";
 import {
   ConnectionManagerOptions,
   EPeersByDiscoveryEvents,
   IConnectionManager,
   IPeersByDiscoveryEvents,
   IRelay,
   KeepAliveOptions,
-  PeersByDiscoveryResult
+  PeersByDiscoveryResult,
+  PubSubTopic,
+  ShardInfo
 } from "@waku/interfaces";
 import { Libp2p, Tags } from "@waku/interfaces";
+import { shardInfoToPubSubTopics } from "@waku/utils";
 import debug from "debug";
 
 import { KeepAliveManager } from "./keep_alive_manager.js";
@@ -40,6 +45,7 @@ export class ConnectionManager
     peerId: string,
     libp2p: Libp2p,
     keepAliveOptions: KeepAliveOptions,
+    pubSubTopics: PubSubTopic[],
     relay?: IRelay,
     options?: ConnectionManagerOptions
   ): ConnectionManager {
@@ -48,6 +54,7 @@ export class ConnectionManager
       instance = new ConnectionManager(
         libp2p,
         keepAliveOptions,
+        pubSubTopics,
         relay,
         options
       );
@@ -104,11 +111,13 @@ export class ConnectionManager
   private constructor(
     libp2p: Libp2p,
     keepAliveOptions: KeepAliveOptions,
+    private configuredPubSubTopics: PubSubTopic[],
     relay?: IRelay,
     options?: Partial<ConnectionManagerOptions>
   ) {
     super();
     this.libp2p = libp2p;
+    this.configuredPubSubTopics = configuredPubSubTopics;
     this.options = {
       maxDialAttemptsForPeer: DEFAULT_MAX_DIAL_ATTEMPTS_FOR_PEER,
       maxBootstrapPeersAllowed: DEFAULT_MAX_BOOTSTRAP_PEERS_ALLOWED,
@@ -314,6 +323,20 @@ export class ConnectionManager
       void (async () => {
         const { id: peerId } = evt.detail;
 
+        if (!(await this.isPeerTopicConfigured(peerId))) {
+          const shardInfo = await this.getPeerShardInfo(
+            peerId,
+            this.libp2p.peerStore
+          );
+          log(
+            `Discovered peer ${peerId.toString()} with ShardInfo ${shardInfo} is not part of any of the configured pubsub topics (${
+              this.configuredPubSubTopics
+            }). 
+            Not dialing.`
+          );
+          return;
+        }
+
         const isBootstrap = (await this.getTagNamesForPeer(peerId)).includes(
           Tags.BOOTSTRAP
         );
@@ -430,4 +453,31 @@ export class ConnectionManager
       return [];
     }
   }
+
+  private async isPeerTopicConfigured(peerId: PeerId): Promise<boolean> {
+    const shardInfo = await this.getPeerShardInfo(
+      peerId,
+      this.libp2p.peerStore
+    );
+
+    // If there's no shard information, simply return true
+    if (!shardInfo) return true;
+
+    const pubSubTopics = shardInfoToPubSubTopics(shardInfo);
+
+    const isTopicConfigured = pubSubTopics.some((topic) =>
+      this.configuredPubSubTopics.includes(topic)
+    );
+    return isTopicConfigured;
+  }
+
+  private async getPeerShardInfo(
+    peerId: PeerId,
+    peerStore: PeerStore
+  ): Promise<ShardInfo | undefined> {
+    const peer = await peerStore.get(peerId);
+    const shardInfoBytes = peer.metadata.get("shardInfo");
+    if (!shardInfoBytes) return undefined;
+    return decodeRelayShard(shardInfoBytes);
+  }
 }

packages/core/src/lib/waku.ts

@@ -7,6 +7,7 @@ import type {
   IRelay,
   IStore,
   Libp2p,
+  PubSubTopic,
   Waku
 } from "@waku/interfaces";
 import { Protocols } from "@waku/interfaces";
@@ -52,6 +53,7 @@ export class WakuNode implements Waku {
 
   constructor(
     options: WakuOptions,
+    public readonly pubSubTopics: PubSubTopic[],
     libp2p: Libp2p,
     store?: (libp2p: Libp2p) => IStore,
     lightPush?: (libp2p: Libp2p) => ILightPush,
@@ -86,6 +88,7 @@ export class WakuNode implements Waku {
       peerId,
       libp2p,
       { pingKeepAlive, relayKeepAlive },
+      pubSubTopics,
       this.relay
     );
 

packages/dns-discovery/src/dns_discovery.ts

@@ -5,6 +5,7 @@ import type {
 } from "@libp2p/interface/peer-discovery";
 import { peerDiscovery as symbol } from "@libp2p/interface/peer-discovery";
 import type { PeerInfo } from "@libp2p/interface/peer-info";
+import { encodeRelayShard } from "@waku/enr";
 import type {
   DnsDiscOptions,
   DnsDiscoveryComponents,
@@ -72,18 +73,16 @@ export class PeerDiscoveryDns
         return;
       }
 
-      const peerInfo = peerEnr.peerInfo;
+      const { peerInfo, shardInfo } = peerEnr;
 
       if (!peerInfo) {
         continue;
       }
 
       const tagsToUpdate = {
-        tags: {
-          [DEFAULT_BOOTSTRAP_TAG_NAME]: {
-            value: this._options.tagValue ?? DEFAULT_BOOTSTRAP_TAG_VALUE,
-            ttl: this._options.tagTTL ?? DEFAULT_BOOTSTRAP_TAG_TTL
-          }
+        [DEFAULT_BOOTSTRAP_TAG_NAME]: {
+          value: this._options.tagValue ?? DEFAULT_BOOTSTRAP_TAG_VALUE,
+          ttl: this._options.tagTTL ?? DEFAULT_BOOTSTRAP_TAG_TTL
         }
       };
 
@@ -96,11 +95,20 @@ export class PeerDiscoveryDns
 
         if (!hasBootstrapTag) {
           isPeerChanged = true;
-          await this._components.peerStore.merge(peerInfo.id, tagsToUpdate);
+          await this._components.peerStore.merge(peerInfo.id, {
+            tags: tagsToUpdate
+          });
         }
       } else {
         isPeerChanged = true;
-        await this._components.peerStore.save(peerInfo.id, tagsToUpdate);
+        await this._components.peerStore.save(peerInfo.id, {
+          tags: tagsToUpdate,
+          ...(shardInfo && {
+            metadata: {
+              shardInfo: encodeRelayShard(shardInfo)
+            }
+          })
+        });
       }
 
       if (isPeerChanged) {

packages/enr/package.json

@@ -71,6 +71,7 @@
     "@waku/interfaces": "0.0.18",
     "chai": "^4.3.7",
     "cspell": "^7.3.2",
+    "fast-check": "^3.13.1",
     "mocha": "^10.2.0",
     "npm-run-all": "^4.1.5",
     "process": "^0.11.10",

packages/enr/src/enr.ts

@@ -6,7 +6,8 @@ import type {
   ENRValue,
   IEnr,
   NodeId,
-  SequenceNumber
+  SequenceNumber,
+  ShardInfo
 } from "@waku/interfaces";
 import debug from "debug";
 
@@ -64,6 +65,13 @@ export class ENR extends RawEnr implements IEnr {
     protocol: TransportProtocol | TransportProtocolPerIpVersion
   ) => Multiaddr | undefined = locationMultiaddrFromEnrFields.bind({}, this);
 
+  get shardInfo(): ShardInfo | undefined {
+    if (this.rs && this.rsv) {
+      log("Warning: ENR contains both `rs` and `rsv` fields.");
+    }
+    return this.rs || this.rsv;
+  }
+
   setLocationMultiaddr(multiaddr: Multiaddr): void {
     const protoNames = multiaddr.protoNames();
     if (

packages/enr/src/index.ts

@@ -5,3 +5,4 @@ export * from "./enr.js";
 export * from "./peer_id.js";
 export * from "./waku2_codec.js";
 export * from "./crypto.js";
+export * from "./relay_shard_codec.js";

packages/enr/src/raw_enr.ts

@@ -3,11 +3,18 @@ import {
   convertToBytes,
   convertToString
 } from "@multiformats/multiaddr/convert";
-import type { ENRKey, ENRValue, SequenceNumber, Waku2 } from "@waku/interfaces";
+import type {
+  ENRKey,
+  ENRValue,
+  SequenceNumber,
+  ShardInfo,
+  Waku2
+} from "@waku/interfaces";
 import { bytesToUtf8 } from "@waku/utils/bytes";
 
 import { ERR_INVALID_ID } from "./constants.js";
 import { decodeMultiaddrs, encodeMultiaddrs } from "./multiaddrs_codec.js";
+import { decodeRelayShard } from "./relay_shard_codec.js";
 import { decodeWaku2, encodeWaku2 } from "./waku2_codec.js";
 
 export class RawEnr extends Map<ENRKey, ENRValue> {
@@ -45,6 +52,18 @@ export class RawEnr extends Map<ENRKey, ENRValue> {
     }
   }
 
+  get rs(): ShardInfo | undefined {
+    const rs = this.get("rs");
+    if (!rs) return undefined;
+    return decodeRelayShard(rs);
+  }
+
+  get rsv(): ShardInfo | undefined {
+    const rsv = this.get("rsv");
+    if (!rsv) return undefined;
+    return decodeRelayShard(rsv);
+  }
+
   get ip(): string | undefined {
     return getStringValue(this, "ip", "ip4");
   }

packages/enr/src/relay_shard_codec.spec.ts

@@ -0,0 +1,68 @@
+import { expect } from "chai";
+import fc from "fast-check";
+
+import { decodeRelayShard, encodeRelayShard } from "./relay_shard_codec.js";
+
+describe("Relay Shard codec", () => {
+  // Boundary test case
+  it("should handle a minimal index list", () => {
+    const shardInfo = { cluster: 0, indexList: [0] };
+    const encoded = encodeRelayShard(shardInfo);
+    const decoded = decodeRelayShard(encoded);
+    expect(decoded).to.deep.equal(
+      shardInfo,
+      "Decoded shard info does not match the original for minimal index list"
+    );
+  });
+
+  // Property-based test for rs format (Index List)
+  it("should correctly encode and decode relay shards using rs format (Index List)", () => {
+    fc.assert(
+      fc.property(
+        fc.nat(65535), // cluster
+        fc
+          .array(fc.nat(1023), { minLength: 1, maxLength: 63 }) // indexList
+          .map((arr) => [...new Set(arr)].sort((a, b) => a - b)),
+        (cluster, indexList) => {
+          const shardInfo = { cluster, indexList };
+          const encoded = encodeRelayShard(shardInfo);
+          const decoded = decodeRelayShard(encoded);
+
+          expect(decoded).to.deep.equal(
+            shardInfo,
+            "Decoded shard info does not match the original for rs format"
+          );
+        }
+      )
+    );
+  });
+
+  // Property-based test for rsv format (Bit Vector)
+  it("should correctly encode and decode relay shards using rsv format (Bit Vector)", () => {
+    fc.assert(
+      fc.property(
+        fc.nat(65535), // cluster
+        fc
+          .array(fc.nat(1023), { minLength: 64, maxLength: 1024 }) // indexList
+          .map((arr) => [...new Set(arr)].sort((a, b) => a - b)),
+        (cluster, indexList) => {
+          const shardInfo = { cluster, indexList };
+          const encoded = encodeRelayShard(shardInfo);
+          const decoded = decodeRelayShard(encoded);
+
+          expect(decoded).to.deep.equal(
+            shardInfo,
+            "Decoded shard info does not match the original for rsv format"
+          );
+        }
+      )
+    );
+  });
+
+  // Error handling test case
+  it("should throw an error for insufficient data", () => {
+    expect(() => decodeRelayShard(new Uint8Array([0, 0]))).to.throw(
+      "Insufficient data"
+    );
+  });
+});