Repository: google/go-safeweb (go)

Issue:
Extend the Interceptor interface with an OnError method
This will run when the `safehttp.ResponseWriter.WriteError` method is called, before the error response is actually written, similarly to how `Commit` is run in the other methods for writing to the `ResponseWriter`.

Pull request:
Added the OnError method to the safehttp.Interceptor interface
Fixes #201 

The `OnError` method of all interceptors associated with a handler will be called in `safehttp.ResponseWriter.WriteError` before the error response is actually written.

Patch:
diff --git a/safehttp/interceptor.go b/safehttp/interceptor.go
index d2bb404..3966cac 100644
--- a/safehttp/interceptor.go
+++ b/safehttp/interceptor.go
@@ -30,6 +30,11 @@ type Interceptor interface {
 	// is written to the ResponseWriter, then the Commit phases from the
 	// remaining interceptors won't execute.
 	Commit(w *ResponseWriter, r *IncomingRequest, resp Response, cfg InterceptorConfig) Result
+
+	// OnError runs when ResponseWriter.WriteError is called, before the
+	// actual error response is written. An attempt to write to the
+	// ResponseWriter in this phase will result in an irrecoverable error.
+	OnError(w *ResponseWriter, r *IncomingRequest, resp Response, cfg InterceptorConfig) Result
 }
 
 // InterceptorConfig is a configuration of an interceptor.
diff --git a/safehttp/mux.go b/safehttp/mux.go
index d95f1bf..ebcd51f 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -289,3 +289,16 @@ func (h handler) commitPhase(w *ResponseWriter, resp Response) {
 		}
 	}
 }
+
+// errorPhase calls the OnError phases of all the interceptors associated with
+// a handler. This stage runs before an error response is written to the
+// ResponseWriter.
+//
+// TODO: BIG WARNING, if an interceptor attempts to write to the ResponseWriter
+// in the OnError phase will result in an irrecoverable error.
+func (h handler) errorPhase(w *ResponseWriter, resp Response) {
+	for i := len(h.interceps) - 1; i >= 0; i-- {
+		it := h.interceps[i]
+		it.Interceptor.OnError(w, w.req, resp, it.Config)
+	}
+}
diff --git a/safehttp/plugins/coop/coop.go b/safehttp/plugins/coop/coop.go
index 6df0d0b..abf1300 100644
--- a/safehttp/plugins/coop/coop.go
+++ b/safehttp/plugins/coop/coop.go
@@ -85,8 +85,16 @@ func (it Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingReq
 	return safehttp.NotWritten()
 }
 
-// Commit does nothing.
-func (it Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+// Commit is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (it Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+//
+// TODO: should OnError take as argument something that notifies it the commit
+// phase was already called?
+func (it Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
 	return safehttp.NotWritten()
 }
 
diff --git a/safehttp/plugins/cors/cors.go b/safehttp/plugins/cors/cors.go
index fce6acf..b0b0f6a 100644
--- a/safehttp/plugins/cors/cors.go
+++ b/safehttp/plugins/cors/cors.go
@@ -168,6 +168,11 @@ func appendToVary(w *safehttp.ResponseWriter, val string) {
 	}
 }
 
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (it Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 // preflight handles requests that have the method OPTIONS.
 func (it *Interceptor) preflight(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.StatusCode {
 	rh := r.Header
diff --git a/safehttp/plugins/csp/csp.go b/safehttp/plugins/csp/csp.go
index b546711..122a01c 100644
--- a/safehttp/plugins/csp/csp.go
+++ b/safehttp/plugins/csp/csp.go
@@ -238,3 +238,8 @@ func (it Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingReq
 	tmplResp.FuncMap["CSPNonce"] = func() string { return nonce }
 	return safehttp.NotWritten()
 }
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (it Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
diff --git a/safehttp/plugins/fetch_metadata/fetch_metadata.go b/safehttp/plugins/fetch_metadata/fetch_metadata.go
index 408abe8..ca3233e 100644
--- a/safehttp/plugins/fetch_metadata/fetch_metadata.go
+++ b/safehttp/plugins/fetch_metadata/fetch_metadata.go
@@ -173,3 +173,8 @@ func (p *Plugin) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest,
 func (p *Plugin) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
 	return safehttp.NotWritten()
 }
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (p *Plugin) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
diff --git a/safehttp/plugins/hostcheck/hostcheck.go b/safehttp/plugins/hostcheck/hostcheck.go
index f82ab2e..75e1978 100644
--- a/safehttp/plugins/hostcheck/hostcheck.go
+++ b/safehttp/plugins/hostcheck/hostcheck.go
@@ -54,3 +54,8 @@ func (it Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingReq
 func (Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
 	return safehttp.NotWritten()
 }
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
diff --git a/safehttp/plugins/hsts/hsts.go b/safehttp/plugins/hsts/hsts.go
index 2f19136..c1d3322 100644
--- a/safehttp/plugins/hsts/hsts.go
+++ b/safehttp/plugins/hsts/hsts.go
@@ -116,3 +116,8 @@ func (it Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingReq
 func (it Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
 	return safehttp.NotWritten()
 }
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (it Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
diff --git a/safehttp/plugins/staticheaders/staticheaders.go b/safehttp/plugins/staticheaders/staticheaders.go
index 9b579f4..c904162 100644
--- a/safehttp/plugins/staticheaders/staticheaders.go
+++ b/safehttp/plugins/staticheaders/staticheaders.go
@@ -54,3 +54,8 @@ func (Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingReques
 func (Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
 	return safehttp.NotWritten()
 }
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
diff --git a/safehttp/plugins/xsrf/xsrf.go b/safehttp/plugins/xsrf/xsrf.go
index 5d9cdc8..82e7271 100644
--- a/safehttp/plugins/xsrf/xsrf.go
+++ b/safehttp/plugins/xsrf/xsrf.go
@@ -161,3 +161,8 @@ func (it *Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRe
 	tmplResp.FuncMap["XSRFToken"] = func() string { return tok }
 	return safehttp.NotWritten()
 }
+
+// OnError is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (it *Interceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
diff --git a/safehttp/response.go b/safehttp/response.go
index d3ef1df..7ebb2d2 100644
--- a/safehttp/response.go
+++ b/safehttp/response.go
@@ -14,7 +14,9 @@
 
 package safehttp
 
-import "io"
+import (
+	"io"
+)
 
 // Response should encapsulate the data passed to the ResponseWriter to be
 // written by the Dispatcher. Any implementation of the interface should be
@@ -49,3 +51,10 @@ type TemplateResponse struct {
 // NoContentResponse is sent to the commit phase when it's initiated from
 // ResponseWriter.NoContent.
 type NoContentResponse struct{}
+
+// ErrorResponse is sent to the on error phase when initiated from
+// ResponseWriter.WriteError, encapsulating the status code with which the
+// method was called.
+type ErrorResponse struct {
+	Code StatusCode
+}
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 0abbcd5..73324ce 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -166,13 +166,15 @@ func (w *ResponseWriter) NoContent() Result {
 	return Result{}
 }
 
-// WriteError writes an error response (400-599) according to the provided status
-// code.
+// WriteError writes an error response (400-599) according to the provided
+// status code.
 //
 // If the ResponseWriter has already been written to, then this method will panic.
 func (w *ResponseWriter) WriteError(code StatusCode) Result {
 	w.markWritten()
-	http.Error(w.rw, http.StatusText(int(code)), int(code))
+	resp := &ErrorResponse{Code: code}
+	w.handler.errorPhase(w, resp)
+	http.Error(w.rw, http.StatusText(int(resp.Code)), int(resp.Code))
 	return Result{}
 }
 


Test patch:
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index c740734..4b287b5 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -177,6 +177,10 @@ func (p setHeaderInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.Inc
 	return safehttp.NotWritten()
 }
 
+func (p setHeaderInterceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 type internalErrorInterceptor struct{}
 
 func (internalErrorInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, cfg safehttp.InterceptorConfig) safehttp.Result {
@@ -188,6 +192,10 @@ func (internalErrorInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.I
 	return safehttp.NotWritten()
 }
 
+func (internalErrorInterceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 type claimHeaderInterceptor struct {
 	headerToClaim string
 }
@@ -204,6 +212,10 @@ func (p *claimHeaderInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.
 	return safehttp.NotWritten()
 }
 
+func (p *claimHeaderInterceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 func claimInterceptorSetHeader(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, value string) {
 	f := r.Context().Value(claimCtxKey{}).(func([]string))
 	f([]string{value})
@@ -219,6 +231,10 @@ func (panickingInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.Incom
 	return safehttp.NotWritten()
 }
 
+func (panickingInterceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 type committerInterceptor struct{}
 
 func (committerInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, cfg safehttp.InterceptorConfig) safehttp.Result {
@@ -230,15 +246,24 @@ func (committerInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.Incom
 	return safehttp.NotWritten()
 }
 
-type commitErroringInterceptor struct{}
-
-func (commitErroringInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, cfg safehttp.InterceptorConfig) safehttp.Result {
-	w.Header().Set("foo", "bar")
+func (committerInterceptor) OnError(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, _ safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
 	return safehttp.NotWritten()
 }
 
-func (commitErroringInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
-	return w.WriteError(safehttp.StatusInternalServerError)
+type setHeaderErroringInterceptor struct{}
+
+func (setHeaderErroringInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return w.WriteError(safehttp.StatusForbidden)
+}
+
+func (setHeaderErroringInterceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	w.Header().Set("name", "foo")
+	return safehttp.Result{}
+}
+
+func (setHeaderErroringInterceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	w.Header().Set("name", "bar")
+	return safehttp.Result{}
 }
 
 func TestMuxInterceptors(t *testing.T) {
@@ -352,6 +377,27 @@ func TestMuxInterceptors(t *testing.T) {
 			},
 			wantBody: "&lt;h1&gt;Hello World!&lt;/h1&gt;",
 		},
+		{
+			name: "OnError phase sets header",
+			mux: func() *safehttp.ServeMux {
+				mb := &safehttp.ServeMuxConfig{}
+				mb.Intercept(setHeaderErroringInterceptor{})
+
+				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
+				})
+				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+
+				return mb.Mux()
+			}(),
+			wantStatus: safehttp.StatusForbidden,
+			wantHeaders: map[string][]string{
+				"Content-Type":           {"text/plain; charset=utf-8"},
+				"X-Content-Type-Options": {"nosniff"},
+				"Name":                   {"bar"},
+			},
+			wantBody: "Forbidden\n",
+		},
 	}
 
 	for _, tt := range tests {
@@ -412,6 +458,10 @@ func (p setHeaderConfigInterceptor) Commit(w *safehttp.ResponseWriter, r *safeht
 	return safehttp.NotWritten()
 }
 
+func (p setHeaderConfigInterceptor) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 type noInterceptorConfig struct{}
 
 func (noInterceptorConfig) Match(i safehttp.Interceptor) bool {
@@ -498,6 +548,10 @@ func (interceptorOne) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingReq
 	return safehttp.NotWritten()
 }
 
+func (interceptorOne) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 type interceptorTwo struct{}
 
 func (interceptorTwo) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, cfg safehttp.InterceptorConfig) safehttp.Result {
@@ -516,6 +570,10 @@ func (interceptorTwo) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingReq
 	return safehttp.NotWritten()
 }
 
+func (interceptorTwo) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 type interceptorThree struct{}
 
 func (interceptorThree) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, cfg safehttp.InterceptorConfig) safehttp.Result {
@@ -534,6 +592,10 @@ func (interceptorThree) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingR
 	return safehttp.NotWritten()
 }
 
+func (interceptorThree) OnError(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, cfg safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}
+
 func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 	mb := &safehttp.ServeMuxConfig{}
 	mb.Intercept(interceptorOne{})
@@ -597,3 +659,41 @@ func TestMuxHandlerReturnsNotWritten(t *testing.T) {
 		t.Errorf(`response body got: %q want: ""`, got)
 	}
 }
+
+type onErrorWriteInterceptor struct{}
+
+func (onErrorWriteInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, _ safehttp.InterceptorConfig) safehttp.Result {
+	return w.WriteError(safehttp.StatusBadRequest)
+}
+
+func (onErrorWriteInterceptor) Commit(_ *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, _ safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.Result{}
+}
+
+func (onErrorWriteInterceptor) OnError(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, _ safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	w.Header().Set("header", "bad")
+	return w.WriteError(safehttp.StatusInsufficientStorage)
+}
+
+func TestMuxInterceptorWriteInOnError(t *testing.T) {
+	b := &strings.Builder{}
+	rw := safehttptest.NewTestResponseWriter(b)
+
+	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
+
+	mb := &safehttp.ServeMuxConfig{}
+	mb.Intercept(onErrorWriteInterceptor{})
+
+	registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
+	})
+	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf(`mb.Mux().ServeHTTP(rw, req): expected panic`)
+		}
+	}()
+
+	mb.Mux().ServeHTTP(rw, req)
+}