Skip to content

Commit

Permalink
chore: refactor PromQL enforcer
Browse files Browse the repository at this point in the history
This change renames the Enforcer struct to PromQLEnforcer to clarify its
purpose. It also removes the error type casting in favor of error
wrapping.

Signed-off-by: Simon Pasquier <[email protected]>
  • Loading branch information
simonpasquier committed Jun 7, 2024
1 parent cf44930 commit 40a5ab4
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 106 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ prom-label-proxy \
`prom-label-proxy` will enforce the `tenant=~"prometheus|alertmanager"` label selector in all requests.
You can match the label value using a regular expression with the `-regex-match` option. For example:
You can match the label value using a regular expression with the `-regex-match` option. For example:
```
prom-label-proxy \
Expand Down
45 changes: 32 additions & 13 deletions injectproxy/enforce.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,40 +14,59 @@
package injectproxy

import (
"errors"
"fmt"

"github.com/prometheus/prometheus/model/labels"
"github.com/prometheus/prometheus/promql/parser"
)

type Enforcer struct {
// PromQLEnforcer can enforce label matchers in PromQL expressions.
type PromQLEnforcer struct {
labelMatchers map[string]*labels.Matcher
errorOnReplace bool
}

func NewEnforcer(errorOnReplace bool, ms ...*labels.Matcher) *Enforcer {
func NewPromQLEnforcer(errorOnReplace bool, ms ...*labels.Matcher) *PromQLEnforcer {
entries := make(map[string]*labels.Matcher)

for _, matcher := range ms {
entries[matcher.Name] = matcher
}

return &Enforcer{
return &PromQLEnforcer{
labelMatchers: entries,
errorOnReplace: errorOnReplace,
}
}

type IllegalLabelMatcherError struct {
msg string
}
var (
// ErrQueryParse is returned when the input query is invalid.
ErrQueryParse = errors.New("failed to parse query string")

// ErrIllegalLabelMatcher is returned when the input query contains a conflicting label matcher.
ErrIllegalLabelMatcher = errors.New("conflicting label matcher")

// ErrEnforceLabel is returned when the label matchers couldn't be enforced.
ErrEnforceLabel = errors.New("failed to enforce label")
)

func (e IllegalLabelMatcherError) Error() string { return e.msg }
// Enforce the label matchers in a PromQL expression.
func (ms *PromQLEnforcer) Enforce(q string) (string, error) {
expr, err := parser.ParseExpr(q)
if err != nil {
return "", fmt.Errorf("%w: %w", ErrQueryParse, err)
}

if err := ms.EnforceNode(expr); err != nil {
if errors.Is(err, ErrIllegalLabelMatcher) {
return "", err
}

func newIllegalLabelMatcherError(existing string, replacement string) IllegalLabelMatcherError {
return IllegalLabelMatcherError{
msg: fmt.Sprintf("label matcher value (%s) conflicts with injected value (%s)", existing, replacement),
return "", fmt.Errorf("%w: %w", ErrEnforceLabel, err)
}

return expr.String(), nil
}

// EnforceNode walks the given node recursively
Expand All @@ -57,7 +76,7 @@ func newIllegalLabelMatcherError(existing string, replacement string) IllegalLab
// their label enforcer is being potentially modified.
// If a node's label matcher has the same name as a label matcher
// of the given enforcer, then it will be replaced.
func (ms Enforcer) EnforceNode(node parser.Node) error {
func (ms PromQLEnforcer) EnforceNode(node parser.Node) error {
switch n := node.(type) {
case *parser.EvalStmt:
if err := ms.EnforceNode(n.Expr); err != nil {
Expand Down Expand Up @@ -140,14 +159,14 @@ func (ms Enforcer) EnforceNode(node parser.Node) error {
// * if errorOnReplace is true, an error is returned,
// * if errorOnReplace is false and the label matcher type is '=', the existing matcher is silently replaced.
// * otherwise the existing matcher is preserved.
func (ms Enforcer) EnforceMatchers(targets []*labels.Matcher) ([]*labels.Matcher, error) {
func (ms PromQLEnforcer) EnforceMatchers(targets []*labels.Matcher) ([]*labels.Matcher, error) {
var res []*labels.Matcher

for _, target := range targets {
if matcher, ok := ms.labelMatchers[target.Name]; ok {
// matcher.String() returns something like "labelfoo=value"
if ms.errorOnReplace && matcher.String() != target.String() {
return res, newIllegalLabelMatcherError(matcher.String(), target.String())
return res, fmt.Errorf("%w: label matcher value %q conflicts with injected value %q", ErrIllegalLabelMatcher, matcher.String(), target.String())
}

// Drop the existing matcher only if the enforced matcher is an
Expand Down
Loading

0 comments on commit 40a5ab4

Please sign in to comment.