Server struct with tests

errors.go

@@ -11,7 +11,6 @@ type (
 	ResponseWriteError   struct{ error }
 	InvalidURLError      string
 	TooManyRequestsError struct{}
-	InvalidSecretError   struct{}
 )
 
 func newResponseWriteError(cause error) *ierrors.Error {
@@ -53,15 +52,3 @@ func newTooManyRequestsError() error {
 }
 
 func (e TooManyRequestsError) Error() string { return "Too many requests" }
-
-func newInvalidSecretError() error {
-	return ierrors.Wrap(
-		InvalidSecretError{},
-		1,
-		ierrors.WithStatusCode(http.StatusForbidden),
-		ierrors.WithPublicMessage("Forbidden"),
-		ierrors.WithShouldReport(false),
-	)
-}
-
-func (e InvalidSecretError) Error() string { return "Invalid secret" }

main.go

@@ -16,6 +16,7 @@ import (
 	"github.com/imgproxy/imgproxy/v3/config/loadenv"
 	"github.com/imgproxy/imgproxy/v3/errorreport"
 	"github.com/imgproxy/imgproxy/v3/gliblog"
+	"github.com/imgproxy/imgproxy/v3/handlers"
 	"github.com/imgproxy/imgproxy/v3/imagedata"
 	"github.com/imgproxy/imgproxy/v3/logger"
 	"github.com/imgproxy/imgproxy/v3/memory"
@@ -23,10 +24,37 @@ import (
 	"github.com/imgproxy/imgproxy/v3/metrics/prometheus"
 	"github.com/imgproxy/imgproxy/v3/options"
 	"github.com/imgproxy/imgproxy/v3/processing"
+	"github.com/imgproxy/imgproxy/v3/server"
 	"github.com/imgproxy/imgproxy/v3/version"
 	"github.com/imgproxy/imgproxy/v3/vips"
 )
 
+const (
+	faviconPath = "/favicon.ico"
+	healthPath  = "/health"
+)
+
+func buildRouter(r *server.Router) *server.Router {
+	r.GET("/", true, handlers.LandingHandler)
+	r.GET("", true, handlers.LandingHandler)
+
+	r.GET(
+		"/", false, handleProcessing,
+		r.WithSecret, r.WithCORS, r.WithPanic, r.WithReportError, r.WithMetrics,
+	)
+
+	r.HEAD("/", false, r.OkHandler, r.WithCORS)
+	r.OPTIONS("/", false, r.OkHandler, r.WithCORS)
+
+	r.GET(faviconPath, true, r.NotFoundHandler).Silent()
+	r.GET(healthPath, true, handlers.HealthHandler).Silent()
+	if config.HealthCheckPath != "" {
+		r.GET(config.HealthCheckPath, true, handlers.HealthHandler).Silent()
+	}
+
+	return r
+}
+
 func initialize() error {
 	if err := loadenv.Load(); err != nil {
 		return err
@@ -103,25 +131,21 @@ func run(ctx context.Context) error {
 		}
 	}()
 
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
 
 	if err := prometheus.StartServer(cancel); err != nil {
 		return err
 	}
 
-	s, err := startServer(cancel)
+	cfg := server.NewConfigFromEnv()
+	r := server.NewRouter(cfg)
+	s, err := server.Start(cancel, buildRouter(r))
 	if err != nil {
 		return err
 	}
-	defer shutdownServer(s)
+	defer s.Shutdown(ctx)
 
-	stop := make(chan os.Signal, 1)
-	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
-
-	select {
-	case <-ctx.Done():
-	case <-stop:
-	}
+	<-ctx.Done()
 
 	return nil
 }

metrics/prometheus/prometheus.go

@@ -161,7 +161,7 @@ func StartServer(cancel context.CancelFunc) error {
 
 	s := http.Server{Handler: promhttp.Handler()}
 
-	l, err := reuseport.Listen("tcp", config.PrometheusBind)
+	l, err := reuseport.Listen("tcp", config.PrometheusBind, config.SoReuseport)
 	if err != nil {
 		return fmt.Errorf("Can't start Prometheus metrics server: %s", err)
 	}

processing_handler_test.go

@@ -48,7 +48,7 @@ func (s *ProcessingHandlerTestSuite) SetupSuite() {
 
 	logrus.SetOutput(io.Discard)
 
-	s.router = buildRouter()
+	s.router = buildRouter(server.NewRouter(server.NewConfigFromEnv()))
 }
 
 func (s *ProcessingHandlerTestSuite) TeardownSuite() {

reuseport/listen_no_reuseport.go

@@ -7,12 +7,10 @@ import (
 	"net"
 
 	log "github.com/sirupsen/logrus"
-
-	"github.com/imgproxy/imgproxy/v3/config"
 )
 
-func Listen(network, address string) (net.Listener, error) {
-	if config.SoReuseport {
+func Listen(network, address string, reuse bool) (net.Listener, error) {
+	if reuse {
 		log.Warning("SO_REUSEPORT support is not implemented for your OS or Go version")
 	}
 

reuseport/listen_reuseport.go

@@ -10,12 +10,10 @@ import (
 	"syscall"
 
 	"golang.org/x/sys/unix"
-
-	"github.com/imgproxy/imgproxy/v3/config"
 )
 
-func Listen(network, address string) (net.Listener, error) {
-	if !config.SoReuseport {
+func Listen(network, address string, reuse bool) (net.Listener, error) {
+	if !reuse {
 		return net.Listen(network, address)
 	}
 

server.go

@@ -1,183 +0,0 @@
-package main
-
-import (
-	"context"
-	"crypto/subtle"
-	"fmt"
-	golog "log"
-	"net/http"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-	"golang.org/x/net/netutil"
-
-	"github.com/imgproxy/imgproxy/v3/config"
-	"github.com/imgproxy/imgproxy/v3/errorreport"
-	"github.com/imgproxy/imgproxy/v3/handlers"
-	"github.com/imgproxy/imgproxy/v3/ierrors"
-	"github.com/imgproxy/imgproxy/v3/metrics"
-	"github.com/imgproxy/imgproxy/v3/reuseport"
-	"github.com/imgproxy/imgproxy/v3/server"
-)
-
-const (
-	faviconPath = "/favicon.ico"
-	healthPath  = "/health"
-)
-
-func buildRouter() *server.Router {
-	r := server.NewRouter(config.PathPrefix)
-
-	r.GET("/", true, handlers.LandingHandler)
-	r.GET("", true, handlers.LandingHandler)
-
-	r.GET("/", false, handleProcessing, withMetrics, withPanicHandler, withCORS, withSecret)
-
-	r.HEAD("/", false, r.OkHandler, withCORS)
-	r.OPTIONS("/", false, r.OkHandler, withCORS)
-
-	r.GET(faviconPath, true, r.NotFoundHandler).Silent()
-	r.GET(healthPath, true, handlers.HealthHandler).Silent()
-	if config.HealthCheckPath != "" {
-		r.GET(config.HealthCheckPath, true, handlers.HealthHandler).Silent()
-	}
-
-	return r
-}
-
-func startServer(cancel context.CancelFunc) (*http.Server, error) {
-	l, err := reuseport.Listen(config.Network, config.Bind)
-	if err != nil {
-		return nil, fmt.Errorf("can't start server: %s", err)
-	}
-
-	if config.MaxClients > 0 {
-		l = netutil.LimitListener(l, config.MaxClients)
-	}
-
-	errLogger := golog.New(
-		log.WithField("source", "http_server").WriterLevel(log.ErrorLevel),
-		"", 0,
-	)
-
-	s := &http.Server{
-		Handler:        buildRouter(),
-		ReadTimeout:    time.Duration(config.ReadRequestTimeout) * time.Second,
-		MaxHeaderBytes: 1 << 20,
-		ErrorLog:       errLogger,
-	}
-
-	if config.KeepAliveTimeout > 0 {
-		s.IdleTimeout = time.Duration(config.KeepAliveTimeout) * time.Second
-	} else {
-		s.SetKeepAlivesEnabled(false)
-	}
-
-	go func() {
-		log.Infof("Starting server at %s", config.Bind)
-		if err := s.Serve(l); err != nil && err != http.ErrServerClosed {
-			log.Error(err)
-		}
-		cancel()
-	}()
-
-	return s, nil
-}
-
-func shutdownServer(s *http.Server) {
-	log.Info("Shutting down the server...")
-
-	ctx, close := context.WithTimeout(context.Background(), 5*time.Second)
-	defer close()
-
-	s.Shutdown(ctx)
-}
-
-func withMetrics(h server.RouteHandler) server.RouteHandler {
-	if !metrics.Enabled() {
-		return h
-	}
-
-	return func(reqID string, rw http.ResponseWriter, r *http.Request) error {
-		ctx, metricsCancel, rw := metrics.StartRequest(r.Context(), rw, r)
-		defer metricsCancel()
-
-		h(reqID, rw, r.WithContext(ctx))
-
-		return nil
-	}
-}
-
-func withCORS(h server.RouteHandler) server.RouteHandler {
-	return func(reqID string, rw http.ResponseWriter, r *http.Request) error {
-		if len(config.AllowOrigin) > 0 {
-			rw.Header().Set("Access-Control-Allow-Origin", config.AllowOrigin)
-			rw.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS")
-		}
-
-		h(reqID, rw, r)
-
-		return nil
-	}
-}
-
-func withSecret(h server.RouteHandler) server.RouteHandler {
-	if len(config.Secret) == 0 {
-		return h
-	}
-
-	authHeader := []byte(fmt.Sprintf("Bearer %s", config.Secret))
-
-	return func(reqID string, rw http.ResponseWriter, r *http.Request) error {
-		if subtle.ConstantTimeCompare([]byte(r.Header.Get("Authorization")), authHeader) == 1 {
-			h(reqID, rw, r)
-		} else {
-			panic(newInvalidSecretError())
-		}
-
-		return nil
-	}
-}
-
-func withPanicHandler(h server.RouteHandler) server.RouteHandler {
-	return func(reqID string, rw http.ResponseWriter, r *http.Request) error {
-		ctx := errorreport.StartRequest(r)
-		r = r.WithContext(ctx)
-
-		errorreport.SetMetadata(r, "Request ID", reqID)
-
-		defer func() {
-			if rerr := recover(); rerr != nil {
-				if rerr == http.ErrAbortHandler {
-					panic(rerr)
-				}
-
-				err, ok := rerr.(error)
-				if !ok {
-					panic(rerr)
-				}
-
-				ierr := ierrors.Wrap(err, 0)
-
-				if ierr.ShouldReport() {
-					errorreport.Report(err, r)
-				}
-
-				server.LogResponse(reqID, r, ierr.StatusCode(), ierr)
-
-				rw.Header().Set("Content-Type", "text/plain")
-				rw.WriteHeader(ierr.StatusCode())
-
-				if config.DevelopmentErrorsMode {
-					rw.Write([]byte(ierr.Error()))
-				} else {
-					rw.Write([]byte(ierr.PublicMessage()))
-				}
-			}
-		}()
-
-		h(reqID, rw, r)
-
-		return nil
-	}
-}

server/config.go

@@ -0,0 +1,49 @@
+package server
+
+import (
+	"time"
+
+	"github.com/imgproxy/imgproxy/v3/config"
+)
+
+const (
+	// gracefulTimeout represents graceful shutdown timeout
+	gracefulTimeout = time.Duration(5 * time.Second)
+)
+
+// Config represents HTTP server config
+type Config struct {
+	Listen                string        // Address to listen on
+	Network               string        // Network type (tcp, unix)
+	Bind                  string        // Bind address
+	PathPrefix            string        // Path prefix for the server
+	MaxClients            int           // Maximum number of concurrent clients
+	ReadRequestTimeout    time.Duration // Timeout for reading requests
+	KeepAliveTimeout      time.Duration // Timeout for keep-alive connections
+	GracefulTimeout       time.Duration // Timeout for graceful shutdown
+	CORSAllowOrigin       string        // CORS allowed origin
+	Secret                string        // Secret for authorization
+	DevelopmentErrorsMode bool          // Enable development mode for detailed error messages
+	SocketReusePort       bool          // Enable SO_REUSEPORT socket option
+	HealthCheckPath       string        // Health check path from config
+	WriteResponseTimeout  time.Duration
+}
+
+// NewConfigFromEnv creates a new Config instance from environment variables
+func NewConfigFromEnv() *Config {
+	return &Config{
+		Network:               config.Network,
+		Bind:                  config.Bind,
+		PathPrefix:            config.PathPrefix,
+		MaxClients:            config.MaxClients,
+		ReadRequestTimeout:    time.Duration(config.ReadRequestTimeout) * time.Second,
+		KeepAliveTimeout:      time.Duration(config.KeepAliveTimeout) * time.Second,
+		GracefulTimeout:       gracefulTimeout,
+		CORSAllowOrigin:       config.AllowOrigin,
+		Secret:                config.Secret,
+		DevelopmentErrorsMode: config.DevelopmentErrorsMode,
+		SocketReusePort:       config.SoReuseport,
+		HealthCheckPath:       config.HealthCheckPath,
+		WriteResponseTimeout:  time.Duration(config.WriteResponseTimeout) * time.Second,
+	}
+}

server/errors.go

@@ -12,6 +12,7 @@ type (
 	RouteNotDefinedError  string
 	RequestCancelledError string
 	RequestTimeoutError   string
+	InvalidSecretError    struct{}
 )
 
 func newRouteNotDefinedError(path string) *ierrors.Error {
@@ -49,3 +50,15 @@ func newRequestTimeoutError(after time.Duration) *ierrors.Error {
 }
 
 func (e RequestTimeoutError) Error() string { return string(e) }
+
+func newInvalidSecretError() error {
+	return ierrors.Wrap(
+		InvalidSecretError{},
+		1,
+		ierrors.WithStatusCode(http.StatusForbidden),
+		ierrors.WithPublicMessage("Forbidden"),
+		ierrors.WithShouldReport(false),
+	)
+}
+
+func (e InvalidSecretError) Error() string { return "Invalid secret" }

server/middlewares.go

@@ -0,0 +1,135 @@
+package server
+
+import (
+	"crypto/subtle"
+	"fmt"
+	"net/http"
+
+	"github.com/imgproxy/imgproxy/v3/errorreport"
+	"github.com/imgproxy/imgproxy/v3/httpheaders"
+	"github.com/imgproxy/imgproxy/v3/ierrors"
+	"github.com/imgproxy/imgproxy/v3/metrics"
+)
+
+// WithMetrics wraps RouteHandler with metrics handling.
+func (r *Router) WithMetrics(h RouteHandler) RouteHandler {
+	if !metrics.Enabled() {
+		return h
+	}
+
+	return func(reqID string, rw http.ResponseWriter, req *http.Request) error {
+		ctx, metricsCancel, rw := metrics.StartRequest(req.Context(), rw, req)
+		defer metricsCancel()
+
+		return h(reqID, rw, req.WithContext(ctx))
+	}
+}
+
+// WithCORS wraps RouteHandler with CORS handling
+func (r *Router) WithCORS(h RouteHandler) RouteHandler {
+	if len(r.config.CORSAllowOrigin) == 0 {
+		return h
+	}
+
+	return func(reqID string, rw http.ResponseWriter, req *http.Request) error {
+		rw.Header().Set(httpheaders.AccessControlAllowOrigin, r.config.CORSAllowOrigin)
+		rw.Header().Set(httpheaders.AccessControlAllowMethods, "GET, OPTIONS")
+
+		return h(reqID, rw, req)
+	}
+}
+
+// WithSecret wraps RouteHandler with secret handling
+func (r *Router) WithSecret(h RouteHandler) RouteHandler {
+	if len(r.config.Secret) == 0 {
+		return h
+	}
+
+	authHeader := fmt.Appendf(nil, "Bearer %s", r.config.Secret)
+
+	return func(reqID string, rw http.ResponseWriter, req *http.Request) error {
+		if subtle.ConstantTimeCompare([]byte(req.Header.Get(httpheaders.Authorization)), authHeader) == 1 {
+			return h(reqID, rw, req)
+		} else {
+			return newInvalidSecretError()
+		}
+	}
+}
+
+// WithReportError handles error reporting.
+// It should be placed after `WithMetrics`, but before `WithPanic`.
+func (r *Router) WithReportError(h RouteHandler) RouteHandler {
+	return func(reqID string, rw http.ResponseWriter, req *http.Request) error {
+		// Open the error context
+		ctx := errorreport.StartRequest(req)
+		req = req.WithContext(ctx)
+		errorreport.SetMetadata(req, "Request ID", reqID)
+
+		// Call the underlying handler passing the context downwards
+		err := h(reqID, rw, req)
+		if err == nil {
+			return nil
+		}
+
+		// Wrap a resulting error into ierrors.Error
+		ierr := ierrors.Wrap(err, 0)
+
+		// Get the error category
+		errCat := ierr.Category()
+
+		// Report error to metrics (if metrics are disabled, it will be a no-op)
+		// Skip context closed error
+		if ierr.StatusCode() != 499 {
+			metrics.SendError(ctx, errCat, err)
+		}
+
+		// Report error to error collectors
+		if ierr.ShouldReport() {
+			errorreport.Report(ierr, req)
+		}
+
+		// Log response and format the error output
+		LogResponse(reqID, req, ierr.StatusCode(), ierr)
+
+		// Error message: either is public message or full development error
+		rw.Header().Set(httpheaders.ContentType, "text/plain")
+		rw.WriteHeader(ierr.StatusCode())
+
+		if r.config.DevelopmentErrorsMode {
+			rw.Write([]byte(ierr.Error()))
+		} else {
+			rw.Write([]byte(ierr.PublicMessage()))
+		}
+
+		return nil
+	}
+}
+
+// WithPanic recovers panic and converts it to normal error
+func (r *Router) WithPanic(h RouteHandler) RouteHandler {
+	return func(reqID string, rw http.ResponseWriter, r *http.Request) (retErr error) {
+		defer func() {
+			// try to recover from panic
+			rerr := recover()
+			if rerr == nil {
+				return
+			}
+
+			// abort handler is an exception of net/http, we should simply repanic it.
+			// it will supress the stack trace
+			if rerr == http.ErrAbortHandler {
+				panic(rerr)
+			}
+
+			// let's recover error value from panic if it has panicked with error
+			err, ok := rerr.(error)
+			if !ok {
+				err = fmt.Errorf("panic: %v", err)
+			}
+
+			retErr = err
+		}()
+
+		return h(reqID, rw, r)
+	}
+}

server/router.go

@@ -9,7 +9,6 @@ import (
 
 	nanoid "github.com/matoous/go-nanoid/v2"
 
-	"github.com/imgproxy/imgproxy/v3/config"
 	"github.com/imgproxy/imgproxy/v3/httpheaders"
 )
 
@@ -40,25 +39,25 @@ type route struct {
 
 // Router is responsible for routing HTTP requests
 type Router struct {
-	// prefix represents global path prefix for all routes
-	prefix string
+	// config represents the server configuration
+	config *Config
 
 	// routes is the collection of all routes
 	routes []*route
 }
 
 // NewRouter creates a new Router instance
-func NewRouter(prefix string) *Router {
-	return &Router{prefix: prefix}
+func NewRouter(config *Config) *Router {
+	return &Router{config: config}
 }
 
 // add adds an abitary route to the router
 func (r *Router) add(method, prefix string, exact bool, handler RouteHandler, middlewares ...Middleware) *route {
-	for i := len(middlewares) - 1; i >= 0; i-- {
-		handler = middlewares[i](handler)
+	for _, m := range middlewares {
+		handler = m(handler)
 	}
 
-	route := &route{method: method, path: r.prefix + prefix, handler: handler, exact: exact}
+	route := &route{method: method, path: r.config.PathPrefix + prefix, handler: handler, exact: exact}
 
 	r.routes = append(
 		r.routes,
@@ -90,7 +89,7 @@ func (r *Router) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	defer timeoutCancel()
 
 	// Create the response writer which times out on write
-	rw = newTimeoutResponse(rw, config.WriteResponseTimeout)
+	rw = newTimeoutResponse(rw, r.config.WriteResponseTimeout)
 
 	// Get/create request ID
 	reqID := r.getRequestID(req)
@@ -138,7 +137,7 @@ func (r *Router) OkHandler(reqID string, rw http.ResponseWriter, req *http.Reque
 
 // getRequestID tries to read request id from headers or from lambda
 // context or generates a new one if nothing found.
-func (rw *Router) getRequestID(req *http.Request) string {
+func (r *Router) getRequestID(req *http.Request) string {
 	// Get request ID from headers (if any)
 	reqID := req.Header.Get(httpheaders.XRequestID)
 
@@ -165,7 +164,7 @@ func (rw *Router) getRequestID(req *http.Request) string {
 }
 
 // replaceRemoteAddr rewrites the req.RemoteAddr property from request headers
-func (rw *Router) replaceRemoteAddr(req *http.Request) {
+func (r *Router) replaceRemoteAddr(req *http.Request) {
 	cfConnectingIP := req.Header.Get(httpheaders.CFConnectingIP)
 	xForwardedFor := req.Header.Get(httpheaders.XForwardedFor)
 	xRealIP := req.Header.Get(httpheaders.XRealIP)
@@ -202,7 +201,8 @@ func (r *route) isMatch(req *http.Request) bool {
 	return methodMatches && (notExactPathMathes || exactPathMatches)
 }
 
-// Silent sets Silent flag which supresses logs to true
+// Silent sets Silent flag which supresses logs to true. We do not need to log
+// requests like /health of /favicon.ico
 func (r *route) Silent() *route {
 	r.silent = true
 	return r

server/router_test.go

@@ -16,7 +16,9 @@ type RouterTestSuite struct {
 }
 
 func (s *RouterTestSuite) SetupTest() {
-	s.router = NewRouter("/api")
+	c := NewConfigFromEnv()
+	c.PathPrefix = "/api"
+	s.router = NewRouter(c)
 }
 
 func TestRouterSuite(t *testing.T) {
@@ -131,7 +133,7 @@ func (s *RouterTestSuite) TestMiddlewareOrder() {
 		return nil
 	}
 
-	s.router.GET("/test", true, handler, middleware1, middleware2)
+	s.router.GET("/test", true, handler, middleware2, middleware1)
 
 	req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
 	rw := httptest.NewRecorder()

server/server.go

@@ -0,0 +1,82 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	golog "log"
+	"net/http"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/net/netutil"
+
+	"github.com/imgproxy/imgproxy/v3/config"
+	"github.com/imgproxy/imgproxy/v3/reuseport"
+)
+
+const (
+	// maxHeaderBytes represents max bytes in request header
+	maxHeaderBytes = 1 << 20
+)
+
+// Server represents the HTTP server wrapper struct
+type Server struct {
+	router *Router
+	server *http.Server
+}
+
+// Start starts the http server. cancel is called in case server failed to start, but it happened
+// asynchronously. It should cancel the upstream context.
+func Start(cancel context.CancelFunc, router *Router) (*Server, error) {
+	l, err := reuseport.Listen(router.config.Network, router.config.Bind, router.config.SocketReusePort)
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("can't start server: %s", err)
+	}
+
+	if router.config.MaxClients > 0 {
+		l = netutil.LimitListener(l, router.config.MaxClients)
+	}
+
+	errLogger := golog.New(
+		log.WithField("source", "http_server").WriterLevel(log.ErrorLevel),
+		"", 0,
+	)
+
+	srv := &http.Server{
+		Handler:        router,
+		ReadTimeout:    router.config.ReadRequestTimeout,
+		MaxHeaderBytes: maxHeaderBytes,
+		ErrorLog:       errLogger,
+	}
+
+	if config.KeepAliveTimeout > 0 {
+		srv.IdleTimeout = router.config.KeepAliveTimeout
+	} else {
+		srv.SetKeepAlivesEnabled(false)
+	}
+
+	go func() {
+		log.Infof("Starting server at %s", router.config.Bind)
+
+		if err := srv.Serve(l); err != nil && err != http.ErrServerClosed {
+			log.Error(err)
+		}
+
+		cancel()
+	}()
+
+	return &Server{
+		router: router,
+		server: srv,
+	}, nil
+}
+
+// Shutdown gracefully shuts down the server
+func (s *Server) Shutdown(ctx context.Context) {
+	log.Info("Shutting down the server...")
+
+	ctx, close := context.WithTimeout(ctx, s.router.config.GracefulTimeout)
+	defer close()
+
+	s.server.Shutdown(ctx)
+}

server/server_test.go

@@ -0,0 +1,268 @@
+package server
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/imgproxy/imgproxy/v3/config"
+	"github.com/imgproxy/imgproxy/v3/httpheaders"
+	"github.com/stretchr/testify/suite"
+)
+
+type ServerTestSuite struct {
+	suite.Suite
+	config      *Config
+	blankRouter *Router
+}
+
+func (s *ServerTestSuite) SetupTest() {
+	config.Reset()
+	s.config = NewConfigFromEnv()
+	s.config.Bind = "127.0.0.1:0" // Use port 0 for auto-assignment
+	s.blankRouter = NewRouter(s.config)
+}
+
+func (s *ServerTestSuite) mockHandler(reqID string, rw http.ResponseWriter, r *http.Request) error {
+	return nil
+}
+
+func (s *ServerTestSuite) TestStartServerWithInvalidBind() {
+	ctx, cancel := context.WithCancel(s.T().Context())
+
+	// Track if cancel was called using atomic
+	var cancelCalled atomic.Bool
+	cancelWrapper := func() {
+		cancel()
+		cancelCalled.Store(true)
+	}
+
+	invalidConfig := &Config{
+		Network: "tcp",
+		Bind:    "invalid-address", // Invalid address
+	}
+
+	r := NewRouter(invalidConfig)
+
+	server, err := Start(cancelWrapper, r)
+
+	s.Require().Error(err)
+	s.Nil(server)
+	s.Contains(err.Error(), "can't start server")
+
+	// Check if cancel was called using Eventually
+	s.Require().Eventually(cancelCalled.Load, 100*time.Millisecond, 10*time.Millisecond)
+
+	// Also verify the context was cancelled
+	s.Require().Eventually(func() bool {
+		select {
+		case <-ctx.Done():
+			return true
+		default:
+			return false
+		}
+	}, 100*time.Millisecond, 10*time.Millisecond)
+}
+
+func (s *ServerTestSuite) TestShutdown() {
+	_, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	server, err := Start(cancel, s.blankRouter)
+	s.Require().NoError(err)
+	s.NotNil(server)
+
+	// Test graceful shutdown
+	shutdownCtx, shutdownCancel := context.WithTimeout(s.T().Context(), 10*time.Second)
+	defer shutdownCancel()
+
+	// Should not panic or hang
+	s.NotPanics(func() {
+		server.Shutdown(shutdownCtx)
+	})
+}
+
+func (s *ServerTestSuite) TestWithCORS() {
+	tests := []struct {
+		name            string
+		corsAllowOrigin string
+		expectedOrigin  string
+		expectedMethods string
+	}{
+		{
+			name:            "WithCORSOrigin",
+			corsAllowOrigin: "https://example.com",
+			expectedOrigin:  "https://example.com",
+			expectedMethods: "GET, OPTIONS",
+		},
+		{
+			name:            "NoCORSOrigin",
+			corsAllowOrigin: "",
+			expectedOrigin:  "",
+			expectedMethods: "",
+		},
+	}
+
+	for _, tt := range tests {
+		s.Run(tt.name, func() {
+			config := &Config{
+				CORSAllowOrigin: tt.corsAllowOrigin,
+			}
+			router := NewRouter(config)
+
+			wrappedHandler := router.WithCORS(s.mockHandler)
+
+			req := httptest.NewRequest("GET", "/test", nil)
+			rw := httptest.NewRecorder()
+
+			wrappedHandler("test-req-id", rw, req)
+
+			s.Equal(tt.expectedOrigin, rw.Header().Get(httpheaders.AccessControlAllowOrigin))
+			s.Equal(tt.expectedMethods, rw.Header().Get(httpheaders.AccessControlAllowMethods))
+		})
+	}
+}
+
+func (s *ServerTestSuite) TestWithSecret() {
+	tests := []struct {
+		name        string
+		secret      string
+		authHeader  string
+		expectError bool
+	}{
+		{
+			name:       "ValidSecret",
+			secret:     "test-secret",
+			authHeader: "Bearer test-secret",
+		},
+		{
+			name:        "InvalidSecret",
+			secret:      "foo-secret",
+			authHeader:  "Bearer wrong-secret",
+			expectError: true,
+		},
+		{
+			name:       "NoSecretConfigured",
+			secret:     "",
+			authHeader: "",
+		},
+	}
+
+	for _, tt := range tests {
+		s.Run(tt.name, func() {
+			config := &Config{
+				Secret: tt.secret,
+			}
+			router := NewRouter(config)
+
+			wrappedHandler := router.WithSecret(s.mockHandler)
+
+			req := httptest.NewRequest("GET", "/test", nil)
+			if tt.authHeader != "" {
+				req.Header.Set(httpheaders.Authorization, tt.authHeader)
+			}
+			rw := httptest.NewRecorder()
+
+			err := wrappedHandler("test-req-id", rw, req)
+
+			if tt.expectError {
+				s.Require().Error(err)
+			} else {
+				s.Require().NoError(err)
+			}
+		})
+	}
+}
+
+func (s *ServerTestSuite) TestIntoSuccess() {
+	mockHandler := func(reqID string, rw http.ResponseWriter, r *http.Request) error {
+		rw.WriteHeader(http.StatusOK)
+		return nil
+	}
+
+	wrappedHandler := s.blankRouter.WithReportError(mockHandler)
+
+	req := httptest.NewRequest("GET", "/test", nil)
+	rw := httptest.NewRecorder()
+
+	wrappedHandler("test-req-id", rw, req)
+
+	s.Equal(http.StatusOK, rw.Code)
+}
+
+func (s *ServerTestSuite) TestIntoWithError() {
+	testError := errors.New("test error")
+	mockHandler := func(reqID string, rw http.ResponseWriter, r *http.Request) error {
+		return testError
+	}
+
+	wrappedHandler := s.blankRouter.WithReportError(mockHandler)
+
+	req := httptest.NewRequest("GET", "/test", nil)
+	rw := httptest.NewRecorder()
+
+	wrappedHandler("test-req-id", rw, req)
+
+	s.Equal(http.StatusInternalServerError, rw.Code)
+	s.Equal("text/plain", rw.Header().Get(httpheaders.ContentType))
+}
+
+func (s *ServerTestSuite) TestIntoPanicWithError() {
+	testError := errors.New("panic error")
+	mockHandler := func(reqID string, rw http.ResponseWriter, r *http.Request) error {
+		panic(testError)
+	}
+
+	wrappedHandler := s.blankRouter.WithPanic(mockHandler)
+
+	req := httptest.NewRequest("GET", "/test", nil)
+	rw := httptest.NewRecorder()
+
+	s.NotPanics(func() {
+		err := wrappedHandler("test-req-id", rw, req)
+		s.Require().Error(err, "panic error")
+	})
+
+	s.Equal(http.StatusOK, rw.Code)
+}
+
+func (s *ServerTestSuite) TestIntoPanicWithAbortHandler() {
+	mockHandler := func(reqID string, rw http.ResponseWriter, r *http.Request) error {
+		panic(http.ErrAbortHandler)
+	}
+
+	wrappedHandler := s.blankRouter.WithPanic(mockHandler)
+
+	req := httptest.NewRequest("GET", "/test", nil)
+	rw := httptest.NewRecorder()
+
+	// Should re-panic with ErrAbortHandler
+	s.Panics(func() {
+		wrappedHandler("test-req-id", rw, req)
+	})
+}
+
+func (s *ServerTestSuite) TestIntoPanicWithNonError() {
+	mockHandler := func(reqID string, rw http.ResponseWriter, r *http.Request) error {
+		panic("string panic")
+	}
+
+	wrappedHandler := s.blankRouter.WithPanic(mockHandler)
+
+	req := httptest.NewRequest("GET", "/test", nil)
+	rw := httptest.NewRecorder()
+
+	// Should re-panic with non-error panics
+	s.NotPanics(func() {
+		err := wrappedHandler("test-req-id", rw, req)
+		s.Require().Error(err, "string panic")
+	})
+}
+
+func TestServerTestSuite(t *testing.T) {
+	suite.Run(t, new(ServerTestSuite))
+}

server/timeout_response.go

@@ -10,11 +10,11 @@ import (
 type timeoutResponse struct {
 	http.ResponseWriter
 	controller *http.ResponseController
-	timeout    int
+	timeout    time.Duration
 }
 
 // newTimeoutResponse creates a new timeoutResponse
-func newTimeoutResponse(rw http.ResponseWriter, timeout int) http.ResponseWriter {
+func newTimeoutResponse(rw http.ResponseWriter, timeout time.Duration) http.ResponseWriter {
 	return &timeoutResponse{
 		ResponseWriter: rw,
 		controller:     http.NewResponseController(rw),
@@ -22,13 +22,6 @@ func newTimeoutResponse(rw http.ResponseWriter, timeout int) http.ResponseWriter
 	}
 }
 
-// WriteHeader implements http.ResponseWriter.WriteHeader
-func (rw *timeoutResponse) WriteHeader(statusCode int) {
-	rw.withWriteDeadline(func() {
-		rw.ResponseWriter.WriteHeader(statusCode)
-	})
-}
-
 // Write implements http.ResponseWriter.Write
 func (rw *timeoutResponse) Write(b []byte) (int, error) {
 	var (
@@ -48,7 +41,7 @@ func (rw *timeoutResponse) Header() http.Header {
 
 // withWriteDeadline executes a Write* function with a deadline
 func (rw *timeoutResponse) withWriteDeadline(f func()) {
-	deadline := time.Now().Add(time.Duration(rw.timeout) * time.Second)
+	deadline := time.Now().Add(rw.timeout)
 
 	// Set write deadline
 	rw.controller.SetWriteDeadline(deadline)