From 1369cf3a30cdb37f9e878f2778f13efeed6dcac9 Mon Sep 17 00:00:00 2001
From: Zahary Karadjov <zahary@gmail.com>
Date: Sat, 22 Feb 2020 01:55:37 +0200
Subject: [PATCH] [discv5] Allow specifying the advertised IP address; Code
 style improvements

---
 eth/p2p/discoveryv5/node.nim     | 22 +++++++++++-----------
 eth/p2p/discoveryv5/protocol.nim | 31 ++++++++++++++++---------------
 tests/p2p/test_discoveryv5.nim   |  4 +++-
 3 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/eth/p2p/discoveryv5/node.nim b/eth/p2p/discoveryv5/node.nim
index 82e7f1f..f77a89b 100644
--- a/eth/p2p/discoveryv5/node.nim
+++ b/eth/p2p/discoveryv5/node.nim
@@ -11,20 +11,20 @@ type
 proc toNodeId*(pk: PublicKey): NodeId =
   readUintBE[256](keccak256.digest(pk.getRaw()).data)
 
-proc newNode*(pk: PublicKey, address: Address): Node =
-  result.new()
-  result.node = initENode(pk, address)
-  result.id = pk.toNodeId()
+proc newNode*(enode: ENode): Node =
+  Node(node: enode,
+       id: enode.pubkey.toNodeId())
+
+proc newNode*(enode: ENode, r: Record): Node =
+  Node(node: enode,
+       id: enode.pubkey.toNodeId(),
+       record: r)
 
 proc newNode*(uriString: string): Node =
-  result.new()
-  result.node = initENode(uriString)
-  result.id = result.node.pubkey.toNodeId()
+  newNode initENode(uriString)
 
-proc newNode*(enode: ENode): Node =
-  result.new()
-  result.node = enode
-  result.id = result.node.pubkey.toNodeId()
+proc newNode*(pk: PublicKey, address: Address): Node =
+  newNode initENode(pk, address)
 
 proc newNode*(r: Record): Node =
   # TODO: Handle IPv6
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 2d4ac91..fc1f79b 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -36,22 +36,23 @@ proc whoareyouMagic(toNode: NodeId): array[32, byte] =
   for i, c in prefix: data[sizeof(toNode) + i] = byte(c)
   sha256.digest(data).data
 
-proc newProtocol*(privKey: PrivateKey,
-                  db: Database,
-                  tcpPort, udpPort: Port): Protocol =
-  result = Protocol(privateKey: privKey, db: db)
-  let a = Address(ip: parseIpAddress("127.0.0.1"),
-                  tcpPort: tcpPort, udpPort: udpPort)
+proc newProtocol*(privKey: PrivateKey, db: Database,
+                  ip: IpAddress, tcpPort, udpPort: Port): Protocol =
+  let
+    a = Address(ip: ip, tcpPort: tcpPort, udpPort: udpPort)
+    enode = initENode(privKey.getPublicKey(), a)
+    enrRec = enr.Record.init(12, privKey, a)
+    node = newNode(enode, enrRec)
 
-  result.localNode = newNode(initENode(result.privateKey.getPublicKey(), a))
-  result.localNode.record = enr.Record.init(12, result.privateKey, a)
+  result = Protocol(
+    privateKey: privKey,
+    db: db,
+    localNode: node,
+    whoareyouMagic: whoareyouMagic(node.id),
+    idHash: sha256.digest(node.id.toByteArrayBE).data,
+    codec: Codec(localNode: node, privKey: privKey, db: db))
 
-  result.whoareyouMagic = whoareyouMagic(result.localNode.id)
-
-  result.idHash = sha256.digest(result.localNode.id.toByteArrayBE).data
-  result.routingTable.init(result.localNode)
-
-  result.codec = Codec(localNode: result.localNode, privKey: result.privateKey, db: result.db)
+  result.routingTable.init(node)
 
 proc start*(p: Protocol) =
   discard
@@ -376,7 +377,7 @@ when isMainModule:
         pk = newPrivateKey()
 
       let d = newProtocol(pk, DiscoveryDB.init(newMemoryDB()),
-                          Port(12001 + i), Port(12001 + i))
+                          parseIpAddress("127.0.0.1"), Port(12001 + i), Port(12001 + i))
       d.open()
       result.add(d)
 
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
index 4399b6b..5159105 100644
--- a/tests/p2p/test_discoveryv5.nim
+++ b/tests/p2p/test_discoveryv5.nim
@@ -8,7 +8,9 @@ import
 proc startDiscoveryv5Node*(privKey: PrivateKey, address: Address,
                            bootnodes: seq[Record]): discv5_protocol.Protocol =
   var db = DiscoveryDB.init(newMemoryDB())
-  result = newProtocol(privKey, db, address.tcpPort, address.udpPort)
+  result = newProtocol(privKey, db,
+                       parseIpAddress("127.0.0.1"),
+                       address.tcpPort, address.udpPort)
 
   for node in bootnodes:
     result.addNode(node)