llmgw/internal/auth/middleware.go

package auth

import (
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"strings"
)

// Middleware authenticates every request before it reaches the router.
//
// It is a wrapper rather than a check inside the handler so there is no path
// into the API that forgets to call it. Adding a route cannot accidentally add
// an unauthenticated one.
func Middleware(a *Authenticator, log *slog.Logger, exempt ...string) func(http.Handler) http.Handler {
	if log == nil {
		log = slog.Default()
	}
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if isExempt(r.URL.Path, exempt) {
				next.ServeHTTP(w, r)
				return
			}

			p, err := a.Authenticate(r.Context(), BearerToken(r))
			if err != nil {
				writeAuthError(w, err)
				// The key prefix, never the key: a rejected credential in a log
				// is still a credential.
				log.Debug("rejected a request", "path", r.URL.Path,
					"reason", err, "remote", r.RemoteAddr)
				return
			}
			next.ServeHTTP(w, r.WithContext(WithPrincipal(r.Context(), p)))
		})
	}
}

// isExempt reports whether a path is served without a key.
//
// Health is the only one by default. Metrics live on the admin listener, which
// is a separate port precisely so it can be firewalled rather than
// authenticated.
func isExempt(path string, exempt []string) bool {
	for _, e := range exempt {
		if path == e || (strings.HasSuffix(e, "/") && strings.HasPrefix(path, e)) {
			return true
		}
	}
	return false
}

func writeAuthError(w http.ResponseWriter, err error) {
	status, kind, msg := http.StatusUnauthorized, "invalid_request_error", "invalid api key"
	switch {
	case errors.Is(err, ErrMissing):
		msg = "no api key was provided"
	case errors.Is(err, ErrExpired):
		msg = "this api key has expired"
	case errors.Is(err, ErrRevoked):
		// Deliberately the same message as an unknown key. Telling a caller
		// that a key exists but was revoked confirms the key was real.
		msg = "invalid api key"
	case errors.Is(err, ErrInvalid):
	default:
		// A store failure is ours, not the caller's, and must not read as a
		// rejected credential — that would send someone hunting for a key
		// problem that does not exist.
		status, kind, msg = http.StatusServiceUnavailable, "server_error",
			"authentication is temporarily unavailable"
	}

	if status == http.StatusUnauthorized {
		// What the OpenAI SDKs look for to stop retrying.
		w.Header().Set("WWW-Authenticate", `Bearer realm="llmgw"`)
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
		"message": msg, "type": kind, "code": nil,
	}})
}

// WriteForbiddenModel reports a model this key may not use.
//
// Separate from authentication because it happens later: the model is only
// known once the body has been parsed.
func WriteForbiddenModel(w http.ResponseWriter, model string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusForbidden)
	_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
		"message": "this api key is not permitted to use model " + model,
		"type":    "invalid_request_error", "code": nil,
	}})
}