From 101a269873b1d91653e43cf36247a1c024b6f087 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 24 Jun 2021 09:53:18 +1000
Subject: [PATCH] Rewrite udp announce to use udp client

---
 tracker/udp.go             | 61 +++++---------------------
 tracker/udp/conn-client.go | 89 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 49 deletions(-)
 create mode 100644 tracker/udp/conn-client.go

diff --git a/tracker/udp.go b/tracker/udp.go
index 1290c7d5..f7afc9ea 100644
--- a/tracker/udp.go
+++ b/tracker/udp.go
@@ -1,13 +1,9 @@
 package tracker
 
 import (
-	"encoding"
 	"encoding/binary"
-	"net"
 	"net/url"
 
-	"github.com/anacrolix/dht/v2/krpc"
-	"github.com/anacrolix/missinggo/v2"
 	trHttp "github.com/anacrolix/torrent/tracker/http"
 	"github.com/anacrolix/torrent/tracker/udp"
 )
@@ -17,56 +13,24 @@ type udpAnnounce struct {
 	a   *Announce
 }
 
-func (c *udpAnnounce) Close() error {
-	return nil
-}
-
-func (c *udpAnnounce) ipv6(conn net.Conn) bool {
-	if c.a.UdpNetwork == "udp6" {
-		return true
-	}
-	rip := missinggo.AddrIP(conn.RemoteAddr())
-	return rip.To16() != nil && rip.To4() == nil
-}
-
 func (c *udpAnnounce) Do(req AnnounceRequest) (res AnnounceResponse, err error) {
-	conn, err := net.Dial(c.dialNetwork(), c.url.Host)
+	cl, err := udp.NewConnClient(udp.NewConnClientOpts{
+		Network: c.dialNetwork(),
+		Host:    c.url.Host,
+		Ipv6:    nil,
+	})
 	if err != nil {
 		return
 	}
-	defer conn.Close()
-	if c.ipv6(conn) {
-		// BEP 15
-		req.IPAddress = 0
-	} else if req.IPAddress == 0 && c.a.ClientIp4.IP != nil {
+	defer cl.Close()
+	if req.IPAddress == 0 && c.a.ClientIp4.IP != nil {
+		// I think we're taking bytes in big-endian order (all IPs), and writing it to a natively
+		// ordered uint32. This will be correctly ordered when written back out by the UDP client
+		// later. I'm ignoring the fact that IPv6 announces shouldn't have an IP address, we have a
+		// perfectly good IPv4 address.
 		req.IPAddress = binary.BigEndian.Uint32(c.a.ClientIp4.IP.To4())
 	}
-	d := udp.Dispatcher{}
-	go func() {
-		for {
-			b := make([]byte, 0x800)
-			n, err := conn.Read(b)
-			if err != nil {
-				break
-			}
-			d.Dispatch(b[:n])
-		}
-	}()
-	cl := udp.Client{
-		Dispatcher: &d,
-		Writer:     conn,
-	}
-	nas := func() interface {
-		encoding.BinaryUnmarshaler
-		NodeAddrs() []krpc.NodeAddr
-	} {
-		if c.ipv6(conn) {
-			return &krpc.CompactIPv6NodeAddrs{}
-		} else {
-			return &krpc.CompactIPv4NodeAddrs{}
-		}
-	}()
-	h, err := cl.Announce(c.a.Context, req, nas, udp.Options{RequestUri: c.url.RequestURI()})
+	h, nas, err := cl.Announce(c.a.Context, req, udp.Options{RequestUri: c.url.RequestURI()})
 	if err != nil {
 		return
 	}
@@ -92,6 +56,5 @@ func announceUDP(opt Announce, _url *url.URL) (AnnounceResponse, error) {
 		url: *_url,
 		a:   &opt,
 	}
-	defer ua.Close()
 	return ua.Do(opt.Request)
 }
diff --git a/tracker/udp/conn-client.go b/tracker/udp/conn-client.go
new file mode 100644
index 00000000..64020569
--- /dev/null
+++ b/tracker/udp/conn-client.go
@@ -0,0 +1,89 @@
+package udp
+
+import (
+	"context"
+	"log"
+	"net"
+
+	"github.com/anacrolix/dht/v2/krpc"
+	"github.com/anacrolix/missinggo/v2"
+)
+
+type NewConnClientOpts struct {
+	Network string
+	Host    string
+	Ipv6    *bool
+}
+
+type ConnClient struct {
+	cl      Client
+	conn    net.Conn
+	d       Dispatcher
+	readErr error
+	ipv6    bool
+}
+
+func (cc *ConnClient) reader() {
+	for {
+		b := make([]byte, 0x800)
+		n, err := cc.conn.Read(b)
+		if err != nil {
+			// TODO: Do bad things to the dispatcher, and incoming calls to the client if we have a
+			// read error.
+			cc.readErr = err
+			break
+		}
+		err = cc.d.Dispatch(b[:n])
+		if err != nil {
+			log.Printf("dispatching packet received on %v: %v", cc.conn, err)
+		}
+	}
+}
+
+func ipv6(opt *bool, network string, conn net.Conn) bool {
+	if opt != nil {
+		return *opt
+	}
+	switch network {
+	case "udp4":
+		return false
+	case "udp6":
+		return true
+	}
+	rip := missinggo.AddrIP(conn.RemoteAddr())
+	return rip.To16() != nil && rip.To4() == nil
+}
+
+func NewConnClient(opts NewConnClientOpts) (cc ConnClient, err error) {
+	cc.conn, err = net.Dial(opts.Network, opts.Host)
+	if err != nil {
+		return
+	}
+	cc.ipv6 = ipv6(opts.Ipv6, opts.Network, cc.conn)
+	go cc.reader()
+	cc.cl = Client{
+		Dispatcher: &cc.d,
+		Writer:     cc.conn,
+	}
+	return
+}
+
+func (c *ConnClient) Close() error {
+	return c.conn.Close()
+}
+
+func (c *ConnClient) Announce(
+	ctx context.Context, req AnnounceRequest, opts Options,
+) (
+	h AnnounceResponseHeader, nas AnnounceResponsePeers, err error,
+) {
+	nas = func() AnnounceResponsePeers {
+		if c.ipv6 {
+			return &krpc.CompactIPv6NodeAddrs{}
+		} else {
+			return &krpc.CompactIPv4NodeAddrs{}
+		}
+	}()
+	h, err = c.cl.Announce(ctx, req, nas, opts)
+	return
+}
-- 
2.48.1