diff --git a/pkg/queryfrontend/protection.go b/pkg/queryfrontend/protection.go index b95f67fa6aa..3ca3b3067db 100644 --- a/pkg/queryfrontend/protection.go +++ b/pkg/queryfrontend/protection.go @@ -5,6 +5,7 @@ package queryfrontend import ( "context" + "regexp" ) // RuleAction represents the action to take when a protection rule is triggered. @@ -50,3 +51,38 @@ func RuleActionToString(action RuleAction) string { return "Unknown" } } + +// thanosQueryReq wraps a query request with actor information. +type thanosQueryReq struct { + actor string +} + +// Protection is the interface that all protection implementations must satisfy. +// It only determines whether a query matches — the action (log/block) is configured in Rule. +type Protection interface { + Name() string + Run(ctx context.Context, req thanosQueryReq) (bool, error) +} + +// ProtectionFactory constructs a Protection from a map of args parsed from config. +type ProtectionFactory func(args map[string]string) (Protection, error) + +// Rule combines a Protection with its configured action, actor filter, and enabled state. +type Rule struct { + name string + protection Protection + action RuleAction + actorRegex *regexp.Regexp + enabled bool +} + +// NewRule creates a Rule. actorRegex may be nil to match all actors. +func NewRule(name string, protection Protection, action RuleAction, actorRegex *regexp.Regexp, enabled bool) *Rule { + return &Rule{ + name: name, + protection: protection, + action: action, + actorRegex: actorRegex, + enabled: enabled, + } +} diff --git a/pkg/queryfrontend/protection_engine.go b/pkg/queryfrontend/protection_engine.go new file mode 100644 index 00000000000..a6b3b012061 --- /dev/null +++ b/pkg/queryfrontend/protection_engine.go @@ -0,0 +1,61 @@ +// Copyright (c) The Thanos Authors. +// Licensed under the Apache License 2.0. + +package queryfrontend + +import ( + "context" + "sync" +) + +// ProtectionEngine evaluates a list of rules against a query request. +type ProtectionEngine struct { + mu sync.Mutex + rules []*Rule +} + +// NewProtectionEngine creates a new ProtectionEngine with the given rules. +func NewProtectionEngine(rules []*Rule) *ProtectionEngine { + return &ProtectionEngine{rules: rules} +} + +// UpdateRules replaces the current rule set. Safe for concurrent use. +func (e *ProtectionEngine) UpdateRules(rules []*Rule) { + e.mu.Lock() + defer e.mu.Unlock() + e.rules = rules +} + +// Evaluate runs all applicable rules against the request. +// Rules are evaluated in order; the first matching rule wins. +// Returns the updated context (with ProtectionResult if a rule triggered), +// the action to take, and any error. +func (e *ProtectionEngine) Evaluate(ctx context.Context, req thanosQueryReq) (*ProtectionResult, error) { + e.mu.Lock() + rules := e.rules + e.mu.Unlock() + + for _, rule := range rules { + // Skip disabled rules. + if !rule.enabled { + continue + } + + // Check actor filter. + if rule.actorRegex != nil && !rule.actorRegex.MatchString(req.actor) { + continue + } + + matched, err := rule.protection.Run(ctx, req) + if !matched { + continue + } + + return &ProtectionResult{ + Triggered: true, + RuleName: rule.name, + Action: rule.action, + }, err + } + return nil, nil +} diff --git a/pkg/queryfrontend/protection_engine_test.go b/pkg/queryfrontend/protection_engine_test.go new file mode 100644 index 00000000000..5383ddd07bb --- /dev/null +++ b/pkg/queryfrontend/protection_engine_test.go @@ -0,0 +1,106 @@ +// Copyright (c) The Thanos Authors. +// Licensed under the Apache License 2.0. + +// Tests for ProtectionEngine.Evaluate: verifies that rules are correctly +// evaluated, filtered, and prioritized against query requests. +package queryfrontend + +import ( + "context" + "regexp" + "testing" + + "github.com/pkg/errors" + "github.com/stretchr/testify/require" +) + +// neverMatchProtection never matches (returns false). +type neverMatchProtection struct{} + +func (p *neverMatchProtection) Name() string { return "never" } +func (p *neverMatchProtection) Run(_ context.Context, _ thanosQueryReq) (bool, error) { + return false, nil +} + +// errorProtection always matches and returns an error. +type errorProtection struct{} + +func (p *errorProtection) Name() string { return "error" } +func (p *errorProtection) Run(_ context.Context, _ thanosQueryReq) (bool, error) { + return true, errors.New("protection error") +} + +func TestProtectionEngine_NoRules(t *testing.T) { + engine := NewProtectionEngine(nil) + protectionResult, err := engine.Evaluate(context.Background(), thanosQueryReq{}) + require.NoError(t, err) + require.Nil(t, protectionResult) +} + +func TestProtectionEngine_DisabledRuleSkipped(t *testing.T) { + engine := NewProtectionEngine([]*Rule{ + NewRule("disabled", &AlwaysMatchProtection{}, RuleActionBlock, nil, false), + }) + protectionResult, err := engine.Evaluate(context.Background(), thanosQueryReq{}) + require.NoError(t, err) + require.Nil(t, protectionResult) +} + +func TestProtectionEngine_ActorRegexNoMatch(t *testing.T) { + engine := NewProtectionEngine([]*Rule{ + NewRule("filtered", &AlwaysMatchProtection{}, RuleActionBlock, regexp.MustCompile("^admin$"), true), + }) + protectionResult, err := engine.Evaluate(context.Background(), thanosQueryReq{actor: "user"}) + require.NoError(t, err) + require.Nil(t, protectionResult) +} + +func TestProtectionEngine_ActorRegexMatch(t *testing.T) { + engine := NewProtectionEngine([]*Rule{ + NewRule("filtered", &AlwaysMatchProtection{}, RuleActionBlock, regexp.MustCompile("^admin$"), true), + }) + protectionResult, err := engine.Evaluate(context.Background(), thanosQueryReq{actor: "admin"}) + require.NoError(t, err) + require.Equal(t, RuleActionBlock, protectionResult.Action) + require.True(t, protectionResult.Triggered) + require.Equal(t, "filtered", protectionResult.RuleName) +} + +func TestProtectionEngine_FirstMatchingRuleWins(t *testing.T) { + engine := NewProtectionEngine([]*Rule{ + NewRule("first", &neverMatchProtection{}, RuleActionBlock, nil, true), + NewRule("second", &AlwaysMatchProtection{}, RuleActionLog, nil, true), + NewRule("third", &AlwaysMatchProtection{}, RuleActionBlock, nil, true), + }) + protectionResult, err := engine.Evaluate(context.Background(), thanosQueryReq{}) + require.NoError(t, err) + require.Equal(t, RuleActionLog, protectionResult.Action) + require.Equal(t, true, protectionResult.Triggered) + require.Equal(t, "second", protectionResult.RuleName) +} + +func TestProtectionEngine_RunError(t *testing.T) { + engine := NewProtectionEngine([]*Rule{ + NewRule("error-rule", &errorProtection{}, RuleActionLog, nil, true), + }) + _, err := engine.Evaluate(context.Background(), thanosQueryReq{}) + require.Error(t, err) + require.Contains(t, err.Error(), "protection error") +} + +func TestProtectionEngine_UpdateRules(t *testing.T) { + engine := NewProtectionEngine([]*Rule{ + NewRule("block-all", &AlwaysMatchProtection{}, RuleActionBlock, nil, true), + }) + + protectionResult, err := engine.Evaluate(context.Background(), thanosQueryReq{}) + require.NoError(t, err) + require.Equal(t, RuleActionBlock, protectionResult.Action) + require.Equal(t, true, protectionResult.Triggered) + + engine.UpdateRules(nil) + + protectionResult, err = engine.Evaluate(context.Background(), thanosQueryReq{}) + require.NoError(t, err) + require.Nil(t, protectionResult) +} diff --git a/pkg/queryfrontend/protection_impls.go b/pkg/queryfrontend/protection_impls.go new file mode 100644 index 00000000000..a012238ee31 --- /dev/null +++ b/pkg/queryfrontend/protection_impls.go @@ -0,0 +1,38 @@ +// Copyright (c) The Thanos Authors. +// Licensed under the Apache License 2.0. + +package queryfrontend + +import ( + "context" + + "github.com/pkg/errors" +) + +// This file contains all protection rule implementations. +// Each protection implements the Protection interface defined in protection.go. + +// AlwaysMatchProtection is a protection that always matches every query. +type AlwaysMatchProtection struct{} + +func (n *AlwaysMatchProtection) Name() string { return "always-match" } + +func (n *AlwaysMatchProtection) Run(_ context.Context, _ thanosQueryReq) (bool, error) { + return true, nil +} + +// protectionRegistry maps protection names (as used in config) to their factory functions. +var protectionRegistry = map[string]ProtectionFactory{ + "always-match": func(_ map[string]string) (Protection, error) { + return &AlwaysMatchProtection{}, nil + }, +} + +// LookupProtection returns the factory for the given protection name. +func LookupProtection(name string) (ProtectionFactory, error) { + factory, ok := protectionRegistry[name] + if !ok { + return nil, errors.Errorf("unknown protection %q", name) + } + return factory, nil +} diff --git a/pkg/queryfrontend/protection_impls_test.go b/pkg/queryfrontend/protection_impls_test.go new file mode 100644 index 00000000000..32d87e4def0 --- /dev/null +++ b/pkg/queryfrontend/protection_impls_test.go @@ -0,0 +1,34 @@ +// Copyright (c) The Thanos Authors. +// Licensed under the Apache License 2.0. + +// Tests for protection implementations in protection_impls.go. +package queryfrontend + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAlwaysMatchProtection_AlwaysMatches(t *testing.T) { + p := &AlwaysMatchProtection{} + matched, err := p.Run(context.Background(), thanosQueryReq{}) + require.NoError(t, err) + require.True(t, matched) +} + +func TestLookupProtection_AlwaysMatch(t *testing.T) { + factory, err := LookupProtection("always-match") + require.NoError(t, err) + require.NotNil(t, factory) + + p, err := factory(nil) + require.NoError(t, err) + require.Equal(t, "always-match", p.Name()) +} + +func TestLookupProtection_Unknown(t *testing.T) { + _, err := LookupProtection("unknown") + require.Error(t, err) +}