ITWorx-Pulse-Public/internal/metricquery/handler.go

package metricquery

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"

	"github.com/itworx/pulse/internal/auth"
	"github.com/itworx/pulse/internal/problem"
	"github.com/itworx/pulse/internal/promqlbinding"
	"github.com/itworx/pulse/internal/queryplan"
)

const maxRequestBody = 64 << 10

type Handler struct{ Service *Service }

func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.NotFound(w, r)
		return
	}
	if _, ok := auth.PrincipalFromContext(r.Context()); !ok {
		problem.Write(w, r, http.StatusUnauthorized, "UNAUTHORIZED", "Authentication required", "Authentication is required to query metrics.", nil)
		return
	}
	if h.Service == nil {
		problem.Write(w, r, http.StatusServiceUnavailable, "METRIC_SOURCE_UNAVAILABLE", "Metric source unavailable", "The configured metric source is not available.", nil)
		return
	}
	var response Response
	var err error
	switch strings.TrimPrefix(r.URL.Path, "/api/v1/metrics/") {
	case "inspect":
		var request queryplan.Request
		if !decodeJSON(w, r, &request) {
			return
		}
		inspector, inspectErr := h.Service.Inspect(r.Context(), request)
		if inspectErr != nil {
			writeQueryError(w, r, inspectErr)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"inspector": inspector})
		return
	case "query-range":
		var request queryplan.Request
		if decodeJSON(w, r, &request) {
			response, err = h.Service.ExecuteRange(r.Context(), request)
		}
	case "query":
		var request InstantRequest
		if decodeJSON(w, r, &request) {
			response, err = h.Service.ExecuteInstant(r.Context(), request)
		}
	default:
		http.NotFound(w, r)
		return
	}
	if err != nil {
		writeQueryError(w, r, err)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(response)
}

func decodeJSON(w http.ResponseWriter, r *http.Request, destination any) bool {
	body, readErr := io.ReadAll(io.LimitReader(r.Body, maxRequestBody+1))
	if readErr != nil || len(body) > maxRequestBody {
		problem.Write(w, r, http.StatusRequestEntityTooLarge, "QUERY_BODY_LIMIT", "Query body too large", "The metric query body exceeds the allowed size.", nil)
		return false
	}
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(destination); err != nil {
		problem.Write(w, r, http.StatusBadRequest, "QUERY_BODY_INVALID", "Invalid query body", "The metric query body is invalid.", nil)
		return false
	}
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		problem.Write(w, r, http.StatusBadRequest, "QUERY_BODY_INVALID", "Invalid query body", "The metric query body contains trailing data.", nil)
		return false
	}
	return true
}

func writeQueryError(w http.ResponseWriter, r *http.Request, err error) {
	if errors.Is(err, ErrInspectorUnauthorized) {
		problem.Write(w, r, http.StatusForbidden, "QUERY_INSPECTOR_FORBIDDEN", "Query inspector forbidden", "The query inspector requires operate permission.", nil)
		return
	}
	if errors.Is(err, context.Canceled) {
		return
	}
	var sourceErr sourceError
	if errors.Is(err, ErrSourceUnavailable) || errors.As(err, &sourceErr) {
		problem.Write(w, r, http.StatusServiceUnavailable, "QUERY_SOURCE_UNAVAILABLE", "Metric source unavailable", "The metric source could not be queried.", nil)
		return
	}
	var planError queryplan.Error
	if errors.As(err, &planError) {
		status := http.StatusUnprocessableEntity
		if planError.Code == queryplan.ErrUnauthorized.Code {
			status = http.StatusUnauthorized
		}
		problem.Write(w, r, status, planError.Code, "Metric query rejected", planError.Detail, map[string]string{"field": planError.Field})
		return
	}
	var bindingError promqlbinding.Error
	if errors.As(err, &bindingError) {
		status := http.StatusUnprocessableEntity
		if bindingError.Code == "PROMQL_RAW_UNAUTHORIZED" {
			status = http.StatusForbidden
		}
		problem.Write(w, r, status, bindingError.Code, "Metric query rejected", bindingError.Detail, map[string]string{"field": bindingError.Field})
		return
	}
	problem.Write(w, r, http.StatusServiceUnavailable, "QUERY_UNAVAILABLE", "Metric query unavailable", "The metric query could not be completed.", nil)
}