From 084f3dfa042c24c02d97ccaba6605b43ccb12493 Mon Sep 17 00:00:00 2001
From: Arnaud <arnaud@status.im>
Date: Thu, 14 May 2026 12:14:41 +0400
Subject: [PATCH] Refactoring to make udp port explicit

---
 storage/discovery.nim               | 10 ++++------
 storage/nat.nim                     | 16 +++++++++++-----
 storage/storage.nim                 |  2 +-
 tests/storage/helpers/nodeutils.nim |  2 +-
 4 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/storage/discovery.nim b/storage/discovery.nim
index 02278af3..70e8d1b2 100644
--- a/storage/discovery.nim
+++ b/storage/discovery.nim
@@ -198,14 +198,12 @@ proc updateSpr(d: Discovery) =
       d.protocol.updateRecord(spr).expect("Should update SPR")
 
 proc updateRecords*(
-    d: Discovery, announceAddrs: openArray[MultiAddress], discoveryPort: Port
+    d: Discovery, announceAddrs: openArray[MultiAddress], udpPort: Port
 ) =
-  ## Update both provider and DHT records from TCP announce addresses.
-  ## Discovery (UDP) addresses are derived by remapping announceAddrs to UDP with discoveryPort.
-  ## Updates the discv5 SPR once with the full set of addresses.
+  # UDP addresses are derived from TCP announce addresses by remapping protocol and port.
   let tcpAddrs = @announceAddrs
   let udpAddrs =
-    tcpAddrs.mapIt(it.remapAddr(protocol = some("udp"), port = some(discoveryPort)))
+    tcpAddrs.mapIt(it.remapAddr(protocol = some("udp"), port = some(udpPort)))
 
   debug "Updating addresses", tcpAddrs, udpAddrs
 
@@ -289,7 +287,7 @@ proc new*(
     key: key, peerId: PeerId.init(key).expect("Should construct PeerId"), store: store
   )
 
-  self.updateRecords(announceAddrs, discoveryPort)
+  self.updateRecords(announceAddrs, udpPort = discoveryPort)
 
   let discoveryConfig =
     DiscoveryConfig(tableIpLimits: tableIpLimits, bitsPerHop: DefaultBitsPerHop)
diff --git a/storage/nat.nim b/storage/nat.nim
index fa9d51b4..0cf06e9c 100644
--- a/storage/nat.nim
+++ b/storage/nat.nim
@@ -148,7 +148,7 @@ method handleNatStatus*(
       else:
         debug "AutoRelayService stopped"
 
-    discovery.updateRecords(@[dialBackAddr.get], discoveryPort)
+    discovery.updateRecords(@[dialBackAddr.get], udpPort = discoveryPort)
     discovery.protocol.clientMode = false
   of NotReachable:
     var hasPortMapping = false
@@ -160,6 +160,9 @@ method handleNatStatus*(
     else:
       debug "Node is not reachable trying UPnP / PMP now"
 
+      # Here we should check first that a mapping exists.
+      # If it does exist but Autonat still report as Not Reachable
+      # we should fallback to relay.
       let maybePorts = await m.mapNatPorts()
 
       if maybePorts.isSome:
@@ -169,16 +172,19 @@ method handleNatStatus*(
 
         let announceAddress = dialBackAddr.get.remapAddr(port = some(tcpPort))
 
-        # TODO: Try a dial me to make sure we are reachable
-
         if autoRelayService.isRunning:
+          # Here we stop the relay because the node *should* be reachable
           if not await autoRelayService.stop(switch):
             debug "AutoRelayService stop method returned false"
           else:
             debug "AutoRelayService stopped"
 
-        discovery.updateRecords(@[announceAddress], udpPort)
-        discovery.protocol.clientMode = false
+        # Note that we update the DHT records but we don't set the client mode
+        # to false because we are not sure the node is reachable.
+        # The client mode will be updated on the next iteration of autonat.
+        # Trying to check manually that the node is reachable is not trivial,
+        # this is exactly what Autonat does.
+        discovery.updateRecords(@[announceAddress], udpPort = udpPort)
         hasPortMapping = true
 
     if not hasPortMapping and not autoRelayService.isRunning:
diff --git a/storage/storage.nim b/storage/storage.nim
index d89c76a9..61e8edc1 100644
--- a/storage/storage.nim
+++ b/storage/storage.nim
@@ -106,7 +106,7 @@ proc start*(s: StorageServer) {.async.} =
     # It will be updated if reachable.
     s.storageNode.discovery.protocol.clientMode = true
 
-  s.storageNode.discovery.updateRecords(announceAddrs, s.config.discoveryPort)
+  s.storageNode.discovery.updateRecords(announceAddrs, udpPort = s.config.discoveryPort)
 
   await s.storageNode.start()
 
diff --git a/tests/storage/helpers/nodeutils.nim b/tests/storage/helpers/nodeutils.nim
index 95d7a8c6..fd9e658f 100644
--- a/tests/storage/helpers/nodeutils.nim
+++ b/tests/storage/helpers/nodeutils.nim
@@ -224,7 +224,7 @@ proc generateNodes*(
 
         if config.enableBootstrap:
           waitFor switch.peerInfo.update()
-          blockDiscovery.updateRecords(switch.peerInfo.addrs, bindPort.Port)
+          blockDiscovery.updateRecords(switch.peerInfo.addrs, udpPort = bindPort.Port)
           if blockDiscovery.getSpr().isSome:
             bootstrapNodes.add !blockDiscovery.getSpr()