Add an AddressMapper to the AutoRelayService

vacp2p · Mar 13, 2023 · 73a46a1 · 73a46a1
1 parent a378b2d
commit 73a46a1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/libp2p/services/autorelayservice.nim b/libp2p/services/autorelayservice.nim
@@ -34,7 +34,12 @@ type
     onReservation: OnReservationHandler
     rng: ref HmacDrbgContext
 
-proc reserveAndUpdate(self: AutoRelayService, relayPid: PeerId, selfPid: PeerId) {.async.} =
+proc addressMapper(
+  self: AutoRelayService,
+  listenAddrs: seq[MultiAddress]): Future[seq[MultiAddress]] {.gcsafe, async.} =
+  return concat(toSeq(self.relayAddresses.values))
+
+proc reserveAndUpdate(self: AutoRelayService, relayPid: PeerId, switch: Switch) {.async.} =
   while self.running:
     let
       rsvp = await self.client.reserve(relayPid).wait(chronos.seconds(5))
@@ -46,11 +51,15 @@ proc reserveAndUpdate(self: AutoRelayService, relayPid: PeerId, selfPid: PeerId)
       break
     if relayPid notin self.relayAddresses or self.relayAddresses[relayPid] != relayedAddr:
       self.relayAddresses[relayPid] = relayedAddr
+      await switch.peerInfo.update()
       if not self.onReservation.isNil():
         self.onReservation(concat(toSeq(self.relayAddresses.values)))
     await sleepAsync chronos.seconds(ttl - 30)
 
 method setup*(self: AutoRelayService, switch: Switch): Future[bool] {.async, gcsafe.} =
+  proc addressMapper(listenAddrs: seq[MultiAddress]): Future[seq[MultiAddress]] {.gcsafe, async.} =
+    return await self.addressMapper(listenAddrs)
+
   let hasBeenSetUp = await procCall Service(self).setup(switch)
   if hasBeenSetUp:
     proc handlePeerJoined(peerId: PeerId, event: PeerEvent) {.async.} =
@@ -63,6 +72,7 @@ method setup*(self: AutoRelayService, switch: Switch): Future[bool] {.async, gcs
         future[].cancel()
     switch.addPeerEventHandler(handlePeerJoined, Joined)
     switch.addPeerEventHandler(handlePeerLeft, Left)
+    switch.peerInfo.addressMappers.add(addressMapper)
     await self.run(switch)
   return hasBeenSetUp
 
@@ -96,7 +106,7 @@ proc innerRun(self: AutoRelayService, switch: Switch) {.async, gcsafe.} =
     for relayPid in connectedPeers:
       if self.relayPeers.len() >= self.numRelays:
         break
-      self.relayPeers[relayPid] = self.reserveAndUpdate(relayPid, switch.peerInfo.peerId)
+      self.relayPeers[relayPid] = self.reserveAndUpdate(relayPid, switch)
 
     if self.relayPeers.len() > 0:
       await one(toSeq(self.relayPeers.values())) or self.peerAvailable.wait()

diff --git a/tests/testautorelay.nim b/tests/testautorelay.nim
@@ -72,8 +72,9 @@ suite "Autorelay":
     await fut.wait(1.seconds)
     let addresses = autorelay.getAddresses()
     check:
-      addresses[0] == buildRelayMA(switchRelay, switchClient)
+      addresses == @[buildRelayMA(switchRelay, switchClient)]
       addresses.len() == 1
+      addresses == switchClient.peerInfo.addrs
     await allFutures(switchClient.stop(), switchRelay.stop())
 
   asyncTest "Three relays connections":