From 2d93fa9e698ccbe4074676f32e7e432a6161ce05 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Mon, 12 Sep 2022 16:22:56 -0600
Subject: [PATCH] fix `updateRecord`

- support incrementing seqNo
- support updating with new record
---
 .../private/eth/p2p/discoveryv5/protocol.nim  | 21 +++++++--
 libp2pdht/private/eth/p2p/discoveryv5/spr.nim | 43 ++++++++++---------
 2 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/libp2pdht/private/eth/p2p/discoveryv5/protocol.nim b/libp2pdht/private/eth/p2p/discoveryv5/protocol.nim
index 58fdd74..c26dae6 100644
--- a/libp2pdht/private/eth/p2p/discoveryv5/protocol.nim
+++ b/libp2pdht/private/eth/p2p/discoveryv5/protocol.nim
@@ -222,16 +222,29 @@ func getRecord*(d: Protocol): SignedPeerRecord =
   ## Get the SPR of the local node.
   d.localNode.record
 
-proc updateRecord*(d: Protocol, newSpr: SignedPeerRecord): DiscResult[void] =
+proc updateRecord*(
+  d: Protocol,
+  spr: Option[SignedPeerRecord] = SignedPeerRecord.none): DiscResult[void] =
   ## Update the ENR of the local node with provided `enrFields` k:v pairs.
+  ##
 
-  info "Updated discovery SPR", uri=toURI(newSpr)
-  d.localNode.record = newSpr
-  ok()
+  if spr.isSome:
+    let
+      newSpr = spr.get()
+      seqNo = d.localNode.record.seqNum
+
+    info "Updated discovery SPR", uri = toURI(newSpr)
+
+    d.localNode.record = newSpr
+    d.localNode.record.data.seqNo = seqNo
+
+  ? d.localNode.record.incSeqNo(d.privateKey)
 
   # TODO: Would it make sense to actively ping ("broadcast") to all the peers
   # we stored a handshake with in order to get that ENR updated?
 
+  ok()
+
 proc sendResponse(d: Protocol, dstId: NodeId, dstAddr: Address,
     message: SomeMessage, reqId: RequestId) =
   ## send Response using the specifid reqId
diff --git a/libp2pdht/private/eth/p2p/discoveryv5/spr.nim b/libp2pdht/private/eth/p2p/discoveryv5/spr.nim
index b55bc19..1e90e22 100644
--- a/libp2pdht/private/eth/p2p/discoveryv5/spr.nim
+++ b/libp2pdht/private/eth/p2p/discoveryv5/spr.nim
@@ -31,7 +31,6 @@ proc seqNum*(r: SignedPeerRecord): uint64 =
     r.data.seqNo
 
 proc fromBytes(r: var SignedPeerRecord, s: openArray[byte]): bool =
-
   let decoded = SignedPeerRecord.decode(@s)
   if decoded.isErr:
     error "Error decoding SignedPeerRecord", error = decoded.error
@@ -47,22 +46,21 @@ proc get*(r: SignedPeerRecord, T: type PublicKey): Option[T] =
   r.envelope.publicKey.some
 
 proc incSeqNo*(
-    r: var SignedPeerRecord,
-    pk: PrivateKey): RecordResult[void] =
+  r: var SignedPeerRecord,
+  pk: PrivateKey): RecordResult[void] =
 
   r.data.seqNo.inc()
   r = ? SignedPeerRecord.init(pk, r.data).mapErr(
         (e: CryptoError) =>
-          ("Error initialising SignedPeerRecord with incremented seqNo: " &
-          $e).cstring
-      )
+          ("Error initializing SignedPeerRecord with incremented seqNo: " & $e).cstring)
   ok()
 
-
-proc update*(r: var SignedPeerRecord, pk: crypto.PrivateKey,
-                            ip: Option[ValidIpAddress],
-                            tcpPort, udpPort: Option[Port] = none[Port]()):
-                            RecordResult[void] =
+proc update*(
+  r: var SignedPeerRecord,
+  pk: crypto.PrivateKey,
+  ip: Option[ValidIpAddress],
+  tcpPort, udpPort: Option[Port] = none[Port]()):
+  RecordResult[void] =
   ## Update a `SignedPeerRecord` with given ip address, tcp port, udp port and optional
   ## custom k:v pairs.
   ##
@@ -110,7 +108,6 @@ proc update*(r: var SignedPeerRecord, pk: crypto.PrivateKey,
       transProtoPort = udpPort.get
 
     updated = MultiAddress.init(ipAddr, transProto, transProtoPort)
-
   else:
     let
       existing = r.data.addresses[0].address
@@ -158,8 +155,7 @@ proc update*(r: var SignedPeerRecord, pk: crypto.PrivateKey,
 
   r = ? SignedPeerRecord.init(pk, r.data)
           .mapErr((e: CryptoError) =>
-            ("Failed to update SignedPeerRecord: " & $e).cstring
-          )
+            ("Failed to update SignedPeerRecord: " & $e).cstring)
 
   return ok()
 
@@ -223,11 +219,13 @@ proc toBase64*(r: SignedPeerRecord): string =
 
 proc toURI*(r: SignedPeerRecord): string = "spr:" & r.toBase64
 
-proc init*(T: type SignedPeerRecord, seqNum: uint64,
-                           pk: PrivateKey,
-                           ip: Option[ValidIpAddress],
-                           tcpPort, udpPort: Option[Port]):
-                           RecordResult[T] =
+proc init*(
+  T: type SignedPeerRecord,
+  seqNum: uint64,
+  pk: PrivateKey,
+  ip: Option[ValidIpAddress],
+  tcpPort, udpPort: Option[Port]):
+  RecordResult[T] =
   ## Initialize a `SignedPeerRecord` with given sequence number, private key, optional
   ## ip address, tcp port, udp port, and optional custom k:v pairs.
   ##
@@ -268,7 +266,9 @@ proc init*(T: type SignedPeerRecord, seqNum: uint64,
   let ma = MultiAddress.init(ipAddr, proto, protoPort)
 
   let pr = PeerRecord.init(peerId, @[ma], seqNum)
-  SignedPeerRecord.init(pk, pr).mapErr((e: CryptoError) => ("Failed to init SignedPeerRecord: " & $e).cstring)
+  SignedPeerRecord.init(pk, pr)
+    .mapErr(
+      (e: CryptoError) => ("Failed to init SignedPeerRecord: " & $e).cstring)
 
 proc contains*(r: SignedPeerRecord, fp: (string, seq[byte])): bool =
   # TODO: use FieldPair for this, but that is a bit cumbersome. Perhaps the
@@ -281,4 +281,5 @@ proc contains*(r: SignedPeerRecord, fp: (string, seq[byte])): bool =
   debugEcho "`contains` is not yet implemented for SignedPeerRecords"
   return false
 
-proc `==`*(a, b: SignedPeerRecord): bool = a.data == b.data
+proc `==`*(a, b: SignedPeerRecord): bool =
+  a.data == b.data