From d4bf547801826200c0705bc08c529e0500c5549f Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 24 Jun 2021 14:38:16 +1000
Subject: [PATCH] Add udp tracker scrape support

---
 tracker/udp/client.go      | 32 +++++++++++++++++++++++++++-----
 tracker/udp/conn-client.go |  8 ++++----
 tracker/udp/protocol.go    | 10 ++++++++++
 tracker/udp/scrape.go      | 11 +++++++++++
 4 files changed, 52 insertions(+), 9 deletions(-)
 create mode 100644 tracker/udp/scrape.go

diff --git a/tracker/udp/client.go b/tracker/udp/client.go
index d66348e1..70df469f 100644
--- a/tracker/udp/client.go
+++ b/tracker/udp/client.go
@@ -22,11 +22,7 @@ func (cl *Client) Announce(
 ) (
 	respHdr AnnounceResponseHeader, err error,
 ) {
-	body, err := marshal(req)
-	if err != nil {
-		return
-	}
-	respBody, err := cl.request(ctx, ActionAnnounce, append(body, opts.Encode()...))
+	respBody, err := cl.request(ctx, ActionAnnounce, append(mustMarshal(req), opts.Encode()...))
 	if err != nil {
 		return
 	}
@@ -43,6 +39,32 @@ func (cl *Client) Announce(
 	return
 }
 
+func (cl *Client) Scrape(
+	ctx context.Context, ihs []InfoHash,
+) (
+	out ScrapeResponse, err error,
+) {
+	// There's no way to pass options in a scrape, since we don't when the request body ends.
+	respBody, err := cl.request(ctx, ActionScrape, mustMarshal(ScrapeRequest(ihs)))
+	if err != nil {
+		return
+	}
+	r := bytes.NewBuffer(respBody)
+	for r.Len() != 0 {
+		var item ScrapeInfohashResult
+		err = Read(r, &item)
+		if err != nil {
+			return
+		}
+		out = append(out, item)
+	}
+	if len(out) > len(ihs) {
+		err = fmt.Errorf("got %v results but expected %v", len(out), len(ihs))
+		return
+	}
+	return
+}
+
 func (cl *Client) connect(ctx context.Context) (err error) {
 	if time.Since(cl.connIdIssued) < time.Minute {
 		return nil
diff --git a/tracker/udp/conn-client.go b/tracker/udp/conn-client.go
index a91cacba..d7e19027 100644
--- a/tracker/udp/conn-client.go
+++ b/tracker/udp/conn-client.go
@@ -15,7 +15,7 @@ type NewConnClientOpts struct {
 }
 
 type ConnClient struct {
-	cl      Client
+	Client  Client
 	conn    net.Conn
 	d       Dispatcher
 	readErr error
@@ -59,13 +59,13 @@ func NewConnClient(opts NewConnClientOpts) (cc *ConnClient, err error) {
 		return
 	}
 	cc = &ConnClient{
-		cl: Client{
+		Client: Client{
 			Writer: conn,
 		},
 		conn: conn,
 		ipv6: ipv6(opts.Ipv6, opts.Network, conn),
 	}
-	cc.cl.Dispatcher = &cc.d
+	cc.Client.Dispatcher = &cc.d
 	go cc.reader()
 	return
 }
@@ -86,6 +86,6 @@ func (c *ConnClient) Announce(
 			return &krpc.CompactIPv4NodeAddrs{}
 		}
 	}()
-	h, err = c.cl.Announce(ctx, req, nas, opts)
+	h, err = c.Client.Announce(ctx, req, nas, opts)
 	return
 }
diff --git a/tracker/udp/protocol.go b/tracker/udp/protocol.go
index 365d3c5c..4a8dc668 100644
--- a/tracker/udp/protocol.go
+++ b/tracker/udp/protocol.go
@@ -53,6 +53,8 @@ type AnnounceResponseHeader struct {
 	Seeders  int32
 }
 
+type InfoHash = [20]byte
+
 func marshal(data interface{}) (b []byte, err error) {
 	var buf bytes.Buffer
 	err = binary.Write(&buf, binary.BigEndian, data)
@@ -60,6 +62,14 @@ func marshal(data interface{}) (b []byte, err error) {
 	return
 }
 
+func mustMarshal(data interface{}) []byte {
+	b, err := marshal(data)
+	if err != nil {
+		panic(err)
+	}
+	return b
+}
+
 func Write(w io.Writer, data interface{}) error {
 	return binary.Write(w, binary.BigEndian, data)
 }
diff --git a/tracker/udp/scrape.go b/tracker/udp/scrape.go
new file mode 100644
index 00000000..331f109e
--- /dev/null
+++ b/tracker/udp/scrape.go
@@ -0,0 +1,11 @@
+package udp
+
+type ScrapeRequest []InfoHash
+
+type ScrapeResponse []ScrapeInfohashResult
+
+type ScrapeInfohashResult struct {
+	Seeders   int32
+	Completed int32
+	Leechers  int32
+}
-- 
2.50.0