From 765edfa7cb0786d6540180be8281c5f4109ebe0c Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 21 Jun 2021 12:54:57 +1000
Subject: [PATCH] Expose DialFirst

---
 client.go | 53 +++++++++++++++++++++++------------------------------
 dialer.go | 10 +++++-----
 socket.go |  5 +++++
 3 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/client.go b/client.go
index 032f253a..c7a5f87c 100644
--- a/client.go
+++ b/client.go
@@ -553,9 +553,9 @@ func (cl *Client) torrent(ih metainfo.Hash) *Torrent {
 	return cl.torrents[ih]
 }
 
-type dialResult struct {
-	Conn    net.Conn
-	Network string
+type DialResult struct {
+	Conn   net.Conn
+	Dialer Dialer
 }
 
 func countDialResult(err error) {
@@ -581,14 +581,19 @@ func (cl *Client) dopplegangerAddr(addr string) bool {
 }
 
 // Returns a connection over UTP or TCP, whichever is first to connect.
-func (cl *Client) dialFirst(ctx context.Context, addr string) (res dialResult) {
+func (cl *Client) dialFirst(ctx context.Context, addr string) (res DialResult) {
+	return DialFirst(ctx, addr, cl.dialers)
+}
+
+// Returns a connection over UTP or TCP, whichever is first to connect.
+func DialFirst(ctx context.Context, addr string, dialers []Dialer) (res DialResult) {
 	{
 		t := perf.NewTimer(perf.CallerName(0))
 		defer func() {
 			if res.Conn == nil {
 				t.Mark(fmt.Sprintf("returned no conn (context: %v)", ctx.Err()))
 			} else {
-				t.Mark("returned conn over " + res.Network)
+				t.Mark("returned conn over " + res.Dialer.DialerNetwork())
 			}
 		}()
 	}
@@ -596,24 +601,17 @@ func (cl *Client) dialFirst(ctx context.Context, addr string) (res dialResult) {
 	// As soon as we return one connection, cancel the others.
 	defer cancel()
 	left := 0
-	resCh := make(chan dialResult, left)
-	func() {
-		cl.lock()
-		defer cl.unlock()
-		cl.eachDialer(func(s Dialer) bool {
-			func() {
-				left++
-				//cl.logger.Printf("dialing %s on %s/%s", addr, s.Addr().Network(), s.Addr())
-				go func() {
-					resCh <- dialResult{
-						cl.dialFromSocket(ctx, s, addr),
-						s.LocalAddr().Network(),
-					}
-				}()
-			}()
-			return true
-		})
-	}()
+	resCh := make(chan DialResult, left)
+	for _, _s := range dialers {
+		left++
+		s := _s
+		go func() {
+			resCh <- DialResult{
+				dialFromSocket(ctx, s, addr),
+				s,
+			}
+		}()
+	}
 	// Wait for a successful connection.
 	func() {
 		defer perf.ScopeTimer()()
@@ -633,15 +631,10 @@ func (cl *Client) dialFirst(ctx context.Context, addr string) (res dialResult) {
 	if res.Conn != nil {
 		go torrent.Add(fmt.Sprintf("network dialed first: %s", res.Conn.RemoteAddr().Network()), 1)
 	}
-	//if res.Conn != nil {
-	//	cl.logger.Printf("first connection for %s from %s/%s", addr, res.Conn.LocalAddr().Network(), res.Conn.LocalAddr().String())
-	//} else {
-	//	cl.logger.Printf("failed to dial %s", addr)
-	//}
 	return res
 }
 
-func (cl *Client) dialFromSocket(ctx context.Context, s Dialer, addr string) net.Conn {
+func dialFromSocket(ctx context.Context, s Dialer, addr string) net.Conn {
 	c, err := s.Dial(ctx, addr)
 	// This is a bit optimistic, but it looks non-trivial to thread this through the proxy code. Set
 	// it now in case we close the connection forthwith.
@@ -711,7 +704,7 @@ func (cl *Client) establishOutgoingConnEx(t *Torrent, addr PeerRemoteAddr, obfus
 		}
 		return nil, errors.New("dial failed")
 	}
-	c, err := cl.initiateProtocolHandshakes(context.Background(), nc, t, true, obfuscatedHeader, addr, dr.Network, regularNetConnPeerConnConnString(nc))
+	c, err := cl.initiateProtocolHandshakes(context.Background(), nc, t, true, obfuscatedHeader, addr, dr.Dialer.DialerNetwork(), regularNetConnPeerConnConnString(nc))
 	if err != nil {
 		nc.Close()
 	}
diff --git a/dialer.go b/dialer.go
index 32ab91f7..e8126bd6 100644
--- a/dialer.go
+++ b/dialer.go
@@ -8,12 +8,8 @@ import (
 )
 
 type Dialer interface {
-	// The network is implied by the instance.
 	Dial(_ context.Context, addr string) (net.Conn, error)
-	// This is required for registering with the connection tracker (router connection table
-	// emulating rate-limiter) before dialing. TODO: What about connections that wouldn't infringe
-	// on routers, like localhost or unix sockets.
-	LocalAddr() net.Addr
+	DialerNetwork() string
 }
 
 type NetDialer struct {
@@ -21,6 +17,10 @@ type NetDialer struct {
 	Dialer  net.Dialer
 }
 
+func (me NetDialer) DialerNetwork() string {
+	return me.Network
+}
+
 func (me NetDialer) Dial(ctx context.Context, addr string) (_ net.Conn, err error) {
 	defer perf.ScopeTimerErr(&err)()
 	return me.Dialer.DialContext(ctx, me.Network, addr)
diff --git a/socket.go b/socket.go
index d9654ace..ba2a091b 100644
--- a/socket.go
+++ b/socket.go
@@ -106,11 +106,16 @@ func listenUtp(network, addr string, fc firewallCallback) (socket, error) {
 	return utpSocketSocket{us, network}, err
 }
 
+// utpSocket wrapper, additionally wrapped for the torrent package's socket interface.
 type utpSocketSocket struct {
 	utpSocket
 	network string
 }
 
+func (me utpSocketSocket) DialerNetwork() string {
+	return me.network
+}
+
 func (me utpSocketSocket) Dial(ctx context.Context, addr string) (conn net.Conn, err error) {
 	defer perf.ScopeTimerErr(&err)()
 	return me.utpSocket.DialContext(ctx, me.network, addr)