From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 20 Dec 2021 03:29:43 +0000 (+1100)
Subject: Apply download rate limiter to webseeds
X-Git-Tag: v1.40.0~9
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=cc0d223a65d0e82a036ed984b01ba2f118207590;p=btrtrc.git

Apply download rate limiter to webseeds
---

diff --git a/torrent.go b/torrent.go
index 0f360374..70609274 100644
--- a/torrent.go
+++ b/torrent.go
@@ -2273,6 +2273,12 @@ func (t *Torrent) addWebSeed(url string) {
 		client: webseed.Client{
 			HttpClient: t.cl.webseedHttpClient,
 			Url:        url,
+			ResponseBodyWrapper: func(r io.Reader) io.Reader {
+				return &rateLimitedReader{
+					l: t.cl.config.DownloadRateLimiter,
+					r: r,
+				}
+			},
 		},
 		activeRequests: make(map[Request]webseed.Request, maxRequests),
 		maxRequests:    maxRequests,
diff --git a/webseed/client.go b/webseed/client.go
index ff246a1c..a04b3430 100644
--- a/webseed/client.go
+++ b/webseed/client.go
@@ -28,6 +28,8 @@ type requestPart struct {
 	e      segments.Extent
 	result chan requestPartResult
 	start  func()
+	// Wrap http response bodies for such things as download rate limiting.
+	responseBodyWrapper ResponseBodyWrapper
 }
 
 type Request struct {
@@ -48,9 +50,12 @@ type Client struct {
 	// given that's how requests are mapped to webseeds, but the torrent.Client works at the piece
 	// level. We can map our file-level adjustments to the pieces here. This probably need to be
 	// private in the future, if Client ever starts removing pieces.
-	Pieces roaring.Bitmap
+	Pieces              roaring.Bitmap
+	ResponseBodyWrapper ResponseBodyWrapper
 }
 
+type ResponseBodyWrapper func(io.Reader) io.Reader
+
 func (me *Client) SetInfo(info *metainfo.Info) {
 	if !strings.HasSuffix(me.Url, "/") && info.IsDir() {
 		// In my experience, this is a non-conforming webseed. For example the
@@ -77,9 +82,10 @@ func (ws *Client) NewRequest(r RequestSpec) Request {
 		}
 		req = req.WithContext(ctx)
 		part := requestPart{
-			req:    req,
-			result: make(chan requestPartResult, 1),
-			e:      e,
+			req:                 req,
+			result:              make(chan requestPartResult, 1),
+			e:                   e,
+			responseBodyWrapper: ws.ResponseBodyWrapper,
 		}
 		part.start = func() {
 			go func() {
@@ -126,12 +132,18 @@ func recvPartResult(ctx context.Context, buf io.Writer, part requestPart) error
 		return result.err
 	}
 	defer result.resp.Body.Close()
+	var body io.Reader = result.resp.Body
+	if part.responseBodyWrapper != nil {
+		body = part.responseBodyWrapper(body)
+	}
+	// Prevent further accidental use
+	result.resp.Body = nil
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}
 	switch result.resp.StatusCode {
 	case http.StatusPartialContent:
-		copied, err := io.Copy(buf, result.resp.Body)
+		copied, err := io.Copy(buf, body)
 		if err != nil {
 			return err
 		}
@@ -154,11 +166,11 @@ func recvPartResult(ctx context.Context, buf io.Writer, part requestPart) error
 			// body. I don't know how one would handle multiple chunk requests resulting in an OK
 			// response for the same file. The request algorithm might be need to be smarter for
 			// that.
-			discarded, _ := io.CopyN(io.Discard, result.resp.Body, part.e.Start)
+			discarded, _ := io.CopyN(io.Discard, body, part.e.Start)
 			if discarded != 0 {
 				log.Printf("discarded %v bytes in webseed request response part", discarded)
 			}
-			_, err := io.CopyN(buf, result.resp.Body, part.e.Length)
+			_, err := io.CopyN(buf, body, part.e.Length)
 			return err
 		} else {
 			return ErrBadResponse{"resp status ok but requested range", result.resp}