From 594d47fe23a03ae316c1954165c8168d72ddac95 Mon Sep 17 00:00:00 2001
From: Arnaud <arnaud@status.im>
Date: Tue, 16 Jun 2026 10:11:22 +0400
Subject: [PATCH] Add more tests

---
 storage/nat.nim                   | 16 +++++++++++
 storage/storage.nim               | 13 +--------
 tests/storage/testnatreaction.nim | 46 +++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/storage/nat.nim b/storage/nat.nim
index b45cba0d..2d984524 100644
--- a/storage/nat.nim
+++ b/storage/nat.nim
@@ -274,6 +274,22 @@ proc findReachableNodes*(bootstrapNodes: seq[SignedPeerRecord]): seq[SignedPeerR
   ## confirmed reachable by AutoNAT could be included.
   bootstrapNodes
 
+proc announceRelayReservation*(
+    discovery: Discovery, addresses: seq[MultiAddress]
+) {.gcsafe.} =
+  ## Announce the publicly dialable circuit addresses from a relay reservation.
+  ## A reservation response can also carry loopback/private addresses, which a
+  ## remote peer can never dial, so they are dropped. If none are public, the
+  ## previous announce is kept untouched.
+  let publicAddrs = addresses.filterIt(it.hasPublicRelayTransport())
+  if publicAddrs.len == 0:
+    warn "Relay reservation has no publicly dialable address, keeping previous announce",
+      addresses
+    return
+  info "Relay reservation updated", addresses = publicAddrs
+  # relay addresses are for download traffic only, not DHT routing
+  discovery.announceRelayAddrs(publicAddrs)
+
 # Hole punching logic below is adapted from libp2p's HPService
 # (libp2p/services/hpservice.nim). HPService cannot be used directly because it
 # depends on AutoNAT v1 and starts the relay immediately on NotReachable,
diff --git a/storage/storage.nim b/storage/storage.nim
index a57705ea..0dd6bd2f 100644
--- a/storage/storage.nim
+++ b/storage/storage.nim
@@ -453,18 +453,7 @@ proc new*(
       maxNumRelays = config.natMaxRelays,
       client = relayClient,
       onReservation = proc(addresses: seq[MultiAddress]) {.gcsafe, raises: [].} =
-        # A relay server is required to have a public extip, so its
-        # circuit addresses always include a public one. The relay's reservation
-        # response can also carry loopback/private addresses:
-        # they are never dialable by a remote peer, so drop them.
-        let publicAddrs = addresses.filterIt(it.hasPublicRelayTransport())
-        if publicAddrs.len == 0:
-          warn "Relay reservation has no publicly dialable address, keeping previous announce",
-            addresses
-          return
-        info "Relay reservation updated", addresses = publicAddrs
-        # relay addresses are for download traffic only, not DHT routing
-        discovery.announceRelayAddrs(publicAddrs),
+        discovery.announceRelayReservation(addresses),
       rng = random.Rng.instance().libp2pRng,
     )
 
diff --git a/tests/storage/testnatreaction.nim b/tests/storage/testnatreaction.nim
index 728bbcbb..1a6d1e31 100644
--- a/tests/storage/testnatreaction.nim
+++ b/tests/storage/testnatreaction.nim
@@ -59,6 +59,13 @@ proc mappingOk(id: cint, port: uint16): Result[MappingResult, string] =
     )
   )
 
+const relayId = "16Uiu2HAmQu456Ae52JqPuqog6wCex47LLvNY8oHMBC4GRRtaStHs"
+
+proc circuitAddr(relayIp: string): MultiAddress =
+  MultiAddress
+    .init("/ip4/" & relayIp & "/tcp/8070/p2p/" & relayId & "/p2p-circuit")
+    .expect("valid")
+
 asyncchecksuite "NAT reaction - port mapping":
   var sw: Switch
   var key: PrivateKey
@@ -261,6 +268,45 @@ asyncchecksuite "NAT reaction - address announcing":
     # Ensure that nothing is injected because there is no active mapping
     check sw.peerInfo.addrs == sw.peerInfo.listenAddrs
 
+  test "handleNatStatus clears the DHT routing addresses when it becomes NotReachable":
+    let dialBack = MultiAddress.init("/ip4/1.2.3.4/tcp/9000").expect("valid")
+    let mapper = MockNatPortMapper(mappedPorts: none((Port, Port, MappingProtocol)))
+
+    autorelayservice.setup(autoRelay, sw)
+
+    # Reachable: the node announces direct addresses, including UDP for the DHT.
+    await mapper.handleNatStatus(
+      Reachable, Opt.some(dialBack), discoveryPort, disc, sw, autoRelay
+    )
+    check disc.dhtAddrs.len > 0
+
+    # NotReachable: the DHT routing addresses are cleared
+    await mapper.handleNatStatus(
+      NotReachable, Opt.some(dialBack), discoveryPort, disc, sw, autoRelay
+    )
+    check disc.dhtAddrs.len == 0
+
+  test "mapped-addr mapper does not inject a non-public mapped address":
+    # Active mapping, but no public observed address: the candidate stays private
+    # and must not be injected.
+    setupMappedAddrMapper(sw, NatPortMapper(activeTcpPort: some(Port(40000))))
+
+    await sw.peerInfo.update()
+
+    check sw.peerInfo.addrs == sw.peerInfo.listenAddrs
+
+  test "announceRelayReservation announces only the publicly dialable circuit address":
+    disc.announceRelayReservation(
+      @[circuitAddr("127.0.0.1"), circuitAddr("204.168.234.45")]
+    )
+
+    check disc.announceAddrs == @[circuitAddr("204.168.234.45")]
+
+  test "announceRelayReservation does not announce a private circuit address":
+    disc.announceRelayReservation(@[circuitAddr("127.0.0.1")])
+
+    check disc.announceAddrs.len == 0
+
 proc mapperWith(protocol: MappingProtocol): Option[NatPortMapper] =
   some(NatPortMapper(activeMappingProtocol: some(protocol)))