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.
This commit is contained in:
Kim De Mey 2023-08-30 17:44:05 +02:00 committed by GitHub
parent 22df101fcf
commit 0d7e7448c4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 23 additions and 16 deletions

View File

@ -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():

View File

@ -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():

View File

@ -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)

View File

@ -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