Enforce Labels on Returned Rules

Makefile

@@ -204,4 +204,4 @@ jsonnet-fmt: | $(JSONNETFMT)
 	PATH=$$PATH:$(BIN_DIR):$(FIRST_GOPATH)/bin echo ${JSONNET_SRC} | xargs -n 1 -- $(JSONNETFMT_CMD) -i
 
 rules/rules.go: $(OAPI_CODEGEN) rules/spec.yaml
-	$(OAPI_CODEGEN) -generate types,client,chi-server -package rules -o $@ rules/spec.yaml
+	$(OAPI_CODEGEN) -generate types,client,chi-server -package rules rules/spec.yaml | sed 's|gopkg.in/yaml.v2|github.com/ghodss/yaml|g' | gofmt -s > $@

api/metrics/v1/http.go

@@ -40,6 +40,7 @@ type handlerConfiguration struct {
 	registry         *prometheus.Registry
 	instrument       handlerInstrumenter
 	spanRoutePrefix  string
+	tenantLabel      string
 	queryMiddlewares []func(http.Handler) http.Handler
 	readMiddlewares  []func(http.Handler) http.Handler
 	uiMiddlewares    []func(http.Handler) http.Handler
@@ -77,6 +78,13 @@ func WithSpanRoutePrefix(spanRoutePrefix string) HandlerOption {
 	}
 }
 
+// WithTenantLabel adds tenant label for the handler to use.
+func WithTenantLabel(tenantLabel string) HandlerOption {
+	return func(h *handlerConfiguration) {
+		h.tenantLabel = tenantLabel
+	}
+}
+
 // WithReadMiddleware adds a middleware for all "matcher based" read operations (series, label names and values).
 func WithReadMiddleware(m func(http.Handler) http.Handler) HandlerOption {
 	return func(h *handlerConfiguration) {
@@ -284,7 +292,7 @@ func NewHandler(read, write, rulesEndpoint *url.URL, upstreamCA []byte, opts ...
 			return r
 		}
 
-		rh := rulesHandler{client: client}
+		rh := rulesHandler{client: client, logger: c.logger, tenantLabel: c.tenantLabel}
 
 		r.Group(func(r chi.Router) {
 			r.Use(c.uiMiddlewares...)

api/metrics/v1/rules.go

@@ -2,14 +2,20 @@ package v1
 
 import (
 	"io"
+	"io/ioutil"
 	"net/http"
 
+	"github.com/ghodss/yaml"
+	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
 	"github.com/observatorium/api/authentication"
 	"github.com/observatorium/api/rules"
 )
 
 type rulesHandler struct {
-	client rules.ClientInterface
+	client      rules.ClientInterface
+	logger      log.Logger
+	tenantLabel string
 }
 
 func (rh *rulesHandler) get(w http.ResponseWriter, r *http.Request) {
@@ -19,8 +25,16 @@ func (rh *rulesHandler) get(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	id, ok := authentication.GetTenantID(r.Context())
+	if !ok {
+		http.Error(w, "error finding tenant ID", http.StatusUnauthorized)
+		return
+	}
+
 	resp, err := rh.client.ListRules(r.Context(), tenant)
 	if err != nil {
+		level.Error(rh.logger).Log("msg", "could not list rules", "err", err.Error())
+
 		sc := http.StatusInternalServerError
 		if resp != nil {
 			sc = resp.StatusCode
@@ -33,8 +47,56 @@ func (rh *rulesHandler) get(w http.ResponseWriter, r *http.Request) {
 
 	defer resp.Body.Close()
 
-	if _, err := io.Copy(w, resp.Body); err != nil {
-		http.Error(w, "error writing rules response", http.StatusInternalServerError)
+	if resp.StatusCode/100 != 2 {
+		http.Error(w, "error listing rules", resp.StatusCode)
+		return
+	}
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		http.Error(w, "error listing rules", http.StatusInternalServerError)
+		return
+	}
+
+	var rawRules rules.Rules
+	if err := yaml.Unmarshal(body, &rawRules); err != nil {
+		level.Error(rh.logger).Log("msg", "could not unmarshal rules", "err", err.Error())
+		http.Error(w, "error unmarshaling rules", http.StatusInternalServerError)
+
+		return
+	}
+
+	for i := range rawRules.Groups {
+		for j := range rawRules.Groups[i].Rules {
+			switch r := rawRules.Groups[i].Rules[j].(type) {
+			case rules.RecordingRule:
+				if r.Labels.AdditionalProperties == nil {
+					r.Labels.AdditionalProperties = make(map[string]string)
+				}
+
+				r.Labels.AdditionalProperties[rh.tenantLabel] = id
+				rawRules.Groups[i].Rules[j] = r
+			case rules.AlertingRule:
+				if r.Labels.AdditionalProperties == nil {
+					r.Labels.AdditionalProperties = make(map[string]string)
+				}
+
+				r.Labels.AdditionalProperties[rh.tenantLabel] = id
+				rawRules.Groups[i].Rules[j] = r
+			}
+		}
+	}
+
+	body, err = yaml.Marshal(rawRules)
+	if err != nil {
+		level.Error(rh.logger).Log("msg", "could not marshal YAML", "err", err.Error())
+		http.Error(w, "error marshaling YAML", http.StatusInternalServerError)
+
+		return
+	}
+
+	if _, err := w.Write(body); err != nil {
+		level.Error(rh.logger).Log("msg", "could not write body", "err", err.Error())
 		return
 	}
 }
@@ -52,6 +114,7 @@ func (rh *rulesHandler) put(w http.ResponseWriter, r *http.Request) {
 			sc = resp.StatusCode
 		}
 
+		level.Error(rh.logger).Log("msg", "could not set rules", "err", err.Error())
 		http.Error(w, "error creating rules", sc)
 
 		return

go.mod

@@ -41,7 +41,6 @@ require (
 	google.golang.org/grpc v1.38.0
 	google.golang.org/protobuf v1.27.1
 	gopkg.in/square/go-jose.v2 v2.4.1 // indirect
-	gopkg.in/yaml.v2 v2.4.0
 	k8s.io/apimachinery v0.21.1
 	k8s.io/apiserver v0.21.1
 	k8s.io/client-go v0.21.1

main.go

@@ -548,6 +548,7 @@ func main() {
 								metricsv1.WithRegistry(reg),
 								metricsv1.WithHandlerInstrumenter(ins),
 								metricsv1.WithSpanRoutePrefix("/api/metrics/v1/{tenant}"),
+								metricsv1.WithTenantLabel(cfg.metrics.tenantLabel),
 								metricsv1.WithQueryMiddleware(authorization.WithAuthorizers(authorizers, rbac.Read, "metrics")),
 								metricsv1.WithQueryMiddleware(metricsv1.WithEnforceTenancyOnQuery(cfg.metrics.tenantLabel)),
 								metricsv1.WithReadMiddleware(authorization.WithAuthorizers(authorizers, rbac.Read, "metrics")),

rules/custom_types.go

@@ -0,0 +1,50 @@
+package rules
+
+import (
+	"encoding/json"
+)
+
+func (rg *RuleGroup) UnmarshalJSON(data []byte) error {
+	raw := struct {
+		Interval string            `json:"interval"`
+		Name     string            `json:"name"`
+		Rules    []json.RawMessage `json:"rules"`
+	}{}
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	rg.Interval = raw.Interval
+	rg.Name = raw.Name
+	rules := make([]interface{}, 0, len(raw.Rules))
+
+	for i := range raw.Rules {
+		rawRule := make(map[string]json.RawMessage)
+		if err := json.Unmarshal(raw.Rules[i], &rawRule); err != nil {
+			return err
+		}
+
+		switch _, ok := rawRule["alert"]; ok {
+		case true:
+			var ar AlertingRule
+			if err := json.Unmarshal(raw.Rules[i], &ar); err != nil {
+				return err
+			}
+
+			rules = append(rules, ar)
+		case false:
+			var rr RecordingRule
+			if err := json.Unmarshal(raw.Rules[i], &rr); err != nil {
+				return err
+			}
+
+			rules = append(rules, rr)
+		}
+	}
+
+	if len(rules) != 0 {
+		rg.Rules = rules
+	}
+
+	return nil
+}