danielgtaylor · danielgtaylor · May 9, 2025 · Apr 29, 2025 · Apr 29, 2025 · Apr 29, 2025
diff --git a/adapters/humamux/humagmux_test.go b/adapters/humamux/humagmux_test.go
@@ -2,18 +2,81 @@ package humamux
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/danielgtaylor/huma/v2"
+	"github.com/danielgtaylor/huma/v2/humatest"
 	"github.com/gorilla/mux"
+	"github.com/stretchr/testify/assert"
 )
 
 var lastModified = time.Now()
 
+type TestInput struct {
+	Group      string `path:"group"`
+	Verbose    bool   `query:"verbose"`
+	Auth       string `header:"Authorization"`
+	TestHeader string `header:"TestHeader"`
+	Body       struct {
+		Name  string `json:"name"`
+		Email string `json:"email"`
+	}
+}
+
+// Test outputs (headers, body).
+type TestOutput struct {
+	MyHeader   string `header:"MyHeader"`
+	TestHeader string `header:"TestHeader"`
+	Body       struct {
+		Message string `json:"message"`
+	}
+}
+
+func testHandler(ctx context.Context, input *TestInput) (*TestOutput, error) {
+	resp := &TestOutput{}
+	resp.MyHeader = "my-value"
+	resp.TestHeader = input.TestHeader
+	resp.Body.Message = fmt.Sprintf("Hello, %s <%s>! (%s, %v, %s)", input.Body.Name, input.Body.Email, input.Group, input.Verbose, input.Auth)
+	return resp, nil
+}
+
+func TestCustomMiddleware(t *testing.T) {
+	mw1 := func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			r.Header.Set("TestHeader", "test-value")
+			next.ServeHTTP(w, r)
+		})
+	}
+
+	r := mux.NewRouter()
+	api := New(r, huma.DefaultConfig("Test", "1.0.0"),
+		WithRouteCustomizer(func(op *huma.Operation, r *mux.Route) {
+			r.Handler(mw1(r.GetHandler()))
+		}))
+
+	huma.Register(api, huma.Operation{
+		OperationID: "test",
+		Method:      http.MethodGet,
+		Path:        "/{group}",
+	}, testHandler)
+
+	testAPI := humatest.Wrap(t, api)
+	resp := testAPI.Do(http.MethodGet, "/foo",
+		"Host: localhost",
+		"Authorization: Bearer abc123",
+		strings.NewReader(`{"name": "Daniel", "email": "daniel@example.com"}`),
+	)
+
+	assert.Equal(t, http.StatusOK, resp.Code)
+	assert.Equal(t, "my-value", resp.Header().Get("MyHeader"))
+	assert.Equal(t, "test-value", resp.Header().Get("TestHeader"))
+}
+
 func BenchmarkHumaGorillaMux(b *testing.B) {
 	type GreetingInput struct {
 		ID          string `path:"id"`

diff --git a/adapters/humamux/humamux.go b/adapters/humamux/humamux.go
@@ -132,23 +132,27 @@ func (c *gmuxContext) BodyWriter() io.Writer {
 }
 
 type gMux struct {
+	options
 	router *mux.Router
 }
 
 func (a *gMux) Handle(op *huma.Operation, handler func(huma.Context)) {
-	a.router.
+	route := a.router.
 		NewRoute().
 		Path(op.Path).
 		Methods(op.Method).
 		HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			handler(&gmuxContext{op: op, r: r, w: w})
 		})
+	if a.routeCustomizer != nil {
+		a.routeCustomizer(op, route)
+	}
 }
 
 func (a *gMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	a.router.ServeHTTP(w, r)
 }
 
-func New(r *mux.Router, config huma.Config) huma.API {
-	return huma.NewAPI(config, &gMux{router: r})
+func New(r *mux.Router, config huma.Config, options ...Option) huma.API {
+	return huma.NewAPI(config, &gMux{router: r, options: parseOptions(options)})
 }
diff --git a/adapters/humamux/options.go b/adapters/humamux/options.go
@@ -0,0 +1,29 @@
+package humamux
+
+import (
+	"github.com/danielgtaylor/huma/v2"
+	"github.com/gorilla/mux"
+)
+
+type Option func(*options)
+
+// WithRouteCustomizer allows customizing a mux route, like adding HTTP middlewares.
+func WithRouteCustomizer(f func(op *huma.Operation, r *mux.Route)) Option {
+	return func(o *options) {
+		o.routeCustomizer = f
+	}
+}
+
+// options
+
+func parseOptions(optionList []Option) options {
+	var optns options
+	for _, opt := range optionList {
+		opt(&optns)
+	}
+	return optns
+}
+
+type options struct {
+	routeCustomizer func(op *huma.Operation, r *mux.Route)
+}