routing/http: feat: limit the resp body payload

routing/http/client/client.go

@@ -69,9 +69,15 @@ func WithProviderInfo(peerID peer.ID, addrs []multiaddr.Multiaddr) option {
 // New creates a content routing API client.
 // The Provider and identity parameters are option. If they are nil, the `Provide` method will not function.
 func New(baseURL string, opts ...option) (*client, error) {
+	defaultHTTPClient := &http.Client{
+		Transport: &ResponseBodyLimitedTransport{
+			RoundTripper: http.DefaultTransport,
+			LimitBytes:   1 << 20,
+		},
+	}
 	client := &client{
 		baseURL:    baseURL,
-		httpClient: http.DefaultClient,
+		httpClient: defaultHTTPClient,
 		validator:  ipns.Validator{},
 		clock:      clock.New(),
 	}

routing/http/client/transport.go

@@ -0,0 +1,38 @@
+package client
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+)
+
+type ResponseBodyLimitedTransport struct {
+	http.RoundTripper
+	LimitBytes int64
+}
+
+func (r *ResponseBodyLimitedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	resp, err := r.RoundTripper.RoundTrip(req)
+	if resp != nil && resp.Body != nil {
+		resp.Body = &limitReadCloser{
+			limit:      r.LimitBytes,
+			ReadCloser: resp.Body,
+		}
+	}
+	return resp, err
+}
+
+type limitReadCloser struct {
+	limit     int64
+	bytesRead int64
+	io.ReadCloser
+}
+
+func (l *limitReadCloser) Read(p []byte) (int, error) {
+	n, err := l.ReadCloser.Read(p)
+	l.bytesRead += int64(n)
+	if l.bytesRead > l.limit {
+		return 0, fmt.Errorf("reached read limit of %d bytes after reading %d bytes", l.limit, l.bytesRead)
+	}
+	return n, err
+}

routing/http/client/transport_test.go

@@ -0,0 +1,44 @@
+package client
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type testServer struct {
+	bytesToWrite int
+}
+
+func (s *testServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	bytes := make([]byte, s.bytesToWrite)
+	for i := 0; i < s.bytesToWrite; i++ {
+		bytes[i] = 'a'
+	}
+	_, err := w.Write(bytes)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func TestResponseBodyLimitedTransport(t *testing.T) {
+	server := httptest.NewServer(&testServer{bytesToWrite: 1 << 21})
+	t.Cleanup(server.Close)
+
+	client := server.Client()
+	client.Transport = &ResponseBodyLimitedTransport{
+		LimitBytes:   1 << 20,
+		RoundTripper: client.Transport,
+	}
+
+	resp, err := client.Get(server.URL)
+	if err != nil {
+		panic(err)
+	}
+	_, err = io.ReadAll(resp.Body)
+
+	assert.Contains(t, err.Error(), "reached read limit of 1048576 bytes after reading")
+}