Add support for discv5 talk protocols (#357)

This commit is contained in:
Kim De Mey 2021-05-20 09:49:46 +02:00 committed by GitHub
parent 1995afb87e
commit 8abe6b7144
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 115 additions and 4 deletions

View File

@ -132,12 +132,16 @@ type
bootstrapRecords*: seq[Record] bootstrapRecords*: seq[Record]
ipVote: IpVote ipVote: IpVote
enrAutoUpdate: bool enrAutoUpdate: bool
talkProtocols: Table[seq[byte], TalkProtocolHandler]
rng*: ref BrHmacDrbgContext rng*: ref BrHmacDrbgContext
PendingRequest = object PendingRequest = object
node: Node node: Node
message: seq[byte] message: seq[byte]
TalkProtocolHandler* = proc(request: seq[byte]): seq[byte]
{.gcsafe, raises: [Defect].}
DiscResult*[T] = Result[T, cstring] DiscResult*[T] = Result[T, cstring]
proc addNode*(d: Protocol, node: Node): bool = proc addNode*(d: Protocol, node: Node): bool =
@ -295,9 +299,15 @@ proc handleFindNode(d: Protocol, fromId: NodeId, fromAddr: Address,
proc handleTalkReq(d: Protocol, fromId: NodeId, fromAddr: Address, proc handleTalkReq(d: Protocol, fromId: NodeId, fromAddr: Address,
talkreq: TalkReqMessage, reqId: RequestId) = talkreq: TalkReqMessage, reqId: RequestId) =
# No support for any protocol yet so an empty response is send as per let protocolHandler = d.talkProtocols.getOrDefault(talkreq.protocol)
# specification.
let talkresp = TalkRespMessage(response: @[]) let talkresp =
if protocolHandler.isNil():
# Protocol identifier that is not registered and thus not supported. An
# empty response is send as per specification.
TalkRespMessage(response: @[])
else:
TalkRespMessage(response: protocolHandler(talkreq.request))
let (data, _) = encodeMessagePacket(d.rng[], d.codec, fromId, fromAddr, let (data, _) = encodeMessagePacket(d.rng[], d.codec, fromId, fromAddr,
encodeMessage(talkresp, reqId)) encodeMessage(talkresp, reqId))
@ -331,6 +341,14 @@ proc handleMessage(d: Protocol, srcId: NodeId, fromAddr: Address,
trace "Timed out or unrequested message", kind = message.kind, trace "Timed out or unrequested message", kind = message.kind,
origin = fromAddr origin = fromAddr
proc registerTalkProtocol*(d: Protocol, protocol: seq[byte],
handler: TalkProtocolHandler): DiscResult[void] =
# Currently allow only for one handler per talk protocol.
if d.talkProtocols.hasKeyOrPut(protocol, handler):
err("Protocol identifier already registered")
else:
ok()
proc sendWhoareyou(d: Protocol, toId: NodeId, a: Address, proc sendWhoareyou(d: Protocol, toId: NodeId, a: Address,
requestNonce: AESGCMNonce, node: Option[Node]) = requestNonce: AESGCMNonce, node: Option[Node]) =
let key = HandShakeKey(nodeId: toId, address: a) let key = HandShakeKey(nodeId: toId, address: a)

View File

@ -2,7 +2,8 @@
import import
std/tables, std/tables,
chronos, chronicles, stint, testutils/unittests, stew/shims/net, bearssl, chronos, chronicles, stint, testutils/unittests, stew/shims/net,
stew/byteutils, bearssl,
../../eth/keys, ../../eth/keys,
../../eth/p2p/discoveryv5/[enr, node, routing_table, encoding, sessions, messages], ../../eth/p2p/discoveryv5/[enr, node, routing_table, encoding, sessions, messages],
../../eth/p2p/discoveryv5/protocol as discv5_protocol, ../../eth/p2p/discoveryv5/protocol as discv5_protocol,
@ -619,3 +620,60 @@ procSuite "Discovery v5 Tests":
firstRequestNonce firstRequestNonce
await receiveNode.closeWait() await receiveNode.closeWait()
asyncTest "Talkreq no protocol":
let
node1 = initDiscoveryNode(
rng, PrivateKey.random(rng[]), localAddress(20302))
node2 = initDiscoveryNode(
rng, PrivateKey.random(rng[]), localAddress(20303))
talkresp = await discv5_protocol.talkreq(node1, node2.localNode,
@[byte 0x01], @[])
check:
talkresp.isOk()
talkresp.get().response.len == 0
await node1.closeWait()
await node2.closeWait()
asyncTest "Talkreq echo protocol":
let
node1 = initDiscoveryNode(
rng, PrivateKey.random(rng[]), localAddress(20302))
node2 = initDiscoveryNode(
rng, PrivateKey.random(rng[]), localAddress(20303))
talkProtocol = "echo".toBytes()
proc handler(request: seq[byte]): seq[byte] {.gcsafe, raises: [Defect].} =
request
check node2.registerTalkProtocol(talkProtocol, handler).isOk()
let talkresp = await discv5_protocol.talkreq(node1, node2.localNode,
talkProtocol, "hello".toBytes())
check:
talkresp.isOk()
talkresp.get().response == "hello".toBytes()
await node1.closeWait()
await node2.closeWait()
asyncTest "Talkreq register protocols":
let
node1 = initDiscoveryNode(
rng, PrivateKey.random(rng[]), localAddress(20302))
node2 = initDiscoveryNode(
rng, PrivateKey.random(rng[]), localAddress(20303))
talkProtocol = "echo".toBytes()
proc handler(request: seq[byte]): seq[byte] {.gcsafe, raises: [Defect].} =
request
check:
node2.registerTalkProtocol(talkProtocol, handler).isOk()
node2.registerTalkProtocol(talkProtocol, handler).isErr()
node2.registerTalkProtocol("test".toBytes(), handler).isOk()
await node1.closeWait()
await node2.closeWait()

View File

@ -110,6 +110,41 @@ suite "Discovery v5.1 Protocol Message Encodings":
message.nodes.enrs[0] == e1 message.nodes.enrs[0] == e1
message.nodes.enrs[1] == e2 message.nodes.enrs[1] == e2
test "Talk Request":
let
tr = TalkReqMessage(protocol: "echo".toBytes(), request: "hi".toBytes())
reqId = RequestId(id: @[1.byte])
let encoded = encodeMessage(tr, reqId)
check encoded.toHex == "05c901846563686f826869"
let decoded = decodeMessage(encoded)
check decoded.isOk()
let message = decoded.get()
check:
message.reqId == reqId
message.kind == talkreq
message.talkreq.protocol == "echo".toBytes()
message.talkreq.request == "hi".toBytes()
test "Talk Response":
let
tr = TalkRespMessage(response: "hi".toBytes())
reqId = RequestId(id: @[1.byte])
let encoded = encodeMessage(tr, reqId)
check encoded.toHex == "06c401826869"
let decoded = decodeMessage(encoded)
check decoded.isOk()
let message = decoded.get()
check:
message.reqId == reqId
message.kind == talkresp
message.talkresp.response == "hi".toBytes()
test "Ping with too large RequestId": test "Ping with too large RequestId":
let let
enrSeq = 1'u64 enrSeq = 1'u64