chore: store self protocols in protobook (#760)

libp2p · Oct 27, 2020 · c9e6757 · c9e6757
1 parent 5246cf8
commit c9e6757
Show file tree

Hide file tree

Showing 6 changed files with 264 additions and 96 deletions.
diff --git a/doc/API.md b/doc/API.md
@@ -37,6 +37,7 @@
   * [`peerStore.protoBook.add`](#peerstoreprotobookadd)
   * [`peerStore.protoBook.delete`](#peerstoreprotobookdelete)
   * [`peerStore.protoBook.get`](#peerstoreprotobookget)
+  * [`peerStore.protoBook.remove`](#peerstoreprotobookremove)
   * [`peerStore.protoBook.set`](#peerstoreprotobookset)
   * [`peerStore.delete`](#peerstoredelete)
   * [`peerStore.get`](#peerstoreget)
@@ -843,32 +844,6 @@ Consider using `addressBook.add()` if you're not sure this is what you want to d
 peerStore.addressBook.add(peerId, multiaddr)
 ```
 
-### peerStore.protoBook.add
-
-Add known `protocols` of a given peer.
-
-`peerStore.protoBook.add(peerId, protocols)`
-
-#### Parameters
-
-| Name | Type | Description |
-|------|------|-------------|
-| peerId | [`PeerId`][peer-id] | peerId to set |
-| protocols | `Array<string>` | protocols to add |
-
-#### Returns
-
-| Type | Description |
-|------|-------------|
-| `ProtoBook` | Returns the Proto Book component |
-
-#### Example
-
-```js
-peerStore.protoBook.add(peerId, protocols)
-```
-
-
 ### peerStore.keyBook.delete
 
 Delete the provided peer from the book.
@@ -1091,6 +1066,31 @@ Set known metadata of a given `peerId`.
 peerStore.metadataBook.set(peerId, 'location', uint8ArrayFromString('Berlin'))
 ```
 
+### peerStore.protoBook.add
+
+Add known `protocols` of a given peer.
+
+`peerStore.protoBook.add(peerId, protocols)`
+
+#### Parameters
+
+| Name | Type | Description |
+|------|------|-------------|
+| peerId | [`PeerId`][peer-id] | peerId to set |
+| protocols | `Array<string>` | protocols to add |
+
+#### Returns
+
+| Type | Description |
+|------|-------------|
+| `ProtoBook` | Returns the Proto Book component |
+
+#### Example
+
+```js
+peerStore.protoBook.add(peerId, protocols)
+```
+
 ### peerStore.protoBook.delete
 
 Delete the provided peer from the book.
@@ -1147,6 +1147,31 @@ peerStore.protoBook.get(peerId)
 // [ '/proto/1.0.0', '/proto/1.1.0' ]
 ```
 
+### peerStore.protoBook.remove
+
+Remove given `protocols` of a given peer.
+
+`peerStore.protoBook.remove(peerId, protocols)`
+
+#### Parameters
+
+| Name | Type | Description |
+|------|------|-------------|
+| peerId | [`PeerId`][peer-id] | peerId to set |
+| protocols | `Array<string>` | protocols to remove |
+
+#### Returns
+
+| Type | Description |
+|------|-------------|
+| `ProtoBook` | Returns the Proto Book component |
+
+#### Example
+
+```js
+peerStore.protoBook.remove(peerId, protocols)
+```
+
 ### peerStore.protoBook.set
 
 Set known `protocols` of a given peer.

diff --git a/src/identify/index.js b/src/identify/index.js
@@ -51,9 +51,8 @@ class IdentifyService {
    * @class
    * @param {object} options
    * @param {Libp2p} options.libp2p
-   * @param {Map<string, handler>} options.protocols - A reference to the protocols we support
    */
-  constructor ({ libp2p, protocols }) {
+  constructor ({ libp2p }) {
     /**
      * @property {PeerStore}
      */
@@ -74,10 +73,9 @@ class IdentifyService {
      */
     this._libp2p = libp2p
 
-    this._protocols = protocols
-
     this.handleMessage = this.handleMessage.bind(this)
 
+    // When a new connection happens, trigger identify
     this.connectionManager.on('peer:connect', (connection) => {
       const peerId = connection.remotePeer
 
@@ -90,6 +88,13 @@ class IdentifyService {
         this.pushToPeerStore()
       }
     })
+
+    // When self protocols change, trigger identify-push
+    this.peerStore.on('change:protocols', ({ peerId }) => {
+      if (peerId.toString() === this.peerId.toString()) {
+        this.pushToPeerStore()
+      }
+    })
   }
 
   /**
@@ -101,7 +106,7 @@ class IdentifyService {
   async push (connections) {
     const signedPeerRecord = await this.peerStore.addressBook.getRawEnvelope(this.peerId)
     const listenAddrs = this._libp2p.multiaddrs.map((ma) => ma.bytes)
-    const protocols = Array.from(this._protocols.keys())
+    const protocols = this.peerStore.protoBook.get(this.peerId) || []
 
     const pushes = connections.map(async connection => {
       try {
@@ -132,6 +137,11 @@ class IdentifyService {
    * @returns {void}
    */
   pushToPeerStore () {
+    // Do not try to push if libp2p node is not running
+    if (!this._libp2p.isStarted()) {
+      return
+    }
+
     const connections = []
     let connection
     for (const peer of this.peerStore.peers.values()) {
@@ -251,6 +261,7 @@ class IdentifyService {
     }
 
     const signedPeerRecord = await this.peerStore.addressBook.getRawEnvelope(this.peerId)
+    const protocols = this.peerStore.protoBook.get(this.peerId) || []
 
     const message = Message.encode({
       protocolVersion: PROTOCOL_VERSION,
@@ -259,7 +270,7 @@ class IdentifyService {
       listenAddrs: this._libp2p.multiaddrs.map((ma) => ma.bytes),
       signedPeerRecord,
       observedAddr: connection.remoteAddr.bytes,
-      protocols: Array.from(this._protocols.keys())
+      protocols
     })
 
     try {

diff --git a/src/index.js b/src/index.js
@@ -157,10 +157,7 @@ class Libp2p extends EventEmitter {
       })
 
       // Add the identify service since we can multiplex
-      this.identifyService = new IdentifyService({
-        libp2p: this,
-        protocols: this.upgrader.protocols
-      })
+      this.identifyService = new IdentifyService({ libp2p: this })
       this.handle(Object.values(IDENTIFY_PROTOCOLS), this.identifyService.handleMessage)
     }
 
@@ -436,10 +433,8 @@ class Libp2p extends EventEmitter {
       this.upgrader.protocols.set(protocol, handler)
     })
 
-    // Only push if libp2p is running
-    if (this.isStarted() && this.identifyService) {
-      this.identifyService.pushToPeerStore()
-    }
+    // Add new protocols to self protocols in the Protobook
+    this.peerStore.protoBook.add(this.peerId, protocols)
   }
 
   /**
@@ -454,10 +449,8 @@ class Libp2p extends EventEmitter {
       this.upgrader.protocols.delete(protocol)
     })
 
-    // Only push if libp2p is running
-    if (this.isStarted() && this.identifyService) {
-      this.identifyService.pushToPeerStore()
-    }
+    // Remove protocols from self protocols in the Protobook
+    this.peerStore.protoBook.remove(this.peerId, protocols)
   }
 
   async _onStarting () {

diff --git a/src/peer-store/proto-book.js b/src/peer-store/proto-book.js
@@ -112,13 +112,50 @@ class ProtoBook extends Book {
       return this
     }
 
-    protocols = [...newSet]
-
     this._setData(peerId, newSet)
     log(`added provided protocols for ${id}`)
 
     return this
   }
+
+  /**
+   * Removes known protocols of a provided peer.
+   * If the protocols did not exist before, nothing will be done.
+   *
+   * @param {PeerId} peerId
+   * @param {Array<string>} protocols
+   * @returns {ProtoBook}
+   */
+  remove (peerId, protocols) {
+    if (!PeerId.isPeerId(peerId)) {
+      log.error('peerId must be an instance of peer-id to store data')
+      throw errcode(new Error('peerId must be an instance of peer-id'), ERR_INVALID_PARAMETERS)
+    }
+
+    if (!protocols) {
+      log.error('protocols must be provided to store data')
+      throw errcode(new Error('protocols must be provided'), ERR_INVALID_PARAMETERS)
+    }
+
+    const id = peerId.toB58String()
+    const recSet = this.data.get(id)
+
+    if (recSet) {
+      const newSet = new Set([
+        ...recSet
+      ].filter((p) => !protocols.includes(p)))
+
+      // Any protocol removed?
+      if (recSet.size === newSet.size) {
+        return this
+      }
+
+      this._setData(peerId, newSet)
+      log(`removed provided protocols for ${id}`)
+    }
+
+    return this
+  }
 }
 
 module.exports = ProtoBook