diff --git a/libp2pdht/private/eth/p2p/discoveryv5/encoding.nim b/libp2pdht/private/eth/p2p/discoveryv5/encoding.nim index 5bdf475..d879365 100644 --- a/libp2pdht/private/eth/p2p/discoveryv5/encoding.nim +++ b/libp2pdht/private/eth/p2p/discoveryv5/encoding.nim @@ -211,8 +211,9 @@ proc encodeStaticHeader*(flag: Flag, nonce: AESGCMNonce, authSize: int): proc encodeMessagePacket*(rng: var HmacDrbgContext, c: var Codec, toId: NodeId, toAddr: Address, message: openArray[byte]): - (seq[byte], AESGCMNonce) = + (seq[byte], AESGCMNonce, bool) = var nonce: AESGCMNonce + var haskey: bool hmacDrbgGenerate(rng, nonce) # Random AESGCM nonce var iv: array[ivSize, byte] hmacDrbgGenerate(rng, iv) # Random IV @@ -230,6 +231,7 @@ proc encodeMessagePacket*(rng: var HmacDrbgContext, c: var Codec, var messageEncrypted: seq[byte] var initiatorKey, recipientKey: AesKey if c.sessions.load(toId, toAddr, recipientKey, initiatorKey): + haskey = true messageEncrypted = encryptGCM(initiatorKey, nonce, message, @iv & header) discovery_session_lru_cache_hits.inc() else: @@ -240,6 +242,7 @@ proc encodeMessagePacket*(rng: var HmacDrbgContext, c: var Codec, # message. 16 bytes for the gcm tag and 4 bytes for ping with requestId of # 1 byte (e.g "01c20101"). Could increase to 27 for 8 bytes requestId in # case this must not look like a random packet. + haskey = false var randomData: array[gcmTagSize + 4, byte] hmacDrbgGenerate(rng, randomData) messageEncrypted.add(randomData) @@ -252,7 +255,7 @@ proc encodeMessagePacket*(rng: var HmacDrbgContext, c: var Codec, packet.add(maskedHeader) packet.add(messageEncrypted) - return (packet, nonce) + return (packet, nonce, haskey) proc encodeWhoareyouPacket*(rng: var HmacDrbgContext, c: var Codec, toId: NodeId, toAddr: Address, requestNonce: AESGCMNonce, recordSeq: uint64, diff --git a/libp2pdht/private/eth/p2p/discoveryv5/transport.nim b/libp2pdht/private/eth/p2p/discoveryv5/transport.nim index 7f47fe6..63f6db6 100644 --- a/libp2pdht/private/eth/p2p/discoveryv5/transport.nim +++ b/libp2pdht/private/eth/p2p/discoveryv5/transport.nim @@ -56,7 +56,7 @@ proc send(t: Transport, n: Node, data: seq[byte]) = t.sendToA(n.address.get(), data) proc sendMessage*(t: Transport, toId: NodeId, toAddr: Address, message: seq[byte]) = - let (data, _) = encodeMessagePacket(t.rng[], t.codec, toId, toAddr, + let (data, _, _) = encodeMessagePacket(t.rng[], t.codec, toId, toAddr, message) t.sendToA(toAddr, data) @@ -74,7 +74,7 @@ proc registerRequest(t: Transport, n: Node, message: seq[byte], proc sendMessage*(t: Transport, toNode: Node, message: seq[byte]) = doAssert(toNode.address.isSome()) let address = toNode.address.get() - let (data, nonce) = encodeMessagePacket(t.rng[], t.codec, + let (data, nonce, haskey) = encodeMessagePacket(t.rng[], t.codec, toNode.id, address, message) t.registerRequest(toNode, message, nonce) diff --git a/tests/discv5/test_discoveryv5.nim b/tests/discv5/test_discoveryv5.nim index ad0066d..9c9a8d9 100644 --- a/tests/discv5/test_discoveryv5.nim +++ b/tests/discv5/test_discoveryv5.nim @@ -629,7 +629,7 @@ suite "Discovery v5 Tests": sendNode = newNode(enrRec).expect("Properly initialized record") var codec = Codec(localNode: sendNode, privKey: privKey, sessions: Sessions.init(5)) - let (packet, _) = encodeMessagePacket(rng[], codec, + let (packet, _, _) = encodeMessagePacket(rng[], codec, receiveNode.localNode.id, receiveNode.localNode.address.get(), @[]) receiveNode.transport.receive(a, packet) @@ -659,7 +659,7 @@ suite "Discovery v5 Tests": var codec = Codec(localNode: sendNode, privKey: privKey, sessions: Sessions.init(5)) for i in 0 ..< 5: let a = localAddress(20303 + i) - let (packet, _) = encodeMessagePacket(rng[], codec, + let (packet, _, _) = encodeMessagePacket(rng[], codec, receiveNode.localNode.id, receiveNode.localNode.address.get(), @[]) receiveNode.transport.receive(a, packet) @@ -691,7 +691,7 @@ suite "Discovery v5 Tests": var firstRequestNonce: AESGCMNonce for i in 0 ..< 5: - let (packet, requestNonce) = encodeMessagePacket(rng[], codec, + let (packet, requestNonce, _) = encodeMessagePacket(rng[], codec, receiveNode.localNode.id, receiveNode.localNode.address.get(), @[]) receiveNode.transport.receive(a, packet) if i == 0: diff --git a/tests/discv5/test_discoveryv5_encoding.nim b/tests/discv5/test_discoveryv5_encoding.nim index 3eb2380..9f00ca1 100644 --- a/tests/discv5/test_discoveryv5_encoding.nim +++ b/tests/discv5/test_discoveryv5_encoding.nim @@ -526,7 +526,7 @@ suite "Discovery v5.1 Additional Encode/Decode": reqId = RequestId.init(rng[]) message = encodeMessage(m, reqId) - let (data, nonce) = encodeMessagePacket(rng[], codecA, nodeB.id, + let (data, nonce, _) = encodeMessagePacket(rng[], codecA, nodeB.id, nodeB.address.get(), message) let decoded = codecB.decodePacket(nodeA.address.get(), data) @@ -642,7 +642,7 @@ suite "Discovery v5.1 Additional Encode/Decode": codecB.sessions.store(nodeA.id, nodeA.address.get(), secrets.initiatorKey, secrets.recipientKey) - let (data, nonce) = encodeMessagePacket(rng[], codecA, nodeB.id, + let (data, nonce, _) = encodeMessagePacket(rng[], codecA, nodeB.id, nodeB.address.get(), message) let decoded = codecB.decodePacket(nodeA.address.get(), data)