From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 16 May 2023 06:35:03 +0000 (+1000)
Subject: Handle more PeerRemoteAddr variants when calculating dial addr
X-Git-Url: http://www.git.stargrave.org/?p=btrtrc.git;a=commitdiff_plain;h=3a92268f10184f5ed4602fa43d476d0ada96dc8e

Handle more PeerRemoteAddr variants when calculating dial addr

This should fix https://github.com/anacrolix/torrent/issues/820.
---

diff --git a/netip-addrport.go b/netip-addrport.go
index cf9edfd5..4e2dc94d 100644
--- a/netip-addrport.go
+++ b/netip-addrport.go
@@ -2,12 +2,15 @@ package torrent
 
 import (
 	"fmt"
-	"net"
 	"net/netip"
 
 	"github.com/anacrolix/dht/v2/krpc"
 )
 
+type addrPorter interface {
+	AddrPort() netip.AddrPort
+}
+
 func ipv4AddrPortFromKrpcNodeAddr(na krpc.NodeAddr) (_ netip.AddrPort, err error) {
 	ip4 := na.IP.To4()
 	if ip4 == nil {
@@ -32,9 +35,7 @@ func ipv6AddrPortFromKrpcNodeAddr(na krpc.NodeAddr) (_ netip.AddrPort, err error
 
 func addrPortFromPeerRemoteAddr(pra PeerRemoteAddr) (netip.AddrPort, error) {
 	switch v := pra.(type) {
-	case *net.TCPAddr:
-		return v.AddrPort(), nil
-	case *net.UDPAddr:
+	case addrPorter:
 		return v.AddrPort(), nil
 	case netip.AddrPort:
 		return v, nil
diff --git a/peerconn.go b/peerconn.go
index 1443a7b9..99157583 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -1056,21 +1056,20 @@ func (c *PeerConn) pexPeerFlags() pp.PexPeerFlags {
 // This returns the address to use if we want to dial the peer again. It incorporates the peer's
 // advertised listen port.
 func (c *PeerConn) dialAddr() PeerRemoteAddr {
-	if !c.outgoing && c.PeerListenPort != 0 {
-		switch addr := c.RemoteAddr.(type) {
-		case *net.TCPAddr:
-			dialAddr := *addr
-			dialAddr.Port = c.PeerListenPort
-			return &dialAddr
-		case *net.UDPAddr:
-			dialAddr := *addr
-			dialAddr.Port = c.PeerListenPort
-			return &dialAddr
-		default:
-			panic(addr)
-		}
+	if c.outgoing || c.PeerListenPort == 0 {
+		return c.RemoteAddr
 	}
-	return c.RemoteAddr
+	addrPort, err := addrPortFromPeerRemoteAddr(c.RemoteAddr)
+	if err != nil {
+		c.logger.Levelf(
+			log.Warning,
+			"error parsing %q for alternate dial port: %v",
+			c.RemoteAddr,
+			err,
+		)
+		return c.RemoteAddr
+	}
+	return netip.AddrPortFrom(addrPort.Addr(), uint16(c.PeerListenPort))
 }
 
 func (c *PeerConn) pexEvent(t pexEventType) pexEvent {
@@ -1095,12 +1094,6 @@ func (pc *PeerConn) remoteIsTransmission() bool {
 	return bytes.HasPrefix(pc.PeerID[:], []byte("-TR")) && pc.PeerID[7] == '-'
 }
 
-func (pc *PeerConn) remoteAddrPort() Option[netip.AddrPort] {
-	return Some(pc.conn.RemoteAddr().(interface {
-		AddrPort() netip.AddrPort
-	}).AddrPort())
-}
-
 func (pc *PeerConn) remoteDialAddrPort() (netip.AddrPort, error) {
 	dialAddr := pc.dialAddr()
 	return addrPortFromPeerRemoteAddr(dialAddr)