diff --git a/README.md b/README.md index d9a9d31..4ef7464 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,8 @@ native performance untouched. Simple, not simplistic. `Options`, `Head`, `Connect`, `Trace`, and `Any` - Routing follows the [net/http.ServeMux](https://pkg.go.dev/net/http#ServeMux) patterns - Route groups under a common prefix +- Customizable 404 and 405 responses with `SetNotFoundHandler` and + `SetMethodNotAllowedHandler` - Standard `net/http` handlers work everywhere — no framework-specific context type to learn - Wrapper composition with `Chain` and `ChainFunc` diff --git a/app.go b/app.go index 4bd915b..f2df4a1 100644 --- a/app.go +++ b/app.go @@ -36,6 +36,12 @@ type App struct { // first, and its response is what the caller ultimately sees — // the same composition as [Chain]. ss []func(http.Handler) http.Handler + + // Handlers served when a request matches no registered pattern + // (404) or matches a pattern but not its method (405); nil keeps + // the mux defaults. See: https://github.com/golang/go/issues/65648 + notFoundHandler http.Handler + methodNotAllowedHandler http.Handler } var ( @@ -82,6 +88,20 @@ func Default() *App { return app } +// SetNotFoundHandler sets the handler served when a request matches +// no registered pattern. When unset, the mux responds with 404 Not +// Found. +func (a *App) SetNotFoundHandler(h http.Handler) { + a.notFoundHandler = h +} + +// SetMethodNotAllowedHandler sets the handler served when a request +// path matches a registered pattern but its method does not. When +// unset, the mux responds with 405 Method Not Allowed. +func (a *App) SetMethodNotAllowedHandler(h http.Handler) { + a.methodNotAllowedHandler = h +} + // Use registers the given wrappers and applies them to every handler // registered after this call. Wrappers run in registration order: // the first is outermost and receives the request first, the same @@ -170,9 +190,11 @@ func (a *App) Group(relativePath string, fn func(r Router)) { return } app := &App{ - basePath: path.Join(cmp.Or(a.basePath, "/"), relativePath), - mux: a.mux, - ss: slices.Clone(a.ss), + basePath: path.Join(cmp.Or(a.basePath, "/"), relativePath), + mux: a.mux, + ss: slices.Clone(a.ss), + notFoundHandler: a.notFoundHandler, + methodNotAllowedHandler: a.methodNotAllowedHandler, } fn(app) } @@ -209,8 +231,71 @@ func parsePattern(s string) (string, string) { return method, rest } +// ServeHTTP dispatches r to the handler registered for the matching +// pattern. When nothing matches, it serves the handlers set by +// [App.SetNotFoundHandler] and [App.SetMethodNotAllowedHandler], +// falling back to the mux defaults when unset. func (a *App) ServeHTTP(w http.ResponseWriter, r *http.Request) { - a.mux.ServeHTTP(w, r) + // Fast path: without custom error handlers the mux serves its + // default 404 and 405 responses directly. + if a.notFoundHandler == nil && a.methodNotAllowedHandler == nil { + a.mux.ServeHTTP(w, r) + return + } + + h, pattern := a.mux.Handler(r) + if pattern != "" { + // Re-dispatch through ServeHTTP instead of serving h: Handler + // does not populate r.pat/r.matches, so r.Pattern and + // r.PathValue would be broken for wildcard patterns. + a.mux.ServeHTTP(w, r) + return + } + + // No pattern matched: h is either the mux's 404 handler or its + // 405 handler (path matched, method did not). Serve it into a + // recorder to tell the two cases apart. + rec := &statusRecorder{header: make(http.Header)} + h.ServeHTTP(rec, r) + + switch rec.code { + case http.StatusNotFound: + if a.notFoundHandler != nil { + a.notFoundHandler.ServeHTTP(w, r) + return + } + case http.StatusMethodNotAllowed: + if a.methodNotAllowedHandler != nil { + // RFC 9110 mandates Allow on 405 responses; carry over the + // methods the mux computed. The custom handler may still + // override or delete it. + w.Header().Set("Allow", rec.header.Get("Allow")) + a.methodNotAllowedHandler.ServeHTTP(w, r) + return + } + } + // Only the other custom handler is set; replay the mux's default + // response, including Allow for 405. + h.ServeHTTP(w, r) +} + +// statusRecorder is an [http.ResponseWriter] that discards the body +// and records only the status code and header. +type statusRecorder struct { + header http.Header + code int +} + +var _ http.ResponseWriter = (*statusRecorder)(nil) + +func (s *statusRecorder) Header() http.Header { return s.header } + +func (s *statusRecorder) Write(b []byte) (int, error) { return len(b), nil } + +func (s *statusRecorder) WriteHeader(code int) { + if s.code == 0 { + s.code = code + } } // shutdownTimeout bounds how long a graceful shutdown may wait for diff --git a/app_test.go b/app_test.go index b05bcfd..bb55719 100644 --- a/app_test.go +++ b/app_test.go @@ -285,6 +285,140 @@ func TestMethodNotAllowed(t *testing.T) { } } +func handle404405App(notFound, notAllowed http.Handler) *sim.App { + app := sim.NewApp() + app.Get("/items", markHandlerFunc("items")) + app.SetNotFoundHandler(notFound) + app.SetMethodNotAllowedHandler(notAllowed) + return app +} + +func statusHandler(code int, body string) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(code) + _, _ = io.WriteString(w, body) + }) +} + +//nolint:funlen // Table‑driven test covering 404 / 405 handler +func TestCustomErrorHandlers(t *testing.T) { + notFound := statusHandler(http.StatusNotFound, "nf") + notAllowed := statusHandler(http.StatusMethodNotAllowed, "mna") + + tests := []struct { + name string + notFound http.Handler + notAllowed http.Handler + method string + target string + wantCode int + wantBody string + }{ + { + name: "custom not found", + notFound: notFound, + method: http.MethodGet, + target: "/nope", + wantCode: http.StatusNotFound, + wantBody: "nf", + }, + { + name: "custom method not allowed", + notAllowed: notAllowed, + method: http.MethodPost, + target: "/items", + wantCode: http.StatusMethodNotAllowed, + wantBody: "mna", + }, + { + name: "default not found when only 405 handler set", + notAllowed: notAllowed, + method: http.MethodGet, + target: "/nope", + wantCode: http.StatusNotFound, + wantBody: "404 page not found\n", + }, + { + name: "default method not allowed when only 404 handler set", + notFound: notFound, + method: http.MethodPost, + target: "/items", + wantCode: http.StatusMethodNotAllowed, + wantBody: "Method Not Allowed\n", + }, + { + name: "both set - custom not found", + notFound: notFound, + notAllowed: notAllowed, + method: http.MethodGet, + target: "/nope", + wantCode: http.StatusNotFound, + wantBody: "nf", + }, + { + name: "both set - custom method not allowed", + notFound: notFound, + notAllowed: notAllowed, + method: http.MethodPost, + target: "/items", + wantCode: http.StatusMethodNotAllowed, + wantBody: "mna", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rec := serve(t, handle404405App(tt.notFound, tt.notAllowed), tt.method, tt.target) + if rec.Code != tt.wantCode { + t.Fatalf("%s %s = %d, want %d", tt.method, tt.target, rec.Code, tt.wantCode) + } + if rec.Body.String() != tt.wantBody { + t.Errorf("%s %s body = %q, want %q", tt.method, tt.target, rec.Body.String(), tt.wantBody) + } + if tt.wantCode == http.StatusMethodNotAllowed && + !strings.Contains(rec.Header().Get("Allow"), http.MethodGet) { + t.Errorf("Allow = %q, want it to contain %q", rec.Header().Get("Allow"), http.MethodGet) + } + }) + } +} + +func TestCustomErrorHandlersDoNotAffectRouting(t *testing.T) { + app := handle404405App( + statusHandler(http.StatusNotFound, "nf"), + statusHandler(http.StatusMethodNotAllowed, "mna"), + ) + app.Get("/tree/", markHandlerFunc("tree")) + + expect(t, app, http.MethodGet, "/items", "items") + expect(t, app, http.MethodGet, "/tree/", "tree") + + rec := serve(t, app, http.MethodGet, "/tree") + if rec.Code != http.StatusTemporaryRedirect { + t.Fatalf("GET /tree = %d, want %d", rec.Code, http.StatusTemporaryRedirect) + } + if loc := rec.Header().Get("Location"); loc != "/tree/" { + t.Errorf("Location = %q, want %q", loc, "/tree/") + } +} + +func TestCustomErrorHandlersPreservePathValues(t *testing.T) { + app := handle404405App( + statusHandler(http.StatusNotFound, "nf"), + statusHandler(http.StatusMethodNotAllowed, "mna"), + ) + app.Get("/items/{id}", func(w http.ResponseWriter, r *http.Request) { + _, _ = fmt.Fprintf(w, "%s:%s", r.Pattern, r.PathValue("id")) + }) + + rec := serve(t, app, http.MethodGet, "/items/42") + if rec.Code != http.StatusOK { + t.Fatalf("GET /items/42 = %d, want %d", rec.Code, http.StatusOK) + } + if want := "GET /items/{id}:42"; rec.Body.String() != want { + t.Errorf("body = %q, want %q", rec.Body.String(), want) + } +} + func TestConflictingPatternPanics(t *testing.T) { app := sim.NewApp() app.Get("/a/{x}", markHandlerFunc("first"))