From f45365fd986e2957843f70ff2c98c7145e82f17e Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 18 May 2023 11:37:46 +1000
Subject: [PATCH] Ignore dial rate limits for holepunch connects

---
 client.go               | 15 ++++++++++++---
 torrent.go              |  1 +
 ut-holepunching_test.go | 10 +++++++---
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/client.go b/client.go
index 21f213b3..70b98482 100644
--- a/client.go
+++ b/client.go
@@ -731,9 +731,18 @@ func doProtocolHandshakeOnDialResult(
 
 // Returns nil connection and nil error if no connection could be established for valid reasons.
 func (cl *Client) dialAndCompleteHandshake(opts outgoingConnOpts) (c *PeerConn, err error) {
-	err = cl.config.DialRateLimiter.Wait(context.Background())
-	if err != nil {
-		return
+	// It would be better if dial rate limiting could be tested when considering to open connections
+	// instead. Doing it here means if the limit is low, and the half-open limit is high, we could
+	// end up with lots of outgoing connection attempts pending that were initiated on stale data.
+	{
+		dialReservation := cl.config.DialRateLimiter.Reserve()
+		if !opts.receivedHolepunchConnect {
+			if !dialReservation.OK() {
+				err = errors.New("can't make dial limit reservation")
+				return
+			}
+			time.Sleep(dialReservation.Delay())
+		}
 	}
 	torrent.Add("establish outgoing connection", 1)
 	addr := opts.peerInfo.Addr
diff --git a/torrent.go b/torrent.go
index 49ae0605..553fa5db 100644
--- a/torrent.go
+++ b/torrent.go
@@ -2801,6 +2801,7 @@ 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)
 		opts := outgoingConnOpts{
 			peerInfo: PeerInfo{
 				Addr:         msg.AddrPort,
diff --git a/ut-holepunching_test.go b/ut-holepunching_test.go
index 5df98bfb..18521195 100644
--- a/ut-holepunching_test.go
+++ b/ut-holepunching_test.go
@@ -18,6 +18,7 @@ import (
 	qt "github.com/frankban/quicktest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/time/rate"
 
 	"github.com/anacrolix/torrent/internal/testutil"
 )
@@ -36,9 +37,12 @@ func TestHolepunchConnect(t *testing.T) {
 	cfg.DisablePEX = true
 	cfg.Debug = true
 	cfg.AcceptPeerConnections = false
-	// Listening, even without accepting, still means the leecher-leecher completes the dial to the seeder, and so it
-	// won't attempt to holepunch.
+	// Listening, even without accepting, still means the leecher-leecher completes the dial to the
+	// seeder, and so it won't attempt to holepunch.
 	cfg.DisableTCP = true
+	// Ensure that responding to holepunch connects don't wait around for the dial limit. We also
+	// have to allow the initial connection to the leecher though, so it can rendezvous for us.
+	cfg.DialRateLimiter = rate.NewLimiter(0, 1)
 	seeder, err := NewClient(cfg)
 	require.NoError(t, err)
 	defer seeder.Close()
@@ -65,7 +69,7 @@ func TestHolepunchConnect(t *testing.T) {
 	cfg.Seed = false
 	cfg.DataDir = t.TempDir()
 	cfg.MaxAllocPeerRequestDataPerConn = 4
-	cfg.Debug = true
+	//cfg.Debug = true
 	cfg.NominalDialTimeout = time.Second
 	//cfg.DisableUTP = true
 	leecherLeecher, _ := NewClient(cfg)
-- 
2.44.0