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"
	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"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
	config           config.Config
	middlewares      []httpx.Middleware
	closers          []func() error
	closeOnce        sync.Once
	closeErr         error
	clientIPResolver *httpx.ClientIPResolver
	responseEncoder  response.Encoder
}

// 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.Response = settings.Response.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...)
	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}
	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 returns the transport contract used by this application.
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.router }

// ServeHTTP implements http.Handler for standard servers and HTTP tests.
func (a *App) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	a.router.ServeHTTP(writer, request)
}

// Config returns the effective application configuration.
func (a *App) Config() config.Config {
	settings := a.config
	settings.Response = settings.Response.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 when Close or RunContext ends.
// 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 releases application resources. It is safe to call more than once.
func (a *App) Close() error {
	a.closeOnce.Do(func() {
		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 around the application.
// Use RunContext or ServeContext to apply the configured transport and TLS.
func (a *App) NewServer(address string) *http.Server {
	if address == "" {
		address = a.config.Address
		if address == "" {
			address = ":8080"
		}
	}
	return &http.Server{Addr: address, Handler: a.router}
}

// RunContext serves until the context is canceled, then performs a graceful shutdown.
func (a *App) RunContext(ctx context.Context, address ...string) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if ctx.Err() != nil {
		return a.Close()
	}
	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) error {
	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()
	}
	return errors.Join(runner.Run(ctx, a, a.listenAddress(address...)), a.Close())
}

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:
		closeErr := a.Close()
		if err == http.ErrServerClosed {
			return closeErr
		}
		return errors.Join(err, closeErr)
	case <-ctx.Done():
		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 {
			// Shutdown leaves active connections open on timeout. Close them so
			// request contexts are canceled 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)
}