From f82a9c7d3d7b7012f076ee33b7b89f40431fc39f Mon Sep 17 00:00:00 2001
From: sukun <sukunrt@gmail.com>
Date: Sat, 20 Jul 2024 17:45:04 +0530
Subject: [PATCH] add certhashes to addresses provided by AddrsFactory

---
 p2p/host/basic/basic_host.go | 100 ++++++++++++++++++++---------------
 1 file changed, 56 insertions(+), 44 deletions(-)

diff --git a/p2p/host/basic/basic_host.go b/p2p/host/basic/basic_host.go
index 6c89440394..ad75cf7d26 100644
--- a/p2p/host/basic/basic_host.go
+++ b/p2p/host/basic/basic_host.go
@@ -284,6 +284,15 @@ func NewHost(n network.Network, opts *HostOpts) (*BasicHost, error) {
 	if opts.AddrsFactory != nil {
 		h.AddrsFactory = opts.AddrsFactory
 	}
+	// This is a terrible hack to ensure that certhashes are added to addresses
+	// provided by AddrsFactory. This ensures that autonat receives addresses
+	// with the correct certhashes. We need to do it here rather and not in
+	// Addrs method as the autorelay package overrides resets the AddrsFactory
+	// so autonat cannot use Addrs directly. Wrapping it here allows us to provide
+	// the wrapped AddrsFactory to autonat before autorelay updates it.
+	h.AddrsFactory = func(addrs []ma.Multiaddr) []ma.Multiaddr {
+		return h.addCertHashes(h.AddrsFactory(addrs))
+	}
 
 	if opts.NATManager != nil {
 		h.natmgr = opts.NATManager(n)
@@ -800,47 +809,9 @@ func (h *BasicHost) ConnManager() connmgr.ConnManager {
 // Addrs returns listening addresses that are safe to announce to the network.
 // The output is the same as AllAddrs, but processed by AddrsFactory.
 func (h *BasicHost) Addrs() []ma.Multiaddr {
-	// This is a temporary workaround/hack that fixes #2233. Once we have a
-	// proper address pipeline, rework this. See the issue for more context.
-	type transportForListeninger interface {
-		TransportForListening(a ma.Multiaddr) transport.Transport
-	}
-
-	type addCertHasher interface {
-		AddCertHashes(m ma.Multiaddr) (ma.Multiaddr, bool)
-	}
-
-	addrs := h.AddrsFactory(h.AllAddrs())
-
-	s, ok := h.Network().(transportForListeninger)
-	if !ok {
-		return addrs
-	}
-
-	// Copy addrs slice since we'll be modifying it.
-	addrsOld := addrs
-	addrs = make([]ma.Multiaddr, len(addrsOld))
-	copy(addrs, addrsOld)
-
-	for i, addr := range addrs {
-		wtOK, wtN := libp2pwebtransport.IsWebtransportMultiaddr(addr)
-		webrtcOK, webrtcN := libp2pwebrtc.IsWebRTCDirectMultiaddr(addr)
-		if (wtOK && wtN == 0) || (webrtcOK && webrtcN == 0) {
-			t := s.TransportForListening(addr)
-			tpt, ok := t.(addCertHasher)
-			if !ok {
-				continue
-			}
-			addrWithCerthash, added := tpt.AddCertHashes(addr)
-			if !added {
-				log.Debugf("Couldn't add certhashes to multiaddr: %s", addr)
-				continue
-			}
-			addrs[i] = addrWithCerthash
-		}
-	}
-
-	return addrs
+	// We don't need to append certhashes here, the user provided addrsFactory was
+	// wrapped with addCertHashes in the constructor.
+	return h.AddrsFactory(h.AllAddrs())
 }
 
 // NormalizeMultiaddr returns a multiaddr suitable for equality checks.
@@ -860,8 +831,9 @@ func (h *BasicHost) NormalizeMultiaddr(addr ma.Multiaddr) ma.Multiaddr {
 	return addr
 }
 
-// AllAddrs returns all the addresses of BasicHost at this moment in time.
-// It's ok to not include addresses if they're not available to be used now.
+// AllAddrs returns all the addresses the host is listening on except circuit addresses.
+// The output has webtransport addresses inferred from quic addresses.
+// All the addresses have the correct
 func (h *BasicHost) AllAddrs() []ma.Multiaddr {
 	listenAddrs := h.Network().ListenAddresses()
 	if len(listenAddrs) == 0 {
@@ -955,10 +927,50 @@ func (h *BasicHost) AllAddrs() []ma.Multiaddr {
 	}
 	finalAddrs = ma.Unique(finalAddrs)
 	finalAddrs = inferWebtransportAddrsFromQuic(finalAddrs)
-
 	return finalAddrs
 }
 
+func (h *BasicHost) addCertHashes(addrs []ma.Multiaddr) []ma.Multiaddr {
+	// This is a temporary workaround/hack that fixes #2233. Once we have a
+	// proper address pipeline, rework this. See the issue for more context.
+	type transportForListeninger interface {
+		TransportForListening(a ma.Multiaddr) transport.Transport
+	}
+
+	type addCertHasher interface {
+		AddCertHashes(m ma.Multiaddr) (ma.Multiaddr, bool)
+	}
+
+	s, ok := h.Network().(transportForListeninger)
+	if !ok {
+		return addrs
+	}
+
+	// Copy addrs slice since we'll be modifying it.
+	addrsOld := addrs
+	addrs = make([]ma.Multiaddr, len(addrsOld))
+	copy(addrs, addrsOld)
+
+	for i, addr := range addrs {
+		wtOK, wtN := libp2pwebtransport.IsWebtransportMultiaddr(addr)
+		webrtcOK, webrtcN := libp2pwebrtc.IsWebRTCDirectMultiaddr(addr)
+		if (wtOK && wtN == 0) || (webrtcOK && webrtcN == 0) {
+			t := s.TransportForListening(addr)
+			tpt, ok := t.(addCertHasher)
+			if !ok {
+				continue
+			}
+			addrWithCerthash, added := tpt.AddCertHashes(addr)
+			if !added {
+				log.Debugf("Couldn't add certhashes to multiaddr: %s", addr)
+				continue
+			}
+			addrs[i] = addrWithCerthash
+		}
+	}
+	return addrs
+}
+
 var wtComponent = ma.StringCast("/webtransport")
 
 // inferWebtransportAddrsFromQuic infers more webtransport addresses from QUIC addresses.