Add query-frontend limit for max length of query expression

CHANGELOG.md

@@ -4,6 +4,7 @@
 * [CHANGE] slo: include request cancellations within SLO [#4355] (https://github.com/grafana/tempo/pull/4355) (@electron0zero)
   request cancellations are exposed under `result` label in `tempo_query_frontend_queries_total` and `tempo_query_frontend_queries_within_slo_total` with `completed` or `canceled` values to differentiate between completed and canceled requests.
 * [CHANGE] update default config values to better align with production workloads [#4340](https://github.com/grafana/tempo/pull/4340) (@electron0zero)
+* [CHANGE] Add query-frontend limit for max length of query expression [##4397](https://github.com/grafana/tempo/pull/4397) (@electron0zero)
 * [CHANGE] fix deprecation warning by switching to DoBatchWithOptions [#4343](https://github.com/grafana/tempo/pull/4343) (@dastrobu)
 * [CHANGE] **BREAKING CHANGE** The Tempo serverless is now deprecated and will be removed in an upcoming release [#4017](https://github.com/grafana/tempo/pull/4017/) @electron0zero
 * [CHANGE] **BREAKING CHANGE** Change the AWS Lambda serverless build tooling output from "main" to "bootstrap". Refer to https://aws.amazon.com/blogs/compute/migrating-aws-lambda-functions-from-the-go1-x-runtime-to-the-custom-runtime-on-amazon-linux-2/ for migration steps [#3852](https://github.com/grafana/tempo/pull/3852) (@zatlodan)

docs/sources/tempo/configuration/_index.md

@@ -593,6 +593,10 @@ query_frontend:
     # A list of regular expressions for refusing matching requests, these will apply for every request regardless of the endpoint.
     [url_deny_list: <list of strings> | default = <empty list>]]
 
+    # Max allowed TraceQL expression size, in bytes. queries bigger then this size will be rejected.
+    # (default: 128 KiB)
+    [max_query_expression_size_bytes: <int> | default = 131072]]
+
     search:
 
         # The number of concurrent jobs to execute when searching the backend.

modules/frontend/config.go

@@ -32,6 +32,9 @@ type Config struct {
 	// A list of regexes for black listing requests, these will apply for every request regardless the endpoint
 	URLDenyList []string `yaml:"url_deny_list,omitempty"`
 
+	// Maximum allowed size of the raw TraceQL Query expression in bytes
+	MaxQueryExpressionSizeBytes int `yaml:"max_query_expression_size_bytes,omitempty"`
+
 	// A list of headers allowed through the HTTP pipeline. Everything else will be stripped.
 	AllowedHeaders []string `yaml:"-"`
 }
@@ -105,6 +108,8 @@ func (cfg *Config) RegisterFlagsAndApplyDefaults(string, *flag.FlagSet) {
 		MaxTraceQLConditions: 4,
 	}
 
+	// set default max query size to 128 KiB, queries larger than this will be rejected
+	cfg.MaxQueryExpressionSizeBytes = 128 * 1024
 	// enable multi tenant queries by default
 	cfg.MultiTenantQueriesEnabled = true
 }

modules/frontend/frontend.go

@@ -95,7 +95,7 @@ func New(cfg Config, next pipeline.RoundTripper, o overrides.Interface, reader t
 	statusCodeWare := pipeline.NewStatusCodeAdjustWare()
 	traceIDStatusCodeWare := pipeline.NewStatusCodeAdjustWareWithAllowedCode(http.StatusNotFound)
 	urlDenyListWare := pipeline.NewURLDenyListWare(cfg.URLDenyList)
-	queryValidatorWare := pipeline.NewQueryValidatorWare()
+	queryValidatorWare := pipeline.NewQueryValidatorWare(cfg.MaxQueryExpressionSizeBytes)
 	headerStripWare := pipeline.NewStripHeadersWare(cfg.AllowedHeaders)
 
 	tracePipeline := pipeline.Build(

modules/frontend/pipeline/async_query_validator_middleware.go

@@ -9,13 +9,15 @@ import (
 )
 
 type queryValidatorWare struct {
-	next AsyncRoundTripper[combiner.PipelineResponse]
+	next              AsyncRoundTripper[combiner.PipelineResponse]
+	maxQuerySizeBytes int
 }
 
-func NewQueryValidatorWare() AsyncMiddleware[combiner.PipelineResponse] {
+func NewQueryValidatorWare(maxQuerySizeBytes int) AsyncMiddleware[combiner.PipelineResponse] {
 	return AsyncMiddlewareFunc[combiner.PipelineResponse](func(next AsyncRoundTripper[combiner.PipelineResponse]) AsyncRoundTripper[combiner.PipelineResponse] {
 		return &queryValidatorWare{
-			next: next,
+			next:              next,
+			maxQuerySizeBytes: maxQuerySizeBytes,
 		}
 	})
 }
@@ -45,6 +47,11 @@ func (c queryValidatorWare) validateTraceQLQuery(queryParams url.Values) error {
 		if err != nil {
 			return fmt.Errorf("invalid TraceQL query: %w", err)
 		}
+
+		// reject query if the query expression size exceeds the maximum allowed size
+		if len(traceQLQuery) > c.maxQuerySizeBytes {
+			return fmt.Errorf("TraceQL expression exceeds the configured maximum size of %d bytes, reduce the query expression size or contact your system administrator", c.maxQuerySizeBytes)
+		}
 	}
 	return nil
 }

modules/frontend/pipeline/async_query_validator_middleware_test.go

@@ -8,45 +8,80 @@ import (
 	"testing"
 
 	"github.com/grafana/tempo/modules/frontend/combiner"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 var nextFunc = AsyncRoundTripperFunc[combiner.PipelineResponse](func(_ Request) (Responses[combiner.PipelineResponse], error) {
 	return NewHTTPToAsyncResponse(&http.Response{
 		StatusCode: 200,
-		Body:       io.NopCloser(bytes.NewReader([]byte{})),
+		Body:       io.NopCloser(bytes.NewReader([]byte("foo"))),
 	}), nil
 })
 
 func TestQueryValidator(t *testing.T) {
-	roundTrip := NewQueryValidatorWare().Wrap(nextFunc)
-	statusCode := doRequest(t, "http://localhost:8080/api/search", roundTrip)
-	assert.Equal(t, 200, statusCode)
-}
+	tests := []struct {
+		name              string
+		url               string
+		statusCode        int
+		maxQuerySizeBytes int
+	}{
+		{
+			name:              "No Query",
+			url:               "http://localhost:8080/api/search",
+			statusCode:        200,
+			maxQuerySizeBytes: 1000,
+		},
+		{
+			name:              "Empty query value",
+			url:               "http://localhost:8080/api/search&q=",
+			statusCode:        200,
+			maxQuerySizeBytes: 1000,
+		},
+		{
+			name:              "Valid query",
+			url:               "http://localhost:8080/api/search&q={}",
+			statusCode:        200,
+			maxQuerySizeBytes: 1000,
+		},
+		{
+			name:              "Invalid TraceQL query",
+			url:               "http://localhost:8080/api/search?q={. hi}",
+			statusCode:        400,
+			maxQuerySizeBytes: 1000,
+		},
+		{
+			name:              "Invalid TraceQL query regex",
+			url:               "http://localhost:8080/api/search?query={span.a =~ \"[\"}",
+			statusCode:        400,
+			maxQuerySizeBytes: 1000,
+		},
+		{
+			name:              "TraceQL query smaller then max size",
+			url:               "http://localhost:8080/api/search?q={ resource.service.name=\"service\" }",
+			statusCode:        200,
+			maxQuerySizeBytes: 1000,
+		},
+		{
+			name:              "TraceQL query bigger then max size",
+			url:               "http://localhost:8080/api/search?q={ resource.service.name=\"service\" }",
+			statusCode:        400,
+			maxQuerySizeBytes: 10,
+		},
+	}
 
-func TestQueryValidatorForAValidQuery(t *testing.T) {
-	roundTrip := NewQueryValidatorWare().Wrap(nextFunc)
-	statusCode := doRequest(t, "http://localhost:8080/api/search&q={}", roundTrip)
-	assert.Equal(t, 200, statusCode)
-}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rt := NewQueryValidatorWare(tt.maxQuerySizeBytes).Wrap(nextFunc)
+			req, _ := http.NewRequest(http.MethodGet, tt.url, nil)
+			resp, _ := rt.RoundTrip(NewHTTPRequest(req))
 
-func TestQueryValidatorForAnInvalidTraceQLQuery(t *testing.T) {
-	roundTrip := NewQueryValidatorWare().Wrap(nextFunc)
-	statusCode := doRequest(t, "http://localhost:8080/api/search?q={. hi}", roundTrip)
-	assert.Equal(t, 400, statusCode)
-}
-
-func TestQueryValidatorForAnInvalidTraceQlQueryRegex(t *testing.T) {
-	roundTrip := NewQueryValidatorWare().Wrap(nextFunc)
-	statusCode := doRequest(t, "http://localhost:8080/api/search?query={span.a =~ \"[\"}", roundTrip)
-	assert.Equal(t, 400, statusCode)
-}
+			httpResponse, _, err := resp.Next(context.Background())
+			require.NoError(t, err)
+			body, err := io.ReadAll(httpResponse.HTTPResponse().Body)
+			require.NoError(t, err)
+			require.NotEmpty(t, body)
 
-func doRequest(t *testing.T, url string, rt AsyncRoundTripper[combiner.PipelineResponse]) int {
-	req, _ := http.NewRequest(http.MethodGet, url, nil)
-	resp, _ := rt.RoundTrip(NewHTTPRequest(req))
-	httpResponse, _, err := resp.Next(context.Background())
-	require.NoError(t, err)
-	return httpResponse.HTTPResponse().StatusCode
+			require.Equal(t, tt.statusCode, httpResponse.HTTPResponse().StatusCode)
+		})
+	}
 }