go-omni/app/app.go

package app

import (
	"context"
	"errors"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/config"
	"git.yuntuo.hk/k2safe/go-omni/cors"
	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
	"git.yuntuo.hk/k2safe/go-omni/middleware"
	"git.yuntuo.hk/k2safe/go-omni/response"
	netrouter "git.yuntuo.hk/k2safe/go-omni/transport/nethttp"
)

// App is the top-level GoOmni application. Its public behavior depends on
// httpx.Router, leaving the network implementation replaceable.
type App struct {
	router           httpx.Router
	handler          http.Handler
	config           config.Config
	middlewares      []httpx.Middleware
	closers          []func() error
	closeOnce        sync.Once
	closeErr         error
	clientIPResolver *httpx.ClientIPResolver
	responseEncoder  response.Encoder
	requests         requestDrain
	bodyLimits       []routeBodyLimit
	validateOnce     sync.Once
	validateErr      error
}

// Runner is the small lifecycle boundary used by an HTTP server
// implementation. App.RunContext selects a configured built-in server; a custom runner
// can own the listener, protocol parser and connection lifecycle.
type Runner interface {
	Run(context.Context, http.Handler, string) error
}

// New creates an application using GoOmni's standard-library router.
func New(options ...config.Option) *App {
	settings := configure(options...)
	return newApp(netrouter.New(), settings)
}

// NewWithRouter creates an application with a custom network layer. A custom
// router only needs to implement httpx.Router; handlers and business code stay
// unchanged when the transport changes.
func NewWithRouter(router httpx.Router, options ...config.Option) *App {
	if router == nil {
		panic("go-omni: nil router")
	}
	return newApp(router, configure(options...))
}

func configure(options ...config.Option) config.Config {
	settings := config.Default()
	for _, option := range options {
		if option != nil {
			option(&settings)
		}
	}
	return settings
}

func newApp(router httpx.Router, settings config.Config) *App {
	settings.CORS = settings.CORS.Clone()
	corsPolicy, corsErr := cors.New(settings.CORS)
	if corsErr != nil {
		panic(config.AtField("cors", corsErr))
	}
	settings.Response = settings.Response.Clone()
	settings.Errors = settings.Errors.Clone()
	encoder, responseErr := response.New(settings.Response)
	if responseErr != nil {
		panic(responseErr)
	}
	var responseEncoder response.Encoder = encoder
	if settings.ResponseEncoder != nil {
		responseEncoder = settings.ResponseEncoder
	}
	settings.TrustedProxies = append([]string(nil), settings.TrustedProxies...)
	settings.Log.RedactKeys = append([]string(nil), settings.Log.RedactKeys...)
	responseEncoder, responseErr = response.WithErrors(settings.Errors, responseEncoder)
	if responseErr != nil {
		panic(responseErr)
	}
	resolver, err := httpx.NewClientIPResolver(settings.TrustedProxies...)
	if err != nil {
		panic("go-omni: " + err.Error())
	}
	middlewares := []httpx.Middleware{middleware.RequestID(settings.RequestIDHeader)}
	if settings.EnableLogger {
		middlewares = append(middlewares, middleware.AccessLog(middleware.WithSlowRequestThreshold(settings.SlowRequestThreshold)))
	}
	middlewares = append(middlewares, middleware.Recovery())
	application := &App{router: router, config: settings, middlewares: middlewares, clientIPResolver: resolver, responseEncoder: responseEncoder}
	application.requests.done = make(chan struct{})
	application.handler = corsPolicy.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.ContentLength > application.hardBodyLimit() {
			application.adapt(func(*httpx.Context) error { return apierrors.New(413, "BODY_TOO_LARGE", "request body exceeds limit") }).ServeHTTP(w, r)
			return
		}
		router.ServeHTTP(w, r)
	}))
	if fallback, ok := router.(httpx.FallbackRouter); ok {
		adaptError := func(status int, code, message string) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				application.adapt(func(*httpx.Context) error { return apierrors.New(status, code, message) }).ServeHTTP(w, r)
			})
		}
		fallback.SetFallbacks(adaptError(404, "NOT_FOUND", "route not found"), adaptError(405, "METHOD_NOT_ALLOWED", "method not allowed"))
	}
	return application
}

// Router exposes the underlying routing contract for advanced registration.
// Serve through App or Engine for CORS, hard body limits and shutdown tracking.
func (a *App) Router() httpx.Router { return a.router }

// Engine is kept as a compatibility name for the old API. It returns a
// standard http.Handler, so callers do not depend on a router type.
func (a *App) Engine() http.Handler { return a }

// ServeHTTP implements http.Handler for standard servers and HTTP tests.
func (a *App) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	ctx, work := a.requests.begin(request.Context())
	if work == nil {
		// Do not invoke encoders, middleware or loggers after their owners close.
		http.Error(writer, "server shutting down", http.StatusServiceUnavailable)
		return
	}
	defer a.requests.end(work)
	a.validateOnce.Do(func() { a.validateErr = a.ValidateHTTP() })
	if a.validateErr != nil {
		http.Error(writer, "invalid HTTP configuration", http.StatusServiceUnavailable)
		return
	}
	request = request.WithContext(ctx)
	if request.Body != nil {
		request.Body = http.MaxBytesReader(writer, request.Body, a.hardBodyLimit())
	}
	a.handler.ServeHTTP(writer, request)
}

// Config returns the effective application configuration.
func (a *App) Config() config.Config {
	settings := a.config
	settings.CORS = settings.CORS.Clone()
	settings.Response = settings.Response.Clone()
	settings.Errors = settings.Errors.Clone()
	settings.TrustedProxies = append([]string(nil), settings.TrustedProxies...)
	settings.Log.RedactKeys = append([]string(nil), settings.Log.RedactKeys...)
	return settings
}

// Use adds middleware to routes registered after this call.
func (a *App) Use(middleware ...httpx.Middleware) *App {
	a.middlewares = append(a.middlewares, middleware...)
	return a
}

// OnShutdown registers a cleanup function owned by the application. Cleanup
// functions run in reverse registration order after all handlers have returned.
// Incomplete draining skips these callbacks and returns ErrResourcesRetained.
// This keeps omni.New as the only public startup entry point while allowing
// projects to attach databases, loggers and other resources to its lifecycle.
func (a *App) OnShutdown(cleanup func() error) *App {
	if cleanup != nil {
		a.closers = append(a.closers, cleanup)
	}
	return a
}

// Close stops admission, cancels active requests and waits for their handlers
// to return before releasing resources. It is idempotent. If cancellation grace
// expires, ErrResourcesRetained is permanent and no cleanup callbacks run.
// Use RunContext/ServeContext for graceful draining before cancellation. Do not
// call Close from a handler or cleanup callback; it would wait for itself.
func (a *App) Close() error {
	a.closeOnce.Do(func() {
		timeout := a.config.HTTP.ShutdownCancelTimeout
		if timeout <= 0 {
			timeout = 5 * time.Second
		}
		if err := a.requests.cancelAndWait(timeout); err != nil {
			a.closeErr = err
			return
		}
		for index := len(a.closers) - 1; index >= 0; index-- {
			if err := a.closers[index](); err != nil {
				a.closeErr = errors.Join(a.closeErr, err)
			}
		}
	})
	return a.closeErr
}

// Run starts the HTTP server and handles SIGINT/SIGTERM. Use RunContext when
// the caller owns cancellation. The default address is :8080.
func (a *App) Run(address ...string) error {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	return a.RunContext(ctx, address...)
}

// NewServer creates a plain standard HTTP server with the configured HTTP options.
// Invalid HTTP options panic; RunContext/ServeContext return validation errors.
// Use RunContext or ServeContext to apply the configured transport and TLS.
func (a *App) NewServer(address string) *http.Server {
	if err := a.ValidateHTTP(); err != nil {
		panic(err)
	}
	if address == "" {
		address = a.config.Address
		if address == "" {
			address = ":8080"
		}
	}
	cfg := a.config.HTTP
	return &http.Server{Addr: address, Handler: a,
		ReadHeaderTimeout: cfg.ReadHeaderTimeout, ReadTimeout: cfg.ReadTimeout,
		WriteTimeout: cfg.WriteTimeout, IdleTimeout: cfg.IdleTimeout,
		MaxHeaderBytes: cfg.MaxHeaderBytes, ErrorLog: cfg.ErrorLog}
}

// RunContext serves until the context is canceled, then performs a graceful shutdown.
func (a *App) RunContext(ctx context.Context, address ...string) (runErr error) {
	defer func() { runErr = shutdownResult(ctx, runErr) }()
	if ctx == nil {
		ctx = context.Background()
	}
	if ctx.Err() != nil {
		return a.Close()
	}
	stopWatching := a.watchShutdown(ctx)
	defer stopWatching()
	serve, err := a.prepareTransport()
	if err != nil {
		return errors.Join(err, a.Close())
	}
	listener, err := net.Listen("tcp", a.listenAddress(address...))
	if err != nil {
		return errors.Join(err, a.Close())
	}
	defer listener.Close()
	return errors.Join(serve(ctx, listener), a.Close())
}

// RunWith starts a custom server runner while keeping App-owned cleanup in the
// same lifecycle. The runner receives the transport-neutral App handler and
// owns transport/TLS settings; built-in transport selection is bypassed.
func (a *App) RunWith(ctx context.Context, runner Runner, address ...string) (runErr error) {
	defer func() { runErr = shutdownResult(ctx, runErr) }()
	if runner == nil {
		return errors.Join(errors.New("go-omni: nil runner"), a.Close())
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if ctx.Err() != nil {
		return a.Close()
	}
	if err := a.ValidateHTTP(); err != nil {
		return errors.Join(err, a.Close())
	}
	stopWatching := a.watchShutdown(ctx)
	defer stopWatching()
	return errors.Join(runner.Run(ctx, a, a.listenAddress(address...)), a.Close())
}

func shutdownResult(ctx context.Context, err error) error {
	if cause := lifecycle.Cause(ctx); cause != nil && !errors.Is(err, cause) {
		return errors.Join(err, cause)
	}
	return err
}

func (a *App) listenAddress(address ...string) string {
	listenOn := a.config.Address
	if listenOn == "" {
		listenOn = ":8080"
	}
	if len(address) > 0 && address[0] != "" {
		listenOn = address[0]
	}
	return listenOn
}

func (a *App) serveContext(ctx context.Context, server *http.Server, listener net.Listener) error {
	serveErr := make(chan error, 1)
	go func() {
		if server.TLSConfig != nil {
			serveErr <- server.ServeTLS(listener, "", "")
		} else {
			serveErr <- server.Serve(listener)
		}
	}()

	select {
	case err := <-serveErr:
		// Serve returning only proves the accept loop ended. Cancel any remaining
		// connections before the application completion barrier and cleanup.
		_ = server.Close()
		closeErr := a.Close()
		if err == http.ErrServerClosed {
			return closeErr
		}
		return errors.Join(err, closeErr)
	case <-ctx.Done():
		a.requests.stop()
		shutdownTimeout := a.config.ShutdownTimeout
		if shutdownTimeout <= 0 {
			shutdownTimeout = 10 * time.Second
		}
		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
		defer cancel()
		shutdownErr := server.Shutdown(shutdownCtx)
		if shutdownErr != nil {
			// Closing sockets requests cancellation; App.Close separately waits
			// for the actual handler return before releasing shared resources.
			shutdownErr = errors.Join(shutdownErr, server.Close())
		}
		if err := <-serveErr; err != nil && err != http.ErrServerClosed {
			shutdownErr = errors.Join(shutdownErr, err)
		}
		return errors.Join(shutdownErr, a.Close())
	}
}

func (a *App) writeError(ctx *httpx.Context, err error) {
	if err == nil || ctx.Responded() {
		return
	}
	ctx.WriteError(err)
}