Repository: google/go-safeweb (go)

Issue:
When mux responds with Method Not Allowed, make the response go through WriteError
https://github.com/google/go-safeweb/blob/407fc4e34c05ad9fea347a5388fb855045791551/safehttp/mux.go#L78

Quick fix: create a dummy handler with `w.WriteError(safehttp.StatusMethodNotAllowed)`.

Pull request:
MethodNotAllowed errors now go through ResponseWriter.WriteError
Fixes #230.

Patch:
diff --git a/safehttp/mux.go b/safehttp/mux.go
index 39bb65e..10d144d 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -71,14 +71,13 @@ type ServeMux struct {
 	mux *http.ServeMux
 }
 
-func registerHandlers(mux *http.ServeMux, handlers map[string]map[string]handlerConfig) {
+func registerHandlers(mux *http.ServeMux, handlers map[string]map[string]handlerConfig, methodNotAllowed handlerConfig) {
 	for pattern, handlersPerMethod := range handlers {
 		handlersPerMethod := handlersPerMethod
 		mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
 			cfg, ok := handlersPerMethod[r.Method]
 			if !ok {
-				http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
-				return
+				cfg = methodNotAllowed
 			}
 			processRequest(cfg, w, r)
 		})
@@ -111,9 +110,10 @@ func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 // ServeMuxConfig is a builder for ServeMux.
 type ServeMuxConfig struct {
-	dispatcher   Dispatcher
-	handlers     []handlerRegistration
-	interceptors []Interceptor
+	dispatcher              Dispatcher
+	handlers                []handlerRegistration
+	interceptors            []Interceptor
+	methodNotAllowedHandler handlerRegistration
 }
 
 // NewServeMuxConfig crates a ServeMuxConfig with the provided Dispatcher. If
@@ -122,7 +122,12 @@ func NewServeMuxConfig(disp Dispatcher) *ServeMuxConfig {
 	if disp == nil {
 		disp = &DefaultDispatcher{}
 	}
-	return &ServeMuxConfig{dispatcher: disp}
+	return &ServeMuxConfig{
+		dispatcher: disp,
+		methodNotAllowedHandler: handlerRegistration{
+			handler: HandlerFunc(defaultMethotNotAllowed),
+		},
+	}
 }
 
 type handlerRegistration struct {
@@ -149,6 +154,19 @@ func (s *ServeMuxConfig) Handle(pattern string, method string, h Handler, cfgs .
 	})
 }
 
+// HandleMethodNotAllowed registers a handler that runs when a given method is
+// not allowed for a registered path.
+func (s *ServeMuxConfig) HandleMethodNotAllowed(h Handler, cfgs ...InterceptorConfig) {
+	s.methodNotAllowedHandler = handlerRegistration{
+		handler: h,
+		cfgs:    cfgs,
+	}
+}
+
+func defaultMethotNotAllowed(w ResponseWriter, req *IncomingRequest) Result {
+	return w.WriteError(StatusMethodNotAllowed)
+}
+
 // Intercept installs an Interceptor.
 //
 // Interceptors order is undetermined and should not be relied on.
@@ -183,8 +201,13 @@ func (s *ServeMuxConfig) Mux() *ServeMux {
 				Interceptors: configureInterceptors(s.interceptors, hr.cfgs),
 			}
 	}
+	methodNotAllowed := handlerConfig{
+		Dispatcher:   s.dispatcher,
+		Handler:      s.methodNotAllowedHandler.handler,
+		Interceptors: configureInterceptors(s.interceptors, s.methodNotAllowedHandler.cfgs),
+	}
 	m := http.NewServeMux()
-	registerHandlers(m, handlers)
+	registerHandlers(m, handlers, methodNotAllowed)
 	return &ServeMux{mux: m}
 }
 


Test patch:
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index c29037b..2a14dba 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -16,6 +16,7 @@ package safehttp_test
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -589,3 +590,114 @@ func TestMuxHandlerReturnsNotWritten(t *testing.T) {
 		t.Errorf(`response body got: %q want: ""`, got)
 	}
 }
+
+func TestMuxMethodNotAllowedDefaults(t *testing.T) {
+	mb := safehttp.NewServeMuxConfig(nil)
+
+	h := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		panic("not tested")
+	})
+	mb.Handle("/", safehttp.MethodGet, h)
+
+	rw := httptest.NewRecorder()
+
+	mux := mb.Mux()
+	mux.ServeHTTP(rw, httptest.NewRequest(safehttp.MethodPost, "http://foo.com/", nil))
+
+	if got, want := rw.Code, int(safehttp.StatusMethodNotAllowed); got != want {
+		t.Errorf("rw.Code: got %v want %v", got, want)
+	}
+
+	wantHeader := map[string][]string{
+		"Content-Type":           {"text/plain; charset=utf-8"},
+		"X-Content-Type-Options": {"nosniff"},
+	}
+	if diff := cmp.Diff(wantHeader, map[string][]string(rw.Header())); diff != "" {
+		t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
+	}
+
+	wantBody := "Method Not Allowed\n"
+	if got := rw.Body.String(); got != wantBody {
+		t.Errorf("response body: got %q want %q", got, wantBody)
+	}
+}
+
+type methodNotAllowedError struct {
+	message string
+}
+
+func (err *methodNotAllowedError) Code() safehttp.StatusCode {
+	return safehttp.StatusMethodNotAllowed
+}
+
+type methodNotAllowedDispatcher struct {
+	safehttp.DefaultDispatcher
+}
+
+func (d *methodNotAllowedDispatcher) Error(rw http.ResponseWriter, resp safehttp.ErrorResponse) error {
+	x := resp.(*methodNotAllowedError)
+	rw.Header().Set("Content-Type", "text/html; charset=utf-8")
+	rw.WriteHeader(int(resp.Code()))
+	_, err := io.WriteString(rw, "<h1>"+http.StatusText(int(resp.Code()))+"</h1>"+"<p>"+x.message+"</p>")
+	return err
+}
+
+type methodNotAllowedInterceptor struct{}
+
+func (ip *methodNotAllowedInterceptor) Before(w safehttp.ResponseWriter, r *safehttp.IncomingRequest, ipcfg safehttp.InterceptorConfig) safehttp.Result {
+	cfg := ipcfg.(methodNotAllowedInterceptorConfig)
+	w.Header().Set("Before-Interceptor", cfg.before)
+	return safehttp.NotWritten()
+}
+
+// Commit runs before the response is written by the Dispatcher. If an error
+// is written to the ResponseWriter, then the Commit phases from the
+// remaining interceptors won't execute.
+func (ip *methodNotAllowedInterceptor) Commit(w safehttp.ResponseHeadersWriter, r *safehttp.IncomingRequest, resp safehttp.Response, ipcfg safehttp.InterceptorConfig) {
+	cfg := ipcfg.(methodNotAllowedInterceptorConfig)
+	w.Header().Set("Commit-Interceptor", cfg.commit)
+}
+
+type methodNotAllowedInterceptorConfig struct {
+	before, commit string
+}
+
+func (methodNotAllowedInterceptorConfig) Match(ip safehttp.Interceptor) bool {
+	_, ok := ip.(*methodNotAllowedInterceptor)
+	return ok
+}
+
+func TestMuxMethodNotAllowedCustom(t *testing.T) {
+	mb := safehttp.NewServeMuxConfig(&methodNotAllowedDispatcher{})
+	mb.Intercept(&methodNotAllowedInterceptor{})
+
+	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		panic("not tested")
+	}))
+	mb.HandleMethodNotAllowed(safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+		return rw.WriteError(&methodNotAllowedError{"custom message"})
+	}), methodNotAllowedInterceptorConfig{before: "foo", commit: "bar"})
+
+	rw := httptest.NewRecorder()
+
+	mux := mb.Mux()
+	mux.ServeHTTP(rw, httptest.NewRequest(safehttp.MethodPost, "http://foo.com/", nil))
+
+	if got, want := rw.Code, int(safehttp.StatusMethodNotAllowed); got != want {
+		t.Errorf("rw.Code: got %v want %v", got, want)
+	}
+
+	wantHeader := map[string][]string{
+		"Content-Type":       {"text/html; charset=utf-8"},
+		"Before-Interceptor": {"foo"},
+		"Commit-Interceptor": {"bar"},
+	}
+	if diff := cmp.Diff(wantHeader, map[string][]string(rw.Header())); diff != "" {
+		t.Errorf("rw.Header() mismatch (-want +got):\n%s", diff)
+	}
+
+	wantBody := "<h1>Method Not Allowed</h1><p>custom message</p>"
+	if got := rw.Body.String(); got != wantBody {
+		t.Errorf("response body: got %q want %q", got, wantBody)
+	}
+}