From 8a63f142a7f28aa11d89f470a9c0a80102876aa9 Mon Sep 17 00:00:00 2001
From: Aleksey Shakhmatov <dev@shakhmatov.com>
Date: Sun, 22 Mar 2026 21:48:13 +0300
Subject: [PATCH] Add WithNotFoundHandler option for custom 404 responses on
 Router

Allows configuring a custom handler for unmatched routes, enabling
consistent JSON error responses instead of ServeMux's default plain
text. NewRouter now accepts RouterOption functional options.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 server/route.go               | 33 ++++++++++++++---
 server/route_notfound_test.go | 70 +++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 5 deletions(-)
 create mode 100644 server/route_notfound_test.go

diff --git a/server/route.go b/server/route.go
index 27599ca..2f32d1c 100644
--- a/server/route.go
+++ b/server/route.go
@@ -9,16 +9,31 @@ import (
 // groups and sub-router mounting. It leverages Go 1.22+ enhanced patterns
 // like "GET /users/{id}".
 type Router struct {
-	mux        *http.ServeMux
-	prefix     string
-	middlewares []Middleware
+	mux             *http.ServeMux
+	prefix          string
+	middlewares     []Middleware
+	notFoundHandler http.Handler
+}
+
+// RouterOption configures a Router.
+type RouterOption func(*Router)
+
+// WithNotFoundHandler sets a custom handler for requests that don't match
+// any registered pattern. This is useful for returning JSON 404/405 responses
+// instead of the default plain text.
+func WithNotFoundHandler(h http.Handler) RouterOption {
+	return func(r *Router) { r.notFoundHandler = h }
 }
 
 // NewRouter creates a new Router backed by a fresh http.ServeMux.
-func NewRouter() *Router {
-	return &Router{
+func NewRouter(opts ...RouterOption) *Router {
+	r := &Router{
 		mux: http.NewServeMux(),
 	}
+	for _, opt := range opts {
+		opt(r)
+	}
+	return r
 }
 
 // Handle registers a handler for the given pattern. The pattern follows
@@ -63,6 +78,14 @@ func (r *Router) Mount(prefix string, handler http.Handler) {
 
 // ServeHTTP implements http.Handler, making Router usable as a handler.
 func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	if r.notFoundHandler != nil {
+		// Use the mux to check for a match. If none, use the custom handler.
+		_, pattern := r.mux.Handler(req)
+		if pattern == "" {
+			r.notFoundHandler.ServeHTTP(w, req)
+			return
+		}
+	}
 	r.mux.ServeHTTP(w, req)
 }
 
diff --git a/server/route_notfound_test.go b/server/route_notfound_test.go
new file mode 100644
index 0000000..134557e
--- /dev/null
+++ b/server/route_notfound_test.go
@@ -0,0 +1,70 @@
+package server_test
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"git.codelab.vc/pkg/httpx/server"
+)
+
+func TestRouter_NotFoundHandler(t *testing.T) {
+	t.Run("custom 404 handler", func(t *testing.T) {
+		notFound := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusNotFound)
+			json.NewEncoder(w).Encode(map[string]string{"error": "not found"})
+		})
+
+		r := server.NewRouter(server.WithNotFoundHandler(notFound))
+		r.HandleFunc("GET /exists", func(w http.ResponseWriter, _ *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		})
+
+		// Matched route works normally.
+		w := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, "/exists", nil)
+		r.ServeHTTP(w, req)
+
+		if w.Code != http.StatusOK {
+			t.Fatalf("matched route: got status %d, want %d", w.Code, http.StatusOK)
+		}
+
+		// Unmatched route uses custom handler.
+		w = httptest.NewRecorder()
+		req = httptest.NewRequest(http.MethodGet, "/nope", nil)
+		r.ServeHTTP(w, req)
+
+		if w.Code != http.StatusNotFound {
+			t.Fatalf("not found: got status %d, want %d", w.Code, http.StatusNotFound)
+		}
+		if ct := w.Header().Get("Content-Type"); ct != "application/json" {
+			t.Fatalf("Content-Type = %q, want %q", ct, "application/json")
+		}
+
+		var body map[string]string
+		if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
+			t.Fatalf("unmarshal: %v", err)
+		}
+		if body["error"] != "not found" {
+			t.Fatalf("error = %q, want %q", body["error"], "not found")
+		}
+	})
+
+	t.Run("default behavior without custom handler", func(t *testing.T) {
+		r := server.NewRouter()
+		r.HandleFunc("GET /exists", func(w http.ResponseWriter, _ *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		})
+
+		w := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, "/nope", nil)
+		r.ServeHTTP(w, req)
+
+		// Default ServeMux returns 404.
+		if w.Code != http.StatusNotFound {
+			t.Fatalf("got status %d, want %d", w.Code, http.StatusNotFound)
+		}
+	})
+}