egress-proxies-operator/internal/discovery/server.go

// Package discovery implements the HTTP API crawler clients use to find
// and lease proxies: list healthy proxies filtered by attributes, acquire a
// TTL-based lease, release it early, and report how a target treated the
// proxy. Reads go through the manager's informer cache; lease state lives
// in the injected store.
package discovery

import (
	"context"
	"crypto/subtle"
	"encoding/json"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-logr/logr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	logf "sigs.k8s.io/controller-runtime/pkg/log"

	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/lease"
	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/tracing"
)

// LeaseStore is what the handlers need from a lease backend. Defined here,
// consumer-side, so a CRD- or Redis-backed store can replace the in-memory
// one (which internal/lease's *Store satisfies) without touching handlers.
type LeaseStore interface {
	Acquire(ctx context.Context, req lease.AcquireRequest) (*lease.Lease, lease.AcquireStats, error)
	Release(ctx context.Context, id string)
	Report(ctx context.Context, id string, result lease.Result, target string) error
	Counts() map[string]int
}

// LeaseMetrics counts lease acquisitions by outcome. Implemented by
// internal/metrics; defined here so this package carries no metrics
// dependency.
type LeaseMetrics interface {
	LeaseRequest(outcome string)
}

const (
	defaultAddr     = ":8090"
	defaultTTL      = 5 * time.Minute
	defaultMaxTTL   = time.Hour
	maxBodyBytes    = 64 << 10
	shutdownGrace   = 10 * time.Second
	readHeadTimeout = 5 * time.Second
)

// Server serves the discovery API as a manager Runnable.
type Server struct {
	// Reader lists Proxies from the manager's cache.
	Reader client.Reader
	// Store is the lease backend.
	Store LeaseStore
	// Addr is the listen address (default ":8090"; --discovery-addr).
	Addr string
	// Token is the static bearer token from DISCOVERY_TOKEN. Empty
	// disables auth — allowed for the prototype, but loudly warned about
	// at startup, because in-cluster that is a silent security hole.
	Token string
	// MaxLeaseTTL caps requested lease TTLs (default 1h; --max-lease-ttl).
	MaxLeaseTTL time.Duration
	// Metrics, when non-nil, counts lease requests by outcome.
	Metrics LeaseMetrics

	log logr.Logger

	mu        sync.Mutex
	boundAddr string
}

// NeedLeaderElection is false, and the deployment ships replicas: 1.
// Verified against controller-runtime's runnable ordering: caches start and
// sync before non-leader-election runnables, so cache reads here are safe.
// If this were leader-elected, non-leader replicas would refuse connections
// while still being Service endpoints. The 1-replica constraint comes from
// lease state being per-process — both facts are README caveats.
func (s *Server) NeedLeaderElection() bool { return false }

// BoundAddr returns the actual listen address once Start has bound it —
// meaningful when Addr uses port 0 (tests).
func (s *Server) BoundAddr() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.boundAddr
}

// Start listens and serves until ctx ends, then shuts down gracefully with
// a 10-second grace period.
func (s *Server) Start(ctx context.Context) error {
	if s.Addr == "" {
		s.Addr = defaultAddr
	}
	if s.MaxLeaseTTL == 0 {
		s.MaxLeaseTTL = defaultMaxTTL
	}
	s.log = logf.FromContext(ctx).WithName("discovery")
	if s.Token == "" {
		s.log.Info("WARNING: DISCOVERY_TOKEN is empty — the discovery API is served without authentication")
	}

	ln, err := net.Listen("tcp", s.Addr)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.boundAddr = ln.Addr().String()
	s.mu.Unlock()

	srv := &http.Server{
		Handler:           s.handler(),
		ReadHeaderTimeout: readHeadTimeout,
		ReadTimeout:       10 * time.Second,
		WriteTimeout:      10 * time.Second,
		IdleTimeout:       60 * time.Second,
	}

	errCh := make(chan error, 1)
	go func() { errCh <- srv.Serve(ln) }()
	s.log.Info("discovery API listening", "addr", s.boundAddr)

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownGrace)
		defer cancel()
		if err := srv.Shutdown(shutdownCtx); err != nil {
			return err
		}
		<-errCh // always http.ErrServerClosed after a clean Shutdown
		return nil
	}
}

// handler assembles the mux and the middleware chain, outermost first:
// recover → tracing (server span + request logger) → request-log →
// body-size cap → bearer auth. Everything inside the tracing layer logs via
// logf.FromContext(r.Context()) and so carries traceID/spanID.
func (s *Server) handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok\n"))
	})
	mux.HandleFunc("GET /v1/proxies", s.handleListProxies)
	mux.HandleFunc("POST /v1/leases", s.handleAcquireLease)
	mux.HandleFunc("DELETE /v1/leases/{id}", s.handleReleaseLease)
	mux.HandleFunc("POST /v1/leases/{id}/report", s.handleReportLease)

	var h http.Handler = mux
	h = s.authMiddleware(h)
	h = maxBytesMiddleware(h)
	h = s.logMiddleware(h)
	h = tracing.HTTPMiddleware("discovery", s.log)(h)
	h = s.recoverMiddleware(h)
	return h
}

func (s *Server) recoverMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if p := recover(); p != nil {
				s.log.Error(nil, "panic in discovery handler", "panic", p, "path", r.URL.Path)
				writeError(w, http.StatusInternalServerError, "internal", "internal server error")
			}
		}()
		next.ServeHTTP(w, r)
	})
}

// statusRecorder captures the response code for the request log.
type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (r *statusRecorder) WriteHeader(code int) {
	r.status = code
	r.ResponseWriter.WriteHeader(code)
}

func (s *Server) logMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/healthz" {
			next.ServeHTTP(w, r) // probes are noise
			return
		}
		rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
		start := time.Now()
		next.ServeHTTP(rec, r)
		logf.FromContext(r.Context()).Info("request",
			"method", r.Method, "path", r.URL.Path,
			"status", rec.status, "duration", time.Since(start).String())
	})
}

func maxBytesMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes)
		next.ServeHTTP(w, r)
	})
}

func (s *Server) authMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.Token == "" || r.URL.Path == "/healthz" {
			next.ServeHTTP(w, r)
			return
		}
		token, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
		if !ok || subtle.ConstantTimeCompare([]byte(token), []byte(s.Token)) != 1 {
			writeError(w, http.StatusUnauthorized, "unauthorized", "missing or invalid bearer token")
			return
		}
		next.ServeHTTP(w, r)
	})
}

// errorBody is the shared error shape:
// {"error":"<machine_code>","message":"<human>"}.
type errorBody struct {
	Error   string `json:"error"`
	Message string `json:"message"`
}

func writeError(w http.ResponseWriter, status int, code, message string) {
	writeJSON(w, status, errorBody{Error: code, Message: message})
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}