From b3c4afbe25aadd0c6d01084d0c1bd097c8690864 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 12 Mar 2015 20:07:10 +1100
Subject: [PATCH] Add BEP 41 support

---
 tracker/tracker.go              |  2 +-
 tracker/udp/udp_tracker.go      | 42 ++++++++++++++++++++++----
 tracker/udp/udp_tracker_test.go | 52 +++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 7 deletions(-)

diff --git a/tracker/tracker.go b/tracker/tracker.go
index b220c292..888228ec 100644
--- a/tracker/tracker.go
+++ b/tracker/tracker.go
@@ -17,7 +17,7 @@ type AnnounceRequest struct {
 	Key        int32
 	NumWant    int32 // How many peer addresses are desired. -1 for default.
 	Port       int16
-}
+} // 82 bytes
 
 type AnnounceResponse struct {
 	Interval int32 // Minimum seconds the local peer should wait before next announce.
diff --git a/tracker/udp/udp_tracker.go b/tracker/udp/udp_tracker.go
index 94bb1d78..3ba30a84 100644
--- a/tracker/udp/udp_tracker.go
+++ b/tracker/udp/udp_tracker.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"io"
 	"math/rand"
 	"net"
@@ -20,6 +21,11 @@ const (
 	Announce
 	Scrape
 	Error
+
+	// BEP 41
+	optionTypeEndOfOptions = 0
+	optionTypeNOP          = 1
+	optionTypeURLData      = 2
 )
 
 type ConnectionRequest struct {
@@ -41,7 +47,7 @@ type RequestHeader struct {
 	ConnectionId  int64
 	Action        Action
 	TransactionId int32
-}
+} // 16 bytes
 
 type AnnounceResponseHeader struct {
 	Interval int32
@@ -100,13 +106,21 @@ func (c *client) Announce(req *tracker.AnnounceRequest) (res tracker.AnnounceRes
 		err = tracker.ErrNotConnected
 		return
 	}
-	b, err := c.request(Announce, req)
+	reqURI := c.url.RequestURI()
+	// Clearly this limits the request URI to 255 bytes. BEP 41 supports
+	// longer but I'm not fussed.
+	options := append([]byte{optionTypeURLData, byte(len(reqURI))}, []byte(reqURI)...)
+	b, err := c.request(Announce, req, options)
 	if err != nil {
 		return
 	}
 	var h AnnounceResponseHeader
 	err = readBody(b, &h)
 	if err != nil {
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
+		}
+		err = fmt.Errorf("error parsing announce response: %s", err)
 		return
 	}
 	res.Interval = h.Interval
@@ -130,7 +144,9 @@ func (c *client) Announce(req *tracker.AnnounceRequest) (res tracker.AnnounceRes
 	}
 }
 
-func (c *client) write(h *RequestHeader, body interface{}) (err error) {
+// body is the binary serializable request body. trailer is optional data
+// following it, such as for BEP 41.
+func (c *client) write(h *RequestHeader, body interface{}, trailer []byte) (err error) {
 	buf := &bytes.Buffer{}
 	err = binary.Write(buf, binary.BigEndian, h)
 	if err != nil {
@@ -142,6 +158,10 @@ func (c *client) write(h *RequestHeader, body interface{}) (err error) {
 			panic(err)
 		}
 	}
+	_, err = buf.Write(trailer)
+	if err != nil {
+		return
+	}
 	n, err := c.socket.Write(buf.Bytes())
 	if err != nil {
 		return
@@ -152,13 +172,23 @@ func (c *client) write(h *RequestHeader, body interface{}) (err error) {
 	return
 }
 
-func (c *client) request(action Action, args interface{}) (responseBody *bytes.Reader, err error) {
+func read(r io.Reader, data interface{}) error {
+	return binary.Read(r, binary.BigEndian, data)
+}
+
+func write(w io.Writer, data interface{}) error {
+	return binary.Write(w, binary.BigEndian, data)
+}
+
+// args is the binary serializable request body. trailer is optional data
+// following it, such as for BEP 41.
+func (c *client) request(action Action, args interface{}, options []byte) (responseBody *bytes.Reader, err error) {
 	tid := newTransactionId()
 	err = c.write(&RequestHeader{
 		ConnectionId:  c.connectionId,
 		Action:        action,
 		TransactionId: tid,
-	}, args)
+	}, args, options)
 	if err != nil {
 		return
 	}
@@ -223,7 +253,7 @@ func (c *client) Connect() (err error) {
 			return
 		}
 	}
-	b, err := c.request(Connect, nil)
+	b, err := c.request(Connect, nil, nil)
 	if err != nil {
 		return
 	}
diff --git a/tracker/udp/udp_tracker_test.go b/tracker/udp/udp_tracker_test.go
index 44ba25fe..488cd49a 100644
--- a/tracker/udp/udp_tracker_test.go
+++ b/tracker/udp/udp_tracker_test.go
@@ -5,7 +5,10 @@ import (
 	"crypto/rand"
 	"encoding/binary"
 	"io"
+	"io/ioutil"
+	"log"
 	"net"
+	"net/url"
 	"sync"
 	"syscall"
 	"testing"
@@ -144,3 +147,52 @@ func TestAnnounceRandomInfoHash(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+// Check that URLPath option is done correctly.
+func TestURLPathOption(t *testing.T) {
+	conn, err := net.ListenUDP("udp", nil)
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close()
+	cl := newClient(&url.URL{
+		Host: conn.LocalAddr().String(),
+		Path: "/announce",
+	})
+	go func() {
+		err = cl.Connect()
+		if err != nil {
+			t.Fatal(err)
+		}
+		log.Print("connected")
+		_, err = cl.Announce(&tracker.AnnounceRequest{})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+	var b [512]byte
+	_, addr, _ := conn.ReadFrom(b[:])
+	r := bytes.NewReader(b[:])
+	var h RequestHeader
+	read(r, &h)
+	w := &bytes.Buffer{}
+	write(w, ResponseHeader{
+		TransactionId: h.TransactionId,
+	})
+	write(w, ConnectionResponse{42})
+	conn.WriteTo(w.Bytes(), addr)
+	n, _, _ := conn.ReadFrom(b[:])
+	r = bytes.NewReader(b[:n])
+	read(r, &h)
+	read(r, &tracker.AnnounceRequest{})
+	all, _ := ioutil.ReadAll(r)
+	if string(all) != "\x02\x09/announce" {
+		t.FailNow()
+	}
+	w = &bytes.Buffer{}
+	write(w, ResponseHeader{
+		TransactionId: h.TransactionId,
+	})
+	write(w, AnnounceResponseHeader{})
+	conn.WriteTo(w.Bytes(), addr)
+}
-- 
2.48.1