Improve limiting HTTP based backends

internal/limiter/limiter.go

@@ -2,6 +2,7 @@ package limiter
 
 import (
 	"io"
+	"net/http"
 )
 
 // Limiter defines an interface that implementors can use to rate limit I/O
@@ -14,4 +15,7 @@ type Limiter interface {
 	// Downstream returns a rate limited reader that is intended to be used
 	// for downloads.
 	Downstream(r io.Reader) io.Reader
+
+	// Transport returns an http.RoundTripper limited with the limiter.
+	Transport(http.RoundTripper) http.RoundTripper
 }

internal/limiter/static_limiter.go

@@ -2,6 +2,7 @@ package limiter
 
 import (
 	"io"
+	"net/http"
 
 	"github.com/juju/ratelimit"
 )
@@ -41,6 +42,39 @@ func (l staticLimiter) Downstream(r io.Reader) io.Reader {
 	return l.limit(r, l.downstream)
 }
 
+type roundTripper func(*http.Request) (*http.Response, error)
+
+func (rt roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	return rt(req)
+}
+
+func (l staticLimiter) roundTripper(rt http.RoundTripper, req *http.Request) (*http.Response, error) {
+	if req.Body != nil {
+		req.Body = limitedReadCloser{
+			limited:  l.Upstream(req.Body),
+			original: req.Body,
+		}
+	}
+
+	res, err := rt.RoundTrip(req)
+
+	if res != nil && res.Body != nil {
+		res.Body = limitedReadCloser{
+			limited:  l.Downstream(res.Body),
+			original: res.Body,
+		}
+	}
+
+	return res, err
+}
+
+// Transport returns an HTTP transport limited with the limiter l.
+func (l staticLimiter) Transport(rt http.RoundTripper) http.RoundTripper {
+	return roundTripper(func(req *http.Request) (*http.Response, error) {
+		return l.roundTripper(rt, req)
+	})
+}
+
 func (l staticLimiter) limit(r io.Reader, b *ratelimit.Bucket) io.Reader {
 	if b == nil {
 		return r