From 0d7e7448c4a2a50136e726be7d4e4a38e094e8e3 Mon Sep 17 00:00:00 2001 From: Kim De Mey Date: Wed, 30 Aug 2023 17:44:05 +0200 Subject: [PATCH] Allow passing along the handshake ENR through talkresp handler (#634) This allows for protocols build on top of discv5 to use the ENR provided in the handshake directly, instead of having to rely on requesting it from the discv5 routing table. --- eth/p2p/discoveryv5/encoding.nim | 6 +++--- eth/p2p/discoveryv5/protocol.nim | 16 +++++++++------- eth/utp/utp_discv5_protocol.nim | 5 +++-- tests/p2p/test_discoveryv5.nim | 12 ++++++++---- 4 files changed, 23 insertions(+), 16 deletions(-) diff --git a/eth/p2p/discoveryv5/encoding.nim b/eth/p2p/discoveryv5/encoding.nim index e92b92b..2a71568 100644 --- a/eth/p2p/discoveryv5/encoding.nim +++ b/eth/p2p/discoveryv5/encoding.nim @@ -94,7 +94,7 @@ type of HandshakeMessage: message*: Message # In a handshake we expect to always be able to decrypt # TODO record or node immediately? - node*: Option[Node] + node*: Opt[Node] srcIdHs*: NodeId HandshakeKey* = object @@ -484,7 +484,7 @@ proc decodeHandshakePacket(c: var Codec, fromAddr: Address, nonce: AESGCMNonce, return err("Invalid encoded ENR") var pubkey: PublicKey - var newNode: Option[Node] + var newNode: Opt[Node] # TODO: Shall we return Node or Record? Record makes more sense, but we do # need the pubkey and the nodeid if record.isSome(): @@ -496,7 +496,7 @@ proc decodeHandshakePacket(c: var Codec, fromAddr: Address, nonce: AESGCMNonce, # Note: Not checking if the record seqNum is higher than the one we might # have stored as it comes from this node directly. pubkey = node.pubkey - newNode = some(node) + newNode = Opt.some(node) else: # TODO: Hmm, should we still verify node id of the ENR of this node? if challenge.pubkey.isSome(): diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim index e213c76..3aa74b8 100644 --- a/eth/p2p/discoveryv5/protocol.nim +++ b/eth/p2p/discoveryv5/protocol.nim @@ -157,7 +157,8 @@ type TalkProtocolHandler* = proc( p: TalkProtocol, request: seq[byte], - fromId: NodeId, fromUdpAddress: Address): seq[byte] + fromId: NodeId, fromUdpAddress: Address, + node: Opt[Node]): seq[byte] {.gcsafe, raises: [].} TalkProtocol* = ref object of RootObj @@ -338,7 +339,7 @@ proc handleFindNode(d: Protocol, fromId: NodeId, fromAddr: Address, d.sendNodes(fromId, fromAddr, reqId, []) proc handleTalkReq(d: Protocol, fromId: NodeId, fromAddr: Address, - talkreq: TalkReqMessage, reqId: RequestId) = + talkreq: TalkReqMessage, reqId: RequestId, node: Opt[Node]) = let talkProtocol = d.talkProtocols.getOrDefault(talkreq.protocol) let talkresp = @@ -348,7 +349,7 @@ proc handleTalkReq(d: Protocol, fromId: NodeId, fromAddr: Address, TalkRespMessage(response: @[]) else: TalkRespMessage(response: talkProtocol.protocolHandler(talkProtocol, - talkreq.request, fromId, fromAddr)) + talkreq.request, fromId, fromAddr, node)) let (data, _) = encodeMessagePacket(d.rng[], d.codec, fromId, fromAddr, encodeMessage(talkresp, reqId)) @@ -356,8 +357,9 @@ proc handleTalkReq(d: Protocol, fromId: NodeId, fromAddr: Address, kind = MessageKind.talkresp d.send(fromAddr, data) -proc handleMessage(d: Protocol, srcId: NodeId, fromAddr: Address, - message: Message) = +proc handleMessage( + d: Protocol, srcId: NodeId, fromAddr: Address, + message: Message, node: Opt[Node] = Opt.none(Node)) = case message.kind of ping: discovery_message_requests_incoming.inc() @@ -367,7 +369,7 @@ proc handleMessage(d: Protocol, srcId: NodeId, fromAddr: Address, d.handleFindNode(srcId, fromAddr, message.findNode, message.reqId) of talkReq: discovery_message_requests_incoming.inc() - d.handleTalkReq(srcId, fromAddr, message.talkReq, message.reqId) + d.handleTalkReq(srcId, fromAddr, message.talkReq, message.reqId, node) of regTopic, topicQuery: discovery_message_requests_incoming.inc() discovery_message_requests_incoming.inc(labelValues = ["no_response"]) @@ -447,7 +449,7 @@ proc receive*(d: Protocol, a: Address, packet: openArray[byte]) = of HandshakeMessage: trace "Received handshake message packet", srcId = packet.srcIdHs, address = a, kind = packet.message.kind - d.handleMessage(packet.srcIdHs, a, packet.message) + d.handleMessage(packet.srcIdHs, a, packet.message, packet.node) # For a handshake message it is possible that we received an newer ENR. # In that case we can add/update it to the routing table. if packet.node.isSome(): diff --git a/eth/utp/utp_discv5_protocol.nim b/eth/utp/utp_discv5_protocol.nim index 5c8d727..111e31e 100644 --- a/eth/utp/utp_discv5_protocol.nim +++ b/eth/utp/utp_discv5_protocol.nim @@ -82,8 +82,9 @@ proc initSendCallback( return fut ) -proc messageHandler(protocol: TalkProtocol, request: seq[byte], - srcId: NodeId, srcUdpAddress: Address): seq[byte] = +proc messageHandler( + protocol: TalkProtocol, request: seq[byte], + srcId: NodeId, srcUdpAddress: Address, node: Opt[Node]): seq[byte] = let p = UtpDiscv5Protocol(protocol) nodeAddress = NodeAddress.init(srcId, srcUdpAddress) diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim index 5df2d0d..cb52972 100644 --- a/tests/p2p/test_discoveryv5.nim +++ b/tests/p2p/test_discoveryv5.nim @@ -792,7 +792,8 @@ suite "Discovery v5 Tests": proc handler( protocol: TalkProtocol, request: seq[byte], - fromId: NodeId, fromUdpAddress: Address): + fromId: NodeId, fromUdpAddress: Address, + node: Opt[Node]): seq[byte] {.gcsafe, raises: [].} = request @@ -819,7 +820,8 @@ suite "Discovery v5 Tests": proc handler( protocol: TalkProtocol, request: seq[byte], - fromId: NodeId, fromUdpAddress: Address): + fromId: NodeId, fromUdpAddress: Address, + node: Opt[Node]): seq[byte] {.gcsafe, raises: [].} = request @@ -843,7 +845,8 @@ suite "Discovery v5 Tests": proc handler( protocol: TalkProtocol, request: seq[byte], - fromId: NodeId, fromUdpAddress: Address): + fromId: NodeId, fromUdpAddress: Address, + node: Opt[Node]): seq[byte] {.gcsafe, raises: [].} = request @@ -882,7 +885,8 @@ suite "Discovery v5 Tests": proc handler( protocol: TalkProtocol, request: seq[byte], - fromId: NodeId, fromUdpAddress: Address): + fromId: NodeId, fromUdpAddress: Address, + node: Opt[Node]): seq[byte] {.gcsafe, raises: [].} = # Return the request + same protocol id + 2 bytes, to make it 1 byte # bigger than the request