save remote peer's shard info after successful connection

danisharora099 · danisharora099 · commit d0c9683c64e4 · 2023-12-05T16:24:24.000+05:30
diff --git a/packages/core/src/lib/metadata/index.ts b/packages/core/src/lib/metadata/index.ts
@@ -1,5 +1,6 @@
 import type { PeerId } from "@libp2p/interface/peer-id";
 import { IncomingStreamData } from "@libp2p/interface/stream-handler";
+import { encodeRelayShard } from "@waku/enr";
 import type { IMetadata, Libp2pComponents, ShardInfo } from "@waku/interfaces";
 import { proto_metadata } from "@waku/proto";
 import { Logger } from "@waku/utils";
@@ -16,8 +17,10 @@ export const MetadataCodec = "/vac/waku/metadata/1.0.0";
 
 class Metadata extends BaseProtocol {
   private readonly shardInfo: ShardInfo;
+  private libp2pComponents: Libp2pComponents;
   constructor(shardInfo: ShardInfo, libp2p: Libp2pComponents) {
     super(MetadataCodec, libp2p.components);
+    this.libp2pComponents = libp2p;
     this.shardInfo = shardInfo;
     void libp2p.registrar.handle(MetadataCodec, (streamData) => {
       void this.onRequest(streamData);
@@ -29,11 +32,28 @@ class Metadata extends BaseProtocol {
    */
   private async onRequest(streamData: IncomingStreamData): Promise<void> {
     try {
-      const encodedResponse = proto_metadata.WakuMetadataResponse.encode(
+      const { stream, connection } = streamData;
+      const encodedShardInfo = proto_metadata.WakuMetadataResponse.encode(
         this.shardInfo
       );
 
-      await pipe([encodedResponse], lp.encode, streamData.stream);
+      const encodedResponse = await pipe(
+        [encodedShardInfo],
+        lp.encode,
+        stream,
+        lp.decode,
+        async (source) => await all(source)
+      );
+
+      const remoteShardInfoResponse =
+        this.decodeMetadataResponse(encodedResponse);
+
+      // add or update the shardInfo to peer store
+      await this.libp2pComponents.peerStore.merge(connection.remotePeer, {
+        metadata: {
+          shardInfo: encodeRelayShard(remoteShardInfoResponse)
+        }
+      });
     } catch (error) {
       log.error("Error handling metadata request", error);
     }
@@ -49,22 +69,32 @@ class Metadata extends BaseProtocol {
 
     const stream = await this.getStream(peer);
 
-    const res = await pipe(
+    const encodedResponse = await pipe(
       [request],
       lp.encode,
       stream,
       lp.decode,
       async (source) => await all(source)
     );
 
+    const decodedResponse = this.decodeMetadataResponse(encodedResponse);
+
+    return decodedResponse;
+  }
+
+  private decodeMetadataResponse(encodedResponse: Uint8ArrayList[]): ShardInfo {
     const bytes = new Uint8ArrayList();
-    res.forEach((chunk) => {
+
+    encodedResponse.forEach((chunk) => {
       bytes.append(chunk);
     });
-    const response = proto_metadata.WakuMetadataResponse.decode(bytes);
+    const response = proto_metadata.WakuMetadataResponse.decode(
+      bytes
+    ) as ShardInfo;
+
     if (!response) log.error("Error decoding metadata response");
 
-    return response as ShardInfo;
+    return response;
   }
 }
 
diff --git a/packages/tests/tests/metadata.spec.ts b/packages/tests/tests/metadata.spec.ts
@@ -1,4 +1,5 @@
 import { MetadataCodec } from "@waku/core";
+import { decodeRelayShard } from "@waku/enr";
 import type { LightNode, ShardInfo } from "@waku/interfaces";
 import { createLightNode } from "@waku/sdk";
 import { shardInfoToPubsubTopics } from "@waku/utils";
@@ -11,7 +12,7 @@ import { NimGoNode } from "../src/node/node.js";
 
 chai.use(chaiAsPromised);
 
-describe.only("Metadata Protocol", () => {
+describe("Metadata Protocol", () => {
   let waku: LightNode;
   let nwaku1: NimGoNode;
 
@@ -24,142 +25,180 @@ describe.only("Metadata Protocol", () => {
     await tearDownNodes([nwaku1], waku);
   });
 
-  it("same cluster, same shard: nodes connect", async function () {
-    this.timeout(55_000);
-
-    const shardInfo: ShardInfo = {
-      clusterId: 1,
-      shards: [1]
-    };
+  describe("connections", function () {
+    it("same cluster, same shard: nodes connect", async function () {
+      this.timeout(55_000);
+
+      const shardInfo: ShardInfo = {
+        clusterId: 1,
+        shards: [1]
+      };
+
+      await nwaku1.start({
+        relay: true,
+        discv5Discovery: true,
+        peerExchange: true,
+        clusterId: shardInfo.clusterId,
+        pubsubTopic: shardInfoToPubsubTopics(shardInfo)
+      });
+
+      const nwaku1Ma = await nwaku1.getMultiaddrWithId();
+      const nwaku1PeerId = await nwaku1.getPeerId();
+
+      waku = await createLightNode({ shardInfo });
+      await waku.start();
+      await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
+
+      const shardInfoRes =
+        await waku.libp2p.services.metadata?.query(nwaku1PeerId);
+      expect(shardInfoRes).to.not.be.undefined;
+      expect(shardInfoRes?.clusterId).to.equal(shardInfo.clusterId);
+      expect(shardInfoRes?.shards).to.deep.equal(shardInfo.shards);
+
+      const activeConnections = waku.libp2p.getConnections();
+      expect(activeConnections.length).to.equal(1);
+    });
 
-    await nwaku1.start({
-      relay: true,
-      discv5Discovery: true,
-      peerExchange: true,
-      clusterId: shardInfo.clusterId,
-      pubsubTopic: shardInfoToPubsubTopics(shardInfo)
+    it("same cluster, different shard: nodes connect", async function () {
+      this.timeout(55_000);
+
+      const shardInfo1: ShardInfo = {
+        clusterId: 1,
+        shards: [1]
+      };
+
+      const shardInfo2: ShardInfo = {
+        clusterId: 1,
+        shards: [2]
+      };
+
+      await nwaku1.start({
+        relay: true,
+        discv5Discovery: true,
+        peerExchange: true,
+        clusterId: shardInfo1.clusterId,
+        pubsubTopic: shardInfoToPubsubTopics(shardInfo1)
+      });
+
+      const nwaku1Ma = await nwaku1.getMultiaddrWithId();
+      const nwaku1PeerId = await nwaku1.getPeerId();
+
+      waku = await createLightNode({ shardInfo: shardInfo2 });
+      await waku.start();
+      await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
+
+      const shardInfoRes =
+        await waku.libp2p.services.metadata?.query(nwaku1PeerId);
+      expect(shardInfoRes).to.not.be.undefined;
+      expect(shardInfoRes?.clusterId).to.equal(shardInfo1.clusterId);
+      expect(shardInfoRes?.shards).to.deep.equal(shardInfo1.shards);
+
+      const activeConnections = waku.libp2p.getConnections();
+      expect(activeConnections.length).to.equal(1);
     });
 
-    const nwaku1Ma = await nwaku1.getMultiaddrWithId();
-    const nwaku1PeerId = await nwaku1.getPeerId();
+    it("different cluster, same shard: nodes don't connect", async function () {
+      this.timeout(55_000);
 
-    waku = await createLightNode({ shardInfo });
-    await waku.start();
-    await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
+      const shardInfo1: ShardInfo = {
+        clusterId: 1,
+        shards: [1]
+      };
 
-    const shardInfoRes =
-      await waku.libp2p.services.metadata?.query(nwaku1PeerId);
-    expect(shardInfoRes).to.not.be.undefined;
-    expect(shardInfoRes?.clusterId).to.equal(shardInfo.clusterId);
-    expect(shardInfoRes?.shards).to.deep.equal(shardInfo.shards);
+      const shardInfo2: ShardInfo = {
+        clusterId: 2,
+        shards: [1]
+      };
 
-    const activeConnections = waku.libp2p.getConnections();
-    expect(activeConnections.length).to.equal(1);
-  });
+      await nwaku1.start({
+        relay: true,
+        discv5Discovery: true,
+        peerExchange: true,
+        clusterId: shardInfo1.clusterId,
+        pubsubTopic: shardInfoToPubsubTopics(shardInfo1)
+      });
 
-  it("same cluster, different shard: nodes connect", async function () {
-    this.timeout(55_000);
+      const nwaku1Ma = await nwaku1.getMultiaddrWithId();
 
-    const shardInfo1: ShardInfo = {
-      clusterId: 1,
-      shards: [1]
-    };
+      waku = await createLightNode({ shardInfo: shardInfo2 });
+      await waku.start();
+      await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
 
-    const shardInfo2: ShardInfo = {
-      clusterId: 1,
-      shards: [2]
-    };
+      // add a delay to make sure the connection is closed from the other side
+      await delay(100);
 
-    await nwaku1.start({
-      relay: true,
-      discv5Discovery: true,
-      peerExchange: true,
-      clusterId: shardInfo1.clusterId,
-      pubsubTopic: shardInfoToPubsubTopics(shardInfo1)
+      const activeConnections = waku.libp2p.getConnections();
+      expect(activeConnections.length).to.equal(0);
     });
 
-    const nwaku1Ma = await nwaku1.getMultiaddrWithId();
-    const nwaku1PeerId = await nwaku1.getPeerId();
+    it("different cluster, different shard: nodes don't connect", async function () {
+      this.timeout(55_000);
 
-    waku = await createLightNode({ shardInfo: shardInfo2 });
-    await waku.start();
-    await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
+      const shardInfo1: ShardInfo = {
+        clusterId: 1,
+        shards: [1]
+      };
 
-    const shardInfoRes =
-      await waku.libp2p.services.metadata?.query(nwaku1PeerId);
-    expect(shardInfoRes).to.not.be.undefined;
-    expect(shardInfoRes?.clusterId).to.equal(shardInfo1.clusterId);
-    expect(shardInfoRes?.shards).to.deep.equal(shardInfo1.shards);
+      const shardInfo2: ShardInfo = {
+        clusterId: 2,
+        shards: [2]
+      };
 
-    const activeConnections = waku.libp2p.getConnections();
-    expect(activeConnections.length).to.equal(1);
-  });
+      await nwaku1.start({
+        relay: true,
+        discv5Discovery: true,
+        peerExchange: true,
+        clusterId: shardInfo1.clusterId,
+        pubsubTopic: shardInfoToPubsubTopics(shardInfo1)
+      });
 
-  it("different cluster, same shard: nodes don't connect", async function () {
-    this.timeout(55_000);
+      const nwaku1Ma = await nwaku1.getMultiaddrWithId();
 
-    const shardInfo1: ShardInfo = {
-      clusterId: 1,
-      shards: [1]
-    };
+      waku = await createLightNode({ shardInfo: shardInfo2 });
+      await waku.start();
+      await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
 
-    const shardInfo2: ShardInfo = {
-      clusterId: 2,
-      shards: [1]
-    };
+      // add a delay to make sure the connection is closed from the other side
+      await delay(100);
 
-    await nwaku1.start({
-      relay: true,
-      discv5Discovery: true,
-      peerExchange: true,
-      clusterId: shardInfo1.clusterId,
-      pubsubTopic: shardInfoToPubsubTopics(shardInfo1)
+      const activeConnections = waku.libp2p.getConnections();
+      expect(activeConnections.length).to.equal(0);
     });
-
-    const nwaku1Ma = await nwaku1.getMultiaddrWithId();
-
-    waku = await createLightNode({ shardInfo: shardInfo2 });
-    await waku.start();
-    await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
-
-    // add a delay to make sure the connection is closed from the other side
-    await delay(100);
-
-    const activeConnections = waku.libp2p.getConnections();
-    expect(activeConnections.length).to.equal(0);
   });
 
-  it("different cluster, different shard: nodes don't connect", async function () {
-    this.timeout(55_000);
-
-    const shardInfo1: ShardInfo = {
+  it("PeerStore has remote peer's shard info after successful connection", async function () {
+    const shardInfo: ShardInfo = {
       clusterId: 1,
       shards: [1]
     };
 
-    const shardInfo2: ShardInfo = {
-      clusterId: 2,
-      shards: [2]
-    };
-
     await nwaku1.start({
       relay: true,
       discv5Discovery: true,
       peerExchange: true,
-      clusterId: shardInfo1.clusterId,
-      pubsubTopic: shardInfoToPubsubTopics(shardInfo1)
+      clusterId: shardInfo.clusterId,
+      pubsubTopic: shardInfoToPubsubTopics(shardInfo)
     });
 
     const nwaku1Ma = await nwaku1.getMultiaddrWithId();
+    const nwaku1PeerId = await nwaku1.getPeerId();
 
-    waku = await createLightNode({ shardInfo: shardInfo2 });
+    waku = await createLightNode({ shardInfo });
     await waku.start();
     await waku.libp2p.dialProtocol(nwaku1Ma, MetadataCodec);
 
-    // add a delay to make sure the connection is closed from the other side
-    await delay(100);
+    // delay to ensure the connection is estabilished and shardInfo is updated
+    await delay(500);
+
+    const encodedShardInfo = (
+      await waku.libp2p.peerStore.get(nwaku1PeerId)
+    ).metadata.get("shardInfo");
+    expect(encodedShardInfo).to.not.be.undefined;
+
+    const metadataShardInfo = decodeRelayShard(encodedShardInfo!);
+    expect(metadataShardInfo).not.be.undefined;
 
-    const activeConnections = waku.libp2p.getConnections();
-    expect(activeConnections.length).to.equal(0);
+    expect(metadataShardInfo!.clusterId).to.eq(shardInfo.clusterId);
+    expect(metadataShardInfo.shards).to.deep.eq(shardInfo.shards);
   });
 });
