From 6b01ada194b1a4c57ea8a7d91465746cba706c02 Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Thu, 27 Feb 2020 13:59:36 +0100
Subject: [PATCH] Add address to handshakes key

---
 eth/p2p/discoveryv5/encoding.nim | 6 +++---
 eth/p2p/discoveryv5/protocol.nim | 6 ++++--
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/eth/p2p/discoveryv5/encoding.nim b/eth/p2p/discoveryv5/encoding.nim
index 373a3ac..5ba323e 100644
--- a/eth/p2p/discoveryv5/encoding.nim
+++ b/eth/p2p/discoveryv5/encoding.nim
@@ -25,7 +25,7 @@ type
     localNode*: Node
     privKey*: PrivateKey
     db*: Database
-    handshakes*: Table[string, Whoareyou] # TODO: Implement hash for NodeID
+    handshakes*: Table[string, Whoareyou] # TODO: Implement type & hash for NodeID + address
 
   HandshakeSecrets = object
     writeKey: AesKey
@@ -247,7 +247,7 @@ proc decodeEncrypted*(c: var Codec,
     auth = r.read(AuthHeader)
     authTag = auth.auth
 
-    let challenge = c.handshakes.getOrDefault($fromId)
+    let challenge = c.handshakes.getOrDefault($fromId & $fromAddr)
     if challenge.isNil:
       trace "Decoding failed (no challenge)"
       return HandshakeError
@@ -260,7 +260,7 @@ proc decodeEncrypted*(c: var Codec,
     if not c.decodeAuthResp(fromId, auth, challenge, sec, newNode):
       trace "Decoding failed (bad auth)"
       return HandshakeError
-    c.handshakes.del($fromId)
+    c.handshakes.del($fromId & $fromAddr)
 
     # Swap keys to match remote
     swap(sec.readKey, sec.writeKey)
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 4979f1a..0f1a92f 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -100,11 +100,13 @@ proc sendWhoareyou(d: Protocol, address: Address, toNode: NodeId, authTag: AuthT
   # will need to be canceled each time.
   # TODO: could also clean up handshakes in a seperate call, e.g. triggered in
   # a loop.
-  if not d.codec.handshakes.hasKeyOrPut($toNode, challenge):
+  # Use toNode + address to make it more difficult for an attacker to occupy
+  # the handshake of another node.
+  if not d.codec.handshakes.hasKeyOrPut($toNode & $address, challenge):
     sleepAsync(handshakeTimeout).addCallback() do(data: pointer):
       # TODO: should we still provide cancellation in case handshake completes
       # correctly?
-      d.codec.handshakes.del($toNode)
+      d.codec.handshakes.del($toNode & $address)
 
     var data = @(whoareyouMagic(toNode))
     data.add(rlp.encode(challenge[]))