story/internal/externalprobe/probe.go

package externalprobe

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

const maxJWKSBodyBytes = 1 << 20

type CheckFunc func(context.Context) error

type Observation struct {
	ServiceName string
	Dependency  string
	Available   bool
	CheckedAt   time.Time
	Latency     time.Duration
	Interval    time.Duration
	ErrorCode   string
}

type Observer func(context.Context, Observation) error

// Runner executes low-cost external dependency checks without making them part
// of readiness. A temporary model or RTC network failure must remain visible to
// operations without evicting otherwise healthy API and Worker pods.
type Runner struct {
	interval  time.Duration
	timeout   time.Duration
	checks    map[string]CheckFunc
	logger    *slog.Logger
	states    map[string]bool
	observers []Observer
	mu        sync.Mutex
}

func NewRunner(interval, timeout time.Duration, checks map[string]CheckFunc, logger *slog.Logger, observers ...Observer) (*Runner, error) {
	if interval <= 0 || timeout <= 0 {
		return nil, errors.New("external probe interval and timeout must be positive")
	}
	if logger == nil {
		logger = slog.Default()
	}
	ownedChecks := make(map[string]CheckFunc, len(checks))
	for name, check := range checks {
		name = strings.TrimSpace(name)
		if name == "" || check == nil {
			return nil, errors.New("external probe name and check are required")
		}
		ownedChecks[name] = check
	}
	for _, observer := range observers {
		if observer == nil {
			return nil, errors.New("external probe observer is required")
		}
	}
	return &Runner{interval: interval, timeout: timeout, checks: ownedChecks, logger: logger, states: map[string]bool{}, observers: append([]Observer(nil), observers...)}, nil
}

func (r *Runner) Run(ctx context.Context) {
	r.runOnce(ctx)
	ticker := time.NewTicker(r.interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			r.runOnce(ctx)
		}
	}
}

func (r *Runner) runOnce(ctx context.Context) {
	for name, check := range r.checks {
		checkContext, cancel := context.WithTimeout(ctx, r.timeout)
		startedAt := time.Now()
		err := check(checkContext)
		checkedAt := time.Now().UTC()
		latency := time.Since(startedAt)
		cancel()
		if latency < 0 {
			latency = 0
		}
		r.mu.Lock()
		previous, observed := r.states[name]
		available := err == nil
		r.states[name] = available
		r.mu.Unlock()
		if !available && (!observed || previous) {
			r.logger.Warn("external dependency probe failed", "dependency", name, "error", err)
		}
		if available && observed && !previous {
			r.logger.Info("external dependency probe recovered", "dependency", name)
		}
		observation := Observation{Dependency: name, Available: available, CheckedAt: checkedAt, Latency: latency, Interval: r.interval, ErrorCode: classifyError(err)}
		for _, observer := range r.observers {
			observerContext, observerCancel := context.WithTimeout(ctx, r.timeout)
			observerErr := observer(observerContext, observation)
			observerCancel()
			if observerErr != nil && ctx.Err() == nil {
				r.logger.Warn("external dependency probe observation failed", "dependency", name, "error", observerErr)
			}
		}
	}
}

func classifyError(err error) string {
	if err == nil {
		return ""
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return "TIMEOUT"
	}
	if errors.Is(err, context.Canceled) {
		return "CANCELED"
	}
	message := strings.ToLower(err.Error())
	if strings.Contains(message, "certificate") || strings.Contains(message, "tls") || strings.Contains(message, "x509") {
		return "TLS_ERROR"
	}
	if strings.Contains(message, "jwks") {
		return "JWKS_INVALID"
	}
	var networkError net.Error
	if errors.As(err, &networkError) {
		return "NETWORK_ERROR"
	}
	return "PROBE_FAILED"
}

// JWKS validates that the configured endpoint returns a bounded JSON Web Key
// Set with at least one key. It does not persist the keys or expose their body.
func JWKS(endpoint string, client *http.Client) (CheckFunc, error) {
	parsed, err := url.Parse(strings.TrimSpace(endpoint))
	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
		return nil, errors.New("JWKS endpoint must be an absolute HTTP(S) URL")
	}
	if client == nil {
		client = &http.Client{}
	}
	return func(ctx context.Context) error {
		request, requestErr := http.NewRequestWithContext(ctx, http.MethodGet, parsed.String(), nil)
		if requestErr != nil {
			return fmt.Errorf("create JWKS request: %w", requestErr)
		}
		request.Header.Set("Accept", "application/json")
		response, requestErr := client.Do(request)
		if requestErr != nil {
			return fmt.Errorf("request JWKS: %w", requestErr)
		}
		defer response.Body.Close()
		if response.StatusCode < 200 || response.StatusCode >= 300 {
			_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4*1024))
			return fmt.Errorf("JWKS endpoint returned HTTP %d", response.StatusCode)
		}
		var document struct {
			Keys []json.RawMessage `json:"keys"`
		}
		decoder := json.NewDecoder(io.LimitReader(response.Body, maxJWKSBodyBytes+1))
		if decodeErr := decoder.Decode(&document); decodeErr != nil {
			return fmt.Errorf("decode JWKS document: %w", decodeErr)
		}
		if len(document.Keys) == 0 {
			return errors.New("JWKS document contains no keys")
		}
		return nil
	}, nil
}

// TLSReachability performs only DNS, TCP and certificate-verified TLS setup.
// It intentionally sends no provider API request and therefore consumes no
// model tokens, voice minutes or provider task quota.
func TLSReachability(endpoint string, rootCAs *x509.CertPool) (CheckFunc, error) {
	parsed, err := url.Parse(strings.TrimSpace(endpoint))
	if err != nil || parsed.Scheme != "https" || parsed.Hostname() == "" {
		return nil, errors.New("TLS probe endpoint must be an absolute HTTPS URL")
	}
	port := parsed.Port()
	if port == "" {
		port = "443"
	}
	address := net.JoinHostPort(parsed.Hostname(), port)
	config := &tls.Config{
		MinVersion: tls.VersionTLS12,
		ServerName: parsed.Hostname(),
		RootCAs:    rootCAs,
	}
	return func(ctx context.Context) error {
		dialer := &tls.Dialer{NetDialer: &net.Dialer{}, Config: config.Clone()}
		connection, dialErr := dialer.DialContext(ctx, "tcp", address)
		if dialErr != nil {
			return fmt.Errorf("establish verified TLS connection: %w", dialErr)
		}
		return connection.Close()
	}, nil
}