ai-gateway-go/internal/shadow/middleware.go

package shadow

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"sync/atomic"

	"aigateway.local/core/internal/gateway"
	"aigateway.local/core/internal/platform/config"
)

type Metrics struct{ eligible, sent, matched, mismatched, errors, skipped atomic.Uint64 }

type Middleware struct {
	config  config.Shadow
	logger  *slog.Logger
	client  *http.Client
	limit   chan struct{}
	metrics Metrics
}

func New(cfg config.Shadow, logger *slog.Logger) *Middleware {
	return &Middleware{config: cfg, logger: logger, client: &http.Client{Timeout: cfg.Timeout, CheckRedirect: func(*http.Request, []*http.Request) error { return fmt.Errorf("shadow redirects are disabled") }}, limit: make(chan struct{}, cfg.MaxConcurrent)}
}

func (m *Middleware) Wrap(next http.Handler) http.Handler {
	if m == nil || m.config.BaseURL == "" || m.config.SampleRate <= 0 {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost || !m.sample(gateway.RequestID(r.Context())) {
			next.ServeHTTP(w, r)
			return
		}
		body, eligible := m.captureRequest(r)
		if !eligible {
			next.ServeHTTP(w, r)
			return
		}
		m.metrics.eligible.Add(1)
		select {
		case m.limit <- struct{}{}:
		default:
			m.metrics.skipped.Add(1)
			next.ServeHTTP(w, r)
			return
		}
		capture := &captureWriter{ResponseWriter: w, limit: m.config.MaxBodyBytes}
		next.ServeHTTP(capture, r)
		primaryStatus := capture.Status()
		primaryBody := append([]byte(nil), capture.body...)
		headers := shadowHeaders(r.Header)
		path := r.URL.RequestURI()
		go func() {
			defer func() { <-m.limit }()
			m.compare(context.WithoutCancel(r.Context()), path, headers, body, primaryStatus, primaryBody)
		}()
	})
}

func (m *Middleware) captureRequest(r *http.Request) ([]byte, bool) {
	if r.Method != http.MethodPost || r.Body == nil || r.Body == http.NoBody {
		return nil, false
	}
	original := r.Body
	raw, err := io.ReadAll(io.LimitReader(original, m.config.MaxBodyBytes+1))
	if err != nil || int64(len(raw)) > m.config.MaxBodyBytes {
		r.Body = &multiReadCloser{Reader: io.MultiReader(bytes.NewReader(raw), original), Closer: original}
		return nil, false
	}
	_ = original.Close()
	r.Body = io.NopCloser(bytes.NewReader(raw))
	r.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(raw)), nil }
	var envelope struct {
		Stream bool `json:"stream"`
	}
	if json.Unmarshal(raw, &envelope) != nil || envelope.Stream {
		return nil, false
	}
	return raw, true
}

type multiReadCloser struct {
	io.Reader
	io.Closer
}

func (m *Middleware) sample(requestID string) bool {
	hash := fnv.New64a()
	_, _ = hash.Write([]byte(requestID))
	return float64(hash.Sum64()%1_000_000)/1_000_000 < m.config.SampleRate
}

func (m *Middleware) compare(ctx context.Context, path string, headers http.Header, body []byte, primaryStatus int, primaryBody []byte) {
	ctx, cancel := context.WithTimeout(ctx, m.config.Timeout)
	defer cancel()
	target, err := url.Parse(m.config.BaseURL + path)
	if err != nil {
		m.fail("invalid shadow URL", err)
		return
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, target.String(), bytes.NewReader(body))
	if err != nil {
		m.fail("create shadow request", err)
		return
	}
	request.Header = headers
	request.Header.Set("Authorization", "Bearer "+m.config.APIKey)
	request.Header.Del("X-Gateway-API-Key")
	request.Header.Set("X-Gateway-Shadow", "true")
	response, err := m.client.Do(request)
	if err != nil {
		m.fail("shadow request failed", err)
		return
	}
	defer response.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(response.Body, m.config.MaxBodyBytes+1))
	if err != nil || int64(len(raw)) > m.config.MaxBodyBytes {
		m.fail("shadow response invalid", err)
		return
	}
	m.metrics.sent.Add(1)
	primarySignature := contractSignature(primaryStatus, primaryBody)
	shadowSignature := contractSignature(response.StatusCode, raw)
	if primarySignature == shadowSignature {
		m.metrics.matched.Add(1)
		return
	}
	m.metrics.mismatched.Add(1)
	if m.logger != nil {
		m.logger.Warn("shadow contract mismatch", "path", path, "primary_status", primaryStatus, "shadow_status", response.StatusCode, "primary_signature", primarySignature, "shadow_signature", shadowSignature)
	}
}
func (m *Middleware) fail(message string, err error) {
	m.metrics.errors.Add(1)
	if m.logger != nil {
		m.logger.Warn(message, "error", err)
	}
}

func contractSignature(status int, raw []byte) string {
	var value any
	if json.Unmarshal(raw, &value) != nil {
		return fmt.Sprintf("%d:non-json", status)
	}
	shape, _ := json.Marshal(jsonShape(value))
	return fmt.Sprintf("%d:%s", status, shape)
}
func jsonShape(value any) any {
	switch item := value.(type) {
	case map[string]any:
		result := make(map[string]any, len(item))
		for key, child := range item {
			switch key {
			case "id", "created", "created_at", "system_fingerprint":
				continue
			}
			result[key] = jsonShape(child)
		}
		return result
	case []any:
		if len(item) == 0 {
			return []any{}
		}
		return []any{jsonShape(item[0])}
	case string:
		return "string"
	case float64:
		return "number"
	case bool:
		return "boolean"
	case nil:
		return nil
	default:
		return fmt.Sprintf("%T", item)
	}
}
func shadowHeaders(source http.Header) http.Header {
	result := make(http.Header)
	for _, key := range []string{"Content-Type", "Accept", "Idempotency-Key", "X-Gateway-Provider"} {
		if value := source.Values(key); len(value) > 0 {
			result[key] = append([]string(nil), value...)
		}
	}
	return result
}
func (m *Middleware) Prometheus() string {
	if m == nil {
		return ""
	}
	return fmt.Sprintf("gateway_shadow_eligible_total %d\ngateway_shadow_sent_total %d\ngateway_shadow_contract_matches_total %d\ngateway_shadow_contract_mismatches_total %d\ngateway_shadow_errors_total %d\ngateway_shadow_skipped_total %d\n", m.metrics.eligible.Load(), m.metrics.sent.Load(), m.metrics.matched.Load(), m.metrics.mismatched.Load(), m.metrics.errors.Load(), m.metrics.skipped.Load())
}

type captureWriter struct {
	http.ResponseWriter
	status int
	body   []byte
	limit  int64
}

func (w *captureWriter) WriteHeader(status int) {
	if w.status == 0 {
		w.status = status
	}
	w.ResponseWriter.WriteHeader(status)
}
func (w *captureWriter) Write(value []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	if int64(len(w.body)) < w.limit {
		remaining := int(w.limit - int64(len(w.body)))
		w.body = append(w.body, value[:min(len(value), remaining)]...)
	}
	return w.ResponseWriter.Write(value)
}
func (w *captureWriter) Status() int {
	if w.status == 0 {
		return http.StatusOK
	}
	return w.status
}
func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *captureWriter) Flush() {
	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
		flusher.Flush()
	}
}