diff --git a/pkg/infra/httpclient/httpclientprovider/http_client_provider.go b/pkg/infra/httpclient/httpclientprovider/http_client_provider.go index 69a9f2a1c73..f5c515a527e 100644 --- a/pkg/infra/httpclient/httpclientprovider/http_client_provider.go +++ b/pkg/infra/httpclient/httpclientprovider/http_client_provider.go @@ -27,7 +27,7 @@ func New(cfg *setting.Cfg, validator validations.PluginRequestValidator, tracer SetUserAgentMiddleware(cfg.DataProxyUserAgent), sdkhttpclient.BasicAuthenticationMiddleware(), sdkhttpclient.CustomHeadersMiddleware(), - ResponseLimitMiddleware(cfg.ResponseLimit), + sdkhttpclient.ResponseLimitMiddleware(cfg.ResponseLimit), RedirectLimitMiddleware(validator), } diff --git a/pkg/infra/httpclient/httpclientprovider/http_client_provider_test.go b/pkg/infra/httpclient/httpclientprovider/http_client_provider_test.go index 208c51fff21..dfbb09826c1 100644 --- a/pkg/infra/httpclient/httpclientprovider/http_client_provider_test.go +++ b/pkg/infra/httpclient/httpclientprovider/http_client_provider_test.go @@ -33,7 +33,7 @@ func TestHTTPClientProvider(t *testing.T) { require.Equal(t, SetUserAgentMiddlewareName, o.Middlewares[3].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, sdkhttpclient.BasicAuthenticationMiddlewareName, o.Middlewares[4].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, sdkhttpclient.CustomHeadersMiddlewareName, o.Middlewares[5].(sdkhttpclient.MiddlewareName).MiddlewareName()) - require.Equal(t, ResponseLimitMiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName()) + require.Equal(t, sdkhttpclient.ResponseLimitMiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName()) }) t.Run("When creating new provider and SigV4 is enabled should apply expected middleware", func(t *testing.T) { @@ -57,7 +57,7 @@ func TestHTTPClientProvider(t *testing.T) { require.Equal(t, SetUserAgentMiddlewareName, o.Middlewares[3].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, sdkhttpclient.BasicAuthenticationMiddlewareName, o.Middlewares[4].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, sdkhttpclient.CustomHeadersMiddlewareName, o.Middlewares[5].(sdkhttpclient.MiddlewareName).MiddlewareName()) - require.Equal(t, ResponseLimitMiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName()) + require.Equal(t, sdkhttpclient.ResponseLimitMiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, SigV4MiddlewareName, o.Middlewares[8].(sdkhttpclient.MiddlewareName).MiddlewareName()) }) @@ -82,7 +82,7 @@ func TestHTTPClientProvider(t *testing.T) { require.Equal(t, SetUserAgentMiddlewareName, o.Middlewares[3].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, sdkhttpclient.BasicAuthenticationMiddlewareName, o.Middlewares[4].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, sdkhttpclient.CustomHeadersMiddlewareName, o.Middlewares[5].(sdkhttpclient.MiddlewareName).MiddlewareName()) - require.Equal(t, ResponseLimitMiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName()) + require.Equal(t, sdkhttpclient.ResponseLimitMiddlewareName, o.Middlewares[6].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, HostRedirectValidationMiddlewareName, o.Middlewares[7].(sdkhttpclient.MiddlewareName).MiddlewareName()) require.Equal(t, HTTPLoggerMiddlewareName, o.Middlewares[8].(sdkhttpclient.MiddlewareName).MiddlewareName()) }) diff --git a/pkg/infra/httpclient/httpclientprovider/response_limit_middleware.go b/pkg/infra/httpclient/httpclientprovider/response_limit_middleware.go deleted file mode 100644 index 97d8285d4cb..00000000000 --- a/pkg/infra/httpclient/httpclientprovider/response_limit_middleware.go +++ /dev/null @@ -1,31 +0,0 @@ -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 - } - - if res != nil && res.StatusCode != http.StatusSwitchingProtocols { - res.Body = httpclient.MaxBytesReader(res.Body, limit) - } - - return res, nil - }) - }) -} diff --git a/pkg/infra/httpclient/httpclientprovider/response_limit_middleware_test.go b/pkg/infra/httpclient/httpclientprovider/response_limit_middleware_test.go deleted file mode 100644 index 88a0d572b86..00000000000 --- a/pkg/infra/httpclient/httpclientprovider/response_limit_middleware_test.go +++ /dev/null @@ -1,60 +0,0 @@ -package httpclientprovider - -import ( - "context" - "errors" - "fmt" - "io" - "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: io.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 := io.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) - }) - } -} diff --git a/pkg/infra/httpclient/max_bytes_reader.go b/pkg/infra/httpclient/max_bytes_reader.go deleted file mode 100644 index 9bbdce1e375..00000000000 --- a/pkg/infra/httpclient/max_bytes_reader.go +++ /dev/null @@ -1,66 +0,0 @@ -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() -} diff --git a/pkg/infra/httpclient/max_bytes_reader_test.go b/pkg/infra/httpclient/max_bytes_reader_test.go deleted file mode 100644 index 1f48be0bbc5..00000000000 --- a/pkg/infra/httpclient/max_bytes_reader_test.go +++ /dev/null @@ -1,40 +0,0 @@ -package httpclient - -import ( - "errors" - "fmt" - "io" - "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 := io.NopCloser(strings.NewReader("dummy")) - readCloser := MaxBytesReader(body, tc.limit) - - bodyBytes, err := io.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) - }) - } -} diff --git a/pkg/registry/apis/datasource/middleware.go b/pkg/registry/apis/datasource/middleware.go new file mode 100644 index 00000000000..fac100c8cac --- /dev/null +++ b/pkg/registry/apis/datasource/middleware.go @@ -0,0 +1,15 @@ +package datasource + +import ( + "context" + + "github.com/grafana/grafana-plugin-sdk-go/backend" + "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient" +) + +func contextualMiddlewares(ctx context.Context) context.Context { + cfg := backend.GrafanaConfigFromContext(ctx) + m := httpclient.ResponseLimitMiddleware(cfg.ResponseLimit()) + + return httpclient.WithContextualMiddleware(ctx, m) +} diff --git a/pkg/registry/apis/datasource/sub_health.go b/pkg/registry/apis/datasource/sub_health.go index b58a540916f..01a8ea82d9e 100644 --- a/pkg/registry/apis/datasource/sub_health.go +++ b/pkg/registry/apis/datasource/sub_health.go @@ -50,6 +50,7 @@ func (r *subHealthREST) Connect(ctx context.Context, name string, opts runtime.O return nil, err } ctx = backend.WithGrafanaConfig(ctx, pluginCtx.GrafanaConfig) + ctx = contextualMiddlewares(ctx) healthResponse, err := r.builder.client.CheckHealth(ctx, &backend.CheckHealthRequest{ PluginContext: pluginCtx, diff --git a/pkg/registry/apis/datasource/sub_query.go b/pkg/registry/apis/datasource/sub_query.go index a68f7e3de4f..1fdd151014e 100644 --- a/pkg/registry/apis/datasource/sub_query.go +++ b/pkg/registry/apis/datasource/sub_query.go @@ -72,6 +72,7 @@ func (r *subQueryREST) Connect(ctx context.Context, name string, opts runtime.Ob } ctx = backend.WithGrafanaConfig(ctx, pluginCtx.GrafanaConfig) + ctx = contextualMiddlewares(ctx) rsp, err := r.builder.client.QueryData(ctx, &backend.QueryDataRequest{ Queries: queries, PluginContext: pluginCtx, diff --git a/pkg/registry/apis/datasource/sub_resource.go b/pkg/registry/apis/datasource/sub_resource.go index 751ddd43214..f8f56baf92e 100644 --- a/pkg/registry/apis/datasource/sub_resource.go +++ b/pkg/registry/apis/datasource/sub_resource.go @@ -51,6 +51,7 @@ func (r *subResourceREST) Connect(ctx context.Context, name string, opts runtime return nil, err } ctx = backend.WithGrafanaConfig(ctx, pluginCtx.GrafanaConfig) + ctx = contextualMiddlewares(ctx) return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { body, err := io.ReadAll(req.Body) diff --git a/pkg/services/pluginsintegration/pluginconfig/config.go b/pkg/services/pluginsintegration/pluginconfig/config.go index 46a5f7b9360..7542cbe1f1b 100644 --- a/pkg/services/pluginsintegration/pluginconfig/config.go +++ b/pkg/services/pluginsintegration/pluginconfig/config.go @@ -63,6 +63,7 @@ type PluginInstanceCfg struct { GrafanaVersion string ConcurrentQueryCount int + ResponseLimit int64 UserFacingDefaultError string @@ -115,6 +116,7 @@ func ProvidePluginInstanceConfig(cfg *setting.Cfg, settingProvider setting.Provi SQLDatasourceMaxOpenConnsDefault: cfg.SqlDatasourceMaxOpenConnsDefault, SQLDatasourceMaxIdleConnsDefault: cfg.SqlDatasourceMaxIdleConnsDefault, SQLDatasourceMaxConnLifetimeDefault: cfg.SqlDatasourceMaxConnLifetimeDefault, + ResponseLimit: cfg.ResponseLimit, }, nil } diff --git a/pkg/services/pluginsintegration/pluginconfig/request.go b/pkg/services/pluginsintegration/pluginconfig/request.go index 92cdd8f1efd..0e969b8f40a 100644 --- a/pkg/services/pluginsintegration/pluginconfig/request.go +++ b/pkg/services/pluginsintegration/pluginconfig/request.go @@ -147,5 +147,9 @@ func (s *RequestConfigProvider) PluginRequestConfig(ctx context.Context, pluginI m[backend.SQLMaxIdleConnsDefault] = strconv.Itoa(s.cfg.SQLDatasourceMaxIdleConnsDefault) m[backend.SQLMaxConnLifetimeSecondsDefault] = strconv.Itoa(s.cfg.SQLDatasourceMaxConnLifetimeDefault) + if s.cfg.ResponseLimit > 0 { + m[backend.ResponseLimit] = strconv.FormatInt(s.cfg.ResponseLimit, 10) + } + return m }