From 4b3bc4f8194731781d93a0fc870d86a85add2896 Mon Sep 17 00:00:00 2001
From: diegomrsantos <diego@status.im>
Date: Mon, 20 Nov 2023 11:06:02 +0100
Subject: [PATCH] Make ObservedAddrManager injectable (#970)

---
 libp2p/builders.nim           | 14 ++++++++++++--
 libp2p/protocols/identify.nim |  5 +++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/libp2p/builders.nim b/libp2p/builders.nim
index 71320826a..63a216d46 100644
--- a/libp2p/builders.nim
+++ b/libp2p/builders.nim
@@ -25,7 +25,7 @@ import
   muxers/[muxer, mplex/mplex, yamux/yamux],
   protocols/[identify, secure/secure, secure/noise, rendezvous],
   protocols/connectivity/[autonat/server, relay/relay, relay/client, relay/rtransport],
-  connmanager, upgrademngrs/muxedupgrade,
+  connmanager, upgrademngrs/muxedupgrade, observedaddrmanager,
   nameresolving/nameresolver,
   errors, utility
 
@@ -59,6 +59,7 @@ type
     circuitRelay: Relay
     rdv: RendezVous
     services: seq[Service]
+    observedAddrManager: ObservedAddrManager
 
 proc new*(T: type[SwitchBuilder]): T {.public.} =
   ## Creates a SwitchBuilder
@@ -201,6 +202,10 @@ proc withServices*(b: SwitchBuilder, services: seq[Service]): SwitchBuilder =
   b.services = services
   b
 
+proc withObservedAddrManager*(b: SwitchBuilder, observedAddrManager: ObservedAddrManager): SwitchBuilder =
+  b.observedAddrManager = observedAddrManager
+  b
+
 proc build*(b: SwitchBuilder): Switch
   {.raises: [LPError], public.} =
 
@@ -223,8 +228,13 @@ proc build*(b: SwitchBuilder): Switch
       protoVersion = b.protoVersion,
       agentVersion = b.agentVersion)
 
+  let identify =
+    if b.observedAddrManager != nil:
+      Identify.new(peerInfo, b.sendSignedPeerRecord, b.observedAddrManager)
+    else:
+      Identify.new(peerInfo, b.sendSignedPeerRecord)
+
   let
-    identify = Identify.new(peerInfo, b.sendSignedPeerRecord)
     connManager = ConnManager.new(b.maxConnsPerPeer, b.maxConnections, b.maxIn, b.maxOut)
     ms = MultistreamSelect.new()
     muxedUpgrade = MuxedUpgrade.new(b.muxers, secureManagerInstances, ms)
diff --git a/libp2p/protocols/identify.nim b/libp2p/protocols/identify.nim
index d677e9729..2432a10e0 100644
--- a/libp2p/protocols/identify.nim
+++ b/libp2p/protocols/identify.nim
@@ -139,12 +139,13 @@ proc decodeMsg*(buf: seq[byte]): Opt[IdentifyInfo] =
 proc new*(
   T: typedesc[Identify],
   peerInfo: PeerInfo,
-  sendSignedPeerRecord = false
+  sendSignedPeerRecord = false,
+  observedAddrManager = ObservedAddrManager.new(),
   ): T =
   let identify = T(
     peerInfo: peerInfo,
     sendSignedPeerRecord: sendSignedPeerRecord,
-    observedAddrManager: ObservedAddrManager.new(),
+    observedAddrManager: observedAddrManager,
   )
   identify.init()
   identify