sift/internal/api/metrics_test.go

package api

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/digest"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
)

func loopbackGET(path string) *http.Request {
	req := httptest.NewRequest(http.MethodGet, path, nil)
	req.Host = "127.0.0.1:8890"
	req.RemoteAddr = "127.0.0.1:9"
	return req
}

func serveMetrics(t *testing.T, s *Server, req *http.Request) *httptest.ResponseRecorder {
	t.Helper()
	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, req)
	return rec
}

func parseProm(t *testing.T, body string) map[string]*dto.MetricFamily {
	t.Helper()
	parser := expfmt.NewDecoder(strings.NewReader(body), expfmt.NewFormat(expfmt.TypeTextPlain))
	out := map[string]*dto.MetricFamily{}
	for {
		var mf dto.MetricFamily
		if err := parser.Decode(&mf); err != nil {
			if err == io.EOF {
				break
			}
			t.Fatalf("expfmt: %v\n%s", err, body)
		}
		out[mf.GetName()] = &mf
	}
	return out
}

func TestMetricsJSONSidePath(t *testing.T) {
	s, key, _, _, _ := searchServer(t)
	postJSON(t, s, "/v1/search", key, map[string]any{"query": "same q", "k": 5, "rerank": false})
	postJSON(t, s, "/v1/search", key, map[string]any{"query": "same q", "k": 5, "rerank": false})
	for _, path := range []string{"/metrics.json", "/metrics?format=json"} {
		rec := serveMetrics(t, s, loopbackGET(path))
		if rec.Code != 200 {
			t.Fatalf("%s %d %s", path, rec.Code, rec.Body.String())
		}
		if !strings.Contains(rec.Header().Get("Content-Type"), "json") {
			t.Fatalf("%s ctype %s", path, rec.Header().Get("Content-Type"))
		}
		var m map[string]any
		if err := json.Unmarshal(rec.Body.Bytes(), &m); err != nil {
			t.Fatal(err)
		}
		if m["search_cache_hit"] == 0 && m["search_cache_miss"] == 0 {
			t.Fatalf("%s %v", path, m)
		}
		if _, ok := m["requests"]; !ok {
			t.Fatalf("missing requests: %v", m)
		}
	}
}

func TestMetricsPrometheusLoopback(t *testing.T) {
	s, key, _, _, _ := searchServer(t)
	postJSON(t, s, "/v1/search", key, map[string]any{"query": "fleet q", "k": 5, "rerank": false})
	rec := serveMetrics(t, s, loopbackGET("/metrics"))
	if rec.Code != 200 {
		t.Fatalf("%d %s", rec.Code, rec.Body.String())
	}
	ct := rec.Header().Get("Content-Type")
	if !strings.Contains(ct, "text/plain") {
		t.Fatalf("ctype %s", ct)
	}
	fams := parseProm(t, rec.Body.String())
	if fams["sift_search_queries_total"] == nil {
		t.Fatalf("missing queries total: %v", rec.Body.String())
	}
	if fams["sift_search_queries_caller_total"] == nil {
		t.Fatalf("missing caller total")
	}
	if !strings.Contains(rec.Body.String(), `caller="test"`) {
		t.Fatal("fleet caller label")
	}
	if strings.Contains(rec.Body.String(), "fleet q") {
		t.Fatal("query leaked")
	}
}

func TestMetricsAuth(t *testing.T) {
	dir := t.TempDir()
	keys := filepath.Join(dir, "keys.toml")
	fleetOut := filepath.Join(dir, "fleet")
	metOut := filepath.Join(dir, "met")
	if err := auth.Add(keys, "test", "hermes-fleet", fleetOut); err != nil {
		t.Fatal(err)
	}
	if err := auth.Add(keys, auth.NameMetrics, "should-be-cleared", metOut); err != nil {
		t.Fatal(err)
	}
	store, err := auth.Load(keys)
	if err != nil {
		t.Fatal(err)
	}
	mc := store.LookupBearer(strings.TrimSpace(mustRead(t, metOut)))
	if mc == nil || mc.Workspace != "" || !mc.MetricsOnly() {
		t.Fatalf("metrics caller %+v", mc)
	}
	s, fleetKey := setupAPI(t, nil)
	s.Auth = store
	fleetKey = strings.TrimSpace(mustRead(t, fleetOut))
	metKey := strings.TrimSpace(mustRead(t, metOut))

	req := httptest.NewRequest(http.MethodGet, "/metrics", nil) // 192.0.2.1 + example.com
	rec := serveMetrics(t, s, req)
	if rec.Code != 401 {
		t.Fatalf("no bearer: %d", rec.Code)
	}
	req = httptest.NewRequest(http.MethodGet, "/metrics", nil)
	req.Header.Set("Authorization", "Bearer "+fleetKey)
	rec = serveMetrics(t, s, req)
	if rec.Code != 403 {
		t.Fatalf("fleet bearer: %d %s", rec.Code, rec.Body.String())
	}
	req = httptest.NewRequest(http.MethodGet, "/metrics", nil)
	req.Header.Set("Authorization", "Bearer "+metKey)
	rec = serveMetrics(t, s, req)
	if rec.Code != 200 {
		t.Fatalf("metrics bearer: %d %s", rec.Code, rec.Body.String())
	}

	fwd := loopbackGET("/metrics")
	fwd.Header.Set("X-Forwarded-For", "100.1.2.3")
	rec = serveMetrics(t, s, fwd)
	if rec.Code != 401 {
		t.Fatalf("forwarded loopback: %d", rec.Code)
	}

	rec = postJSON(t, s, "/v1/fetch", metKey, map[string]string{"url": "https://example.com/"})
	if rec.Code != 403 {
		t.Fatalf("metrics fetch: %d", rec.Code)
	}
	rec = postJSON(t, s, "/v1/search", metKey, map[string]any{"query": "x"})
	if rec.Code != 403 {
		t.Fatalf("metrics search: %d", rec.Code)
	}
}

func TestMetricsRestrictedNoCallerLabel(t *testing.T) {
	s, fleetKey, gsKey, _, _ := searchServer(t)
	q := "secret-restricted-query-42"
	postJSON(t, s, "/v1/search", fleetKey, map[string]any{"query": "fleet-only-q", "k": 5, "rerank": false})
	postJSON(t, s, "/v1/search", gsKey, map[string]any{"query": q, "k": 5, "rerank": false})
	rec := serveMetrics(t, s, loopbackGET("/metrics"))
	body := rec.Body.String()
	if rec.Code != 200 {
		t.Fatal(body)
	}
	if strings.Contains(body, q) || strings.Contains(body, "fleet-only-q") {
		t.Fatal("query in prometheus")
	}
	if strings.Contains(body, "example.com") {
		t.Fatal("host in prometheus")
	}
	if strings.Contains(body, `caller="straylight-gs"`) {
		t.Fatal("restricted caller series")
	}
	if strings.Contains(body, "tavily:restricted") {
		t.Fatal("restricted cap partition")
	}
	if !strings.Contains(body, `caller="test"`) {
		t.Fatal("fleet caller missing")
	}
	js := serveMetrics(t, s, loopbackGET("/metrics.json"))
	var m map[string]any
	_ = json.Unmarshal(js.Body.Bytes(), &m)
	queries, _ := m["search_queries"].(map[string]any)
	if queries["straylight-gs"] == nil {
		t.Fatalf("json still has restricted caller: %v", m["search_queries"])
	}
}

func TestMetricsDigestFallbackAndGauges(t *testing.T) {
	s, _ := setupAPI(t, nil)
	g := digest.NewGate(1, 0, time.Second)
	s.Digest = &digest.Engine{Gate: g}
	ok, _ := g.Acquire(context.Background())
	if !ok {
		t.Fatal("hold permit")
	}
	s.Digest.Run(context.Background(), "doc", "q", "https://ex.test/p", "fleet")
	rec := serveMetrics(t, s, loopbackGET("/metrics"))
	body := rec.Body.String()
	if !strings.Contains(body, `sift_digest_fallbacks_total{reason="queue_full"}`) &&
		!strings.Contains(body, `reason="queue_full"`) {
		t.Fatal(body)
	}
	if !strings.Contains(body, "sift_dixie_inflight") || !strings.Contains(body, "sift_dixie_queue_waiters") {
		t.Fatal(body)
	}
}

func TestMetricsDoesNotCountScrapes(t *testing.T) {
	s, key := setupAPI(t, nil)
	before := serveMetrics(t, s, loopbackGET("/metrics"))
	serveMetrics(t, s, loopbackGET("/metrics"))
	after := serveMetrics(t, s, loopbackGET("/metrics"))
	if before.Body.String() != after.Body.String() && strings.Contains(after.Body.String(), `endpoint="/metrics"`) {
		t.Fatal(after.Body.String())
	}
	postJSON(t, s, "/v2/scrape", key, map[string]string{"url": "http://127.0.0.1/"})
	rec := serveMetrics(t, s, loopbackGET("/metrics"))
	if !strings.Contains(rec.Body.String(), `sift_requests_total{code="403",endpoint="/v2/scrape"}`) &&
		!strings.Contains(rec.Body.String(), `endpoint="/v2/scrape"`) {
		t.Fatal(rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "sift_request_duration_seconds") {
		t.Fatal("histogram missing")
	}
}

func TestMetricsCacheHitMiss(t *testing.T) {
	s, key, _, _, _ := searchServer(t)
	postJSON(t, s, "/v1/search", key, map[string]any{"query": "same q", "k": 5, "rerank": false})
	postJSON(t, s, "/v1/search", key, map[string]any{"query": "same q", "k": 5, "rerank": false})
	req := loopbackGET("/metrics.json")
	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, req)
	var m map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &m)
	if m["search_cache_hit"] == 0 && m["search_cache_miss"] == 0 {
		t.Fatalf("%v", m)
	}
}

func mustRead(t *testing.T, path string) string {
	t.Helper()
	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	return string(b)
}