key-cape/src/cmd/keycape/main.go

// keycape is the main server binary for the KeyCape IAM profile service.
// It orchestrates Authelia, LLDAP, and privacyIDEA to implement the
// NetKingdom IAM Profile (OIDC/PKCE Authorization Code Flow).
package main

import (
	"context"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"keycape/internal/server/readiness"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/rs/zerolog"

	"keycape/internal/adapters/authelia"
	"keycape/internal/adapters/lldap"
	"keycape/internal/adapters/privacyidea"
	"keycape/internal/adapters/tenantengine"
	"keycape/internal/authclient"
	"keycape/internal/config"
	"keycape/internal/domain"
	servererrors "keycape/internal/server/errors"
	"keycape/internal/server/oidc"
	"keycape/internal/server/telemetry"
)

const version = "0.1.0"

func main() {
	if len(os.Args) > 1 && os.Args[1] == "probe-upstream-issuer" {
		if err := authelia.RunIssuerProbe(context.Background(), os.Args[2:], os.Stdout); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}

	if len(os.Args) > 1 && (os.Args[1] == "login" || os.Args[1] == "service-token" || os.Args[1] == "verify-client") {
		if err := authclient.Run(context.Background(), os.Args[1:], os.Stderr); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}

	log := zerolog.New(os.Stdout).With().Timestamp().Logger()

	// -----------------------------------------------------------------
	// 1. Parse flags and load config.
	// -----------------------------------------------------------------
	var cfgPath string
	flag.StringVar(&cfgPath, "config", "", "path to YAML config file (env: KEYCAPE_CONFIG)")
	flag.Parse()

	cfg, err := config.Load(cfgPath)
	if err != nil {
		log.Error().Err(err).Msg("failed to load config")
		os.Exit(1)
	}

	// -----------------------------------------------------------------
	// 2. Validate config.
	// -----------------------------------------------------------------
	errs := config.ValidateConfig(cfg)
	if len(errs) > 0 {
		log.Error().Strs("errors", errs).Msg("config validation failed")
		os.Exit(1)
	}

	// -----------------------------------------------------------------
	// 3. Load RSA private key.
	// -----------------------------------------------------------------
	privateKey, err := loadPrivateKey(cfg.PrivateKeyPEM)
	if err != nil {
		log.Error().Err(err).Str("path", cfg.PrivateKeyPEM).Msg("failed to load private key")
		os.Exit(1)
	}

	// -----------------------------------------------------------------
	// 4. Build JWKS from public key.
	// -----------------------------------------------------------------
	ks := oidc.NewKeySet()
	ks.AddKey("key-1", &privateKey.PublicKey)

	// -----------------------------------------------------------------
	// 5. Build client registry.
	// -----------------------------------------------------------------
	clients, err := buildClientRegistry(cfg.Clients)
	if err != nil {
		log.Error().Err(err).Msg("failed to build client registry")
		os.Exit(1)
	}

	// -----------------------------------------------------------------
	// 6. Create adapters.
	// -----------------------------------------------------------------
	lldapAdapter := lldap.New(cfg.LLDAP)
	autheliaAdapter := authelia.New(cfg.Authelia, nil)
	privacyIDEAAdapter := privacyidea.New(cfg.PrivacyIDEA, nil)

	// -----------------------------------------------------------------
	// 7. Create telemetry emitter.
	// -----------------------------------------------------------------
	emitter := telemetry.NewLogEmitter(log)

	// -----------------------------------------------------------------
	// 8. Create enforcement registry.
	// -----------------------------------------------------------------
	enforcement := servererrors.DefaultRegistry()

	// -----------------------------------------------------------------
	// 9. Create session store.
	// -----------------------------------------------------------------
	sessions := oidc.NewSessionStore()

	// -----------------------------------------------------------------
	// 10. Parse token lifetime.
	// -----------------------------------------------------------------
	tokenLifetime := 15 * time.Minute
	if cfg.TokenLifetime != "" {
		d, err := time.ParseDuration(cfg.TokenLifetime)
		if err != nil {
			log.Error().Err(err).Str("tokenLifetime", cfg.TokenLifetime).Msg("invalid tokenLifetime")
			os.Exit(1)
		}
		tokenLifetime = d
	}

	// -----------------------------------------------------------------
	// 11. Build issuer base URL.
	// -----------------------------------------------------------------
	issuer := strings.TrimRight(cfg.Issuer, "/")

	// -----------------------------------------------------------------
	// 12. Register HTTP handlers.
	// -----------------------------------------------------------------
	mux := http.NewServeMux()

	// Discovery.
	mux.Handle("/.well-known/openid-configuration", oidc.NewDiscoveryHandler(oidc.DiscoveryConfig{
		Issuer:                issuer,
		AuthorizationEndpoint: issuer + "/authorize",
		TokenEndpoint:         issuer + "/token",
		JWKSUri:               issuer + "/jwks",
		UserinfoEndpoint:      issuer + "/userinfo",
		EndSessionEndpoint:    issuer + "/logout",
		Clients:               clients,
	}))

	// JWKS.
	mux.Handle("/jwks", oidc.NewJWKSHandler(ks))

	// Authorize handler (with enforcement middleware).
	logins := oidc.NewLoginSessionStore()
	authorizeHandler := &oidc.AuthorizeHandler{
		ClientConfig: clients,
		Auth:         autheliaAdapter,
		MFA:          privacyIDEAAdapter,
		Sessions:     sessions,
		Logins:       logins,
		Handoffs:     oidc.NewHandoffStore(),
		Issuer:       issuer,
		Emitter:      emitter,
	}
	mux.Handle("/authorize", enforcement.Middleware(authorizeHandler))
	mux.Handle("/authorize/callback", authorizeHandler)
	mux.Handle("/authorize/return", authorizeHandler)
	mux.Handle("/authorize/register", authorizeHandler)
	mux.Handle("/logout", &oidc.LogoutHandler{
		ClientConfig: clients,
		Logins:       logins,
		SecureCookie: strings.HasPrefix(strings.ToLower(issuer), "https://"),
	})

	// Token handler (with enforcement middleware).
	tokenHandler := &oidc.TokenHandler{
		ClientConfig:  clients,
		Sessions:      sessions,
		Users:         lldapAdapter,
		SigningKey:    privateKey,
		Issuer:        issuer,
		TokenLifetime: tokenLifetime,
		Emitter:       emitter,
		// Opt-in: nil unless tenantEngine.baseURL is configured, which is what
		// leaves tenant_roles off the stock server (KEY-WP-0024). The adapter
		// fails open, so a configured-but-unreachable cache omits the claim
		// rather than failing issuance.
		TenantEngine: buildTenantEngine(cfg.TenantEngine),
	}
	mux.Handle("/token", enforcement.Middleware(tokenHandler))

	// Userinfo handler.
	mux.Handle("/userinfo", &oidc.UserinfoHandler{
		Users:      lldapAdapter,
		SigningKey: &privateKey.PublicKey,
		Issuer:     issuer,
		Emitter:    emitter,
	})

	// Healthz.
	// /readyz gates traffic on the dependencies a login needs; /healthz stays a
	// liveness signal (KEY-WP-0025). Restarting KeyCape because LLDAP blinked
	// would convert a dependency blip into an outage of our own, so the two must
	// not be the same endpoint.
	mux.Handle("/readyz", readiness.New(2*time.Second, 3*time.Second,
		readiness.Check{Name: "lldap", Probe: func(ctx context.Context) error {
			return lldapAdapter.Ping(ctx)
		}},
		readiness.Check{Name: "authelia", Probe: httpReachable(cfg.Authelia.TokenBaseURL, cfg.Authelia.BaseURL)},
		readiness.Check{Name: "privacyidea", Probe: httpReachable(cfg.PrivacyIDEA.BaseURL)},
	))

	// Liveness only: this says the process is up, and deliberately probes
	// nothing. See /readyz.
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]string{
			"status":  "ok",
			"version": version,
		})
	})

	// Inject emitter into request context.
	handler := withEmitter(mux, emitter)

	// -----------------------------------------------------------------
	// 13. Start HTTP server.
	// -----------------------------------------------------------------
	addr := fmt.Sprintf(":%d", cfg.Port)
	if cfg.Port == 0 {
		addr = ":8080"
	}

	log.Info().
		Str("issuer", issuer).
		Str("addr", addr).
		Str("environment", cfg.Environment).
		Str("version", version).
		Msg("starting keycape server")

	srv := &http.Server{
		Addr:         addr,
		Handler:      handler,
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 30 * time.Second,
		IdleTimeout:  120 * time.Second,
	}

	// Serve in the background so shutdown can be driven by a signal. Without
	// this the process died on SIGTERM mid-request, so every rolling restart
	// returned errors to whoever was logging in at that moment.
	serveErr := make(chan error, 1)
	go func() {
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			serveErr <- err
			return
		}
		serveErr <- nil
	}()

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)

	select {
	case err := <-serveErr:
		if err != nil {
			log.Error().Err(err).Msg("server error")
			os.Exit(1)
		}
	case sig := <-signals:
		log.Info().Str("signal", sig.String()).Msg("shutting down")
		// In-flight requests get a bounded window to finish. In-memory login and
		// authorization state does not survive this and is not meant to: a
		// browser mid-login must start again after a restart. That is a
		// documented limit of the single-replica topology, not a bug to paper
		// over here -- see docs/operations.md.
		ctx, cancel := context.WithTimeout(context.Background(), shutdownGrace)
		defer cancel()
		if err := srv.Shutdown(ctx); err != nil {
			log.Error().Err(err).Msg("graceful shutdown failed; exiting anyway")
			os.Exit(1)
		}
		log.Info().Msg("shutdown complete")
	}
}

// shutdownGrace bounds how long in-flight requests may finish after a signal.
// Kept under the 30s read/write timeouts so a stuck request cannot outlive the
// window an orchestrator typically allows before sending SIGKILL.
const shutdownGrace = 15 * time.Second

// httpReachable probes the first non-empty URL with a GET, reporting whether the
// dependency answered at all. It is a reachability check, not a functional one:
// any HTTP response means something is listening and speaking HTTP there.
func httpReachable(candidates ...string) func(context.Context) error {
	target := ""
	for _, candidate := range candidates {
		if candidate != "" {
			target = candidate
			break
		}
	}
	return func(ctx context.Context) error {
		if target == "" {
			return fmt.Errorf("no endpoint configured")
		}
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
		if err != nil {
			return err
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		return nil
	}
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

// loadPrivateKey reads a PEM file and parses the RSA private key.
// Supports both PKCS#1 ("RSA PRIVATE KEY") and PKCS#8 ("PRIVATE KEY") PEM blocks.
func loadPrivateKey(path string) (*rsa.PrivateKey, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read key file: %w", err)
	}

	block, _ := pem.Decode(data)
	if block == nil {
		return nil, fmt.Errorf("no PEM block found in %q", path)
	}

	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("parse PKCS1 private key: %w", err)
		}
		return key, nil
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("parse PKCS8 private key: %w", err)
		}
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("private key is not an RSA key")
		}
		return rsaKey, nil
	default:
		return nil, fmt.Errorf("unexpected PEM block type %q; expected RSA PRIVATE KEY or PRIVATE KEY", block.Type)
	}
}

// buildClientRegistry converts []ClientConfig into the map used by handlers.
func buildClientRegistry(cfgClients []config.ClientConfig) (map[string]*domain.Client, error) {
	m := make(map[string]*domain.Client, len(cfgClients))
	for i := range cfgClients {
		c := &cfgClients[i]
		clientSecret := ""
		var clientTokenLifetime time.Duration
		if strings.HasPrefix(c.SecretRef, "env:") {
			envName := strings.TrimPrefix(c.SecretRef, "env:")
			clientSecret = os.Getenv(envName)
			if clientSecret == "" {
				return nil, fmt.Errorf("client %q secret environment variable %q is empty", c.ClientID, envName)
			}
		}
		if c.TokenLifetime != "" {
			parsedLifetime, parseErr := time.ParseDuration(c.TokenLifetime)
			if parseErr != nil {
				return nil, fmt.Errorf("client %q tokenLifetime is invalid: %w", c.ClientID, parseErr)
			}
			clientTokenLifetime = parsedLifetime
		}
		m[c.ClientID] = &domain.Client{
			ClientID:        c.ClientID,
			DisplayName:     c.DisplayName,
			RedirectURIs:    c.RedirectURIs,
			AllowedScopes:   c.AllowedScopes,
			GrantTypes:      c.GrantTypes,
			ClientType:      c.ClientType,
			SecretRef:       c.SecretRef,
			ClientSecret:    clientSecret,
			Audience:        c.Audience,
			ServiceSubject:  c.ServiceSubject,
			Tenant:          c.Tenant,
			Roles:           c.Roles,
			TokenLifetime:   clientTokenLifetime,
			MFARequired:     c.MFARequired,
			RegistrationURL: c.RegistrationURL,
			EnrollmentURL:   c.EnrollmentURL,
		}
	}
	return m, nil
}

// withEmitter wraps a handler to inject the telemetry emitter into every request context.
func withEmitter(next http.Handler, e telemetry.Emitter) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := telemetry.WithEmitter(context.Background(), e)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// buildTenantEngine returns the tenant_roles cache client, or nil when the
// claim is not configured. Validation has already accepted the URL and timeout.
func buildTenantEngine(cfg config.TenantEngineConfig) *tenantengine.Client {
	if cfg.BaseURL == "" {
		return nil
	}
	var httpClient *http.Client
	if cfg.Timeout != "" {
		if timeout, err := time.ParseDuration(cfg.Timeout); err == nil {
			httpClient = &http.Client{Timeout: timeout}
		}
	}
	return tenantengine.New(cfg.BaseURL, httpClient)
}