feat: client mode (#115)

* Add client mode in findNode

* Remove clientMode checking in addProvider

* Add clientMode check over TalkProtocol

* Fix client mode checking

* Add register talk protocol unicity check

* Add future callback to prevent unbound growth

* Fix PR reviews

* Use DhtMode

* Add clientMode in the Message

* Fix minor issues

* Add logs

* Switch to Nim 2.2.10 on CI

* Add client mode even when the value is false

* Move the client mode check at the beginning of the proc

* Increment version

* Replace uint64 by uint32

.github/workflows/ci.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        nim: [2.2.8]
+        nim: [2.2.10]
         os: [ubuntu-latest, macos-latest, windows-latest]
     steps:
       - name: Checkout

codexdht.nimble

@@ -1,6 +1,6 @@
 # Package
 
-version       = "0.6.2"
+version       = "0.6.3"
 author        = "Status Research & Development GmbH"
 description   = "DHT based on Eth discv5 implementation"
 license       = "MIT"

codexdht/private/eth/p2p/discoveryv5/messages.nim

@@ -83,6 +83,7 @@ type
 
   Message* = object
     reqId*: RequestId
+    clientMode*: bool
     case kind*: MessageKind
     of ping:
       ping*: PingMessage

codexdht/private/eth/p2p/discoveryv5/messages_encoding.nim

@@ -312,7 +312,7 @@ proc encode*(msg: TalkRespMessage): seq[byte] =
   pb.finish()
   pb.buffer
 
-proc encodeMessage*[T: SomeMessage](p: T, reqId: RequestId): seq[byte] =
+proc encodeMessage*[T: SomeMessage](p: T, reqId: RequestId, clientMode: bool = false): seq[byte] =
   result = newSeqOfCap[byte](64)
   result.add(messageKind(T).ord)
 
@@ -324,6 +324,8 @@ proc encodeMessage*[T: SomeMessage](p: T, reqId: RequestId): seq[byte] =
   var pb = initProtoBuffer()
   pb.write(1, reqId)
   pb.write(2, encoded)
+  pb.write(3, clientMode.uint32)
+
   pb.finish()
   result.add(pb.buffer)
   trace "Encoded protobuf message", typ = $T
@@ -344,6 +346,7 @@ proc decodeMessage*(body: openArray[byte]): DecodeResult[Message] =
   var
     reqId: RequestId
     encoded: EncodedMessage
+    clientModeField: uint32
 
   if pb.getRequiredField(1, reqId).isErr:
     return err("Invalid request-id")
@@ -353,6 +356,11 @@ proc decodeMessage*(body: openArray[byte]): DecodeResult[Message] =
   if pb.getRequiredField(2, encoded).isErr:
     return err("Invalid message encoding")
 
+  if pb.getField(3, clientModeField).isErr:
+    return err("Invalid clientMode field")
+
+  message.clientMode = clientModeField == 1
+
   case kind
   of unused: return err("Invalid message type")
 

codexdht/private/eth/p2p/discoveryv5/protocol.nim

@@ -137,7 +137,6 @@ const
   LookupSeenThreshold = 0.0 ## threshold used for lookup nodeset selection
   QuerySeenThreshold = 0.0 ## threshold used for query nodeset selection
   NoreplyRemoveThreshold = 0.5 ## remove node on no reply if 'seen' is below this value
-
 func shortLog*(record: SignedPeerRecord): string =
   ## Returns compact string representation of ``SignedPeerRecord``.
   ##
@@ -180,6 +179,7 @@ type
     talkProtocols*: Table[seq[byte], TalkProtocol] # TODO: Table is a bit of
     rng*: ref HmacDrbgContext
     providers: ProvidersManager
+    clientMode*: bool
 
   TalkProtocolHandler* = proc(p: TalkProtocol, request: seq[byte], fromId: NodeId, fromUdpAddress: Address): seq[byte]
     {.gcsafe, raises: [Defect].}
@@ -297,7 +297,7 @@ proc updateRecord*(
 proc sendResponse(d: Protocol, dstId: NodeId, dstAddr: Address,
     message: SomeMessage, reqId: RequestId) =
   ## send Response using the specifid reqId
-  d.transport.sendMessage(dstId, dstAddr, encodeMessage(message, reqId))
+  d.transport.sendMessage(dstId, dstAddr, encodeMessage(message, reqId, d.clientMode))
 
 proc sendNodes(d: Protocol, toId: NodeId, toAddr: Address, reqId: RequestId,
     nodes: openArray[Node]) =
@@ -411,6 +411,13 @@ proc handleGetProviders(
 
 proc handleMessage(d: Protocol, srcId: NodeId, fromAddr: Address,
     message: Message) =
+
+  if message.clientMode:
+    let node = d.routingTable.getNode(srcId)
+    if node.isSome:
+      d.routingTable.removeNode(node.get)
+      trace "Node removed from routing table after handling message", srcId
+
   case message.kind
   of ping:
     dht_message_requests_incoming.inc()
@@ -466,7 +473,7 @@ proc sendRequest*[T: SomeMessage](d: Protocol, toNode: Node, m: T,
     reqId: RequestId) =
   doAssert(toNode.address.isSome())
   let
-    message = encodeMessage(m, reqId)
+    message = encodeMessage(m, reqId, d.clientMode)
 
   trace "Send message packet", dstId = toNode.id,
     address = toNode.address, kind = messageKind(T)

codexdht/private/eth/p2p/discoveryv5/transport.nim

@@ -232,8 +232,14 @@ proc receive*(t: Transport, a: Address, packet: openArray[byte]) =
           # sending the 'whoareyou' message to. In that case, we can set 'seen'
           # TODO: verify how this works with restrictive NAT and firewall scenarios.
           node.registerSeen()
-          if t.client.addNode(node):
-            trace "Added new node to routing table after handshake", node, tablesize=t.client.nodesDiscovered()
+
+          if packet.message.clientMode:
+            t.client.routingTable.removeNode(node)
+            trace "Removed node from the routing table after handshake", node, tablesize=t.client.nodesDiscovered()
+          else:
+            if t.client.addNode(node):
+              trace "Added new node to routing table after handshake", node, tablesize=t.client.nodesDiscovered()
+
           discard t.sendPending(node)
         else:
           trace "address mismatch, not adding seen flag", node, address = a, nodeAddress = node.address

tests/discv5/test_discoveryv5.nim

@@ -777,3 +777,66 @@ suite "Discovery v5 Tests":
 
     await node1.closeWait()
     await node2.closeWait()
+
+  test "Node is added to routing table when clientMode is not enabled":
+    let
+      node1 = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20310))
+      node2 = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20311))
+
+    discard await discv5_protocol.ping(node1, node2.localNode)
+
+    check node2.routingTable.len() == 1
+
+    await node1.closeWait()
+    await node2.closeWait()
+
+  test "Node is not added to routing table when clientMode is enabled":
+    let
+      clientNode = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20314))
+      serverNode = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20315))
+
+    clientNode.clientMode = true
+
+    discard await discv5_protocol.ping(clientNode, serverNode.localNode)
+
+    check serverNode.routingTable.len() == 0
+
+    await clientNode.closeWait()
+    await serverNode.closeWait()
+
+  test "Node is removed from routing table when clientMode is enabled after session is established":
+    let
+      node1 = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20318))
+      node2 = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20319))
+
+    discard await discv5_protocol.ping(node1, node2.localNode)
+
+    check node2.routingTable.len() == 1
+
+    node1.clientMode = true
+
+    discard await discv5_protocol.ping(node1, node2.localNode)
+
+    check node2.routingTable.len() == 0
+
+    await node1.closeWait()
+    await node2.closeWait()
+
+  test "Node is removed from routing table when clientMode is enabled during re-validation":
+    let
+      clientNode = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20316))
+      serverNode = initDiscoveryNode(rng, PrivateKey.example(rng), localAddress(20317))
+
+    # Add client node directly to routing table
+    check serverNode.addNode(clientNode.localNode)
+    
+    check serverNode.routingTable.len() == 1
+
+    clientNode.clientMode = true
+
+    await serverNode.revalidateNode(clientNode.localNode)
+
+    check serverNode.routingTable.len() == 0
+
+    await clientNode.closeWait()
+    await serverNode.closeWait()

tests/discv5/test_discoveryv5_encoding.nim

@@ -22,7 +22,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(p, reqId)
-    check byteutils.toHex(encoded) == "010a010112020801"
+    check byteutils.toHex(encoded) == "010a0101120208011800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -42,7 +42,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(p, reqId)
-    check byteutils.toHex(encoded) == "020a01011211080112090a010112047f0000011a021388"
+    check byteutils.toHex(encoded) == "020a01011211080112090a010112047f0000011a0213881800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -62,7 +62,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(fn, reqId)
-    check byteutils.toHex(encoded) == "030a010112040a020100"
+    check byteutils.toHex(encoded) == "030a010112040a0201001800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -80,7 +80,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(n, reqId)
-    check byteutils.toHex(encoded) == "040a010112020801"
+    check byteutils.toHex(encoded) == "040a0101120208011800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -102,7 +102,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(n, reqId)
-    check byteutils.toHex(encoded) == "040a0101129f03080112cc010a250802122102339d487ed237392d83791950dc891f0636de698c1fa051ea01ae3fa58bd78580120203011a560a2700250802122102339d487ed237392d83791950dc891f0636de698c1fa051ea01ae3fa58bd78580109cfd8992061a0b0a090400000000910200011a0b0a090400000000910200021a0b0a090400000000910200032a4730450221008cc77fd265e33c955174b9f49628048b2d72a6395acb30f0ba9d90536fa1a5d502207fa8e5bab8e8ddee9884a8e244b0990228e3546b5a9b6848632abd924796e57612cb010a2508021221026beda5cfddf1cd89130e7b5bb6092bac23db4a044bf847328aa0310dd123a445120203011a560a27002508021221026beda5cfddf1cd89130e7b5bb6092bac23db4a044bf847328aa0310dd123a445109cfd8992061a0b0a090400000000910200011a0b0a090400000000910200021a0b0a090400000000910200032a46304402203d41b1a78c5e6d98c9b4f3fcb213dc16ae4de50a1c8715ab29c516afe6488b4e02205841d09e92b3d2f1ad72c7bc066e561dab57320886f3fbbf272d2cf1732ca259"
+    check byteutils.toHex(encoded) == "040a0101129f03080112cc010a250802122102339d487ed237392d83791950dc891f0636de698c1fa051ea01ae3fa58bd78580120203011a560a2700250802122102339d487ed237392d83791950dc891f0636de698c1fa051ea01ae3fa58bd78580109cfd8992061a0b0a090400000000910200011a0b0a090400000000910200021a0b0a090400000000910200032a4730450221008cc77fd265e33c955174b9f49628048b2d72a6395acb30f0ba9d90536fa1a5d502207fa8e5bab8e8ddee9884a8e244b0990228e3546b5a9b6848632abd924796e57612cb010a2508021221026beda5cfddf1cd89130e7b5bb6092bac23db4a044bf847328aa0310dd123a445120203011a560a27002508021221026beda5cfddf1cd89130e7b5bb6092bac23db4a044bf847328aa0310dd123a445109cfd8992061a0b0a090400000000910200011a0b0a090400000000910200021a0b0a090400000000910200032a46304402203d41b1a78c5e6d98c9b4f3fcb213dc16ae4de50a1c8715ab29c516afe6488b4e02205841d09e92b3d2f1ad72c7bc066e561dab57320886f3fbbf272d2cf1732ca2591800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -122,7 +122,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(tr, reqId)
-    check byteutils.toHex(encoded) == "050a0101120a0a046563686f12026869"
+    check byteutils.toHex(encoded) == "050a0101120a0a046563686f120268691800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -140,7 +140,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       reqId = RequestId(id: @[1.byte])
 
     let encoded = encodeMessage(tr, reqId)
-    check byteutils.toHex(encoded) == "060a0101120412026869"
+    check byteutils.toHex(encoded) == "060a01011204120268691800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
@@ -158,7 +158,7 @@ suite "Discovery v5.1 Protocol Message Encodings":
       # 1 byte too large
       reqId = RequestId(id: @[0.byte, 1, 2, 3, 4, 5, 6, 7, 8])
     let encoded = encodeMessage(p, reqId)
-    check byteutils.toHex(encoded) == "010a0900010203040506070812020801"
+    check byteutils.toHex(encoded) == "010a09000102030405060708120208011800"
 
     let decoded = decodeMessage(encoded)
     check decoded.isErr()
@@ -170,6 +170,27 @@ suite "Discovery v5.1 Protocol Message Encodings":
     let decoded = decodeMessage(hexToSeqByte(encodedPong))
     check decoded.isErr()
 
+  test "clientMode flag is correctly encoded and decoded":
+    let
+      p = PingMessage(sprSeq: 1'u64)
+      reqId = RequestId(id: @[1.byte])
+
+    let encodedClient = encodeMessage(p, reqId, clientMode = true)
+    let decodedClient = decodeMessage(encodedClient)
+    check decodedClient.isOk()
+    check decodedClient.get().clientMode == true
+
+    let encodedServer = encodeMessage(p, reqId, clientMode = false)
+    let decodedServer = decodeMessage(encodedServer)
+    check decodedServer.isOk()
+    check decodedServer.get().clientMode == false
+
+  test "Message without clientMode field decodes as server mode":
+    # "010a010112020801" is a ping from a legacy node (clientMode field)
+    let decoded = decodeMessage(hexToSeqByte("010a010112020801"))
+    check decoded.isOk()
+    check decoded.get().clientMode == false
+
 # According to test vectors:
 # https://github.com/ethereum/devp2p/blob/master/discv5/discv5-wire-test-vectors.md#cryptographic-primitives
 suite "Discovery v5.1 Cryptographic Primitives Test Vectors":