Add udp tracker scrape support

diff --git a/tracker/udp/client.go b/tracker/udp/client.go
index d66348e16928d4bd48a2499cf44340093461154c..70df469f91dc33b0bcd23691383597db8343adbb 100644 (file)
--- 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 a91cacba86cbe2ecd1b92ecee8b61462d3255aa9..d7e19027a8a9aac8e3d4a4fb788ebf2139ff4872 100644 (file)
--- 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 365d3c5ca59a996b140d618b2ebca2e38f9c11f8..4a8dc6681512f14710023f5992647b001b124ec4 100644 (file)
--- 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 (file)
index 0000000..331f109
--- /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
+}