From 6b18583a0642929aa0c6e844e36c85d689657b4b Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 18 May 2023 10:41:51 +1000
Subject: [PATCH] Add ClientConfig.DialRateLimiter, handle dial rate limiting
 errors

---
 client.go | 12 ++++++------
 config.go |  3 +++
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/client.go b/client.go
index bc9bafbb..21f213b3 100644
--- a/client.go
+++ b/client.go
@@ -35,7 +35,6 @@ import (
 	"github.com/dustin/go-humanize"
 	gbtree "github.com/google/btree"
 	"github.com/pion/datachannel"
-	"golang.org/x/time/rate"
 
 	"github.com/anacrolix/torrent/bencode"
 	"github.com/anacrolix/torrent/internal/check"
@@ -82,9 +81,8 @@ type Client struct {
 	torrents          map[InfoHash]*Torrent
 	pieceRequestOrder map[interface{}]*request_strategy.PieceRequestOrder
 
-	acceptLimiter   map[ipStr]int
-	dialRateLimiter *rate.Limiter
-	numHalfOpen     int
+	acceptLimiter map[ipStr]int
+	numHalfOpen   int
 
 	websocketTrackers websocketTrackers
 
@@ -201,7 +199,6 @@ func (cl *Client) init(cfg *ClientConfig) {
 	cl.config = cfg
 	g.MakeMap(&cl.dopplegangerAddrs)
 	cl.torrents = make(map[metainfo.Hash]*Torrent)
-	cl.dialRateLimiter = rate.NewLimiter(10, 10)
 	cl.activeAnnounceLimiter.SlotsPerKey = 2
 	cl.event.L = cl.locker()
 	cl.ipBlockList = cfg.IPBlocklist
@@ -734,6 +731,10 @@ 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
+	}
 	torrent.Add("establish outgoing connection", 1)
 	addr := opts.peerInfo.Addr
 	dialPool := dialPool{
@@ -843,7 +844,6 @@ func (cl *Client) outgoingConnection(
 	opts outgoingConnOpts,
 	attemptKey outgoingConnAttemptKey,
 ) {
-	cl.dialRateLimiter.Wait(context.Background())
 	c, err := cl.dialAndCompleteHandshake(opts)
 	if err == nil {
 		c.conn.SetWriteDeadline(time.Time{})
diff --git a/config.go b/config.go
index 17641689..fca37e63 100644
--- a/config.go
+++ b/config.go
@@ -180,6 +180,8 @@ type ClientConfig struct {
 	DisableWebseeds   bool
 
 	Callbacks Callbacks
+
+	DialRateLimiter *rate.Limiter
 }
 
 func (cfg *ClientConfig) SetListenAddr(addr string) *ClientConfig {
@@ -223,6 +225,7 @@ func NewDefaultClientConfig() *ClientConfig {
 		Extensions:            defaultPeerExtensionBytes(),
 		AcceptPeerConnections: true,
 		MaxUnverifiedBytes:    64 << 20,
+		DialRateLimiter:       rate.NewLimiter(10, 10),
 	}
 	cc.DhtStartingNodes = func(network string) dht.StartingNodesGetter {
 		return func() ([]dht.Addr, error) { return dht.GlobalBootstrapAddrs(network) }