Merge pull request libp2p#45 from diasdavid/feat/id-on-conns

attach peerId to the conn
jacobheun · May 4, 2016 · 352876c · 352876c
2 parents eea7e91 + 41b700f
commit 352876c
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 4 deletions.
diff --git a/src/index.js b/src/index.js
@@ -148,7 +148,13 @@ function Swarm (peerInfo) {
     // for listening
     this.handle(muxer.multicodec, (conn) => {
       const muxedConn = muxer(conn, true)
+
+      var peerIdForConn
+
       muxedConn.on('stream', (conn) => {
+        if (peerIdForConn) {
+          conn.peerId = peerIdForConn
+        }
         connHandler(conn)
       })
 
@@ -158,6 +164,8 @@ function Swarm (peerInfo) {
           if (err) {
             return console.log('Identify exec failed', err)
           }
+
+          peerIdForConn = pi.id
           this.muxedConns[pi.id.toB58String()] = {}
           this.muxedConns[pi.id.toB58String()].muxer = muxedConn
           this.muxedConns[pi.id.toB58String()].conn = conn // to be able to extract addrs
@@ -194,15 +202,19 @@ function Swarm (peerInfo) {
 
   // higher level (public) API
   this.dial = (pi, protocol, callback) => {
-    var pt = null
     if (typeof protocol === 'function') {
       callback = protocol
       protocol = null
-    } else {
-      pt = new DuplexPassThrough()
     }
 
+    if (!callback) {
+      callback = function noop () {}
+    }
+
+    const pt = new DuplexPassThrough()
+
     const b58Id = pi.id.toB58String()
+
     if (!this.muxedConns[b58Id]) {
       if (!this.conns[b58Id]) {
         attemptDial(pi, (err, conn) => {
@@ -303,13 +315,17 @@ function Swarm (peerInfo) {
             })
 
             // in case identify is on
-            muxedConn.on('stream', connHandler)
+            muxedConn.on('stream', (conn) => {
+              conn.peerId = pi.id
+              connHandler(conn)
+            })
 
             cb(null, muxedConn)
           })
         })
       }
     }
+
     function openConnInMuxedConn (muxer, cb) {
       cb(muxer.newStream())
     }
@@ -323,6 +339,7 @@ function Swarm (peerInfo) {
           }
 
           pt.wrapStream(conn)
+          pt.peerId = pi.id
           callback(null, pt)
         })
       })

diff --git a/test/09-swarm-with-muxing.node.js b/test/09-swarm-with-muxing.node.js
@@ -195,6 +195,7 @@ describe('high level API - with everything mixed all together!', function () {
 
     swarmA.dial(peerC, '/mamao/1.0.0', (err, conn) => {
       expect(err).to.not.exist
+      expect(conn.peerId).to.exist
       expect(Object.keys(swarmA.muxedConns).length).to.equal(2)
       conn.end()
 
@@ -203,6 +204,21 @@ describe('high level API - with everything mixed all together!', function () {
     })
   })
 
+  it('again, so that identify had time', (done) => {
+    swarmC.handle('/mamao/1.0.0', (conn) => {
+      expect(conn.peerId).to.exist
+      conn.pipe(conn)
+    })
+
+    swarmA.dial(peerC, '/mamao/1.0.0', (err, conn) => {
+      expect(err).to.not.exist
+      expect(conn.peerId).to.exist
+      conn.end()
+      conn.on('data', () => {}) // let it flow.. let it flooooow
+      conn.on('end', done)
+    })
+  })
+
   it('close a muxer emits event', (done) => {
     swarmC.close(() => {})
     swarmA.once('peer-mux-closed', (peerInfo) => {