feat: add _acceptFrom filter

ChainSafe · Jun 28, 2020 · 1ff5816 · 1ff5816
1 parent b2bdb13
commit 1ff5816
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/ts/index.ts b/ts/index.ts
@@ -254,13 +254,16 @@ class Gossipsub extends BasicPubsub {
    * @param {String} idB58Str
    * @param {Peer} peer
    * @param {RPC} rpc
-   * @returns {void}
+   * @returns {boolean}
    */
-  _processRpc (idB58Str: string, peer: Peer, rpc: RPC): void {
-    super._processRpc(idB58Str, peer, rpc)
-    if (rpc.control) {
-      this._processRpcControlMessage(peer, rpc.control)
+  _processRpc (idB58Str: string, peer: Peer, rpc: RPC): boolean {
+    if (super._processRpc(idB58Str, peer, rpc)) {
+      if (rpc.control) {
+        this._processRpcControlMessage(peer, rpc.control)
+      }
+      return true
     }
+    return false
   }
 
   /**

diff --git a/ts/pubsub.js b/ts/pubsub.js
@@ -142,7 +142,7 @@ class BasicPubSub extends Pubsub {
    * @param {String} idB58Str
    * @param {Peer} peer
    * @param {RPC} rpc
-   * @returns {void}
+   * @returns {boolean}
    */
   _processRpc (idB58Str, peer, rpc) {
     this.log('rpc from', idB58Str)
@@ -156,12 +156,28 @@ class BasicPubSub extends Pubsub {
       this.emit('pubsub:subscription-change', peer.id, peer.topics, subs)
     }
 
+    if (!this._acceptFrom(idB58Str)) {
+      this.log('received message from unacceptable peer %s', idB58Str)
+      return false
+    }
+
     if (msgs.length) {
       msgs.forEach(async message => {
         const msg = utils.normalizeInRpcMessage(message)
         this._processRpcMessage(peer, msg)
       })
     }
+    return true
+  }
+
+  /**
+   * Whether to accept a message from a peer
+   * Override to create a graylist
+   * @param {string} id
+   * @returns {boolean}
+   */
+  _acceptFrom (id: string): boolean {
+    return true
   }
 
   /**