Datasources: Introduce response_limit for datasource responses (#38962)

* Introduce response_limit for datasource responses

* Fix lint

* Fix tests

* Add case where limit <= 0 - added parametrized tests

* Add max_bytes_reader.go

* Use new httpclient.MaxBytesReader instead of net/http one

* Fixes according to reviewer's comments

* Add tests for max_bytes_reader

* Add small piece in configuration.md

* Further fixes according to reviewer's comments

* Fix linting - fix test

conf/defaults.ini

@@ -172,6 +172,9 @@ idle_conn_timeout_seconds = 90
 # If enabled and user is not anonymous, data proxy will add X-Grafana-User header with username into the request.
 send_user_header = false
 
+# Limit the amount of bytes that will be read/accepted from responses of outgoing HTTP requests.
+response_limit = 0
+
 #################################### Analytics ###########################
 [analytics]
 # Server reporting, sends usage counters to stats.grafana.org every 24 hours.

conf/sample.ini

@@ -178,6 +178,9 @@
 # If enabled and user is not anonymous, data proxy will add X-Grafana-User header with username into the request, default is false.
 ;send_user_header = false
 
+# Limit the amount of bytes that will be read/accepted from responses of outgoing HTTP requests.
+;response_limit = 0
+
 #################################### Analytics ####################################
 [analytics]
 # Server reporting, sends usage counters to stats.grafana.org every 24 hours.

docs/sources/administration/configuration.md

@@ -435,6 +435,10 @@ The length of time that Grafana maintains idle connections before closing them.
 
 If enabled and user is not anonymous, data proxy will add X-Grafana-User header with username into the request. Default is `false`.
 
+### response_limit
+
+Limits the amount of bytes that will be read/accepted from responses of outgoing HTTP requests. Default is `0` which means disabled.
+
 <hr />
 
 ## [analytics]

pkg/infra/httpclient/httpclientprovider/http_client_provider.go

@@ -25,6 +25,7 @@ func New(cfg *setting.Cfg) *sdkhttpclient.Provider {
 		SetUserAgentMiddleware(userAgent),
 		sdkhttpclient.BasicAuthenticationMiddleware(),
 		sdkhttpclient.CustomHeadersMiddleware(),
+		ResponseLimitMiddleware(cfg.ResponseLimit),
 	}
 
 	if cfg.SigV4AuthEnabled {

pkg/infra/httpclient/httpclientprovider/http_client_provider_test.go

@@ -22,12 +22,13 @@ func TestHTTPClientProvider(t *testing.T) {
 		_ = New(&setting.Cfg{SigV4AuthEnabled: false})
 		require.Len(t, providerOpts, 1)
 		o := providerOpts[0]
-		require.Len(t, o.Middlewares, 5)
+		require.Len(t, o.Middlewares, 6)
 		require.Equal(t, TracingMiddlewareName, o.Middlewares[0].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, DataSourceMetricsMiddlewareName, o.Middlewares[1].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, SetUserAgentMiddlewareName, o.Middlewares[2].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, sdkhttpclient.BasicAuthenticationMiddlewareName, o.Middlewares[3].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, sdkhttpclient.CustomHeadersMiddlewareName, o.Middlewares[4].(sdkhttpclient.MiddlewareName).MiddlewareName())
+		require.Equal(t, ResponseLimitMiddlewareName, o.Middlewares[5].(sdkhttpclient.MiddlewareName).MiddlewareName())
 	})
 
 	t.Run("When creating new provider and SigV4 is enabled should apply expected middleware", func(t *testing.T) {
@@ -43,12 +44,13 @@ func TestHTTPClientProvider(t *testing.T) {
 		_ = New(&setting.Cfg{SigV4AuthEnabled: true})
 		require.Len(t, providerOpts, 1)
 		o := providerOpts[0]
-		require.Len(t, o.Middlewares, 6)
+		require.Len(t, o.Middlewares, 7)
 		require.Equal(t, TracingMiddlewareName, o.Middlewares[0].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, DataSourceMetricsMiddlewareName, o.Middlewares[1].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, SetUserAgentMiddlewareName, o.Middlewares[2].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, sdkhttpclient.BasicAuthenticationMiddlewareName, o.Middlewares[3].(sdkhttpclient.MiddlewareName).MiddlewareName())
 		require.Equal(t, sdkhttpclient.CustomHeadersMiddlewareName, o.Middlewares[4].(sdkhttpclient.MiddlewareName).MiddlewareName())
-		require.Equal(t, SigV4MiddlewareName, o.Middlewares[5].(sdkhttpclient.MiddlewareName).MiddlewareName())
+		require.Equal(t, ResponseLimitMiddlewareName, o.Middlewares[5].(sdkhttpclient.MiddlewareName).MiddlewareName())
+		require.Equal(t, SigV4MiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName())
 	})
 }

pkg/infra/httpclient/httpclientprovider/response_limit_middleware.go

@@ -0,0 +1,28 @@
+package httpclientprovider
+
+import (
+	"net/http"
+
+	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/infra/httpclient"
+)
+
+// ResponseLimitMiddlewareName is the middleware name used by ResponseLimitMiddleware.
+const ResponseLimitMiddlewareName = "response-limit"
+
+func ResponseLimitMiddleware(limit int64) sdkhttpclient.Middleware {
+	return sdkhttpclient.NamedMiddlewareFunc(ResponseLimitMiddlewareName, func(opts sdkhttpclient.Options, next http.RoundTripper) http.RoundTripper {
+		if limit <= 0 {
+			return next
+		}
+		return sdkhttpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+			res, err := next.RoundTrip(req)
+			if err != nil {
+				return nil, err
+			}
+
+			res.Body = httpclient.MaxBytesReader(res.Body, limit)
+			return res, nil
+		})
+	})
+}

pkg/infra/httpclient/httpclientprovider/response_limit_middleware_test.go

@@ -0,0 +1,60 @@
+package httpclientprovider
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResponseLimitMiddleware(t *testing.T) {
+	tcs := []struct {
+		limit      int64
+		bodyLength int
+		body       string
+		err        error
+	}{
+		{limit: 1, bodyLength: 1, body: "d", err: errors.New("error: http: response body too large, response limit is set to: 1")},
+		{limit: 1000000, bodyLength: 5, body: "dummy", err: nil},
+		{limit: 0, bodyLength: 5, body: "dummy", err: nil},
+	}
+	for _, tc := range tcs {
+		t.Run(fmt.Sprintf("Test ResponseLimitMiddleware with limit: %d", tc.limit), func(t *testing.T) {
+			finalRoundTripper := httpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+				return &http.Response{StatusCode: http.StatusOK, Request: req, Body: ioutil.NopCloser(strings.NewReader("dummy"))}, nil
+			})
+
+			mw := ResponseLimitMiddleware(tc.limit)
+			rt := mw.CreateMiddleware(httpclient.Options{}, finalRoundTripper)
+			require.NotNil(t, rt)
+			middlewareName, ok := mw.(httpclient.MiddlewareName)
+			require.True(t, ok)
+			require.Equal(t, ResponseLimitMiddlewareName, middlewareName.MiddlewareName())
+
+			ctx := context.Background()
+			req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://test.com/query", nil)
+			require.NoError(t, err)
+			res, err := rt.RoundTrip(req)
+			require.NoError(t, err)
+			require.NotNil(t, res)
+			require.NotNil(t, res.Body)
+			require.NoError(t, res.Body.Close())
+
+			bodyBytes, err := ioutil.ReadAll(res.Body)
+			if err != nil {
+				require.EqualError(t, tc.err, err.Error())
+			} else {
+				require.NoError(t, tc.err)
+			}
+
+			require.Len(t, bodyBytes, tc.bodyLength)
+			require.Equal(t, string(bodyBytes), tc.body)
+		})
+	}
+}

pkg/infra/httpclient/max_bytes_reader.go

@@ -0,0 +1,66 @@
+package httpclient
+
+import (
+	"errors"
+	"fmt"
+	"io"
+)
+
+// Similar implementation to http/net MaxBytesReader
+// https://pkg.go.dev/net/http#MaxBytesReader
+// What's happening differently here, is that the field that
+// is limited is the response and not the request, thus
+// the error handling/message needed to be accurate.
+
+// ErrResponseBodyTooLarge indicates response body is too large
+var ErrResponseBodyTooLarge = errors.New("http: response body too large")
+
+// MaxBytesReader is similar to io.LimitReader but is intended for
+// limiting the size of incoming request bodies. In contrast to
+// io.LimitReader, MaxBytesReader's result is a ReadCloser, returns a
+// non-EOF error for a Read beyond the limit, and closes the
+// underlying reader when its Close method is called.
+//
+// MaxBytesReader prevents clients from accidentally or maliciously
+// sending a large request and wasting server resources.
+func MaxBytesReader(r io.ReadCloser, n int64) io.ReadCloser {
+	return &maxBytesReader{r: r, n: n}
+}
+
+type maxBytesReader struct {
+	r   io.ReadCloser // underlying reader
+	n   int64         // max bytes remaining
+	err error         // sticky error
+}
+
+func (l *maxBytesReader) Read(p []byte) (n int, err error) {
+	if l.err != nil {
+		return 0, l.err
+	}
+	if len(p) == 0 {
+		return 0, nil
+	}
+	// If they asked for a 32KB byte read but only 5 bytes are
+	// remaining, no need to read 32KB. 6 bytes will answer the
+	// question of the whether we hit the limit or go past it.
+	if int64(len(p)) > l.n+1 {
+		p = p[:l.n+1]
+	}
+	n, err = l.r.Read(p)
+
+	if int64(n) <= l.n {
+		l.n -= int64(n)
+		l.err = err
+		return n, err
+	}
+
+	n = int(l.n)
+	l.n = 0
+
+	l.err = fmt.Errorf("error: %w, response limit is set to: %d", ErrResponseBodyTooLarge, n)
+	return n, l.err
+}
+
+func (l *maxBytesReader) Close() error {
+	return l.r.Close()
+}

pkg/infra/httpclient/max_bytes_reader_test.go

@@ -0,0 +1,40 @@
+package httpclient
+
+import (
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMaxBytesReader(t *testing.T) {
+	tcs := []struct {
+		limit      int64
+		bodyLength int
+		body       string
+		err        error
+	}{
+		{limit: 1, bodyLength: 1, body: "d", err: errors.New("error: http: response body too large, response limit is set to: 1")},
+		{limit: 1000000, bodyLength: 5, body: "dummy", err: nil},
+		{limit: 0, bodyLength: 0, body: "", err: errors.New("error: http: response body too large, response limit is set to: 0")},
+	}
+	for _, tc := range tcs {
+		t.Run(fmt.Sprintf("Test MaxBytesReader with limit: %d", tc.limit), func(t *testing.T) {
+			body := ioutil.NopCloser(strings.NewReader("dummy"))
+			readCloser := MaxBytesReader(body, tc.limit)
+
+			bodyBytes, err := ioutil.ReadAll(readCloser)
+			if err != nil {
+				require.EqualError(t, tc.err, err.Error())
+			} else {
+				require.NoError(t, tc.err)
+			}
+
+			require.Len(t, bodyBytes, tc.bodyLength)
+			require.Equal(t, string(bodyBytes), tc.body)
+		})
+	}
+}

pkg/setting/setting.go

@@ -326,6 +326,7 @@ type Cfg struct {
 	DataProxyMaxIdleConns          int
 	DataProxyKeepAlive             int
 	DataProxyIdleConnTimeout       int
+	ResponseLimit                  int64
 
 	// DistributedCache
 	RemoteCacheOptions *RemoteCacheOptions

pkg/setting/setting_data_proxy.go

@@ -14,6 +14,7 @@ func readDataProxySettings(iniFile *ini.File, cfg *Cfg) error {
 	cfg.DataProxyMaxConnsPerHost = dataproxy.Key("max_conns_per_host").MustInt(0)
 	cfg.DataProxyMaxIdleConns = dataproxy.Key("max_idle_connections").MustInt()
 	cfg.DataProxyIdleConnTimeout = dataproxy.Key("idle_conn_timeout_seconds").MustInt(90)
+	cfg.ResponseLimit = dataproxy.Key("response_limit").MustInt64(0)
 
 	if val, err := dataproxy.Key("max_idle_connections_per_host").Int(); err == nil {
 		cfg.Logger.Warn("[Deprecated] the configuration setting 'max_idle_connections_per_host' is deprecated, please use 'max_idle_connections' instead")