chore: rendezvous improvements (vacp2p#1319)

Author: gmelodie

libp2p/discovery/rendezvousinterface.nim

@@ -26,14 +26,14 @@ proc `==`*(a, b: RdvNamespace): bool {.borrow.}
 method request*(
     self: RendezVousInterface, pa: PeerAttributes
 ) {.async: (raises: [DiscoveryError, CancelledError]).} =
-  var namespace = ""
+  var namespace = Opt.none(string)
   for attr in pa:
     if attr.ofType(RdvNamespace):
-      namespace = string attr.to(RdvNamespace)
+      namespace = Opt.some(string attr.to(RdvNamespace))
     elif attr.ofType(DiscoveryService):
-      namespace = string attr.to(DiscoveryService)
+      namespace = Opt.some(string attr.to(DiscoveryService))
     elif attr.ofType(PeerId):
-      namespace = $attr.to(PeerId)
+      namespace = Opt.some($attr.to(PeerId))
     else:
       # unhandled type
       return
@@ -44,8 +44,8 @@ method request*(
       for address in pr.addresses:
         peer.add(address.address)
 
-      peer.add(DiscoveryService(namespace))
-      peer.add(RdvNamespace(namespace))
+      peer.add(DiscoveryService(namespace.get()))
+      peer.add(RdvNamespace(namespace.get()))
       self.onPeerFound(peer)
 
     await sleepAsync(self.timeToRequest)

libp2p/protocols/rendezvous.nim

@@ -37,6 +37,7 @@ const
   RendezVousCodec* = "/rendezvous/1.0.0"
   MinimumDuration* = 2.hours
   MaximumDuration = 72.hours
+  MaximumMessageLen = 1 shl 22 # 4MB
   MinimumNamespaceLen = 1
   MaximumNamespaceLen = 255
   RegistrationLimitPerPeer = 1000
@@ -63,7 +64,7 @@ type
 
   Cookie = object
     offset: uint64
-    ns: string
+    ns: Opt[string]
 
   Register = object
     ns: string
@@ -79,7 +80,7 @@ type
     ns: string
 
   Discover = object
-    ns: string
+    ns: Opt[string]
     limit: Opt[uint64]
     cookie: Opt[seq[byte]]
 
@@ -100,7 +101,8 @@ type
 proc encode(c: Cookie): ProtoBuffer =
   result = initProtoBuffer()
   result.write(1, c.offset)
-  result.write(2, c.ns)
+  if c.ns.isSome():
+    result.write(2, c.ns.get())
   result.finish()
 
 proc encode(r: Register): ProtoBuffer =
@@ -127,7 +129,8 @@ proc encode(u: Unregister): ProtoBuffer =
 
 proc encode(d: Discover): ProtoBuffer =
   result = initProtoBuffer()
-  result.write(1, d.ns)
+  if d.ns.isSome():
+    result.write(1, d.ns.get())
   d.limit.withValue(limit):
     result.write(2, limit)
   d.cookie.withValue(cookie):
@@ -161,13 +164,17 @@ proc encode(msg: Message): ProtoBuffer =
   result.finish()
 
 proc decode(_: typedesc[Cookie], buf: seq[byte]): Opt[Cookie] =
-  var c: Cookie
+  var
+    c: Cookie
+    ns: string
   let
     pb = initProtoBuffer(buf)
     r1 = pb.getRequiredField(1, c.offset)
-    r2 = pb.getRequiredField(2, c.ns)
+    r2 = pb.getField(2, ns)
   if r1.isErr() or r2.isErr():
     return Opt.none(Cookie)
+  if r2.get(false):
+    c.ns = Opt.some(ns)
   Opt.some(c)
 
 proc decode(_: typedesc[Register], buf: seq[byte]): Opt[Register] =
@@ -219,13 +226,16 @@ proc decode(_: typedesc[Discover], buf: seq[byte]): Opt[Discover] =
     d: Discover
     limit: uint64
     cookie: seq[byte]
+    ns: string
   let
     pb = initProtoBuffer(buf)
-    r1 = pb.getRequiredField(1, d.ns)
+    r1 = pb.getField(1, ns)
     r2 = pb.getField(2, limit)
     r3 = pb.getField(3, cookie)
   if r1.isErr() or r2.isErr() or r3.isErr:
     return Opt.none(Discover)
+  if r1.get(false):
+    d.ns = Opt.some(ns)
   if r2.get(false):
     d.limit = Opt.some(limit)
   if r3.get(false):
@@ -446,7 +456,7 @@ proc discover(
 ) {.async: (raises: [CancelledError, LPStreamError]).} =
   trace "Received Discover", peerId = conn.peerId, ns = d.ns
   libp2p_rendezvous_discover.inc()
-  if d.ns.len > MaximumNamespaceLen:
+  if d.ns.isSome() and d.ns.get().len > MaximumNamespaceLen:
     await conn.sendDiscoverResponseError(InvalidNamespace)
     return
   var limit = min(DiscoverLimit, d.limit.get(DiscoverLimit))
@@ -459,20 +469,19 @@ proc discover(
         return
     else:
       Cookie(offset: rdv.registered.low().uint64 - 1)
-  if cookie.ns != d.ns or cookie.offset < rdv.registered.low().uint64 or
+  if d.ns.isSome() and cookie.ns.isSome() and cookie.ns.get() != d.ns.get() or
+      cookie.offset < rdv.registered.low().uint64 or
       cookie.offset > rdv.registered.high().uint64:
     cookie = Cookie(offset: rdv.registered.low().uint64 - 1)
-  let
-    nsSalted = d.ns & rdv.salt
-    namespaces =
-      if d.ns != "":
-        try:
-          rdv.namespaces[nsSalted]
-        except KeyError:
-          await conn.sendDiscoverResponseError(InvalidNamespace)
-          return
-      else:
-        toSeq(cookie.offset.int .. rdv.registered.high())
+  let namespaces =
+    if d.ns.isSome():
+      try:
+        rdv.namespaces[d.ns.get() & rdv.salt]
+      except KeyError:
+        await conn.sendDiscoverResponseError(InvalidNamespace)
+        return
+    else:
+      toSeq(max(cookie.offset.int, rdv.registered.offset) .. rdv.registered.high())
   if namespaces.len() == 0:
     await conn.sendDiscoverResponse(@[], Cookie())
     return
@@ -516,7 +525,7 @@ proc advertisePeer(
       rdv.sema.release()
 
   await rdv.sema.acquire()
-  discard await advertiseWrap().withTimeout(5.seconds)
+  await advertiseWrap()
 
 proc advertise*(
     rdv: RendezVous, ns: string, ttl: Duration, peers: seq[PeerId]
@@ -539,7 +548,7 @@ proc advertise*(
   let futs = collect(newSeq()):
     for peer in peers:
       trace "Send Advertise", peerId = peer, ns
-      rdv.advertisePeer(peer, msg.buffer)
+      rdv.advertisePeer(peer, msg.buffer).withTimeout(5.seconds)
 
   await allFutures(futs)
 
@@ -563,7 +572,7 @@ proc requestLocally*(rdv: RendezVous, ns: string): seq[PeerRecord] =
     @[]
 
 proc request*(
-    rdv: RendezVous, ns: string, l: int = DiscoverLimit.int, peers: seq[PeerId]
+    rdv: RendezVous, ns: Opt[string], l: int = DiscoverLimit.int, peers: seq[PeerId]
 ): Future[seq[PeerRecord]] {.async: (raises: [DiscoveryError, CancelledError]).} =
   var
     s: Table[PeerId, (PeerRecord, Register)]
@@ -572,7 +581,7 @@ proc request*(
 
   if l <= 0 or l > DiscoverLimit.int:
     raise newException(AdvertiseError, "Invalid limit")
-  if ns.len > MaximumNamespaceLen:
+  if ns.isSome() and ns.get().len > MaximumNamespaceLen:
     raise newException(AdvertiseError, "Invalid namespace")
 
   limit = l.uint64
@@ -584,15 +593,18 @@ proc request*(
       await conn.close()
     d.limit = Opt.some(limit)
     d.cookie =
-      try:
-        Opt.some(rdv.cookiesSaved[peer][ns])
-      except KeyError as exc:
+      if ns.isSome():
+        try:
+          Opt.some(rdv.cookiesSaved[peer][ns.get()])
+        except KeyError, CatchableError:
+          Opt.none(seq[byte])
+      else:
         Opt.none(seq[byte])
     await conn.writeLp(
       encode(Message(msgType: MessageType.Discover, discover: Opt.some(d))).buffer
     )
     let
-      buf = await conn.readLp(65536)
+      buf = await conn.readLp(MaximumMessageLen)
       msgRcv = Message.decode(buf).valueOr:
         debug "Message undecodable"
         return
@@ -606,12 +618,14 @@ proc request*(
       trace "Cannot discover", ns, status = resp.status, text = resp.text
       return
     resp.cookie.withValue(cookie):
-      if cookie.len() < 1000 and
-          rdv.cookiesSaved.hasKeyOrPut(peer, {ns: cookie}.toTable()):
-        try:
-          rdv.cookiesSaved[peer][ns] = cookie
-        except KeyError:
-          raiseAssert "checked with hasKeyOrPut"
+      if ns.isSome:
+        let namespace = ns.get()
+        if cookie.len() < 1000 and
+            rdv.cookiesSaved.hasKeyOrPut(peer, {namespace: cookie}.toTable()):
+          try:
+            rdv.cookiesSaved[peer][namespace] = cookie
+          except KeyError:
+            raiseAssert "checked with hasKeyOrPut"
     for r in resp.registrations:
       if limit == 0:
         return
@@ -634,8 +648,9 @@ proc request*(
       else:
         s[pr.peerId] = (pr, r)
       limit.dec()
-    for (_, r) in s.values():
-      rdv.save(ns, peer, r, false)
+    if ns.isSome():
+      for (_, r) in s.values():
+        rdv.save(ns.get(), peer, r, false)
 
   for peer in peers:
     if limit == 0:
@@ -654,10 +669,15 @@ proc request*(
   return toSeq(s.values()).mapIt(it[0])
 
 proc request*(
-    rdv: RendezVous, ns: string, l: int = DiscoverLimit.int
+    rdv: RendezVous, ns: Opt[string], l: int = DiscoverLimit.int
 ): Future[seq[PeerRecord]] {.async: (raises: [DiscoveryError, CancelledError]).} =
   await rdv.request(ns, l, rdv.peers)
 
+proc request*(
+    rdv: RendezVous, l: int = DiscoverLimit.int
+): Future[seq[PeerRecord]] {.async: (raises: [DiscoveryError, CancelledError]).} =
+  await rdv.request(Opt.none(string), l, rdv.peers)
+
 proc unsubscribeLocally*(rdv: RendezVous, ns: string) =
   let nsSalted = ns & rdv.salt
   try:
@@ -690,7 +710,7 @@ proc unsubscribe*(
     for peer in peerIds:
       unsubscribePeer(peer)
 
-  discard await allFutures(futs).withTimeout(5.seconds)
+  await allFutures(futs)
 
 proc unsubscribe*(
     rdv: RendezVous, ns: string
@@ -786,8 +806,10 @@ proc new*(
   rdv.setup(switch)
   return rdv
 
-proc deletesRegister(rdv: RendezVous) {.async: (raises: [CancelledError]).} =
-  heartbeat "Register timeout", 1.minutes:
+proc deletesRegister(
+    rdv: RendezVous, interval = 1.minutes
+) {.async: (raises: [CancelledError]).} =
+  heartbeat "Register timeout", interval:
     let n = Moment.now()
     var total = 0
     rdv.registered.flushIfIt(it.expiration < n)

tests/testrendezvous.nim

@@ -58,18 +58,22 @@ suite "RendezVous":
     await client.start()
     await remoteSwitch.start()
     await client.connect(remoteSwitch.peerInfo.peerId, remoteSwitch.peerInfo.addrs)
-    let res0 = await rdv.request("empty")
+    let res0 = await rdv.request(Opt.some("empty"))
     check res0.len == 0
+
     await rdv.advertise("foo")
-    let res1 = await rdv.request("foo")
+    let res1 = await rdv.request(Opt.some("foo"))
     check:
       res1.len == 1
       res1[0] == client.peerInfo.signedPeerRecord.data
-    let res2 = await rdv.request("bar")
+
+    let res2 = await rdv.request(Opt.some("bar"))
     check res2.len == 0
+
     await rdv.unsubscribe("foo")
-    let res3 = await rdv.request("foo")
+    let res3 = await rdv.request(Opt.some("foo"))
     check res3.len == 0
+
     await allFutures(client.stop(), remoteSwitch.stop())
 
   asyncTest "Harder remote test":
@@ -88,17 +92,21 @@ suite "RendezVous":
     )
     await allFutures(rdvSeq.mapIt(it.advertise("foo")))
     var data = clientSeq.mapIt(it.peerInfo.signedPeerRecord.data)
-    let res1 = await rdvSeq[0].request("foo", 5)
+    let res1 = await rdvSeq[0].request(Opt.some("foo"), 5)
     check res1.len == 5
     for d in res1:
       check d in data
     data.keepItIf(it notin res1)
-    let res2 = await rdvSeq[0].request("foo")
+    let res2 = await rdvSeq[0].request(Opt.some("foo"))
     check res2.len == 5
     for d in res2:
       check d in data
-    let res3 = await rdvSeq[0].request("foo")
+    let res3 = await rdvSeq[0].request(Opt.some("foo"))
     check res3.len == 0
+    let res4 = await rdvSeq[0].request()
+    check res4.len == 11
+    let res5 = await rdvSeq[0].request(Opt.none(string))
+    check res5.len == 11
     await remoteSwitch.stop()
     await allFutures(clientSeq.mapIt(it.stop()))
 
@@ -116,9 +124,9 @@ suite "RendezVous":
     await clientA.connect(remoteSwitch.peerInfo.peerId, remoteSwitch.peerInfo.addrs)
     await clientB.connect(remoteSwitch.peerInfo.peerId, remoteSwitch.peerInfo.addrs)
     await rdvA.advertise("foo")
-    let res1 = await rdvA.request("foo")
+    let res1 = await rdvA.request(Opt.some("foo"))
     await rdvB.advertise("foo")
-    let res2 = await rdvA.request("foo")
+    let res2 = await rdvA.request(Opt.some("foo"))
     check:
       res2.len == 1
       res2[0] == clientB.peerInfo.signedPeerRecord.data
@@ -129,11 +137,11 @@ suite "RendezVous":
       rdv = RendezVous.new(minDuration = 1.minutes, maxDuration = 72.hours)
       switch = createSwitch(rdv)
     expect AdvertiseError:
-      discard await rdv.request("A".repeat(300))
+      discard await rdv.request(Opt.some("A".repeat(300)))
     expect AdvertiseError:
-      discard await rdv.request("A", -1)
+      discard await rdv.request(Opt.some("A"), -1)
     expect AdvertiseError:
-      discard await rdv.request("A", 3000)
+      discard await rdv.request(Opt.some("A"), 3000)
     expect AdvertiseError:
       await rdv.advertise("A".repeat(300))
     expect AdvertiseError: