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: if k == key:
return some(i) return some(i)
proc insertFieldPair*(record: var Record, pk: PrivateKey, fieldPair: FieldPair): proc insertFieldPairs*(record: var Record, pk: PrivateKey,
EnrResult[bool] = fieldPairs: openarray[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 var r = record
let pubkey = r.get(PublicKey) let pubkey = r.get(PublicKey)
if pubkey.isNone() or pubkey.get() != pk.toPublicKey(): if pubkey.isNone() or pubkey.get() != pk.toPublicKey():
return err("Public key does not correspond with given private key") return err("Public key does not correspond with given private key")
var updated = false
for fieldPair in fieldPairs:
let index = r.find(fieldPair[0]) let index = r.find(fieldPair[0])
if(index.isSome()): if(index.isSome()):
if r.pairs[index.get()][1] == fieldPair[1]: if r.pairs[index.get()][1] == fieldPair[1]:
# Exact k:v pair is already in record, nothing to do here. # Exact k:v pair is already in record, nothing to do here.
return ok(true) continue
else: else:
# Need to update the value. # Need to update the value.
r.pairs[index.get()] = fieldPair r.pairs[index.get()] = fieldPair
updated = true
else: else:
# Add new k:v pair. # Add new k:v pair.
r.pairs.insert(fieldPair, r.pairs.insert(fieldPair,
lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0])) lowerBound(r.pairs, fieldPair) do(a, b: FieldPair) -> int: cmp(a[0], b[0]))
updated = true
if updated:
r.seqNum.inc() r.seqNum.inc()
r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs) r.raw = ? makeEnrRaw(r.seqNum, pk, r.pairs)
record = r record = r
ok(true) 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] = proc tryGet*(r: Record, key: string, T: type): Option[T] =
try: try:

View File

@ -95,14 +95,14 @@ suite "ENR":
var r = Record.init(1, pk, none(ValidIpAddress), Port(9000), Port(9000))[] var r = Record.init(1, pk, none(ValidIpAddress), Port(9000), Port(9000))[]
block: # Insert new k:v pair, update of seqNum should occur. block: # Insert new k:v pair, update of seqNum should occur.
let res = r.insertFieldPair(pk, newField) let res = r.insertFieldPairs(pk, [newField])
check: check:
res.isOk() res.isOk()
r.get("test", uint) == 123 r.get("test", uint) == 123
r.seqNum == 2 r.seqNum == 2
block: # Insert same k:v pair, no update of seqNum should occur. 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: check:
res.isOk() res.isOk()
r.get("test", uint) == 123 r.get("test", uint) == 123
@ -110,7 +110,7 @@ suite "ENR":
block: # Insert k:v pair with changed value, update of seqNum should occur. block: # Insert k:v pair with changed value, update of seqNum should occur.
let updatedField = toFieldPair("test", 1234'u) let updatedField = toFieldPair("test", 1234'u)
let res = r.insertFieldPair(pk, updatedField) let res = r.insertFieldPairs(pk, [updatedField])
check: check:
res.isOk() res.isOk()
r.get("test", uint) == 1234 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)""" check $r == """(123: "abc", a12: 1, abc: 1234, id: "v4", secp256k1: 0x02E51EFA66628CE09F689BC2B82F165A75A9DDECBB6A804BE15AC3FDF41F3B34E7, z: 0x00)"""
let newField = toFieldPair("test", 123'u) 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 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": test "ENR insert size too big":
let pk = PrivateKey.fromHex( let pk = PrivateKey.fromHex(
@ -138,7 +139,7 @@ suite "ENR":
check r.isOk() check r.isOk()
let newField = toFieldPair("test", 123'u) let newField = toFieldPair("test", 123'u)
let res = r[].insertFieldPair(pk, newField) let res = r[].insertFieldPairs(pk, [newField])
check res.isErr() check res.isErr()
test "ENR insert invalid key": test "ENR insert invalid key":
@ -151,5 +152,5 @@ suite "ENR":
let let
wrongPk = PrivateKey.random(rng[]) wrongPk = PrivateKey.random(rng[])
newField = toFieldPair("test", 123'u) newField = toFieldPair("test", 123'u)
res = r[].insertFieldPair(wrongPk, newField) res = r[].insertFieldPairs(wrongPk, [newField])
check res.isErr() check res.isErr()