logs: re-introduce rules filtering for Prometheus API as this is still not supported by Loki (#805)

Author: JoaoBraveCoding

api/logs/v1/http.go

@@ -58,6 +58,7 @@ type handlerConfiguration struct {
 	registry              *prometheus.Registry
 	instrument            handlerInstrumenter
 	spanRoutePrefix       string
+	rulesLabelFilters     map[string][]string
 	readMiddlewares       []func(http.Handler) http.Handler
 	writeMiddlewares      []func(http.Handler) http.Handler
 	rulesReadMiddlewares  []func(http.Handler) http.Handler
@@ -131,6 +132,13 @@ func WithGlobalMiddleware(m ...func(http.Handler) http.Handler) HandlerOption {
 	}
 }
 
+// WithRulesLabelFilters adds the slice of rule labels filters to the handler configuration.
+func WithRulesLabelFilters(f map[string][]string) HandlerOption {
+	return func(h *handlerConfiguration) {
+		h.rulesLabelFilters = f
+	}
+}
+
 type handlerInstrumenter interface {
 	NewHandler(labels prometheus.Labels, handler http.Handler) http.HandlerFunc
 }
@@ -231,7 +239,7 @@ func NewHandler(read, tail, write, rules *url.URL, rulesReadOnly bool, tlsOption
 	}
 
 	if rules != nil {
-		var proxyRules http.Handler
+		var proxyRules, proxyPrometheusReadRules http.Handler
 		{
 			middlewares := proxy.Middlewares(
 				proxy.MiddlewareSetUpstream(rules),
@@ -240,18 +248,27 @@ func NewHandler(read, tail, write, rules *url.URL, rulesReadOnly bool, tlsOption
 				proxy.MiddlewareMetrics(c.registry, prometheus.Labels{"proxy": "logsv1-rules"}),
 			)
 
+			logger := proxy.Logger(c.logger)
 			t := &http.Transport{
 				DialContext: (&net.Dialer{
 					Timeout: dialTimeout,
 				}).DialContext,
 				TLSClientConfig: tlsOptions.NewClientConfig(),
 			}
+			transport := otelhttp.NewTransport(t)
 
+			proxyPrometheusReadRules = &httputil.ReverseProxy{
+				Director:       middlewares,
+				ErrorLog:       logger,
+				Transport:      transport,
+				ModifyResponse: newModifyResponseProm(c.logger, c.rulesLabelFilters),
+			}
 			proxyRules = &httputil.ReverseProxy{
 				Director:  middlewares,
-				ErrorLog:  proxy.Logger(c.logger),
-				Transport: otelhttp.NewTransport(t),
+				ErrorLog:  logger,
+				Transport: transport,
 			}
+
 		}
 		r.Group(func(r chi.Router) {
 			r.Use(c.readMiddlewares...)
@@ -270,11 +287,11 @@ func NewHandler(read, tail, write, rules *url.URL, rulesReadOnly bool, tlsOption
 			))
 			r.Get(prometheusRulesRoute, c.instrument.NewHandler(
 				prometheus.Labels{"group": "logsv1", "handler": "rules"},
-				otelhttp.WithRouteTag(c.spanRoutePrefix+prometheusRulesRoute, proxyRules),
+				otelhttp.WithRouteTag(c.spanRoutePrefix+prometheusRulesRoute, proxyPrometheusReadRules),
 			))
 			r.Get(prometheusAlertsRoute, c.instrument.NewHandler(
 				prometheus.Labels{"group": "logsv1", "handler": "alerts"},
-				otelhttp.WithRouteTag(c.spanRoutePrefix+prometheusAlertsRoute, proxyRules),
+				otelhttp.WithRouteTag(c.spanRoutePrefix+prometheusAlertsRoute, proxyPrometheusReadRules),
 			))
 			r.Get(promRulesRoute, c.instrument.NewHandler(
 				prometheus.Labels{"group": "logsv1", "handler": "rules"},

api/logs/v1/rules_labels_enforcer.go

@@ -99,6 +99,12 @@ func WithEnforceRulesLabelFilters(labelKeys map[string][]string) func(http.Handl
 func WithParametersAsLabelsFilterRules(labelKeys map[string][]string) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Prometheus rules & alert endpoints do not support filtering using the `labels` query parameter.
+			if strings.Contains(r.URL.Path, prometheusRulesRoute) || strings.Contains(r.URL.Path, prometheusAlertsRoute) {
+				next.ServeHTTP(w, r)
+				return
+			}
+
 			tenant, ok := authentication.GetTenant(r.Context())
 			if !ok {
 				httperr.PrometheusAPIError(w, "missing tenant id", http.StatusBadRequest)

api/logs/v1/rules_prometheus_labels_enforcer.go

@@ -0,0 +1,280 @@
+package http
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/observatorium/api/authentication"
+	"github.com/observatorium/api/authorization"
+	"github.com/prometheus/prometheus/model/labels"
+)
+
+const contentTypeApplicationJSON = "application/json"
+
+var (
+	errUnknownTenantKey        = errors.New("unknown tenant key")
+	errUnknownRulesContentType = errors.New("unknown rules response content type")
+)
+
+type alert struct {
+	Labels      labels.Labels `json:"labels"`
+	Annotations labels.Labels `json:"annotations"`
+	State       string        `json:"state"`
+	ActiveAt    *time.Time    `json:"activeAt,omitempty"`
+	Value       string        `json:"value"`
+}
+
+func (a *alert) GetLabels() labels.Labels { return a.Labels }
+
+type alertingRule struct {
+	State          string        `json:"state"`
+	Name           string        `json:"name"`
+	Query          string        `json:"query"`
+	Duration       float64       `json:"duration"`
+	Labels         labels.Labels `json:"labels"`
+	Annotations    labels.Labels `json:"annotations"`
+	Alerts         []*alert      `json:"alerts"`
+	Health         string        `json:"health"`
+	LastError      string        `json:"lastError"`
+	LastEvaluation string        `json:"lastEvaluation"`
+	EvaluationTime float64       `json:"evaluationTime"`
+	// Type of an alertingRule is always "alerting".
+	Type string `json:"type"`
+}
+
+type recordingRule struct {
+	Name           string        `json:"name"`
+	Query          string        `json:"query"`
+	Labels         labels.Labels `json:"labels"`
+	Health         string        `json:"health"`
+	LastError      string        `json:"lastError"`
+	LastEvaluation string        `json:"lastEvaluation"`
+	EvaluationTime float64       `json:"evaluationTime"`
+	// Type of a recordingRule is always "recording".
+	Type string `json:"type"`
+}
+
+type ruleGroup struct {
+	Name           string  `json:"name"`
+	File           string  `json:"file"`
+	Rules          []rule  `json:"rules"`
+	Interval       float64 `json:"interval"`
+	LastEvaluation string  `json:"lastEvaluation"`
+	EvaluationTime float64 `json:"evaluationTime"`
+}
+
+type rule struct {
+	*alertingRule
+	*recordingRule
+}
+
+func (r *rule) GetLabels() labels.Labels {
+	if r.alertingRule != nil {
+		return r.alertingRule.Labels
+	}
+	return r.recordingRule.Labels
+}
+
+// MarshalJSON implements the json.Marshaler interface for rule.
+func (r *rule) MarshalJSON() ([]byte, error) {
+	if r.alertingRule != nil {
+		return json.Marshal(r.alertingRule)
+	}
+	return json.Marshal(r.recordingRule)
+}
+
+// UnmarshalJSON implements the json.Unmarshaler interface for rule.
+func (r *rule) UnmarshalJSON(b []byte) error {
+	var ruleType struct {
+		Type string `json:"type"`
+	}
+	if err := json.Unmarshal(b, &ruleType); err != nil {
+		return err
+	}
+	switch ruleType.Type {
+	case "alerting":
+		var alertingr alertingRule
+		if err := json.Unmarshal(b, &alertingr); err != nil {
+			return err
+		}
+		r.alertingRule = &alertingr
+	case "recording":
+		var recordingr recordingRule
+		if err := json.Unmarshal(b, &recordingr); err != nil {
+			return err
+		}
+		r.recordingRule = &recordingr
+	default:
+		return fmt.Errorf("failed to unmarshal rule: unknown type %q", ruleType.Type)
+	}
+
+	return nil
+}
+
+type rulesData struct {
+	RuleGroups []*ruleGroup `json:"groups,omitempty"`
+	Alerts     []*alert     `json:"alerts,omitempty"`
+}
+
+type prometheusRulesResponse struct {
+	Status    string    `json:"status"`
+	Data      rulesData `json:"data"`
+	Error     string    `json:"error"`
+	ErrorType string    `json:"errorType"`
+}
+
+func newModifyResponseProm(logger log.Logger, labelKeys map[string][]string) func(*http.Response) error {
+	return func(res *http.Response) error {
+		tenant, ok := authentication.GetTenant(res.Request.Context())
+		if !ok {
+			return errUnknownTenantKey
+		}
+
+		keys, ok := labelKeys[tenant]
+		if !ok {
+			level.Debug(logger).Log("msg", "Skip applying rule label filters", "tenant", tenant)
+			return nil
+		}
+
+		var (
+			matchers    = extractMatchers(res.Request, keys)
+			contentType = res.Header.Get("Content-Type")
+		)
+
+		data, ok := authorization.GetData(res.Request.Context())
+
+		var matchersInfo AuthzResponseData
+		if ok && data != "" {
+			if err := json.Unmarshal([]byte(data), &matchersInfo); err != nil {
+				return nil
+			}
+		}
+
+		strictMode := len(matchersInfo.Matchers) != 0
+
+		matcherStr := fmt.Sprintf("%s", matchers)
+		level.Debug(logger).Log("msg", "filtering using matchers", "tenant", tenant, "matchers", matcherStr)
+
+		body, err := io.ReadAll(res.Body)
+		if err != nil {
+			level.Error(logger).Log("msg", err)
+			return err
+		}
+		res.Body.Close()
+
+		b, err := filterRules(body, contentType, matchers, strictMode)
+		if err != nil {
+			level.Error(logger).Log("msg", err)
+			return err
+		}
+
+		res.Body = io.NopCloser(bytes.NewReader(b))
+		res.ContentLength = int64(len(b))
+		res.Header.Set("Content-Length", strconv.FormatInt(res.ContentLength, 10))
+
+		return nil
+	}
+}
+
+func extractMatchers(r *http.Request, l []string) map[string]string {
+	queryParams := r.URL.Query()
+	matchers := map[string]string{}
+	for _, name := range l {
+		value := queryParams.Get(name)
+		if value != "" {
+			matchers[name] = value
+		}
+	}
+
+	return matchers
+}
+
+func filterRules(body []byte, contentType string, matchers map[string]string, strictMode bool) ([]byte, error) {
+	switch contentType {
+	case contentTypeApplicationJSON:
+		var res prometheusRulesResponse
+		err := json.Unmarshal(body, &res)
+		if err != nil {
+			return nil, err
+		}
+		return json.Marshal(filterPrometheusResponse(res, matchers, strictMode))
+	default:
+		return nil, errUnknownRulesContentType
+	}
+}
+
+func filterPrometheusResponse(res prometheusRulesResponse, matchers map[string]string, strictEnforce bool) prometheusRulesResponse {
+	if len(matchers) == 0 {
+		if strictEnforce {
+			res.Data = rulesData{}
+		}
+
+		return res
+	}
+
+	if len(res.Data.RuleGroups) > 0 {
+		filtered := filterPrometheusRuleGroups(res.Data.RuleGroups, matchers)
+		res.Data = rulesData{RuleGroups: filtered}
+	}
+
+	if len(res.Data.Alerts) > 0 {
+		filtered := filterPrometheusAlerts(res.Data.Alerts, matchers)
+		res.Data = rulesData{Alerts: filtered}
+	}
+
+	return res
+}
+
+type labeledRule interface {
+	GetLabels() labels.Labels
+}
+
+func hasMatchingLabels(rule labeledRule, matchers map[string]string) bool {
+	for key, value := range matchers {
+		labels := rule.GetLabels().Map()
+		val, ok := labels[key]
+		if !ok || val != value {
+			return false
+		}
+	}
+	return true
+}
+
+func filterPrometheusRuleGroups(groups []*ruleGroup, matchers map[string]string) []*ruleGroup {
+	var filtered []*ruleGroup
+
+	for _, group := range groups {
+		var filteredRules []rule
+		for _, r := range group.Rules {
+			if hasMatchingLabels(&r, matchers) {
+				filteredRules = append(filteredRules, r)
+			}
+		}
+
+		if len(filteredRules) > 0 {
+			group.Rules = filteredRules
+			filtered = append(filtered, group)
+		}
+	}
+
+	return filtered
+}
+
+func filterPrometheusAlerts(alerts []*alert, matchers map[string]string) []*alert {
+	var filtered []*alert
+	for _, a := range alerts {
+		if hasMatchingLabels(a, matchers) {
+			filtered = append(filtered, a)
+		}
+	}
+
+	return filtered
+}