feat: add custom route and middleware options

Add options for embedders to register Fiber routes and middleware before the S3 route table is initialized. WithRoute registers a top-level route with explicit method and path matching. WithMiddleware registers prefix middleware that can handle a request or call ctx.Next() to continue into the S3 stack. Add coverage for route registration order when a top-level route and catch-all middleware are both configured.
2026-08-17 12:46:23 +00:00 · 2026-05-04 08:41:16 -07:00
parent 40c02dd31e
commit d1fba07fe6
2 changed files with 277 additions and 0 deletions
@@ -55,9 +55,22 @@ type S3ApiServer struct {
 	maxRequests      int
 	webuiMountPrefix string
 	webuiSrvCfg      *webui.ServerConfig
+	routes           []routeMount
+	middlewares      []middlewareMount
 	socketPerm       os.FileMode
 }

+type routeMount struct {
+	method   string
+	path     string
+	handlers []fiber.Handler
+}
+
+type middlewareMount struct {
+	prefix  string
+	handler fiber.Handler
+}
+
 func New(
 	be backend.Backend,
 	root middlewares.RootUserConfig,
@@ -133,6 +146,21 @@ func New(
 	// initialize total requests cap limiter middleware
 	app.Use(middlewares.RateLimiter(server.maxRequests, mm, l))

+	for _, route := range server.routes {
+		method, err := validateRouteMount(route)
+		if err != nil {
+			return nil, err
+		}
+		app.Add(method, route.path, route.handlers...)
+	}
+
+	for _, mount := range server.middlewares {
+		if err := validateMiddlewareMount(mount); err != nil {
+			return nil, err
+		}
+		app.Use(mount.prefix, mount.handler)
+	}
+
 	// initilaze the default value setter middleware
 	app.Use(middlewares.SetDefaultValues(root, region))

@@ -150,6 +178,48 @@ func New(
 	return server, nil
 }

+func validateRouteMount(route routeMount) (string, error) {
+	if route.method == "" {
+		return "", fmt.Errorf("invalid route for path %q: empty method", route.path)
+	}
+	method := strings.ToUpper(route.method)
+	if !isStandardHTTPMethod(method) {
+		return "", fmt.Errorf("invalid HTTP method %q for route path %q: must be one of %s",
+			route.method, route.path, strings.Join(fiber.DefaultMethods, ", "))
+	}
+	if route.path == "" || route.path[0] != '/' {
+		return "", fmt.Errorf("invalid route path %q: must start with /", route.path)
+	}
+	if len(route.handlers) == 0 {
+		return "", fmt.Errorf("invalid route for %s %s: no handlers", method, route.path)
+	}
+	for i, handler := range route.handlers {
+		if handler == nil {
+			return "", fmt.Errorf("invalid route for %s %s: nil handler at index %d", method, route.path, i)
+		}
+	}
+	return method, nil
+}
+
+func isStandardHTTPMethod(method string) bool {
+	for _, valid := range fiber.DefaultMethods {
+		if method == valid {
+			return true
+		}
+	}
+	return false
+}
+
+func validateMiddlewareMount(mount middlewareMount) error {
+	if mount.prefix == "" || mount.prefix[0] != '/' {
+		return fmt.Errorf("invalid middleware prefix %q: must start with /", mount.prefix)
+	}
+	if mount.handler == nil {
+		return fmt.Errorf("invalid middleware for prefix %q: nil handler", mount.prefix)
+	}
+	return nil
+}
+
 // Option sets various options for New()
 type Option func(*S3ApiServer)

@@ -210,6 +280,32 @@ func WithWebUI(prefix string, cfg *webui.ServerConfig) Option {
 	}
 }

+// WithRoute registers a top-level Fiber route after the gateway rate limiter
+// and before S3 routes are registered. Handlers are terminal unless they
+// explicitly call ctx.Next().
+func WithRoute(method, path string, handlers ...fiber.Handler) Option {
+	return func(s *S3ApiServer) {
+		copied := append([]fiber.Handler(nil), handlers...)
+		s.routes = append(s.routes, routeMount{
+			method:   method,
+			path:     path,
+			handlers: copied,
+		})
+	}
+}
+
+// WithMiddleware mounts a Fiber middleware after the gateway rate limiter and
+// before the S3 route table is registered. The middleware must call ctx.Next()
+// for requests it does not fully handle.
+func WithMiddleware(prefix string, handler fiber.Handler) Option {
+	return func(s *S3ApiServer) {
+		s.middlewares = append(s.middlewares, middlewareMount{
+			prefix:  prefix,
+			handler: handler,
+		})
+	}
+}
+
 // WithConcurrencyLimiter sets the server's maximum connection limit
 // and the hard limit for in-flight requests.
 func WithConcurrencyLimiter(maxConnections, maxRequests int) Option {
@@ -15,13 +15,35 @@
 package s3api

 import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
 	"testing"

 	"github.com/gofiber/fiber/v2"
+	"github.com/versity/versitygw/auth"
 	"github.com/versity/versitygw/backend"
+	"github.com/versity/versitygw/s3api/middlewares"
 	"github.com/versity/versitygw/s3api/utils"
 )

+func newTestS3ApiServer(opts ...Option) (*S3ApiServer, error) {
+	allOpts := append([]Option{WithConcurrencyLimiter(10, 10)}, opts...)
+
+	return New(
+		backend.BackendUnsupported{},
+		middlewares.RootUserConfig{Access: "access", Secret: "secret"},
+		"us-east-1",
+		auth.NewIAMServiceSingle(auth.Account{Access: "access", Secret: "secret"}),
+		nil,
+		nil,
+		nil,
+		nil,
+		allOpts...,
+	)
+}
+
 func TestS3ApiServer_Serve(t *testing.T) {
 	tests := []struct {
 		name    string
@@ -59,3 +81,162 @@ func TestS3ApiServer_Serve(t *testing.T) {
 		})
 	}
 }
+
+func TestWithRouteRegistersBeforeMiddleware(t *testing.T) {
+	const routePath = "/custom/route"
+
+	middlewareCalled := false
+	server, err := newTestS3ApiServer(
+		WithRoute(http.MethodGet, routePath, func(ctx *fiber.Ctx) error {
+			return ctx.SendStatus(http.StatusNoContent)
+		}),
+		WithMiddleware("/", func(ctx *fiber.Ctx) error {
+			middlewareCalled = true
+			return ctx.SendStatus(http.StatusMisdirectedRequest)
+		}),
+	)
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+
+	resp, err := server.app.Test(httptest.NewRequest(http.MethodGet, routePath, nil))
+	if err != nil {
+		t.Fatalf("app.Test() error = %v", err)
+	}
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			t.Fatalf("response close error = %v", err)
+		}
+	}()
+	if resp.StatusCode != http.StatusNoContent {
+		t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusNoContent)
+	}
+	if middlewareCalled {
+		t.Fatal("middleware was called for top-level route")
+	}
+}
+
+func TestWithRouteRegistersAfterRateLimiter(t *testing.T) {
+	const routePath = "/custom/limited"
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+	firstDone := make(chan error, 1)
+	var once sync.Once
+
+	server, err := newTestS3ApiServer(
+		WithConcurrencyLimiter(10, 1),
+		WithRoute(http.MethodGet, routePath, func(ctx *fiber.Ctx) error {
+			once.Do(func() {
+				close(started)
+			})
+			<-release
+			return ctx.SendStatus(http.StatusNoContent)
+		}),
+	)
+	if err != nil {
+		t.Fatalf("New() error = %v", err)
+	}
+
+	go func() {
+		resp, err := server.app.Test(httptest.NewRequest(http.MethodGet, routePath, nil), -1)
+		if err != nil {
+			firstDone <- err
+			return
+		}
+		defer resp.Body.Close()
+		if resp.StatusCode != http.StatusNoContent {
+			firstDone <- fiber.NewError(resp.StatusCode)
+			return
+		}
+		firstDone <- nil
+	}()
+
+	<-started
+
+	resp, err := server.app.Test(httptest.NewRequest(http.MethodGet, routePath, nil), 100)
+	if err != nil {
+		close(release)
+		t.Fatalf("second app.Test() error = %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		close(release)
+		t.Fatalf("second status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable)
+	}
+
+	close(release)
+	if err := <-firstDone; err != nil {
+		t.Fatalf("first request error = %v", err)
+	}
+}
+
+func TestCustomMountValidation(t *testing.T) {
+	validHandler := func(ctx *fiber.Ctx) error {
+		return ctx.SendStatus(http.StatusNoContent)
+	}
+
+	tests := []struct {
+		name    string
+		opt     Option
+		wantErr string
+	}{
+		{
+			name:    "route empty method",
+			opt:     WithRoute("", "/custom", validHandler),
+			wantErr: "empty method",
+		},
+		{
+			name:    "route unsupported HTTP method",
+			opt:     WithRoute("BREW", "/custom", validHandler),
+			wantErr: "invalid HTTP method",
+		},
+		{
+			name:    "route empty path",
+			opt:     WithRoute(http.MethodGet, "", validHandler),
+			wantErr: "must start with /",
+		},
+		{
+			name:    "route relative path",
+			opt:     WithRoute(http.MethodGet, "custom", validHandler),
+			wantErr: "must start with /",
+		},
+		{
+			name:    "route no handlers",
+			opt:     WithRoute(http.MethodGet, "/custom"),
+			wantErr: "no handlers",
+		},
+		{
+			name:    "route nil handler",
+			opt:     WithRoute(http.MethodGet, "/custom", fiber.Handler(nil)),
+			wantErr: "nil handler",
+		},
+		{
+			name:    "middleware empty prefix",
+			opt:     WithMiddleware("", validHandler),
+			wantErr: "must start with /",
+		},
+		{
+			name:    "middleware relative prefix",
+			opt:     WithMiddleware("custom", validHandler),
+			wantErr: "must start with /",
+		},
+		{
+			name:    "middleware nil handler",
+			opt:     WithMiddleware("/custom", nil),
+			wantErr: "nil handler",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := newTestS3ApiServer(tt.opt)
+			if err == nil {
+				t.Fatal("New() error = nil, want error")
+			}
+			if !strings.Contains(err.Error(), tt.wantErr) {
+				t.Fatalf("New() error = %v, want substring %q", err, tt.wantErr)
+			}
+		})
+	}
+}