easyai-ai-gateway/apps/api/internal/securityevents/metrics.go

package securityevents

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type MetricsSnapshotProvider interface {
	SecurityEventMetrics(context.Context, string, string) (string, *time.Time, error)
}

type DynamicMetricsSnapshotProvider interface {
	MetricsSnapshotProvider
	SecurityEventConnection(context.Context) (store.SecurityEventConnection, error)
}

type Metrics struct {
	accepted              atomic.Uint64
	rejected              atomic.Uint64
	duplicate             atomic.Uint64
	sessionsDeleted       atomic.Uint64
	watermarkRejected     atomic.Uint64
	verificationAccepted  atomic.Uint64
	heartbeatAccepted     atomic.Uint64
	heartbeatFailed       atomic.Uint64
	introspectionActive   atomic.Uint64
	introspectionInactive atomic.Uint64
	introspectionFailed   atomic.Uint64
	jwksSSFFailed         atomic.Uint64
	jwksOIDCFailed        atomic.Uint64
	processingCount       atomic.Uint64
	processingNanos       atomic.Uint64
	processingBuckets     [6]atomic.Uint64
}

var processingDurationBounds = [...]time.Duration{
	10 * time.Millisecond, 50 * time.Millisecond, 100 * time.Millisecond,
	500 * time.Millisecond, time.Second, 3 * time.Second,
}

func (m *Metrics) ObserveReceipt(outcome string, duration time.Duration, sessionsDeleted int64) {
	switch outcome {
	case "accepted":
		m.accepted.Add(1)
	case "duplicate":
		m.duplicate.Add(1)
	default:
		m.rejected.Add(1)
	}
	if sessionsDeleted > 0 {
		m.sessionsDeleted.Add(uint64(sessionsDeleted))
	}
	if duration < 0 {
		duration = 0
	}
	m.processingCount.Add(1)
	m.processingNanos.Add(uint64(duration))
	for index, bound := range processingDurationBounds {
		if duration <= bound {
			m.processingBuckets[index].Add(1)
		}
	}
}

func (m *Metrics) ObserveWatermarkRejection() { m.watermarkRejected.Add(1) }

func (m *Metrics) ObserveVerificationAccepted() { m.verificationAccepted.Add(1) }

func (m *Metrics) ObserveHeartbeat(outcome string) {
	if outcome == "accepted" {
		m.heartbeatAccepted.Add(1)
		return
	}
	m.heartbeatFailed.Add(1)
}

func (m *Metrics) ObserveIntrospection(outcome string) {
	switch outcome {
	case "active":
		m.introspectionActive.Add(1)
	case "inactive":
		m.introspectionInactive.Add(1)
	default:
		m.introspectionFailed.Add(1)
	}
}

func (m *Metrics) ObserveJWKSRefreshFailure(source string) {
	if source == "ssf" {
		m.jwksSSFFailed.Add(1)
		return
	}
	m.jwksOIDCFailed.Add(1)
}

func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience string, enabled bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mode := "disabled"
		var lastVerificationAt *time.Time
		if enabled {
			ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
			defer cancel()
			var err error
			mode, lastVerificationAt, err = provider.SecurityEventMetrics(ctx, issuer, audience)
			if err != nil {
				http.Error(w, "metrics unavailable", http.StatusServiceUnavailable)
				return
			}
		}
		w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
		outcomeCounters(w, "easyai_gateway_ssf_receipts_total", "Received SETs by bounded outcome.", []outcomeValue{
			{"accepted", m.accepted.Load()}, {"rejected", m.rejected.Load()}, {"duplicate", m.duplicate.Load()},
		})
		plainCounter(w, "easyai_gateway_ssf_sessions_deleted_total", "OIDC browser sessions deleted by accepted session-revoked events.", m.sessionsDeleted.Load())
		plainCounter(w, "easyai_gateway_ssf_watermark_rejections_total", "OIDC access tokens rejected by a local revocation watermark.", m.watermarkRejected.Load())
		plainCounter(w, "easyai_gateway_ssf_verifications_total", "Valid verification SETs accepted by the receiver.", m.verificationAccepted.Load())
		outcomeCounters(w, "easyai_gateway_ssf_heartbeat_requests_total", "Verification requests by bounded outcome.", []outcomeValue{
			{"accepted", m.heartbeatAccepted.Load()}, {"failed", m.heartbeatFailed.Load()},
		})
		outcomeCounters(w, "easyai_gateway_oidc_introspection_total", "OIDC introspection requests by bounded outcome.", []outcomeValue{
			{"active", m.introspectionActive.Load()}, {"inactive", m.introspectionInactive.Load()}, {"failed", m.introspectionFailed.Load()},
		})
		outcomeCounters(w, "easyai_gateway_jwks_refresh_failures_total", "JWKS refresh failures by bounded source.", []outcomeValue{
			{"ssf", m.jwksSSFFailed.Load()}, {"oidc", m.jwksOIDCFailed.Load()},
		})
		fmt.Fprintln(w, "# HELP easyai_gateway_ssf_processing_duration_seconds SSF receiver transaction processing time.")
		fmt.Fprintln(w, "# TYPE easyai_gateway_ssf_processing_duration_seconds histogram")
		for index, bound := range processingDurationBounds {
			fmt.Fprintf(w, "easyai_gateway_ssf_processing_duration_seconds_bucket{le=\"%.2f\"} %d\n", bound.Seconds(), m.processingBuckets[index].Load())
		}
		fmt.Fprintf(w, "easyai_gateway_ssf_processing_duration_seconds_bucket{le=\"+Inf\"} %d\n", m.processingCount.Load())
		fmt.Fprintf(w, "easyai_gateway_ssf_processing_duration_seconds_sum %.6f\n", float64(m.processingNanos.Load())/float64(time.Second))
		fmt.Fprintf(w, "easyai_gateway_ssf_processing_duration_seconds_count %d\n", m.processingCount.Load())
		verificationAge := float64(0)
		if lastVerificationAt != nil {
			verificationAge = time.Since(*lastVerificationAt).Seconds()
			if verificationAge < 0 {
				verificationAge = 0
			}
		}
		fmt.Fprintln(w, "# HELP easyai_gateway_ssf_verification_age_seconds Seconds since the last matched verification SET.")
		fmt.Fprintln(w, "# TYPE easyai_gateway_ssf_verification_age_seconds gauge")
		fmt.Fprintf(w, "easyai_gateway_ssf_verification_age_seconds %.6f\n", verificationAge)
		fmt.Fprintln(w, "# HELP easyai_gateway_ssf_mode Current SSF receiver mode as a one-hot gauge.")
		fmt.Fprintln(w, "# TYPE easyai_gateway_ssf_mode gauge")
		for _, modeName := range []string{"disabled", "bootstrap", "push_healthy", "introspection_fallback"} {
			value := 0
			if mode == modeName {
				value = 1
			}
			fmt.Fprintf(w, "easyai_gateway_ssf_mode{mode=\"%s\"} %d\n", modeName, value)
		}
	})
}

func (m *Metrics) DynamicHandler(provider DynamicMetricsSnapshotProvider) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		connection, err := provider.SecurityEventConnection(r.Context())
		if errors.Is(err, store.ErrSecurityEventConnectionNotFound) {
			m.Handler(provider, "", "", false).ServeHTTP(w, r)
			return
		}
		if err != nil {
			http.Error(w, "metrics unavailable", http.StatusServiceUnavailable)
			return
		}
		if connection.Audience == nil {
			m.Handler(provider, "", "", false).ServeHTTP(w, r)
			return
		}
		m.Handler(provider, connection.TransmitterIssuer, *connection.Audience, true).ServeHTTP(w, r)
	})
}

type outcomeValue struct {
	name  string
	value uint64
}

func outcomeCounters(w http.ResponseWriter, name, help string, values []outcomeValue) {
	fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s counter\n", name, help, name)
	for _, value := range values {
		fmt.Fprintf(w, "%s{outcome=\"%s\"} %d\n", name, value.name, value.value)
	}
}

func plainCounter(w http.ResponseWriter, name, help string, value uint64) {
	fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s counter\n%s %d\n", name, help, name, name, value)
}