From 916af6e38ad80485dcf7c7b0197cbf8a14da1dbd Mon Sep 17 00:00:00 2001
From: Marco Vidonis <31407403+marcovidonis@users.noreply.github.com>
Date: Mon, 28 Nov 2022 23:35:36 +0000
Subject: [PATCH] Modify HTTP request before sending (#787)

* set up custom request headers

* apply headers defined in torrent client config

* add error handling

* provide better name for method

* update error message

* only apply HTTPRequestDirector if not nil
---
 config.go            |  3 +++
 tracker/http/http.go | 18 ++++++++++++++----
 tracker/tracker.go   | 28 +++++++++++++++-------------
 tracker_scraper.go   | 27 ++++++++++++++-------------
 4 files changed, 46 insertions(+), 30 deletions(-)

diff --git a/config.go b/config.go
index 6f9f6827..f640e1a4 100644
--- a/config.go
+++ b/config.go
@@ -102,6 +102,9 @@ type ClientConfig struct {
 	LookupTrackerIp func(*url.URL) ([]net.IP, error)
 	// HTTPUserAgent changes default UserAgent for HTTP requests
 	HTTPUserAgent string
+	// HTTPRequestDirector modifies the request before it's sent.
+	// Useful for adding authentication headers, for example
+	HTTPRequestDirector func(*http.Request) error
 	// Updated occasionally to when there's been some changes to client
 	// behaviour in case other clients are assuming anything of us. See also
 	// `bep20`.
diff --git a/tracker/http/http.go b/tracker/http/http.go
index 72a95c38..3c4d9052 100644
--- a/tracker/http/http.go
+++ b/tracker/http/http.go
@@ -76,10 +76,11 @@ func setAnnounceParams(_url *url.URL, ar *AnnounceRequest, opts AnnounceOpt) {
 }
 
 type AnnounceOpt struct {
-	UserAgent  string
-	HostHeader string
-	ClientIp4  net.IP
-	ClientIp6  net.IP
+	UserAgent           string
+	HostHeader          string
+	ClientIp4           net.IP
+	ClientIp6           net.IP
+	HTTPRequestDirector func(*http.Request) error
 }
 
 type AnnounceRequest = udp.AnnounceRequest
@@ -95,6 +96,15 @@ func (cl Client) Announce(ctx context.Context, ar AnnounceRequest, opt AnnounceO
 	if userAgent != "" {
 		req.Header.Set("User-Agent", userAgent)
 	}
+
+	if opt.HTTPRequestDirector != nil {
+		err = opt.HTTPRequestDirector(req)
+		if err != nil {
+			err = fmt.Errorf("error modifying HTTP request: %s", err)
+			return
+		}
+	}
+
 	req.Host = opt.HostHeader
 	resp, err := cl.hc.Do(req)
 	if err != nil {
diff --git a/tracker/tracker.go b/tracker/tracker.go
index 122278b6..c169c6f5 100644
--- a/tracker/tracker.go
+++ b/tracker/tracker.go
@@ -35,15 +35,16 @@ type AnnounceEvent = udp.AnnounceEvent
 var ErrBadScheme = errors.New("unknown scheme")
 
 type Announce struct {
-	TrackerUrl   string
-	Request      AnnounceRequest
-	HostHeader   string
-	HTTPProxy    func(*http.Request) (*url.URL, error)
-	DialContext  func(ctx context.Context, network, addr string) (net.Conn, error)
-	ListenPacket func(network, addr string) (net.PacketConn, error)
-	ServerName   string
-	UserAgent    string
-	UdpNetwork   string
+	TrackerUrl          string
+	Request             AnnounceRequest
+	HostHeader          string
+	HTTPProxy           func(*http.Request) (*url.URL, error)
+	HTTPRequestDirector func(*http.Request) error
+	DialContext         func(ctx context.Context, network, addr string) (net.Conn, error)
+	ListenPacket        func(network, addr string) (net.PacketConn, error)
+	ServerName          string
+	UserAgent           string
+	UdpNetwork          string
 	// If the port is zero, it's assumed to be the same as the Request.Port.
 	ClientIp4 krpc.NodeAddr
 	// If the port is zero, it's assumed to be the same as the Request.Port.
@@ -79,9 +80,10 @@ func (me Announce) Do() (res AnnounceResponse, err error) {
 		me.Context = ctx
 	}
 	return cl.Announce(me.Context, me.Request, trHttp.AnnounceOpt{
-		UserAgent:  me.UserAgent,
-		HostHeader: me.HostHeader,
-		ClientIp4:  me.ClientIp4.IP,
-		ClientIp6:  me.ClientIp6.IP,
+		UserAgent:           me.UserAgent,
+		HostHeader:          me.HostHeader,
+		ClientIp4:           me.ClientIp4.IP,
+		ClientIp6:           me.ClientIp6.IP,
+		HTTPRequestDirector: me.HTTPRequestDirector,
 	})
 }
diff --git a/tracker_scraper.go b/tracker_scraper.go
index de65d9fb..3f20cbd9 100644
--- a/tracker_scraper.go
+++ b/tracker_scraper.go
@@ -156,19 +156,20 @@ func (me *trackerScraper) announce(ctx context.Context, event tracker.AnnounceEv
 	defer cancel()
 	me.t.logger.WithDefaultLevel(log.Debug).Printf("announcing to %q: %#v", me.u.String(), req)
 	res, err := tracker.Announce{
-		Context:      ctx,
-		HTTPProxy:    me.t.cl.config.HTTPProxy,
-		DialContext:  me.t.cl.config.TrackerDialContext,
-		ListenPacket: me.t.cl.config.TrackerListenPacket,
-		UserAgent:    me.t.cl.config.HTTPUserAgent,
-		TrackerUrl:   me.trackerUrl(ip),
-		Request:      req,
-		HostHeader:   me.u.Host,
-		ServerName:   me.u.Hostname(),
-		UdpNetwork:   me.u.Scheme,
-		ClientIp4:    krpc.NodeAddr{IP: me.t.cl.config.PublicIp4},
-		ClientIp6:    krpc.NodeAddr{IP: me.t.cl.config.PublicIp6},
-		Logger:       me.t.logger,
+		Context:             ctx,
+		HTTPProxy:           me.t.cl.config.HTTPProxy,
+		HTTPRequestDirector: me.t.cl.config.HTTPRequestDirector,
+		DialContext:         me.t.cl.config.TrackerDialContext,
+		ListenPacket:        me.t.cl.config.TrackerListenPacket,
+		UserAgent:           me.t.cl.config.HTTPUserAgent,
+		TrackerUrl:          me.trackerUrl(ip),
+		Request:             req,
+		HostHeader:          me.u.Host,
+		ServerName:          me.u.Hostname(),
+		UdpNetwork:          me.u.Scheme,
+		ClientIp4:           krpc.NodeAddr{IP: me.t.cl.config.PublicIp4},
+		ClientIp6:           krpc.NodeAddr{IP: me.t.cl.config.PublicIp6},
+		Logger:              me.t.logger,
 	}.Do()
 	me.t.logger.WithDefaultLevel(log.Debug).Printf("announce to %q returned %#v: %v", me.u.String(), res, err)
 	if err != nil {
-- 
2.44.0