Add record size check on init + add insertFieldPair call + tests

This commit is contained in:
kdeme 2020-07-07 17:19:15 +02:00
parent 484fbcab1b
commit 3bbe757fe3
No known key found for this signature in database
GPG Key ID: 4E8DD21420AF43F5
3 changed files with 189 additions and 35 deletions

View File

@ -1,4 +1,4 @@
# ENR implemetation according to spec:
# ENR implementation according to specification in EIP-778:
# https://github.com/ethereum/EIPs/blob/master/EIPS/eip-778.md
import
@ -11,8 +11,9 @@ export options
{.push raises: [Defect].}
const
maxEnrSize = 300
minRlpListLen = 4 # for signature, seqId, "id" key, id
maxEnrSize = 300 ## Maximum size of an encoded node record, in bytes.
minRlpListLen = 4 ## Minimum node record RLP list has: signature, seqId,
## "id" key and value.
type
FieldPair* = (string, Field)
@ -63,8 +64,49 @@ template toField[T](v: T): Field =
else:
{.error: "Unsupported field type".}
proc `==`(a, b: Field): bool =
if a.kind == b.kind:
case a.kind
of kString:
return a.str == b.str
of kNum:
return a.num == b.num
of kBytes:
return a.bytes == b.bytes
else:
return false
proc makeEnrRaw(seqNum: uint64, pk: PrivateKey,
pairs: openarray[FieldPair]): EnrResult[seq[byte]] =
proc append(w: var RlpWriter, seqNum: uint64,
pairs: openarray[FieldPair]): seq[byte] =
w.append(seqNum)
for (k, v) in pairs:
w.append(k)
case v.kind
of kString: w.append(v.str)
of kNum: w.append(v.num)
of kBytes: w.append(v.bytes)
w.finish()
let toSign = block:
var w = initRlpList(pairs.len * 2 + 1)
w.append(seqNum, pairs)
let sig = signNR(pk, toSign)
var raw = block:
var w = initRlpList(pairs.len * 2 + 2)
w.append(sig.toRaw())
w.append(seqNum, pairs)
if raw.len > maxEnrSize:
err("Record exceeds maximum size")
else:
ok(raw)
proc makeEnrAux(seqNum: uint64, pk: PrivateKey,
pairs: openarray[(string, Field)]): EnrResult[Record] =
pairs: openarray[FieldPair]): EnrResult[Record] =
var record: Record
record.pairs = @pairs
record.seqNum = seqNum
@ -76,35 +118,18 @@ proc makeEnrAux(seqNum: uint64, pk: PrivateKey,
Field(kind: kBytes, bytes: @(pubkey.toRawCompressed()))))
# Sort by key
record.pairs.sort() do(a, b: (string, Field)) -> int:
record.pairs.sort() do(a, b: FieldPair) -> int:
cmp(a[0], b[0])
proc append(w: var RlpWriter, seqNum: uint64,
pairs: openarray[(string, Field)]): seq[byte] =
w.append(seqNum)
for (k, v) in pairs:
w.append(k)
case v.kind
of kString: w.append(v.str)
of kNum: w.append(v.num)
of kBytes: w.append(v.bytes)
w.finish()
let toSign = block:
var w = initRlpList(record.pairs.len * 2 + 1)
w.append(seqNum, record.pairs)
let sig = signNR(pk, toSign)
record.raw = block:
var w = initRlpList(record.pairs.len * 2 + 2)
w.append(sig.toRaw())
w.append(seqNum, record.pairs)
record.raw = ? makeEnrRaw(seqNum, pk, record.pairs)
ok(record)
macro initRecord*(seqNum: uint64, pk: PrivateKey,
pairs: untyped{nkTableConstr}): untyped =
## Initialize a `Record` with given sequence number, private key and k:v
## pairs.
##
## Can fail in case the record exceeds the `maxEnrSize`.
for c in pairs:
c.expectKind(nnkExprColonExpr)
c[1] = newCall(bindSym"toField", c[1])
@ -121,6 +146,10 @@ proc init*(T: type Record, seqNum: uint64,
tcpPort, udpPort: Port,
extraFields: openarray[FieldPair] = []):
EnrResult[T] =
## Initialize a `Record` with given sequence number, private key, optional
## ip-address, tcp port, udp port, and optional custom k:v pairs.
##
## Can fail in case the record exceeds the `maxEnrSize`.
var fields = newSeq[FieldPair]()
if ip.isSome():
@ -190,6 +219,50 @@ proc get*(r: Record, T: type PublicKey): Option[T] =
if pk.isOk:
return some pk[]
proc find(r: Record, key: string): Option[int] =
## Search for key in record key:value pairs.
##
## Returns some(index of key) if key is found in record. Else return none.
for i, (k, v) in r.pairs:
if k == key:
return some(i)
proc insertFieldPair*(record: var Record, pk: PrivateKey, fieldPair: FieldPair):
EnrResult[bool] =
## Insert or modify a k:v pair to the `Record`.
##
## If k:v pair doesn't exist yet, it will be inserted. If it does exist, it
## will be updated if the value is different, else nothing is done. In case of
## an insert or update the seqNum will be be incremented and a new signature
## will be applied.
##
## Can fail in case of wrong PrivateKey or if the size of the resulting record
## is > maxEnrSize. `record` will not be altered in these cases.
var r = record
let pubkey = r.get(PublicKey)
if pubkey.isNone() or pubkey.get() != pk.toPublicKey():
return err("Public key does not correspond with given private key")
let index = r.find(fieldPair[0])
if(index.isSome()):
if r.pairs[index.get()][1] == fieldPair[1]:
# Exact k:v pair is already in record, nothing to do here.
return ok(true)
else:
# Need to update the value.
r.pairs[index.get()] = fieldPair
else:
# Add new k:v pair.
r.pairs.insert(fieldPair,
lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
r.seqNum.inc()
r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
record = r
ok(true)
proc tryGet*(r: Record, key: string, T: type): Option[T] =
try:
return some get(r, key, T)

View File

@ -701,8 +701,9 @@ proc newProtocol*(privKey: PrivateKey, db: Database,
# remapping through NAT and this API is also subject to change once we
# introduce support for ipv4 + ipv6 binding/listening.
let
# TODO: Defect or return a result on this call?
enrRec = enr.Record.init(1, privKey, externalIp, tcpPort, udpPort,
localEnrFields).expect("Properly intialized private key")
localEnrFields).expect("Record within size limits")
node = newNode(enrRec).expect("Properly initialized node")
# TODO Consider whether this should be a Defect

View File

@ -1,5 +1,5 @@
import
unittest, options,
unittest, options, sequtils,
nimcrypto/utils, stew/shims/net,
eth/p2p/enode, eth/p2p/discoveryv5/enr, eth/keys
@ -35,14 +35,15 @@ suite "ENR":
test "Create from ENode address":
let
keys = KeyPair.random(rng[])
keypair = KeyPair.random(rng[])
ip = ValidIpAddress.init("10.20.30.40")
enr = Record.init(100, keys.seckey, some(ip), Port(9000), Port(9000), @[])[]
enr = Record.init(
100, keypair.seckey, some(ip), Port(9000), Port(9000),@[])[]
typedEnr = get enr.toTypedRecord()
check:
typedEnr.secp256k1.isSome()
typedEnr.secp256k1.get == keys.pubkey.toRawCompressed()
typedEnr.secp256k1.get == keypair.pubkey.toRawCompressed()
typedEnr.ip.isSome()
typedEnr.ip.get() == [byte 10, 20, 30, 40]
@ -55,13 +56,14 @@ suite "ENR":
test "ENR without address":
let
keys = KeyPair.random(rng[])
enr = Record.init(100, keys.seckey, none(ValidIpAddress), Port(9000), Port(9000))[]
keypair = KeyPair.random(rng[])
enr = Record.init(
100, keypair.seckey, none(ValidIpAddress), Port(9000), Port(9000))[]
typedEnr = get enr.toTypedRecord()
check:
typedEnr.secp256k1.isSome()
typedEnr.secp256k1.get() == keys.pubkey.toRawCompressed()
typedEnr.secp256k1.get() == keypair.pubkey.toRawCompressed()
typedEnr.ip.isNone()
typedEnr.tcp.isSome()
@ -73,3 +75,81 @@ suite "ENR":
typedEnr.ip6.isNone()
typedEnr.tcp6.isNone()
typedEnr.udp6.isNone()
test "ENR init size too big":
let pk = PrivateKey.fromHex(
"5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
block: # This gives ENR of 300 bytes encoded
let r = initRecord(1, pk, {"maxvalue": repeat(byte 2, 169),})
check r.isOk()
block: # This gives ENR of 301 bytes encoded
let r = initRecord(1, pk, {"maxplus1": repeat(byte 2, 170),})
check r.isErr()
test "ENR insert":
let
pk = PrivateKey.fromHex(
"5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
newField = toFieldPair("test", 123'u)
var r = Record.init(1, pk, none(ValidIpAddress), Port(9000), Port(9000))[]
block: # Insert new k:v pair, update of seqNum should occur.
let res = r.insertFieldPair(pk, newField)
check:
res.isOk()
r.get("test", uint) == 123
r.seqNum == 2
block: # Insert same k:v pair, no update of seqNum should occur.
let res = r.insertFieldPair(pk, newField)
check:
res.isOk()
r.get("test", uint) == 123
r.seqNum == 2
block: # Insert k:v pair with changed value, update of seqNum should occur.
let updatedField = toFieldPair("test", 1234'u)
let res = r.insertFieldPair(pk, updatedField)
check:
res.isOk()
r.get("test", uint) == 1234
r.seqNum == 3
test "ENR insert sorted":
let pk = PrivateKey.fromHex(
"5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
var r = initRecord(123, pk, {"abc": 1234'u,
"z": [byte 0],
"123": "abc",
"a12": 1'u})[]
check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, z: 0x00)"""
let newField = toFieldPair("test", 123'u)
let res = r.insertFieldPair(pk, newField)
check res.isOk()
check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00)"""
test "ENR insert size too big":
let pk = PrivateKey.fromHex(
"5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
var r = initRecord(1, pk, {"maxvalue": repeat(byte 2, 169),})
check r.isOk()
let newField = toFieldPair("test", 123'u)
let res = r[].insertFieldPair(pk, newField)
check res.isErr()
test "ENR insert invalid key":
let pk = PrivateKey.fromHex(
"5d2908f3f09ea1ff2e327c3f623159639b00af406e9009de5fd4b910fc34049d")[]
var r = initRecord(1, pk, {"abc": 1'u,})
check r.isOk()
let
wrongPk = PrivateKey.random(rng[])
newField = toFieldPair("test", 123'u)
res = r[].insertFieldPair(wrongPk, newField)
check res.isErr()