From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 28 May 2023 03:34:58 +0000 (+1000)
Subject: Attribute accepted connection to holepunching when connect message is late
X-Git-Url: http://www.git.stargrave.org/?p=btrtrc.git;a=commitdiff_plain;h=5efb4dd9410e28bb2d6320268af4f98366be6508

Attribute accepted connection to holepunching when connect message is late

Also perform holepunch metric adjustments sooner to reduce timing issues in tests.
---

diff --git a/client-stats.go b/client-stats.go
index 1a800970..bfa6994e 100644
--- a/client-stats.go
+++ b/client-stats.go
@@ -2,14 +2,21 @@ package torrent
 
 import (
 	"net/netip"
+
+	g "github.com/anacrolix/generics"
 )
 
+func setAdd[K comparable](m *map[K]struct{}, elem K) {
+	g.MakeMapIfNilAndSet(m, elem, struct{}{})
+}
+
 type clientHolepunchAddrSets struct {
 	undialableWithoutHolepunch                            map[netip.AddrPort]struct{}
 	undialableWithoutHolepunchDialedAfterHolepunchConnect map[netip.AddrPort]struct{}
 	dialableOnlyAfterHolepunch                            map[netip.AddrPort]struct{}
 	dialedSuccessfullyAfterHolepunchConnect               map[netip.AddrPort]struct{}
 	probablyOnlyConnectedDueToHolepunch                   map[netip.AddrPort]struct{}
+	accepted                                              map[netip.AddrPort]struct{}
 }
 
 type ClientStats struct {
diff --git a/client.go b/client.go
index a87e4439..f1a054ab 100644
--- a/client.go
+++ b/client.go
@@ -498,6 +498,22 @@ func (cl *Client) acceptConnections(l Listener) {
 	for {
 		conn, err := l.Accept()
 		torrent.Add("client listener accepts", 1)
+		if err == nil {
+			holepunchAddr, holepunchErr := addrPortFromPeerRemoteAddr(conn.RemoteAddr())
+			if holepunchErr == nil {
+				cl.lock()
+				if g.MapContains(cl.undialableWithoutHolepunch, holepunchAddr) {
+					setAdd(&cl.accepted, holepunchAddr)
+				}
+				if g.MapContains(
+					cl.undialableWithoutHolepunchDialedAfterHolepunchConnect,
+					holepunchAddr,
+				) {
+					setAdd(&cl.probablyOnlyConnectedDueToHolepunch, holepunchAddr)
+				}
+				cl.unlock()
+			}
+		}
 		conn = pproffd.WrapNetConn(conn)
 		cl.rLock()
 		closed := cl.closed.IsSet()
@@ -516,20 +532,6 @@ func (cl *Client) acceptConnections(l Listener) {
 			log.Fmsg("error accepting connection: %s", err).LogLevel(log.Debug, cl.logger)
 			continue
 		}
-		{
-			holepunchAddr, holepunchErr := addrPortFromPeerRemoteAddr(conn.RemoteAddr())
-			if holepunchErr == nil {
-				cl.lock()
-				if g.MapContains(
-					cl.undialableWithoutHolepunchDialedAfterHolepunchConnect,
-					holepunchAddr,
-				) {
-					g.MakeMapIfNil(&cl.probablyOnlyConnectedDueToHolepunch)
-					g.MapInsert(cl.probablyOnlyConnectedDueToHolepunch, holepunchAddr, struct{}{})
-				}
-				cl.unlock()
-			}
-		}
 		go func() {
 			if reject != nil {
 				torrent.Add("rejected accepted connections", 1)
@@ -761,17 +763,6 @@ func (cl *Client) dialAndCompleteHandshake(opts outgoingConnOpts) (c *PeerConn,
 		}
 	}
 	holepunchAddr, holepunchAddrErr := addrPortFromPeerRemoteAddr(addr)
-	if holepunchAddrErr == nil && opts.receivedHolepunchConnect {
-		cl.lock()
-		if g.MapContains(cl.undialableWithoutHolepunch, holepunchAddr) {
-			g.MakeMapIfNilAndSet(
-				&cl.undialableWithoutHolepunchDialedAfterHolepunchConnect,
-				holepunchAddr,
-				struct{}{},
-			)
-		}
-		cl.unlock()
-	}
 	headerObfuscationPolicy := opts.HeaderObfuscationPolicy
 	obfuscatedHeaderFirst := headerObfuscationPolicy.Preferred
 	firstDialResult := dialPool.getFirst()
diff --git a/torrent.go b/torrent.go
index 6cc2e87f..42a62fc1 100644
--- a/torrent.go
+++ b/torrent.go
@@ -2825,7 +2825,14 @@ func (t *Torrent) handleReceivedUtHolepunchMsg(msg utHolepunch.Msg, sender *Peer
 		}
 		return nil
 	case utHolepunch.Connect:
-		t.logger.Printf("got holepunch connect request for %v from %p", msg.AddrPort, sender)
+		holepunchAddr := msg.AddrPort
+		t.logger.Printf("got holepunch connect request for %v from %p", holepunchAddr, sender)
+		if g.MapContains(t.cl.undialableWithoutHolepunch, holepunchAddr) {
+			setAdd(&t.cl.undialableWithoutHolepunchDialedAfterHolepunchConnect, holepunchAddr)
+			if g.MapContains(t.cl.accepted, holepunchAddr) {
+				setAdd(&t.cl.probablyOnlyConnectedDueToHolepunch, holepunchAddr)
+			}
+		}
 		opts := outgoingConnOpts{
 			peerInfo: PeerInfo{
 				Addr:         msg.AddrPort,