sift/internal/metrics/metrics.go

// Package metrics is sift's Prometheus registry and the loopback-scrape predicate.
package metrics

import (
	"net"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var durationBuckets = []float64{0.1, 0.5, 1, 2, 5, 10, 15, 20, 30, 45}

var fallbackOK = map[string]struct{}{
	"queue_full": {}, "queue_wait": {}, "canceled": {},
	"embed": {}, "dixie": {}, "timeout": {}, "empty": {}, "unset": {},
}

// Snapshot is process state gathered on each Prometheus scrape.
type Snapshot struct {
	Blocks, Scrapes, SelfIPFailures uint64
	DigestCalls                     uint64
	DigestFallbacks                 map[string]uint64
	DixieInflight, DixieWaiters     int
	DigestHit, DigestMiss           uint64
	EmbedHit, EmbedMiss             uint64
	SearchHit, SearchMiss           uint64
	SearchQueries                   map[string]map[string]uint64 // caller → source
	SearchCredits                   map[string]map[string]uint64
	SearchDaily                     map[string]uint64
	SearchCapped                    map[string]uint64
	Restricted                      map[string]struct{}
}

// Source is implemented by api.Server.
type Source interface {
	Snapshot() Snapshot
}

// Set is a private registry of sift_* series.
type Set struct {
	Reg      *prometheus.Registry
	Requests *prometheus.CounterVec
	Duration *prometheus.HistogramVec
	handler  http.Handler
}

func New(src Source) *Set {
	reg := prometheus.NewRegistry()
	requests := prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "sift_requests_total",
		Help: "HTTP API requests by endpoint and status.",
	}, []string{"endpoint", "code"})
	duration := prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Name:    "sift_request_duration_seconds",
		Help:    "API wall time (fetch/search) or digest time (digest).",
		Buckets: durationBuckets,
	}, []string{"op"})
	reg.MustRegister(requests, duration, &collector{src: src})
	return &Set{
		Reg:      reg,
		Requests: requests,
		Duration: duration,
		handler:  promhttp.HandlerFor(reg, promhttp.HandlerOpts{EnableOpenMetrics: false}),
	}
}

func (s *Set) Observe(endpoint string, dur time.Duration, code int) {
	if s == nil {
		return
	}
	s.Requests.WithLabelValues(endpoint, strconv.Itoa(code)).Inc()
	switch endpoint {
	case "/v1/fetch", "/v2/scrape":
		s.Duration.WithLabelValues("fetch").Observe(dur.Seconds())
	case "/v1/search", "/v2/search":
		s.Duration.WithLabelValues("search").Observe(dur.Seconds())
	}
}

func (s *Set) ObserveDigest(d time.Duration) {
	if s == nil {
		return
	}
	s.Duration.WithLabelValues("digest").Observe(d.Seconds())
}

func (s *Set) Handler() http.Handler {
	if s == nil {
		return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			http.Error(w, "metrics unavailable", http.StatusServiceUnavailable)
		})
	}
	return s.handler
}

// LoopbackOpen is curl http://127.0.0.1:8890/metrics. tailscale serve on :8443
// is not open: it sets forwarded/Tailscale headers and a non-loopback Host.
func LoopbackOpen(r *http.Request) bool {
	if r == nil {
		return false
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	addr, err := netip.ParseAddr(host)
	if err != nil || !addr.IsLoopback() {
		return false
	}
	for _, h := range []string{
		"X-Forwarded-For", "X-Forwarded-Proto", "X-Forwarded-Host",
		"Tailscale-User-Login", "Tailscale-Headers-Info",
	} {
		if r.Header.Get(h) != "" {
			return false
		}
	}
	h := r.Host
	if hh, _, err := net.SplitHostPort(h); err == nil {
		h = hh
	}
	if strings.EqualFold(h, "localhost") {
		return true
	}
	ip, err := netip.ParseAddr(h)
	return err == nil && ip.IsLoopback()
}

func CoerceFallback(reason string) string {
	if _, ok := fallbackOK[reason]; ok {
		return reason
	}
	return "other"
}

func skipRestrictedSource(source string) bool {
	return strings.HasSuffix(source, ":restricted")
}

type collector struct {
	src Source
}

func (c *collector) Describe(ch chan<- *prometheus.Desc) {
	ch <- dBlocks
	ch <- dSelfIP
	ch <- dDigestCalls
	ch <- dDigestFB
	ch <- dDixieIn
	ch <- dDixieWait
	ch <- dCache
	ch <- dQTotal
	ch <- dQCaller
	ch <- dCTotal
	ch <- dCCaller
	ch <- dDaily
	ch <- dCapped
}

func (c *collector) Collect(ch chan<- prometheus.Metric) {
	if c.src == nil {
		return
	}
	s := c.src.Snapshot()
	ch <- prometheus.MustNewConstMetric(dBlocks, prometheus.CounterValue, float64(s.Blocks))
	ch <- prometheus.MustNewConstMetric(dSelfIP, prometheus.CounterValue, float64(s.SelfIPFailures))
	ch <- prometheus.MustNewConstMetric(dDigestCalls, prometheus.CounterValue, float64(s.DigestCalls))
	fb := map[string]uint64{}
	for reason, n := range s.DigestFallbacks {
		fb[CoerceFallback(reason)] += n
	}
	for reason, n := range fb {
		ch <- prometheus.MustNewConstMetric(dDigestFB, prometheus.CounterValue, float64(n), reason)
	}
	ch <- prometheus.MustNewConstMetric(dDixieIn, prometheus.GaugeValue, float64(s.DixieInflight))
	ch <- prometheus.MustNewConstMetric(dDixieWait, prometheus.GaugeValue, float64(s.DixieWaiters))
	emitCache := func(cache string, hit, miss uint64) {
		ch <- prometheus.MustNewConstMetric(dCache, prometheus.CounterValue, float64(hit), cache, "hit")
		ch <- prometheus.MustNewConstMetric(dCache, prometheus.CounterValue, float64(miss), cache, "miss")
	}
	emitCache("digest", s.DigestHit, s.DigestMiss)
	emitCache("embed", s.EmbedHit, s.EmbedMiss)
	emitCache("search", s.SearchHit, s.SearchMiss)

	sumQ := map[string]uint64{}
	sumC := map[string]uint64{}
	for caller, bySrc := range s.SearchQueries {
		restricted := false
		if s.Restricted != nil {
			_, restricted = s.Restricted[caller]
		}
		for src, n := range bySrc {
			if skipRestrictedSource(src) {
				continue
			}
			sumQ[src] += n
			if !restricted {
				ch <- prometheus.MustNewConstMetric(dQCaller, prometheus.CounterValue, float64(n), caller, src)
			}
		}
	}
	for caller, bySrc := range s.SearchCredits {
		restricted := false
		if s.Restricted != nil {
			_, restricted = s.Restricted[caller]
		}
		for src, n := range bySrc {
			if skipRestrictedSource(src) {
				continue
			}
			sumC[src] += n
			if !restricted {
				ch <- prometheus.MustNewConstMetric(dCCaller, prometheus.CounterValue, float64(n), caller, src)
			}
		}
	}
	for src, n := range sumQ {
		ch <- prometheus.MustNewConstMetric(dQTotal, prometheus.CounterValue, float64(n), src)
	}
	for src, n := range sumC {
		ch <- prometheus.MustNewConstMetric(dCTotal, prometheus.CounterValue, float64(n), src)
	}
	for src, n := range s.SearchDaily {
		if skipRestrictedSource(src) {
			continue
		}
		ch <- prometheus.MustNewConstMetric(dDaily, prometheus.GaugeValue, float64(n), src)
	}
	for src, n := range s.SearchCapped {
		if skipRestrictedSource(src) {
			continue
		}
		ch <- prometheus.MustNewConstMetric(dCapped, prometheus.CounterValue, float64(n), src)
	}
}

var (
	dBlocks      = prometheus.NewDesc("sift_guard_blocks_total", "SSRF/guard rejections.", nil, nil)
	dSelfIP      = prometheus.NewDesc("sift_self_ip_failures_total", "Self-IP discovery failures.", nil, nil)
	dDigestCalls = prometheus.NewDesc("sift_digest_calls_total", "Dixie digest calls started (excludes cache hits).", nil, nil)
	dDigestFB    = prometheus.NewDesc("sift_digest_fallbacks_total", "Digest fallbacks by reason.", []string{"reason"}, nil)
	dDixieIn     = prometheus.NewDesc("sift_dixie_inflight", "Sift-side Dixie permits in use.", nil, nil)
	dDixieWait   = prometheus.NewDesc("sift_dixie_queue_waiters", "Waiters in the sift-side Dixie FIFO.", nil, nil)
	dCache       = prometheus.NewDesc("sift_cache_requests_total", "Cache lookups.", []string{"cache", "result"}, nil)
	dQTotal      = prometheus.NewDesc("sift_search_queries_total", "Search queries by source (all callers).", []string{"source"}, nil)
	dQCaller     = prometheus.NewDesc("sift_search_queries_caller_total", "Search queries by fleet caller and source.", []string{"caller", "source"}, nil)
	dCTotal      = prometheus.NewDesc("sift_search_credits_total", "Search credits by source (all callers).", []string{"source"}, nil)
	dCCaller     = prometheus.NewDesc("sift_search_credits_caller_total", "Search credits by fleet caller and source.", []string{"caller", "source"}, nil)
	dDaily       = prometheus.NewDesc("sift_search_daily_credits", "UTC-day paid-source credits (fleet partition).", []string{"source"}, nil)
	dCapped      = prometheus.NewDesc("sift_search_capped_total", "Times a paid source was skipped for the daily cap.", []string{"source"}, nil)
)