From 35a8de306b4e812704d33794e4fb6cad0c57e6ef Mon Sep 17 00:00:00 2001
From: Danish Arora <danisharora099@gmail.com>
Date: Mon, 16 Dec 2024 20:08:30 +0530
Subject: [PATCH] chore: use protocol-level class for getting peer

---
 packages/core/src/lib/base_protocol.ts    | 17 +------------
 packages/sdk/src/protocols/store/index.ts | 30 +++++++++++++++++------
 2 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/packages/core/src/lib/base_protocol.ts b/packages/core/src/lib/base_protocol.ts
index 45ae85de39..8f8cc9405b 100644
--- a/packages/core/src/lib/base_protocol.ts
+++ b/packages/core/src/lib/base_protocol.ts
@@ -76,30 +76,15 @@ export class BaseProtocol implements IBaseProtocolCore {
   public async getPeers(
     {
       numPeers,
-      maxBootstrapPeers,
-      peerIdStr
+      maxBootstrapPeers
     }: {
       numPeers: number;
       maxBootstrapPeers: number;
-      peerIdStr?: PeerIdStr;
     } = {
       maxBootstrapPeers: 0,
       numPeers: 0
     }
   ): Promise<Peer[]> {
-    if (peerIdStr) {
-      const peer = (await this.connectedPeers()).find(
-        (p) => p.id.toString() === peerIdStr
-      );
-      if (peer) {
-        return [peer];
-      }
-      this.log.warn(
-        `Passed node to use for ${this.multicodec} not found: ${peerIdStr}. Attempting to use random peers.`
-      );
-      return this.getPeers({ numPeers, maxBootstrapPeers });
-    }
-
     // Retrieve all connected peers that support the protocol & shard (if configured)
     const allAvailableConnectedPeers = await this.connectedPeers();
 
diff --git a/packages/sdk/src/protocols/store/index.ts b/packages/sdk/src/protocols/store/index.ts
index b12958184a..655c090722 100644
--- a/packages/sdk/src/protocols/store/index.ts
+++ b/packages/sdk/src/protocols/store/index.ts
@@ -1,3 +1,4 @@
+import type { Peer } from "@libp2p/interface-peer-id";
 import { ConnectionManager, StoreCore } from "@waku/core";
 import {
   IDecodedMessage,
@@ -62,13 +63,14 @@ export class Store extends BaseProtocolSDK implements IStore {
       ...options
     };
 
-    const peer = (
-      await this.protocol.getPeers({
-        numPeers: this.numPeersToUse,
-        maxBootstrapPeers: 1,
-        peerIdStr: this.peerIdStrToUse
-      })
-    )[0];
+    const peer =
+      (await this.getPeerToUse()) ??
+      (
+        await this.protocol.getPeers({
+          numPeers: this.numPeersToUse,
+          maxBootstrapPeers: 1
+        })
+      )[0];
 
     if (!peer) {
       log.error("No peers available to query");
@@ -234,6 +236,20 @@ export class Store extends BaseProtocolSDK implements IStore {
       decodersAsMap
     };
   }
+
+  private async getPeerToUse(): Promise<Peer | null> {
+    const peer = this.connectedPeers.find(
+      (p) => p.id.toString() === this.peerIdStrToUse
+    );
+    if (peer) {
+      return peer;
+    }
+
+    log.warn(
+      `Passed node to use for Store not found: ${this.peerIdStrToUse}. Attempting to use random peers.`
+    );
+    return null;
+  }
 }
 
 /**