vps-control/control-plane/backend/observability.go

package main

import (
	"context"
	"crypto/subtle"
	"errors"
	"fmt"
	"math"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
)

var metricSelector = regexp.MustCompile(`^([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{([^}]*)\})?$`)
var logSelector = regexp.MustCompile(`^([a-z_][a-z0-9_]*)\s*(=|=~)\s*"([^"]*)"$`)

var metricColumns = map[string]string{
	"vps_control_cpu_usage_percent":    "cpu",
	"vps_control_memory_usage_percent": "memory",
	"vps_control_disk_usage_percent":   "disk",
	"vps_control_network_rx_bytes":     "rx",
	"vps_control_network_tx_bytes":     "tx",
	"vps_control_connections":          "connections",
	"vps_control_bans":                 "bans",
}

type metricQuery struct {
	Name       string
	ServerID   string
	ServerName string
}

type logMatcher struct {
	Label string
	Exact string
	Regex *regexp.Regexp
}

func readIntegrationToken(path string) (string, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return "", fmt.Errorf("read integration token metadata: %w", err)
	}
	if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 || info.Mode().Perm()&0o022 != 0 {
		return "", errors.New("integration token must be a non-writable regular file")
	}
	if info.Size() > 4096 {
		return "", errors.New("integration token file is too large")
	}
	value, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("read integration token: %w", err)
	}
	token := strings.TrimSpace(string(value))
	if len(token) < 32 || len(token) > 256 || strings.IndexFunc(token, func(character rune) bool {
		return character < 0x21 || character == 0x7f
	}) >= 0 {
		return "", errors.New("integration token must contain 32-256 visible bytes")
	}
	return token, nil
}

func (app *App) integration(writer http.ResponseWriter, request *http.Request) {
	if len(app.integrationTokenHash) == 0 {
		writeError(writer, http.StatusNotFound, "integration API is disabled", "not_found")
		return
	}
	if !app.integrationLimiter.Allow(app.clientIP(request), time.Now()) {
		writeError(writer, http.StatusTooManyRequests, "integration request rate exceeded", "rate_limited")
		return
	}
	ctx, cancel := context.WithTimeout(request.Context(), 10*time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	provided := strings.TrimPrefix(request.Header.Get("Authorization"), "Bearer ")
	if len(provided) > 256 || subtle.ConstantTimeCompare(tokenHash(provided), app.integrationTokenHash) != 1 {
		writeError(writer, http.StatusUnauthorized, "integration authentication required", "unauthorized")
		return
	}
	if strings.HasPrefix(request.URL.Path, "/integrations/prometheus/") {
		app.prometheusAPI(writer, request)
		return
	}
	if strings.HasPrefix(request.URL.Path, "/integrations/loki/") {
		app.lokiAPI(writer, request)
		return
	}
	writeError(writer, http.StatusNotFound, "not found", "not_found")
}

func (app *App) prometheusAPI(writer http.ResponseWriter, request *http.Request) {
	path := strings.TrimPrefix(request.URL.Path, "/integrations/prometheus")
	if request.Method != http.MethodGet && request.Method != http.MethodPost {
		writePrometheusError(writer, "bad_data", "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	if request.Method == http.MethodPost {
		request.Body = http.MaxBytesReader(writer, request.Body, maxRequestBody)
	}
	switch path {
	case "/-/healthy", "/-/ready":
		writeJSON(writer, 200, map[string]string{"status": "success"})
	case "/api/v1/query":
		app.prometheusQuery(writer, request, false)
	case "/api/v1/query_range":
		app.prometheusQuery(writer, request, true)
	case "/api/v1/labels":
		writeJSON(writer, 200, map[string]any{"status": "success", "data": []string{"__name__", "server_id", "server_name"}})
	case "/api/v1/label/__name__/values":
		names := make([]string, 0, len(metricColumns))
		for name := range metricColumns {
			names = append(names, name)
		}
		sort.Strings(names)
		writeJSON(writer, 200, map[string]any{"status": "success", "data": names})
	case "/api/v1/label/server_id/values", "/api/v1/label/server_name/values":
		label := strings.TrimSuffix(strings.TrimPrefix(path, "/api/v1/label/"), "/values")
		values, err := app.store.MetricLabelValues(request.Context(), label)
		if err != nil {
			writePrometheusError(writer, "internal", "query failed", 500)
			return
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": values})
	case "/api/v1/metadata":
		metadata := map[string][]map[string]string{}
		for name := range metricColumns {
			metadata[name] = []map[string]string{{"type": "gauge", "help": "VPS Control sampled metric", "unit": ""}}
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": metadata})
	case "/api/v1/series":
		series, err := app.store.MetricSeries(request.Context())
		if err != nil {
			writePrometheusError(writer, "internal", "query failed", 500)
			return
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": series})
	case "/api/v1/status/buildinfo":
		writeJSON(writer, 200, map[string]any{"status": "success", "data": map[string]string{"version": "vps-control-compat-v1"}})
	default:
		writePrometheusError(writer, "bad_data", "unsupported Prometheus API operation", 404)
	}
}

func (app *App) prometheusQuery(writer http.ResponseWriter, request *http.Request, ranged bool) {
	if err := request.ParseForm(); err != nil {
		writePrometheusError(writer, "bad_data", "invalid form", 400)
		return
	}
	rawQuery := request.Form.Get("query")
	if !ranged && rawQuery == "1+1" {
		at, err := parseTimestamp(request.Form.Get("time"), time.Now())
		if err != nil {
			writePrometheusError(writer, "bad_data", "invalid time", 400)
			return
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": map[string]any{
			"resultType": "scalar", "result": []any{float64(at.UnixNano()) / 1e9, "2"},
		}})
		return
	}
	query, err := parseMetricQuery(rawQuery)
	if err != nil {
		writePrometheusError(writer, "bad_data", err.Error(), 400)
		return
	}
	timeParameter := "time"
	if ranged {
		timeParameter = "end"
	}
	end, err := parseTimestamp(request.Form.Get(timeParameter), time.Now())
	if err != nil {
		writePrometheusError(writer, "bad_data", "invalid "+timeParameter, 400)
		return
	}
	if !ranged {
		points, queryErr := app.store.InstantMetric(request.Context(), query, end)
		if queryErr != nil {
			writePrometheusError(writer, "internal", "query failed", 500)
			return
		}
		writeJSON(writer, 200, prometheusVector(query.Name, points))
		return
	}
	start, err := parseTimestamp(request.Form.Get("start"), end.Add(-time.Hour))
	if err != nil || start.After(end) || end.Sub(start) > 31*24*time.Hour {
		writePrometheusError(writer, "bad_data", "range must be positive and at most 31 days", 400)
		return
	}
	step, err := parseStep(request.Form.Get("step"))
	if err != nil {
		writePrometheusError(writer, "bad_data", "invalid step", 400)
		return
	}
	if int64(end.Sub(start)/time.Second)/int64(step) > 10_000 {
		writePrometheusError(writer, "bad_data", "range contains more than 10000 points per series", 400)
		return
	}
	points, queryErr := app.store.RangeMetric(request.Context(), query, start, end, step)
	if queryErr != nil {
		writePrometheusError(writer, "internal", "query failed", 500)
		return
	}
	writeJSON(writer, 200, prometheusMatrix(query.Name, points))
}

func (app *App) lokiAPI(writer http.ResponseWriter, request *http.Request) {
	path := strings.TrimPrefix(request.URL.Path, "/integrations/loki")
	if request.Method != http.MethodGet {
		writeError(writer, 405, "method not allowed", "method_not_allowed")
		return
	}
	switch path {
	case "/ready":
		writer.WriteHeader(http.StatusOK)
	case "/loki/api/v1/query_range":
		app.lokiQueryRange(writer, request)
	case "/loki/api/v1/query":
		if request.URL.Query().Get("query") != "vector(1)+vector(1)" {
			writeError(writer, 400, "only the datasource health expression is supported for instant queries", "invalid_query")
			return
		}
		at, err := parseLokiTimestamp(request.URL.Query().Get("time"), time.Now())
		if err != nil {
			writeError(writer, 400, "invalid time", "invalid_query")
			return
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": map[string]any{
			"resultType": "vector", "result": []any{map[string]any{
				"metric": map[string]string{}, "value": []any{float64(at.UnixNano()) / 1e9, "2"},
			}},
		}})
	case "/loki/api/v1/labels":
		writeJSON(writer, 200, map[string]any{"status": "success", "data": []string{"server_id", "server_name", "unit", "priority"}})
	case "/loki/api/v1/label/server_id/values", "/loki/api/v1/label/server_name/values", "/loki/api/v1/label/unit/values", "/loki/api/v1/label/priority/values":
		label := strings.TrimSuffix(strings.TrimPrefix(path, "/loki/api/v1/label/"), "/values")
		values, err := app.store.LogLabelValues(request.Context(), label)
		if err != nil {
			writeError(writer, 500, "query failed", "internal_error")
			return
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": values})
	case "/loki/api/v1/series":
		series, err := app.store.LogSeries(request.Context())
		if err != nil {
			writeError(writer, 500, "query failed", "internal_error")
			return
		}
		writeJSON(writer, 200, map[string]any{"status": "success", "data": series})
	default:
		writeError(writer, 404, "unsupported Loki API operation", "not_found")
	}
}

func (app *App) lokiQueryRange(writer http.ResponseWriter, request *http.Request) {
	matchers, err := parseLogQuery(request.URL.Query().Get("query"))
	if err != nil {
		writeError(writer, 400, err.Error(), "invalid_query")
		return
	}
	end, err := parseLokiTimestamp(request.URL.Query().Get("end"), time.Now())
	if err != nil {
		writeError(writer, 400, "invalid end", "invalid_query")
		return
	}
	start, err := parseLokiTimestamp(request.URL.Query().Get("start"), end.Add(-time.Hour))
	if err != nil || start.After(end) || end.Sub(start) > 31*24*time.Hour {
		writeError(writer, 400, "range must be positive and at most 31 days", "invalid_query")
		return
	}
	limit := parseInt(request.URL.Query().Get("limit"), 1000)
	if limit < 1 {
		limit = 1
	}
	if limit > 5000 {
		limit = 5000
	}
	direction := request.URL.Query().Get("direction")
	if direction == "" {
		direction = "backward"
	}
	if direction != "backward" && direction != "forward" {
		writeError(writer, 400, "invalid direction", "invalid_query")
		return
	}
	entries, queryErr := app.store.IntegrationLogs(request.Context(), start, end, 20_000, direction)
	if queryErr != nil {
		writeError(writer, 500, "query failed", "internal_error")
		return
	}
	streams := buildLokiStreams(entries, matchers, limit, direction)
	writeJSON(writer, 200, map[string]any{"status": "success", "data": map[string]any{
		"resultType": "streams", "result": streams, "stats": map[string]any{},
	}})
}

func parseMetricQuery(raw string) (metricQuery, error) {
	if len(raw) > 512 {
		return metricQuery{}, errors.New("query is too long")
	}
	match := metricSelector.FindStringSubmatch(strings.TrimSpace(raw))
	if match == nil {
		return metricQuery{}, errors.New("only direct metric selectors are supported")
	}
	if _, found := metricColumns[match[1]]; !found {
		return metricQuery{}, errors.New("metric is not supported")
	}
	result := metricQuery{Name: match[1]}
	if match[2] == "" {
		return result, nil
	}
	for _, selector := range strings.Split(match[2], ",") {
		parts := strings.SplitN(strings.TrimSpace(selector), "=", 2)
		if len(parts) != 2 {
			return metricQuery{}, errors.New("invalid label selector")
		}
		value, err := strconv.Unquote(strings.TrimSpace(parts[1]))
		if err != nil || len(value) > 128 {
			return metricQuery{}, errors.New("invalid label value")
		}
		switch strings.TrimSpace(parts[0]) {
		case "server_id":
			result.ServerID = value
		case "server_name":
			result.ServerName = value
		default:
			return metricQuery{}, errors.New("unsupported label")
		}
	}
	return result, nil
}

func parseLogQuery(raw string) ([]logMatcher, error) {
	if len(raw) < 2 || len(raw) > 512 || raw[0] != '{' || raw[len(raw)-1] != '}' {
		return nil, errors.New("a bounded stream selector is required")
	}
	content := strings.TrimSpace(raw[1 : len(raw)-1])
	if content == "" {
		return nil, errors.New("at least one stream matcher is required")
	}
	matchers := []logMatcher{}
	for _, part := range strings.Split(content, ",") {
		match := logSelector.FindStringSubmatch(strings.TrimSpace(part))
		if match == nil || !allowedLogLabel(match[1]) || len(match[3]) > 128 {
			return nil, errors.New("unsupported stream matcher")
		}
		item := logMatcher{Label: match[1]}
		if match[2] == "=" {
			item.Exact = match[3]
		} else {
			compiled, err := regexp.Compile("^(?:" + match[3] + ")$")
			if err != nil {
				return nil, errors.New("invalid matcher regex")
			}
			item.Regex = compiled
		}
		matchers = append(matchers, item)
	}
	return matchers, nil
}

func allowedLogLabel(label string) bool {
	return label == "server_id" || label == "server_name" || label == "unit" || label == "priority"
}

func parseTimestamp(raw string, fallback time.Time) (time.Time, error) {
	if raw == "" {
		return fallback, nil
	}
	value, err := strconv.ParseFloat(raw, 64)
	if err != nil || math.IsNaN(value) || math.IsInf(value, 0) {
		return time.Time{}, errors.New("invalid timestamp")
	}
	seconds, fraction := math.Modf(value)
	parsed := time.Unix(int64(seconds), int64(fraction*1e9))
	if !validIntegrationTime(parsed) {
		return time.Time{}, errors.New("timestamp out of range")
	}
	return parsed, nil
}

func parseLokiTimestamp(raw string, fallback time.Time) (time.Time, error) {
	if raw == "" {
		return fallback, nil
	}
	if value, err := strconv.ParseInt(raw, 10, 64); err == nil {
		var parsed time.Time
		if value > 1_000_000_000_000 {
			parsed = time.Unix(0, value)
		} else {
			parsed = time.Unix(value, 0)
		}
		if !validIntegrationTime(parsed) {
			return time.Time{}, errors.New("timestamp out of range")
		}
		return parsed, nil
	}
	parsed, err := time.Parse(time.RFC3339Nano, raw)
	if err != nil || !validIntegrationTime(parsed) {
		return time.Time{}, errors.New("invalid timestamp")
	}
	return parsed, nil
}

func validIntegrationTime(value time.Time) bool {
	return !value.Before(time.Unix(0, 0)) && value.Before(time.Date(2100, 1, 1, 0, 0, 0, 0, time.UTC))
}

func parseStep(raw string) (int, error) {
	if raw == "" {
		return 60, nil
	}
	if duration, err := time.ParseDuration(raw); err == nil {
		seconds := int(duration.Seconds())
		if seconds >= 5 && seconds <= 86400 {
			return seconds, nil
		}
	}
	value, err := strconv.Atoi(raw)
	if err != nil || value < 5 || value > 86400 {
		return 0, errors.New("step out of range")
	}
	return value, nil
}

func writePrometheusError(writer http.ResponseWriter, kind, message string, status int) {
	writeJSON(writer, status, map[string]any{"status": "error", "errorType": kind, "error": message})
}

func prometheusVector(name string, points []MetricPoint) map[string]any {
	result := make([]map[string]any, 0, len(points))
	for _, point := range points {
		result = append(result, map[string]any{
			"metric": metricLabels(name, point),
			"value":  []any{float64(point.Timestamp.UnixNano()) / 1e9, strconv.FormatFloat(point.Value, 'f', -1, 64)},
		})
	}
	return map[string]any{"status": "success", "data": map[string]any{"resultType": "vector", "result": result}}
}

func prometheusMatrix(name string, points []MetricPoint) map[string]any {
	type series struct {
		Metric map[string]string `json:"metric"`
		Values [][]any           `json:"values"`
	}
	byServer := map[string]*series{}
	order := []string{}
	for _, point := range points {
		item := byServer[point.ServerID]
		if item == nil {
			item = &series{Metric: metricLabels(name, point)}
			byServer[point.ServerID] = item
			order = append(order, point.ServerID)
		}
		item.Values = append(item.Values, []any{
			float64(point.Timestamp.UnixNano()) / 1e9,
			strconv.FormatFloat(point.Value, 'f', -1, 64),
		})
	}
	result := make([]*series, 0, len(order))
	for _, id := range order {
		result = append(result, byServer[id])
	}
	return map[string]any{"status": "success", "data": map[string]any{"resultType": "matrix", "result": result}}
}

func metricLabels(name string, point MetricPoint) map[string]string {
	return map[string]string{"__name__": name, "server_id": point.ServerID, "server_name": point.ServerName}
}

type lokiStream struct {
	Stream map[string]string `json:"stream"`
	Values [][2]string       `json:"values"`
}

func buildLokiStreams(entries []IntegrationLog, matchers []logMatcher, limit int, direction string) []lokiStream {
	groups := map[string]*lokiStream{}
	count := 0
	for _, entry := range entries {
		labels := map[string]string{
			"server_id": entry.ServerID, "server_name": entry.ServerName,
			"unit": entry.Unit, "priority": strconv.Itoa(entry.Priority),
		}
		if !logMatches(labels, matchers) {
			continue
		}
		key := entry.ServerID + "\x00" + entry.Unit + "\x00" + strconv.Itoa(entry.Priority)
		stream := groups[key]
		if stream == nil {
			stream = &lokiStream{Stream: labels}
			groups[key] = stream
		}
		stream.Values = append(stream.Values, [2]string{strconv.FormatInt(entry.Timestamp.UnixNano(), 10), entry.Message})
		count++
		if count >= limit {
			break
		}
	}
	result := make([]lokiStream, 0, len(groups))
	for _, stream := range groups {
		if direction != "backward" {
			sort.Slice(stream.Values, func(i, j int) bool { return stream.Values[i][0] < stream.Values[j][0] })
		}
		result = append(result, *stream)
	}
	sort.Slice(result, func(i, j int) bool {
		return result[i].Stream["server_id"]+result[i].Stream["unit"] < result[j].Stream["server_id"]+result[j].Stream["unit"]
	})
	return result
}

func logMatches(labels map[string]string, matchers []logMatcher) bool {
	for _, matcher := range matchers {
		value := labels[matcher.Label]
		if matcher.Regex != nil {
			if !matcher.Regex.MatchString(value) {
				return false
			}
		} else if value != matcher.Exact {
			return false
		}
	}
	return true
}