pocket-id-pocket-id/backend/internal/bootstrap/router_bootstrap.go

package bootstrap

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/fsnotify/fsnotify"
	sloggin "github.com/gin-contrib/slog"
	"github.com/gin-gonic/gin"
	"github.com/italypaleale/francis/builtin/ratelimit"
	"github.com/italypaleale/go-kit/servicerunner"
	"go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin"
	"gorm.io/gorm"

	"github.com/pocket-id/pocket-id/backend/frontend"
	"github.com/pocket-id/pocket-id/backend/internal/common"
	"github.com/pocket-id/pocket-id/backend/internal/controller"
	"github.com/pocket-id/pocket-id/backend/internal/middleware"
	"github.com/pocket-id/pocket-id/backend/internal/tracing"
	"github.com/pocket-id/pocket-id/backend/internal/utils/systemd"
)

// This is used to register additional controllers for tests
var registerTestControllers []func(apiGroup *gin.RouterGroup, db *gorm.DB, svc *services)

func initRouter(db *gorm.DB, svc *services, rateLimitServices map[string]*ratelimit.RateLimitService) (servicerunner.Service, error) {
	r, err := initEngine()
	if err != nil {
		return nil, err
	}
	err = registerRoutes(r, db, svc, rateLimitServices)
	if err != nil {
		return nil, err
	}

	serverConfig, err := initServer(r)
	if err != nil {
		return nil, err
	}

	runFn := func(ctx context.Context) error {
		return runServer(ctx, serverConfig)
	}

	return runFn, nil
}

type serverConfig struct {
	addr         string
	certProvider *tlsCertProvider
	listener     net.Listener
	server       *http.Server
	tlsConfig    *tls.Config
}

func initEngine() (*gin.Engine, error) {
	setGinMode()

	r := gin.New()
	initLogger(r)
	configureEngine(r)
	registerGlobalMiddleware(r)

	return r, nil
}

func setGinMode() {
	// Set the appropriate Gin mode based on the environment
	switch common.EnvConfig.AppEnv {
	case common.AppEnvProduction:
		gin.SetMode(gin.ReleaseMode)
	case common.AppEnvDevelopment:
		gin.SetMode(gin.DebugMode)
	case common.AppEnvTest:
		gin.SetMode(gin.TestMode)
	}
}

func configureEngine(r *gin.Engine) {
	if !common.EnvConfig.TrustProxy {
		_ = r.SetTrustedProxies(nil)
	}

	if common.EnvConfig.TrustedPlatform != "" {
		r.TrustedPlatform = common.EnvConfig.TrustedPlatform
	}

	r.Use(otelgin.Middleware(
		common.Name,
		otelgin.WithFilter(shouldTraceRequest)),
	)
}

// shouldTraceRequest reports whether an incoming request should be traced.
// It traces only requests handled by real backend routes (the API, the OIDC/OAuth endpoints, and the well-known documents).
// Everything else falls through to the frontend NoRoute handler, which serves the SPA shell and static assets; tracing those would produce noisy, unparented spans named just "GET" with an empty http.route.
func shouldTraceRequest(r *http.Request) bool {
	p := r.URL.Path
	switch {
	case strings.HasPrefix(p, "/api/"),
		strings.HasPrefix(p, "/.well-known/"),
		p == "/authorize":
		return true
	default:
		return false
	}
}

func registerGlobalMiddleware(r *gin.Engine) {
	r.Use(middleware.HeadMiddleware())
	r.Use(middleware.NewCacheControlMiddleware().Add())
	r.Use(middleware.NewCorsMiddleware().Add())
	r.Use(middleware.NewCspMiddleware().Add())
	r.Use(middleware.NewErrorHandlerMiddleware().Add())
}

func registerRoutes(r *gin.Engine, db *gorm.DB, svc *services, rateLimitServices map[string]*ratelimit.RateLimitService) error {

	err := frontend.RegisterFrontend(r)
	if errors.Is(err, frontend.ErrFrontendNotIncluded) {
		slog.Warn("Frontend is not included in the build. Skipping frontend registration.")
	} else if err != nil {
		return fmt.Errorf("failed to register frontend: %w", err)
	}

	// Initialize middleware for specific routes
	authMiddleware := middleware.NewAuthMiddleware(svc.apiKeyModule, svc.userService, svc.jwtService)
	fileSizeLimitMiddleware := middleware.NewFileSizeLimitMiddleware()
	rateLimitMiddleware := middleware.NewRateLimitMiddleware(rateLimitServices)
	apiRateLimitMiddleware := rateLimitMiddleware.Add(middleware.RateLimitAPI)

	apiGroup := r.Group("/api", apiRateLimitMiddleware)
	baseGroup := r.Group("/", apiRateLimitMiddleware)

	svc.apiKeyModule.RegisterRoutes(apiGroup,
		authMiddleware.WithAdminNotRequired().Add(),
		authMiddleware.WithAdminNotRequired().WithApiKeyAuthDisabled().Add(),
	)
	svc.webauthnModule.RegisterRoutes(apiGroup,
		authMiddleware.WithAdminNotRequired().Add(),
		rateLimitMiddleware.Add(middleware.RateLimitWebauthnLogin),
		rateLimitMiddleware.Add(middleware.RateLimitWebauthnReauthenticate),
	)
	controller.NewOidcController(apiGroup, authMiddleware, fileSizeLimitMiddleware, svc.oidcService)
	controller.NewUserController(apiGroup, authMiddleware, rateLimitMiddleware, svc.userService, svc.oneTimeAccessService, svc.webauthnModule, svc.appConfigService)
	controller.NewAppConfigController(apiGroup, authMiddleware, svc.appConfigService, svc.emailService, svc.ldapService)
	controller.NewAppImagesController(apiGroup, authMiddleware, svc.appImagesService)
	controller.NewAuditLogController(apiGroup, svc.auditLogService, authMiddleware)
	controller.NewUserGroupController(apiGroup, authMiddleware, svc.userGroupService)
	controller.NewCustomClaimController(apiGroup, authMiddleware, svc.customClaimService)
	controller.NewVersionController(apiGroup, authMiddleware, svc.versionService)
	controller.NewScimController(apiGroup, authMiddleware, svc.scimService)
	svc.userSignUpModule.RegisterRoutes(apiGroup,
		authMiddleware.Add(),
		rateLimitMiddleware.Add(middleware.RateLimitSignup),
	)

	optionalBrowserAuth := authMiddleware.WithAdminNotRequired().WithSuccessOptional().WithApiKeyAuthDisabled().Add()
	browserAuth := authMiddleware.WithAdminNotRequired().WithApiKeyAuthDisabled().Add()
	svc.oidcModule.RegisterRoutes(baseGroup, apiGroup, optionalBrowserAuth, browserAuth)

	registerTestRoutes(apiGroup, db, svc)

	controller.NewWellKnownController(baseGroup, svc.jwtService)

	// These are not rate-limited.
	controller.NewHealthzController(r)

	// Receives OTLP trace payloads from the browser SPA (POST /internal/telemetry/traces) and forwards them to the collector, when trace export is enabled.
	// Outside /api, so it's unauthenticated and not traced, but it is rate-limited.
	tracing.NewTelemetryController(r, rateLimitMiddleware.Add(middleware.RateLimitInternal))

	return nil
}

func registerTestRoutes(apiGroup *gin.RouterGroup, db *gorm.DB, svc *services) {
	if common.EnvConfig.AppEnv.IsProduction() {
		return
	}

	for _, f := range registerTestControllers {
		f(apiGroup, db, svc)
	}
}

func initServer(r *gin.Engine) (*serverConfig, error) {
	protocols, tlsConfig, certProvider, err := initServerProtocols()
	if err != nil {
		return nil, err
	}

	var socketFn func() (*socket, error)
	switch {
	case common.EnvConfig.SystemdSocket:
		socketFn = systemdSocket
	case common.EnvConfig.UnixSocket != "":
		socketFn = unixSocket
	default:
		socketFn = tcpSocket
	}

	socket, err := socketFn()
	if err != nil {
		return nil, err
	}

	addr := socket.addr
	listener := socket.listener
	server := newHTTPServer(r, protocols)

	return &serverConfig{addr, certProvider, listener, server, tlsConfig}, nil
}

func initServerProtocols() (*http.Protocols, *tls.Config, *tlsCertProvider, error) {
	protocols := new(http.Protocols)
	protocols.SetHTTP1(true)

	if common.EnvConfig.TLSCertFile == "" || common.EnvConfig.TLSKeyFile == "" {
		protocols.SetUnencryptedHTTP2(true)
		return protocols, nil, nil, nil
	}

	protocols.SetHTTP2(true)
	certProvider, err := newCertProvider(common.EnvConfig.TLSCertFile, common.EnvConfig.TLSKeyFile)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to load TLS certificate: %w", err)
	}

	tlsConfig := &tls.Config{
		GetCertificate: certProvider.GetCertificate,
		MinVersion:     tls.VersionTLS13,
		NextProtos:     []string{"h2"},
	}

	slog.Info("TLS enabled")
	return protocols, tlsConfig, certProvider, nil
}

func newHTTPServer(r *gin.Engine, protocols *http.Protocols) *http.Server {
	return &http.Server{
		MaxHeaderBytes:    1 << 20,
		ReadHeaderTimeout: 10 * time.Second,
		Protocols:         protocols,
		Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			// HEAD requests don't get matched by Gin routes, so we convert them to GET
			// middleware.HeadMiddleware will convert them back to HEAD later
			if req.Method == http.MethodHead {
				req.Method = http.MethodGet
				ctx := context.WithValue(req.Context(), middleware.IsHeadRequestCtxKey{}, true)
				req = req.WithContext(ctx)
			}

			r.ServeHTTP(w, req)
		}),
	}
}

func runServer(ctx context.Context, config *serverConfig) error {
	slog.Info("Server listening", slog.String("addr", config.addr), slog.Bool("tls", config.tlsConfig != nil))

	certWatcher, err := startCertWatcher(ctx, config.certProvider)
	if err != nil {
		return err
	}
	defer closeCertWatcher(certWatcher)

	startHTTPServer(config)
	notifySystemdReady()

	<-ctx.Done()

	// We do not pass the context because it's already been canceled
	//nolint:contextcheck
	return shutdownServer(config.server)
}

func startCertWatcher(ctx context.Context, certProvider *tlsCertProvider) (*fsnotify.Watcher, error) {
	if certProvider == nil {
		return nil, nil
	}

	certWatcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("failed to create certificate watcher: %w", err)
	}

	if err := certWatcher.Add(common.EnvConfig.TLSCertFile); err != nil {
		certWatcher.Close()
		return nil, fmt.Errorf("failed to watch TLS certificate: %w", err)
	}
	if err := certWatcher.Add(common.EnvConfig.TLSKeyFile); err != nil {
		certWatcher.Close()
		return nil, fmt.Errorf("failed to watch TLS key: %w", err)
	}

	go certProvider.StartWatching(ctx, certWatcher)
	return certWatcher, nil
}

func closeCertWatcher(certWatcher *fsnotify.Watcher) {
	if certWatcher != nil {
		certWatcher.Close()
	}
}

func startHTTPServer(config *serverConfig) {
	go func() {
		defer config.listener.Close()

		listener := config.listener
		if config.tlsConfig != nil {
			listener = tls.NewListener(config.listener, config.tlsConfig)
		}
		srvErr := config.server.Serve(listener)

		if srvErr != http.ErrServerClosed {
			slog.Error("Error starting app server", "error", srvErr)
			os.Exit(1)
		}
	}()
}

func notifySystemdReady() {
	err := systemd.SdNotifyReady()
	if err != nil {
		// Log the error only
		slog.Warn("Unable to notify systemd that the service is ready", "error", err)
	}
}

func shutdownServer(srv *http.Server) error {
	// Note we use the background context here as ctx has been canceled already
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
	shutdownErr := srv.Shutdown(shutdownCtx) //nolint:contextcheck
	shutdownCancel()
	if shutdownErr != nil {
		// Log the error only (could be context canceled)
		slog.Warn("App server shutdown error", "error", shutdownErr)
	}

	return nil
}

func initLogger(r *gin.Engine) {
	loggerSkipPathsPrefix := []string{
		"GET /api/application-images/logo",
		"GET /api/application-images/background",
		"GET /api/application-images/favicon",
		"GET /api/application-images/email",
		"GET /_app",
		"GET /fonts",
		"GET /healthz",
		"HEAD /healthz",
	}

	r.Use(sloggin.SetLogger(
		sloggin.WithLogger(func(_ *gin.Context, _ *slog.Logger) *slog.Logger {
			return slog.Default()
		}),
		sloggin.WithSkipper(func(c *gin.Context) bool {
			for _, prefix := range loggerSkipPathsPrefix {
				if strings.HasPrefix(c.Request.Method+" "+c.Request.URL.String(), prefix) {
					return true
				}
			}
			return false
		}),
	))
}

// tlsCertProvider holds certificates that can be dynamically reloaded
type tlsCertProvider struct {
	certMutex   sync.RWMutex
	cert        *tls.Certificate
	certFile    string
	keyFile     string
	forceReload atomic.Bool
}

// GetCertificate implements tls.GetCertificate interface for dynamic certificate loading
func (p *tlsCertProvider) GetCertificate(_ *tls.ClientHelloInfo) (*tls.Certificate, error) {
	if p.forceReload.Load() {
		p.certMutex.Lock()
		p.forceReload.Store(false)
		p.certMutex.Unlock()
	}

	p.certMutex.RLock()
	defer p.certMutex.RUnlock()
	return p.cert, nil
}

// newCertProvider creates a new certificate provider with initial certificates loaded
func newCertProvider(certFile, keyFile string) (*tlsCertProvider, error) {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, err
	}

	return &tlsCertProvider{
		cert:     &cert,
		certFile: certFile,
		keyFile:  keyFile,
	}, nil
}

// reloadCertificate reloads the certificate from disk
func (p *tlsCertProvider) reloadCertificate() error {
	cert, err := tls.LoadX509KeyPair(p.certFile, p.keyFile)
	if err != nil {
		return fmt.Errorf("failed to reload TLS certificate: %w", err)
	}

	p.certMutex.Lock()
	p.cert = &cert
	p.certMutex.Unlock()

	return nil
}

// StartWatching begins monitoring the certificate files for changes with debouncing
func (p *tlsCertProvider) StartWatching(ctx context.Context, watcher *fsnotify.Watcher) {
	debounceDuration := 1 * time.Second
	reloadTimer := time.NewTimer(debounceDuration)
	reloadTimer.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case event, ok := <-watcher.Events:
			if !ok {
				return
			}
			// Only process write/rename events for certificate/key files
			if event.Has(fsnotify.Write | fsnotify.Rename) {
				// Reset the debounce timer whenever we get a relevant event
				reloadTimer.Stop()
				// Drain the channel if there's a pending value
				select {
				case <-reloadTimer.C:
				default:
				}
				reloadTimer.Reset(debounceDuration)
				slog.Debug("TLS file change detected, debouncing", slog.String("path", event.Name))
			}
		case <-reloadTimer.C:
			// Timer fired - no more events in 500ms, so reload
			slog.Info("Reloading TLS certificate")

			if err := p.reloadCertificate(); err != nil {
				slog.Error("Failed to reload TLS certificate", "error", err)
				continue
			}

			p.forceReload.Store(true)
			slog.Info("TLS certificate reloaded successfully")
		case err, ok := <-watcher.Errors:
			if !ok {
				return
			}
			slog.Error("Certificate watcher error", "error", err)
		}
	}
}