From 8880eefa8ffeff1203cdf5053a17dbf45f43cc3d Mon Sep 17 00:00:00 2001
From: Cayman <caymannava@gmail.com>
Date: Thu, 11 Aug 2022 11:48:16 -0400
Subject: [PATCH] fix: add successful stream peer to protobook (#1341)

* fix: add successful stream peer to protobook

If a protocol stream has been successfully negotiated and is to be
psased to the application, the peerstore should ensure that the peer is
registerd with that protocol.

* fix: upgrader test fix
---
 src/upgrader.ts                 |  8 ++++++++
 test/upgrading/upgrader.spec.ts | 10 ++++++++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/upgrader.ts b/src/upgrader.ts
index 8d314bb1..bc0a8d96 100644
--- a/src/upgrader.ts
+++ b/src/upgrader.ts
@@ -376,6 +376,10 @@ export class DefaultUpgrader extends EventEmitter<UpgraderEvents> implements Upg
 
               muxedStream.stat.protocol = protocol
 
+              // If a protocol stream has been successfully negotiated and is to be passed to the application,
+              // the peerstore should ensure that the peer is registered with that protocol
+              this.components.getPeerStore().protoBook.add(remotePeer, [protocol]).catch(err => log.error(err))
+
               connection.addStream(muxedStream)
               this._onStream({ connection, stream: { ...muxedStream, ...stream }, protocol })
             })
@@ -438,6 +442,10 @@ export class DefaultUpgrader extends EventEmitter<UpgraderEvents> implements Upg
 
           muxedStream.stat.protocol = protocol
 
+          // If a protocol stream has been successfully negotiated and is to be passed to the application,
+          // the peerstore should ensure that the peer is registered with that protocol
+          this.components.getPeerStore().protoBook.add(remotePeer, [protocol]).catch(err => log.error(err))
+
           return {
             ...muxedStream,
             ...stream,
diff --git a/test/upgrading/upgrader.spec.ts b/test/upgrading/upgrader.spec.ts
index 48ca372e..24027209 100644
--- a/test/upgrading/upgrader.spec.ts
+++ b/test/upgrading/upgrader.spec.ts
@@ -30,6 +30,8 @@ import { TimeoutController } from 'timeout-abort-controller'
 import delay from 'delay'
 import drain from 'it-drain'
 import { Uint8ArrayList } from 'uint8arraylist'
+import { PersistentPeerStore } from '@libp2p/peer-store'
+import { MemoryDatastore } from 'datastore-core'
 
 const addrs = [
   new Multiaddr('/ip4/127.0.0.1/tcp/0'),
@@ -57,7 +59,9 @@ describe('Upgrader', () => {
     localComponents = new Components({
       peerId: localPeer,
       connectionGater: mockConnectionGater(),
-      registrar: mockRegistrar()
+      registrar: mockRegistrar(),
+      peerStore: new PersistentPeerStore(),
+      datastore: new MemoryDatastore()
     })
     localMuxerFactory = new Mplex()
     localUpgrader = new DefaultUpgrader(localComponents, {
@@ -73,7 +77,9 @@ describe('Upgrader', () => {
     remoteComponents = new Components({
       peerId: remotePeer,
       connectionGater: mockConnectionGater(),
-      registrar: mockRegistrar()
+      registrar: mockRegistrar(),
+      peerStore: new PersistentPeerStore(),
+      datastore: new MemoryDatastore()
     })
     remoteUpgrader = new DefaultUpgrader(remoteComponents, {
       connectionEncryption: [