torrent.go | 6 ++++++ webseed/client.go | 26 +++++++++++++++++++------- diff --git a/torrent.go b/torrent.go index 0f3603749337f64f1606d758d5f602c7dae1352f..706092740830a19abad6ba5e2db431a156e4c417 100644 --- a/torrent.go +++ b/torrent.go @@ -2273,6 +2273,12 @@ }, 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 --- a/webseed/client.go +++ b/webseed/client.go @@ -28,6 +28,8 @@ req *http.Request 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,8 +50,11 @@ // The pieces we can request with the Url. We're more likely to ban/block at the file-level // 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() { @@ -77,9 +82,10 @@ panic(err) } 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 @@ if result.err != nil { 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 @@ // Instead of discarding, we could try receiving all the chunks present in the response // 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}