ITWorx-Pulse-Public/internal/promqlbinding/binding.go

package promqlbinding

import (
	"context"
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/itworx/pulse/internal/auth"
	"github.com/itworx/pulse/internal/datasource"
	"github.com/itworx/pulse/internal/metriccatalog"
	"github.com/itworx/pulse/internal/queryplan"
)

const MaxCompiledQueryLength = 4096

var placeholderPattern = regexp.MustCompile(`\{\{([a-zA-Z_][a-zA-Z0-9_]*)\}\}`)
var durationPattern = regexp.MustCompile(`^(?:[0-9]+(?:ms|s|m|h|d|w|y))+$`)
var labelValuePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}$`)
var rawQueryLimit = 4096

type Error struct {
	Code   string
	Field  string
	Detail string
}

func (e Error) Error() string {
	if e.Field == "" {
		return e.Code + ": " + e.Detail
	}
	return e.Code + " (" + e.Field + "): " + e.Detail
}

func Compile(definition metriccatalog.Definition, values map[string]string) (string, error) {
	if err := definition.Validate(); err != nil {
		return "", fmt.Errorf("invalid metric binding: %w", err)
	}
	matches := placeholderPattern.FindAllStringSubmatchIndex(definition.QueryTemplate, -1)
	known := make(map[string]struct{}, len(matches))
	for _, match := range matches {
		name := definition.QueryTemplate[match[2]:match[3]]
		known[name] = struct{}{}
		if _, ok := values[name]; !ok {
			return "", Error{Code: "PROMQL_BINDING_VALUE_REQUIRED", Field: name, Detail: "template value is required"}
		}
		if name == "window" {
			if !durationPattern.MatchString(values[name]) {
				return "", Error{Code: "PROMQL_WINDOW_INVALID", Field: name, Detail: "window must be a Prometheus duration"}
			}
		} else {
			if !contains(definition.AllowedLabels, name) {
				return "", Error{Code: "PROMQL_LABEL_NOT_ALLOWED", Field: name, Detail: "template label is not in the metric contract"}
			}
			if !labelValuePattern.MatchString(values[name]) {
				return "", Error{Code: "PROMQL_LABEL_VALUE_INVALID", Field: name, Detail: "label value contains unsupported characters"}
			}
		}
	}
	for name := range values {
		if _, ok := known[name]; !ok {
			return "", Error{Code: "PROMQL_BINDING_VALUE_UNKNOWN", Field: name, Detail: "binding value is not used by the approved template"}
		}
	}
	var builder strings.Builder
	last := 0
	for _, match := range matches {
		builder.WriteString(definition.QueryTemplate[last:match[0]])
		name := definition.QueryTemplate[match[2]:match[3]]
		value := values[name]
		if name != "window" {
			value = strconv.Quote(value)
		}
		builder.WriteString(value)
		last = match[1]
	}
	builder.WriteString(definition.QueryTemplate[last:])
	query := builder.String()
	if len(query) == 0 || len(query) > MaxCompiledQueryLength || strings.Contains(query, "{{") {
		return "", Error{Code: "PROMQL_QUERY_LIMIT", Field: "queryTemplate", Detail: "compiled query exceeds bounds"}
	}
	return query, nil
}

func CompilePlan(plan queryplan.Plan) (string, error) {
	values := make(map[string]string, len(plan.ResolvedScope)+1)
	for label, value := range plan.ResolvedScope {
		values[label] = value
	}
	// Gauges such as memory utilization have no range selector and therefore no
	// approved {{window}} placeholder. Compile rejects every surplus binding as
	// a safety boundary, so only supply the derived duration when the catalog
	// template explicitly consumes it.
	if strings.Contains(plan.Metric.QueryTemplate, "{{window}}") {
		values["window"] = WindowForStep(plan.Request.Range.StepSeconds)
	}
	return Compile(plan.Metric, values)
}
func WindowForStep(stepSeconds int) string {
	// A range-vector needs at least two source samples. The dashboard step is
	// allowed to be 15 seconds, equal to a common Prometheus scrape interval;
	// using that value verbatim makes rate()/increase() intermittently empty.
	// Keep graph resolution independent and use a conservative bounded lookback.
	if stepSeconds < 60 {
		stepSeconds = 60
	}
	if stepSeconds%3600 == 0 {
		return fmt.Sprintf("%dh", stepSeconds/3600)
	}
	if stepSeconds%60 == 0 {
		return fmt.Sprintf("%dm", stepSeconds/60)
	}
	return fmt.Sprintf("%ds", stepSeconds)
}

type Binding struct {
	Definition           metriccatalog.Definition
	RequiredCapabilities []string
	Priority             int
	ID                   string
}

func SelectBinding(candidates []Binding, capabilities datasource.CapabilitySet) (Binding, error) {
	if len(candidates) == 0 {
		return Binding{}, Error{Code: "PROMQL_BINDING_UNAVAILABLE", Detail: "no approved binding exists"}
	}
	ordered := append([]Binding(nil), candidates...)
	sort.SliceStable(ordered, func(i, j int) bool {
		if ordered[i].Priority != ordered[j].Priority {
			return ordered[i].Priority > ordered[j].Priority
		}
		if ordered[i].Definition.Version != ordered[j].Definition.Version {
			return ordered[i].Definition.Version > ordered[j].Definition.Version
		}
		return ordered[i].ID < ordered[j].ID
	})
	for _, candidate := range ordered {
		if candidate.Definition.Validate() != nil {
			continue
		}
		if allCapabilitiesEnabled(candidate.RequiredCapabilities, capabilities) {
			return candidate, nil
		}
	}
	return ordered[0], Error{Code: "PROMQL_BINDING_UNAVAILABLE", Detail: "no binding has all required capabilities"}
}
func allCapabilitiesEnabled(required []string, capabilities datasource.CapabilitySet) bool {
	for _, requiredID := range required {
		capability, ok := capabilities.Find(requiredID)
		if !ok || capability.State != datasource.CapabilityEnabled {
			return false
		}
	}
	return true
}

func ValidateRawQuery(ctx context.Context, query string) (string, error) {
	if err := ctx.Err(); err != nil {
		return "", err
	}
	principal, ok := auth.PrincipalFromContext(ctx)
	if !ok || !auth.Allows(principal.Role, auth.PermissionAdmin) {
		return "", Error{Code: "PROMQL_RAW_UNAUTHORIZED", Detail: "raw PromQL requires administrator permission"}
	}
	query = strings.TrimSpace(query)
	if query == "" || len(query) > rawQueryLimit || strings.ContainsRune(query, '\x00') {
		return "", Error{Code: "PROMQL_RAW_INVALID", Detail: "raw PromQL is empty or exceeds bounds"}
	}
	return query, nil
}
func contains(values []string, target string) bool {
	for _, value := range values {
		if value == target {
			return true
		}
	}
	return false
}