From: Matt Joiner <anacrolix@gmail.com>
Date: Wed, 21 Nov 2018 06:02:22 +0000 (+1100)
Subject: Add initial connection tracking
X-Git-Tag: v1.0.0~20
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=f4e30656d6afb336b441041fa1a2986835e7589e;p=btrtrc.git

Add initial connection tracking
---

diff --git a/client.go b/client.go
index 81f620ca..99edceee 100644
--- a/client.go
+++ b/client.go
@@ -21,6 +21,7 @@ import (
 	"github.com/anacrolix/log"
 	"github.com/anacrolix/missinggo"
 	"github.com/anacrolix/missinggo/bitmap"
+	"github.com/anacrolix/missinggo/conntrack"
 	"github.com/anacrolix/missinggo/perf"
 	"github.com/anacrolix/missinggo/pproffd"
 	"github.com/anacrolix/missinggo/pubsub"
@@ -291,7 +292,8 @@ func (cl *Client) newDhtServer(conn net.PacketConn) (s *dht.Server, err error) {
 			}
 			return cl.config.PublicIp4
 		}(),
-		StartingNodes: cl.config.DhtStartingNodes,
+		StartingNodes:      cl.config.DhtStartingNodes,
+		ConnectionTracking: cl.config.ConnTracker,
 	}
 	s, err = dht.NewServer(&cfg)
 	if err == nil {
@@ -538,6 +540,9 @@ func (cl *Client) dialFirst(ctx context.Context, addr string) dialResult {
 			if peerNetworkEnabled(network, cl.config) {
 				left++
 				go func() {
+					cte := cl.config.ConnTracker.Wait(
+						conntrack.Entry{network, s.Addr().String(), addr},
+						"dial torrent client")
 					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
@@ -546,7 +551,17 @@ func (cl *Client) dialFirst(ctx context.Context, addr string) dialResult {
 						tc.SetLinger(0)
 					}
 					countDialResult(err)
-					resCh <- dialResult{c, network}
+					dr := dialResult{c, network}
+					if c == nil {
+						cte.Done()
+					} else {
+						dr.Conn = closeWrapper{c, func() error {
+							err := c.Close()
+							cte.Done()
+							return err
+						}}
+					}
+					resCh <- dr
 				}()
 			}
 			return true
diff --git a/closewrapper.go b/closewrapper.go
new file mode 100644
index 00000000..54494189
--- /dev/null
+++ b/closewrapper.go
@@ -0,0 +1,12 @@
+package torrent
+
+import "net"
+
+type closeWrapper struct {
+	net.Conn
+	closer func() error
+}
+
+func (me closeWrapper) Close() error {
+	return me.closer()
+}
diff --git a/config.go b/config.go
index 31f22513..7e1bcdd0 100644
--- a/config.go
+++ b/config.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/anacrolix/dht"
 	"github.com/anacrolix/missinggo"
+	"github.com/anacrolix/missinggo/conntrack"
 	"github.com/anacrolix/missinggo/expect"
 	"github.com/anacrolix/torrent/iplist"
 	"github.com/anacrolix/torrent/storage"
@@ -121,6 +122,8 @@ type ClientConfig struct {
 	// Don't add connections that have the same peer ID as an existing
 	// connection for a given Torrent.
 	dropDuplicatePeerIds bool
+
+	ConnTracker *conntrack.Instance
 }
 
 func (cfg *ClientConfig) SetListenAddr(addr string) *ClientConfig {
@@ -147,6 +150,7 @@ func NewDefaultClientConfig() *ClientConfig {
 		ListenHost:                 func(string) string { return "" },
 		UploadRateLimiter:          unlimited,
 		DownloadRateLimiter:        unlimited,
+		ConnTracker:                conntrack.NewInstance(),
 	}
 }