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 +}