diff --git a/pkg/queryfrontend/protection.go b/pkg/queryfrontend/protection.go index 3ca3b3067db..2577ef80a46 100644 --- a/pkg/queryfrontend/protection.go +++ b/pkg/queryfrontend/protection.go @@ -6,6 +6,9 @@ package queryfrontend import ( "context" "regexp" + + "github.com/prometheus/prometheus/promql/parser" + "github.com/thanos-io/thanos/internal/cortex/querier/queryrange" ) // RuleAction represents the action to take when a protection rule is triggered. @@ -52,9 +55,11 @@ func RuleActionToString(action RuleAction) string { } } -// thanosQueryReq wraps a query request with actor information. +// thanosQueryReq wraps a query request with parsed PromQL and actor information. type thanosQueryReq struct { - actor string + inner queryrange.Request + parsed parser.Expr + actor string } // Protection is the interface that all protection implementations must satisfy. diff --git a/pkg/queryfrontend/protection_middleware.go b/pkg/queryfrontend/protection_middleware.go new file mode 100644 index 00000000000..8f6e9ec3234 --- /dev/null +++ b/pkg/queryfrontend/protection_middleware.go @@ -0,0 +1,111 @@ +// Copyright (c) The Thanos Authors. +// Licensed under the Apache License 2.0. + +package queryfrontend + +import ( + "context" + "net/http" + "time" + + "github.com/go-kit/log" + "github.com/go-kit/log/level" + "github.com/pkg/errors" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "github.com/thanos-io/thanos/internal/cortex/querier/queryrange" + "github.com/thanos-io/thanos/pkg/extpromql" + "github.com/weaveworks/common/httpgrpc" +) + +type protectionMiddleware struct { + next queryrange.Handler + engine *ProtectionEngine + logger log.Logger + totalLatency prometheus.Histogram +} + +// NewProtectionMiddleware creates a new middleware that applies protection rules to queries. +func NewProtectionMiddleware(engine *ProtectionEngine, logger log.Logger, reg prometheus.Registerer) queryrange.Middleware { + totalLatency := promauto.With(reg).NewHistogram(prometheus.HistogramOpts{ + Namespace: "thanos", + Subsystem: "query_frontend", + Name: "protection_duration_seconds", + Help: "Total duration of the protection middleware, including parsing and evaluation.", + Buckets: []float64{0.0001, 0.00025, 0.0005, 0.001, 0.005, 0.01, 0.1, 1.0, 10.0}, + }) + + return queryrange.MiddlewareFunc(func(next queryrange.Handler) queryrange.Handler { + return &protectionMiddleware{ + next: next, + engine: engine, + logger: logger, + totalLatency: totalLatency, + } + }) +} + +func (m *protectionMiddleware) Do(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + query := r.GetQuery() + if query == "" { + return m.next.Do(ctx, r) + } + + totalStart := time.Now() + defer func() { m.totalLatency.Observe(time.Since(totalStart).Seconds()) }() + + // Parse PromQL. + parsed, err := extpromql.ParseExpr(query) + if err != nil { + // Malformed query: let downstream handle it, don't block. + level.Debug(m.logger).Log("msg", "protection middleware: failed to parse query, skipping", "err", err) + return m.next.Do(ctx, r) + } + + // Extract actor from X-Source header. + actor := extractActor(r) + + req := thanosQueryReq{ + inner: r, + parsed: parsed, + actor: actor, + } + + // Evaluate all protection rules. + protectionResult, err := m.engine.Evaluate(ctx, req) + + if err != nil { + return nil, errors.Wrap(err, "protection engine evaluation failed") + } + + if protectionResult != nil { + ctx = WithProtectionResult(ctx, protectionResult) + } + + // Return 400 Bad Request error code if a protection rule is triggered. + if protectionResult != nil && protectionResult.Action == RuleActionBlock { + return nil, httpgrpc.Errorf(http.StatusBadRequest, "query blocked by protection rule: %s", protectionResult.RuleName) + } + + return m.next.Do(ctx, r) +} + +// extractActor returns the actor identifier from the request headers. +// Uses X-Source header as the actor identifier. +func extractActor(r queryrange.Request) string { + headers := getHeaders(r) + userInfo := ExtractUserInfoFromHeaders(headers) + return userInfo.Source +} + +// getHeaders extracts headers from a queryrange.Request. +func getHeaders(r queryrange.Request) []*RequestHeader { + switch req := r.(type) { + case *ThanosQueryRangeRequest: + return req.Headers + case *ThanosQueryInstantRequest: + return req.Headers + default: + return nil + } +} diff --git a/pkg/queryfrontend/protection_middleware_test.go b/pkg/queryfrontend/protection_middleware_test.go new file mode 100644 index 00000000000..817f317218d --- /dev/null +++ b/pkg/queryfrontend/protection_middleware_test.go @@ -0,0 +1,167 @@ +// Copyright (c) The Thanos Authors. +// Licensed under the Apache License 2.0. + +// Tests for protectionMiddleware.Do: verifies routing behavior based on +// query validity, rule evaluation results, and actor filtering. +package queryfrontend + +import ( + "context" + "net/http" + "regexp" + "testing" + + "github.com/go-kit/log" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/require" + "github.com/thanos-io/thanos/internal/cortex/querier/queryrange" + "github.com/weaveworks/common/httpgrpc" +) + +// makeRequest builds a ThanosQueryRangeRequest with the given query and optional X-Source header. +func makeRequest(query, source string) *ThanosQueryRangeRequest { + r := &ThanosQueryRangeRequest{ + Path: "/api/v1/query_range", + Start: 0, + End: 3600000, + Step: 60000, + Query: query, + } + if source != "" { + r.Headers = []*RequestHeader{ + {Name: "x-source", Values: []string{source}}, + } + } + return r +} + +func newTestMiddleware(engine *ProtectionEngine) queryrange.Middleware { + return NewProtectionMiddleware(engine, log.NewNopLogger(), prometheus.NewRegistry()) +} + +func TestProtectionMiddleware_EmptyQuery_PassesThrough(t *testing.T) { + called := false + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + called = true + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("block-all", &alwaysMatchProtection{}, RuleActionBlock, nil, true), + })).Wrap(next) + + _, err := handler.Do(context.Background(), makeRequest("", "")) + require.NoError(t, err) + require.True(t, called, "next should be called for empty query") +} + +func TestProtectionMiddleware_InvalidPromQL_PassesThrough(t *testing.T) { + called := false + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + called = true + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("block-all", &alwaysMatchProtection{}, RuleActionBlock, nil, true), + })).Wrap(next) + + _, err := handler.Do(context.Background(), makeRequest("this is not valid promql!!!", "")) + require.NoError(t, err) + require.True(t, called, "next should be called when PromQL parsing fails") +} + +func TestProtectionMiddleware_RuleActionLog_PassesThrough(t *testing.T) { + called := false + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + called = true + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("log-all", &alwaysMatchProtection{}, RuleActionLog, nil, true), + })).Wrap(next) + + _, err := handler.Do(context.Background(), makeRequest("up", "")) + require.NoError(t, err) + require.True(t, called) +} + +func TestProtectionMiddleware_RuleActionBlock_Returns400(t *testing.T) { + called := false + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + called = true + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("block-all", &alwaysMatchProtection{}, RuleActionBlock, nil, true), + })).Wrap(next) + + _, err := handler.Do(context.Background(), makeRequest("up", "")) + require.Error(t, err) + require.False(t, called, "next should NOT be called when query is blocked") + + httpErr, ok := httpgrpc.HTTPResponseFromError(err) + require.True(t, ok) + require.Equal(t, int32(http.StatusBadRequest), httpErr.Code) +} + +func TestProtectionMiddleware_ProtectionResultInContext(t *testing.T) { + var capturedCtx context.Context + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + capturedCtx = ctx + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("log-rule", &alwaysMatchProtection{}, RuleActionLog, nil, true), + })).Wrap(next) + + _, err := handler.Do(context.Background(), makeRequest("up", "actor1")) + require.NoError(t, err) + + result := GetProtectionResult(capturedCtx) + require.NotNil(t, result) + require.True(t, result.Triggered) + require.Equal(t, "log-rule", result.RuleName) + require.Equal(t, RuleActionLog, result.Action) +} + +func TestProtectionMiddleware_EvaluationError_ReturnsError(t *testing.T) { + called := false + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + called = true + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("error-rule", &errorProtection{}, RuleActionLog, nil, true), + })).Wrap(next) + + _, err := handler.Do(context.Background(), makeRequest("up", "")) + require.Error(t, err) + require.False(t, called, "next should NOT be called when engine evaluation fails") +} + +func TestProtectionMiddleware_ActorFiltering(t *testing.T) { + called := 0 + next := queryrange.HandlerFunc(func(ctx context.Context, r queryrange.Request) (queryrange.Response, error) { + called++ + return &queryrange.PrometheusResponse{Status: "success"}, nil + }) + + handler := newTestMiddleware(NewProtectionEngine([]*Rule{ + NewRule("block-admin", &alwaysMatchProtection{}, RuleActionBlock, regexp.MustCompile("^admin$"), true), + })).Wrap(next) + + // Non-admin should pass through. + _, err := handler.Do(context.Background(), makeRequest("up", "user")) + require.NoError(t, err) + require.Equal(t, 1, called) + + // Admin should be blocked. + _, err = handler.Do(context.Background(), makeRequest("up", "admin")) + require.Error(t, err) + require.Equal(t, 1, called, "next should not be called for blocked admin request") +}