From a7e5fde6f72b395b29014a077a81073c34fcb0cf Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Wed, 4 Sep 2019 11:36:07 -0600
Subject: [PATCH] check identify request against known pubkey

---
 libp2p/identify.nim    | 18 +++++++++++------
 tests/testidentify.nim | 46 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/libp2p/identify.nim b/libp2p/identify.nim
index 7e806a892..3413323b9 100644
--- a/libp2p/identify.nim
+++ b/libp2p/identify.nim
@@ -7,10 +7,12 @@
 ## This file may not be copied, modified, or distributed except according to
 ## those terms.
 
+import options
 import chronos
 import protobuf/minprotobuf, peerinfo, 
        protocol as proto, connection,
-       peer, crypto/crypto, multiaddress
+       peer, crypto/crypto, multiaddress,
+       crypto/crypto
 
 const IdentifyCodec* = "/ipfs/id/1.0.0"
 const IdentifyPushCodec* = "/ipfs/id/push/1.0.0"
@@ -20,10 +22,7 @@ const AgentVersion* = "nim-libp2p/0.0.1"
 #TODO: implment push identify, leaving out for now as it is not essential
 
 type
-  # TODO: we're doing protobuf manualy, this is only temporary
-  ProtoField[T] = object
-    index: int
-    field: T
+  IdentityNoMatchError* = object of CatchableError
 
   IdentifyInfo* = object
     pubKey*: PublicKey
@@ -92,11 +91,18 @@ method init*(p: Identify) =
 
   p.handler = handle
 
-proc identify*(p: Identify, conn: Connection): Future[IdentifyInfo] {.async.} = 
+proc identify*(p: Identify, 
+               conn: Connection, 
+               remotePeerInfo: Option[PeerInfo] = none(PeerInfo)): 
+               Future[IdentifyInfo] {.async.} = 
   var message = await conn.readLp()
   if len(message) == 0:
     raise newException(CatchableError, "Incorrect or empty message received!")
   result = decodeMsg(message)
+  if remotePeerInfo.isSome and 
+      remotePeerInfo.get().peerId.publicKey != result.pubKey:
+    raise newException(IdentityNoMatchError, 
+      "Peer's remote public key doesn't match")
 
 proc push*(p: Identify, conn: Connection) {.async.} =
   await conn.write(IdentifyPushCodec)
diff --git a/tests/testidentify.nim b/tests/testidentify.nim
index bdb2251d5..043c728bd 100644
--- a/tests/testidentify.nim
+++ b/tests/testidentify.nim
@@ -1,4 +1,4 @@
-import unittest
+import unittest, options
 import chronos, strutils, sequtils
 import ../libp2p/identify, ../libp2p/multiaddress, 
        ../libp2p/peerinfo, ../libp2p/peer, 
@@ -17,6 +17,8 @@ suite "Identify":
         var peerInfo: PeerInfo
         peerInfo.peerId = PeerID.init(remoteSeckey)
         peerInfo.addrs.add(ma)
+        peerInfo.protocols.add("/test/proto1/1.0.0")
+        peerInfo.protocols.add("/test/proto2/1.0.0")
 
         let identifyProto = newIdentify(peerInfo)
         let msListen = newMultistream()
@@ -48,10 +50,52 @@ suite "Identify":
         check id.addrs[0] == ma
         check id.protoVersion == ProtoVersion
         check id.agentVersion == AgentVersion
+        check id.protos == @["/test/proto1/1.0.0", "/test/proto2/1.0.0"]
 
       await allFutures(receiver(), sender())
       result = true
 
     check:
       waitFor(testHandle()) == true
+  
+  test "handle failed identify":
+    proc testHandleError() {.async.} =
+      let ma: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53361")
+
+      let remoteSeckey = PrivateKey.random(RSA)
+      var remotePeerInfo: PeerInfo
+      remotePeerInfo.peerId = PeerID.init(remoteSeckey)
+      remotePeerInfo.addrs.add(ma)
+
+      let identifyProto1 = newIdentify(remotePeerInfo)
+      let msListen = newMultistream()
+
+      msListen.addHandler(IdentifyCodec, identifyProto1)
+      proc connHandler(conn: Connection): Future[void] {.async, gcsafe.} =
+        await msListen.handle(conn)
+
+      let transport1: TcpTransport = newTransport(TcpTransport)
+      await transport1.listen(ma, connHandler)
+
+      let msDial = newMultistream()
+      let transport2: TcpTransport = newTransport(TcpTransport)
+      let conn = await transport2.dial(ma)
+
+      let seckey = PrivateKey.random(RSA)
+      var localPeerInfo: PeerInfo
+      localPeerInfo.peerId = PeerID.init(seckey)
+      localPeerInfo.addrs.add(ma)
+
+      let identifyProto2 = newIdentify(localPeerInfo)
+      let res = await msDial.select(conn, IdentifyCodec)
+
+      let wrongSec = PrivateKey.random(RSA)
+      var wrongRemotePeer: PeerInfo
+      wrongRemotePeer.peerId = PeerID.init(wrongSec)
+
+      let id = await identifyProto2.identify(conn, some(wrongRemotePeer))
+      await conn.close()
+
+    expect IdentityNoMatchError:
+      waitFor(testHandleError())
   
\ No newline at end of file