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.51.0