From ff288123e7633b875e0f0145adb2f29b0cc6b1c7 Mon Sep 17 00:00:00 2001
From: kartik <kartik.rajput622001@gmail.com>
Date: Mon, 30 Mar 2026 16:07:27 +0530
Subject: [PATCH] fix router prob writer interception

---
 CHANGELOG.md     |  6 ++++++
 router/group.go  |  4 ++--
 router/router.go | 24 ++++++++++++++++++++----
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 13765ef..8378e53 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.12.1] - 2026-03-30
+
+### Fixed
+
+- **router** — `probeWriter` no longer intercepts intentional 404/405 responses from user handlers; only unmatched routes (ServeMux's own 404/405) are routed through the `ErrorHandler`
+
 ## [0.12.0] - 2026-03-30
 
 ### Added
diff --git a/router/group.go b/router/group.go
index b519ee5..d342c3c 100644
--- a/router/group.go
+++ b/router/group.go
@@ -99,7 +99,7 @@ func (g *Group) Handle(pattern string, handler http.Handler) {
 	if len(chain) > 0 {
 		handler = middleware.Chain(chain...)(handler)
 	}
-	g.router.mux.Handle(fullPattern, handler)
+	g.router.mux.Handle(fullPattern, markMatched(handler))
 }
 
 // HandleFunc registers an http.HandlerFunc for the given pattern.
@@ -133,7 +133,7 @@ func (g *Group) register(method, pattern string, handler http.Handler) {
 	if len(chain) > 0 {
 		handler = middleware.Chain(chain...)(handler)
 	}
-	g.router.mux.Handle(fullPattern, handler)
+	g.router.mux.Handle(fullPattern, markMatched(handler))
 }
 
 // collectMiddleware walks the parent chain from root to current group
diff --git a/router/router.go b/router/router.go
index 9f6e69a..7867756 100644
--- a/router/router.go
+++ b/router/router.go
@@ -93,19 +93,23 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 }
 
 // probeWriter intercepts WriteHeader calls to detect 404/405 from ServeMux.
-// If the status is 404 or 405 and no user handler has written to the body,
+// If the status is 404 or 405 and no user handler has been matched,
 // it suppresses the write so the router's ErrorHandler can produce a consistent response.
+//
+// Registered handlers set pw.matched = true via a thin wrapper so that intentional
+// 404/405 responses from user handlers are forwarded correctly.
 type probeWriter struct {
 	http.ResponseWriter
 	code        int
 	intercepted bool
 	wroteBody   bool
+	matched     bool // true when a registered handler was matched by ServeMux
 }
 
 func (pw *probeWriter) WriteHeader(code int) {
-	// Only intercept if this is the first write (ServeMux's own 404/405).
-	// If a user handler already wrote body bytes, the handler owns the response.
-	if !pw.wroteBody && (code == http.StatusNotFound || code == http.StatusMethodNotAllowed) {
+	// Only intercept unmatched routes (ServeMux's own 404/405).
+	// If a registered handler was matched, forward the status as-is.
+	if !pw.matched && !pw.wroteBody && (code == http.StatusNotFound || code == http.StatusMethodNotAllowed) {
 		pw.code = code
 		pw.intercepted = true
 		return
@@ -122,6 +126,18 @@ func (pw *probeWriter) Write(b []byte) (int, error) {
 	return pw.ResponseWriter.Write(b)
 }
 
+// markMatched wraps a handler to set the matched flag on the probeWriter.
+// This lets probeWriter distinguish ServeMux's default 404/405 from
+// intentional 404/405 responses returned by user handlers.
+func markMatched(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if pw, ok := w.(*probeWriter); ok {
+			pw.matched = true
+		}
+		h.ServeHTTP(w, r)
+	})
+}
+
 // Get registers an error-returning handler for GET requests.
 func (r *Router) Get(pattern string, fn HandlerFunc) { r.group.Get(pattern, fn) }