Allow for multiple enr fields to be inserted + enr update proc

This commit is contained in:
kdeme 2020-07-07 22:48:26 +02:00
parent 3bbe757fe3
commit 57302fcf52
No known key found for this signature in database
GPG Key ID: 4E8DD21420AF43F5
2 changed files with 54 additions and 32 deletions

View File

@ -227,41 +227,62 @@ proc find(r: Record, key: string): Option[int] =
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.
proc insertFieldPairs*(record: var Record, pk: PrivateKey,
fieldPairs: openarray[FieldPair]): EnrResult[bool] =
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")
var updated = false
for fieldPair in fieldPairs:
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)
continue
else:
# Need to update the value.
r.pairs[index.get()] = fieldPair
updated = true
else:
# Add new k:v pair.
r.pairs.insert(fieldPair,
lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
updated = true
if updated:
r.seqNum.inc()
r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
record = r
ok(true)
else:
ok(true)
proc update*(r: var Record, pk: PrivateKey,
ip: Option[ValidIpAddress],
tcpPort, udpPort: Port,
extraFields: openarray[FieldPair]):
EnrResult[bool] =
var fields = newSeq[FieldPair]()
if ip.isSome():
let
ipExt = ip.get()
isV6 = ipExt.family == IPv6
fields.add(if isV6: ("ip6", ipExt.address_v6.toField)
else: ("ip", ipExt.address_v4.toField))
fields.add(((if isV6: "tcp6" else: "tcp"), tcpPort.uint16.toField))
fields.add(((if isV6: "udp6" else: "udp"), udpPort.uint16.toField))
else:
fields.add(("tcp", tcpPort.uint16.toField))
fields.add(("udp", udpPort.uint16.toField))
fields.add extraFields
r.insertFieldPairs(pk, fields)
proc tryGet*(r: Record, key: string, T: type): Option[T] =
try:

View File

@ -95,14 +95,14 @@ suite "ENR":
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)
let res = r.insertFieldPairs(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)
let res = r.insertFieldPairs(pk, [newField])
check:
res.isOk()
r.get("test", uint) == 123
@ -110,7 +110,7 @@ suite "ENR":
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)
let res = r.insertFieldPairs(pk, [updatedField])
check:
res.isOk()
r.get("test", uint) == 1234
@ -126,9 +126,10 @@ suite "ENR":
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)
let newField2 = toFieldPair("zzz", 123'u)
let res = r.insertFieldPairs(pk, [newField, newField2])
check res.isOk()
check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00)"""
check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, test: 123, z: 0x00, zzz: 123)"""
test "ENR insert size too big":
let pk = PrivateKey.fromHex(
@ -138,7 +139,7 @@ suite "ENR":
check r.isOk()
let newField = toFieldPair("test", 123'u)
let res = r[].insertFieldPair(pk, newField)
let res = r[].insertFieldPairs(pk, [newField])
check res.isErr()
test "ENR insert invalid key":
@ -151,5 +152,5 @@ suite "ENR":
let
wrongPk = PrivateKey.random(rng[])
newField = toFieldPair("test", 123'u)
res = r[].insertFieldPair(wrongPk, newField)
res = r[].insertFieldPairs(wrongPk, [newField])
check res.isErr()