131 lines
3.0 KiB
Go
131 lines
3.0 KiB
Go
// Copyright 2020 The Prometheus Authors
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package injectproxy
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/prometheus/prometheus/pkg/labels"
|
|
"github.com/prometheus/prometheus/promql/parser"
|
|
)
|
|
|
|
type Enforcer struct {
|
|
labelMatchers map[string]*labels.Matcher
|
|
}
|
|
|
|
func NewEnforcer(ms ...*labels.Matcher) *Enforcer {
|
|
entries := make(map[string]*labels.Matcher)
|
|
|
|
for _, matcher := range ms {
|
|
entries[matcher.Name] = matcher
|
|
}
|
|
|
|
return &Enforcer{
|
|
labelMatchers: entries,
|
|
}
|
|
}
|
|
|
|
// EnforceNode walks the given node recursively
|
|
// and enforces the given label enforcer on it.
|
|
//
|
|
// Whenever a parser.MatrixSelector or parser.VectorSelector AST node is found,
|
|
// 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 {
|
|
switch n := node.(type) {
|
|
case *parser.EvalStmt:
|
|
if err := ms.EnforceNode(n.Expr); err != nil {
|
|
return err
|
|
}
|
|
|
|
case parser.Expressions:
|
|
for _, e := range n {
|
|
if err := ms.EnforceNode(e); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
case *parser.AggregateExpr:
|
|
if err := ms.EnforceNode(n.Expr); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *parser.BinaryExpr:
|
|
if err := ms.EnforceNode(n.LHS); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := ms.EnforceNode(n.RHS); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *parser.Call:
|
|
if err := ms.EnforceNode(n.Args); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *parser.SubqueryExpr:
|
|
if err := ms.EnforceNode(n.Expr); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *parser.ParenExpr:
|
|
if err := ms.EnforceNode(n.Expr); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *parser.UnaryExpr:
|
|
if err := ms.EnforceNode(n.Expr); err != nil {
|
|
return err
|
|
}
|
|
|
|
case *parser.NumberLiteral, *parser.StringLiteral:
|
|
// nothing to do
|
|
|
|
case *parser.MatrixSelector:
|
|
// inject labelselector
|
|
if vs, ok := n.VectorSelector.(*parser.VectorSelector); ok {
|
|
vs.LabelMatchers = ms.enforceMatchers(vs.LabelMatchers)
|
|
}
|
|
|
|
case *parser.VectorSelector:
|
|
// inject labelselector
|
|
n.LabelMatchers = ms.enforceMatchers(n.LabelMatchers)
|
|
|
|
default:
|
|
panic(fmt.Errorf("parser.Walk: unhandled node type %T", n))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (ms Enforcer) enforceMatchers(targets []*labels.Matcher) []*labels.Matcher {
|
|
var res []*labels.Matcher
|
|
|
|
for _, target := range targets {
|
|
if _, ok := ms.labelMatchers[target.Name]; ok {
|
|
continue
|
|
}
|
|
|
|
res = append(res, target)
|
|
}
|
|
|
|
for _, enforcedMatcher := range ms.labelMatchers {
|
|
res = append(res, enforcedMatcher)
|
|
}
|
|
|
|
return res
|
|
}
|