Apply download rate limiter to webseeds

diff --git a/torrent.go b/torrent.go
index 0f3603749337f64f1606d758d5f602c7dae1352f..706092740830a19abad6ba5e2db431a156e4c417 100644 (file)
--- 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 ff246a1cdc5c79a19bc4f15c84587166a61f844e..a04b34300ef025ba08719699591eb9fdd413af06 100644 (file)
--- 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}