middleware: GetHead advertises HEAD in the 405 Allow header

middleware/get_head.go

@@ -2,11 +2,16 @@ package middleware
 
 import (
 	"net/http"
+	"strings"
 
 	"github.com/go-chi/chi/v5"
 )
 
 // GetHead automatically route undefined HEAD requests to GET handlers.
+//
+// When a 405 Method Not Allowed response is emitted for a route whose
+// Allow header lists GET, HEAD is also added to that header to reflect
+// the implicit GET-to-HEAD routing this middleware provides.
 func GetHead(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.Method == "HEAD" {
@@ -29,11 +34,51 @@ func GetHead(next http.Handler) http.Handler {
 			if !rctx.Routes.Match(tctx, "HEAD", routePath) {
 				rctx.RouteMethod = "GET"
 				rctx.RoutePath = routePath
-				next.ServeHTTP(w, r)
+				next.ServeHTTP(&getHeadAllowWriter{ResponseWriter: w}, r)
 				return
 			}
 		}
 
-		next.ServeHTTP(w, r)
+		next.ServeHTTP(&getHeadAllowWriter{ResponseWriter: w}, r)
 	})
 }
+
+// getHeadAllowWriter advertises HEAD support alongside GET in any
+// 405 Method Not Allowed response. The methodNotAllowedHandler builds
+// the Allow header from the route's registered methods, which doesn't
+// include HEAD when only GET is registered, so the response would
+// otherwise omit it.
+type getHeadAllowWriter struct {
+	http.ResponseWriter
+	wroteHeader bool
+}
+
+func (w *getHeadAllowWriter) WriteHeader(status int) {
+	if !w.wroteHeader && status == http.StatusMethodNotAllowed {
+		hdr := w.ResponseWriter.Header()
+		allow := hdr.Values("Allow")
+		hasGet, hasHead := false, false
+		for _, a := range allow {
+			for _, m := range strings.Split(a, ",") {
+				switch strings.TrimSpace(m) {
+				case http.MethodGet:
+					hasGet = true
+				case http.MethodHead:
+					hasHead = true
+				}
+			}
+		}
+		if hasGet && !hasHead {
+			hdr.Add("Allow", http.MethodHead)
+		}
+	}
+	w.wroteHeader = true
+	w.ResponseWriter.WriteHeader(status)
+}
+
+func (w *getHeadAllowWriter) Write(b []byte) (int, error) {
+	if !w.wroteHeader {
+		w.WriteHeader(http.StatusOK)
+	}
+	return w.ResponseWriter.Write(b)
+}

middleware/get_head_test.go

@@ -3,6 +3,7 @@ package middleware
 import (
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/go-chi/chi/v5"
@@ -64,3 +65,33 @@ func TestGetHead(t *testing.T) {
 		t.Fatalf("expecting X-User header '-' but got '%s'", req.Header.Get("X-User"))
 	}
 }
+
+// Regression for #1030. GetHead implicitly serves HEAD via the GET
+// handler, so a 405 response for a route that has GET should advertise
+// HEAD in the Allow header alongside GET.
+func TestGetHead_AllowHeaderIncludesHead(t *testing.T) {
+	r := chi.NewRouter()
+	r.Use(GetHead)
+	r.Get("/only-get", func(w http.ResponseWriter, _ *http.Request) {
+		w.Write([]byte("ok"))
+	})
+
+	ts := httptest.NewServer(r)
+	defer ts.Close()
+
+	req, _ := testRequest(t, ts, "POST", "/only-get", nil)
+	if req.StatusCode != http.StatusMethodNotAllowed {
+		t.Fatalf("status = %d, want 405", req.StatusCode)
+	}
+	allow := req.Header.Values("Allow")
+	joined := ""
+	for _, v := range allow {
+		joined += "," + v
+	}
+	if !strings.Contains(joined, "GET") {
+		t.Fatalf("Allow header missing GET: %v", allow)
+	}
+	if !strings.Contains(joined, "HEAD") {
+		t.Fatalf("Allow header missing HEAD: %v", allow)
+	}
+}