Repository: google/go-safeweb
Pull request #289 (2021-05-12 13:33:12)

Remove NoContent special-casing.
Support NoContent as a normal Response.


Patch:
diff --git a/safehttp/default_dispatcher.go b/safehttp/default_dispatcher.go
index 87fbfd9..ceb5cce 100644
--- a/safehttp/default_dispatcher.go
+++ b/safehttp/default_dispatcher.go
@@ -80,6 +80,9 @@ func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
 	case RedirectResponse:
 		http.Redirect(rw, x.Request.req, x.Location, int(x.Code))
 		return nil
+	case NoContentResponse:
+		rw.WriteHeader(int(StatusNoContent))
+		return nil
 	default:
 		return fmt.Errorf("%T is not a safe response type and it cannot be written", resp)
 	}
diff --git a/safehttp/flight.go b/safehttp/flight.go
index 3011ad1..05300ad 100644
--- a/safehttp/flight.go
+++ b/safehttp/flight.go
@@ -70,7 +70,7 @@ func processRequest(cfg handlerConfig, rw http.ResponseWriter, req *http.Request
 	}
 	f.cfg.Handler.ServeHTTP(f, f.req)
 	if !f.written {
-		f.NoContent()
+		cfg.Dispatcher.Write(rw, NoContentResponse{})
 	}
 }
 
@@ -89,21 +89,6 @@ func (f *flight) Write(resp Response) Result {
 	return Result{}
 }
 
-// NoContent responds with a 204 No Content response.
-//
-// If the ResponseWriter has already been written to, then this method will panic.
-func (f *flight) NoContent() Result {
-	// TODO: should NoContent call Write under the hood? Should the dispatcher
-	// handle this too?
-	if f.written {
-		panic("ResponseWriter was already written to")
-	}
-	f.written = true
-	f.commitPhase(NoContentResponse{})
-	f.rw.WriteHeader(int(StatusNoContent))
-	return Result{}
-}
-
 // WriteError writes an error response (400-599) according to the provided
 // status code.
 //
diff --git a/safehttp/plugins/collector/collector.go b/safehttp/plugins/collector/collector.go
index 3b2595e..c87c639 100644
--- a/safehttp/plugins/collector/collector.go
+++ b/safehttp/plugins/collector/collector.go
@@ -180,7 +180,7 @@ func handleDeprecatedCSPReports(h func(CSPReport), w safehttp.ResponseWriter, b
 	}
 	h(rr)
 
-	return w.NoContent()
+	return w.Write(safehttp.NoContentResponse{})
 }
 
 var reportHandlers = map[string]func(json.RawMessage) (body interface{}, ok bool){
@@ -231,7 +231,7 @@ func handleReport(h func(Report), w safehttp.ResponseWriter, b []byte) safehttp.
 		return w.WriteError(safehttp.StatusBadRequest)
 	}
 
-	return w.NoContent()
+	return w.Write(safehttp.NoContentResponse{})
 }
 
 // cspViolationHandler parses reports of type csp-violation and returns a CSPReport.
diff --git a/safehttp/plugins/cors/cors.go b/safehttp/plugins/cors/cors.go
index 76bb92c..fe6dd26 100644
--- a/safehttp/plugins/cors/cors.go
+++ b/safehttp/plugins/cors/cors.go
@@ -149,7 +149,7 @@ func (it *Interceptor) Before(w safehttp.ResponseWriter, r *safehttp.IncomingReq
 	}
 
 	if status == safehttp.StatusNoContent {
-		return w.NoContent()
+		return w.Write(safehttp.NoContentResponse{})
 	}
 	return safehttp.NotWritten()
 }
diff --git a/safehttp/response.go b/safehttp/response.go
index f24d004..578dec4 100644
--- a/safehttp/response.go
+++ b/safehttp/response.go
@@ -102,8 +102,7 @@ func ExecuteNamedTemplateWithFuncs(w ResponseWriter, t Template, name string, da
 	return w.Write(&TemplateResponse{Template: t, Name: name, Data: data, FuncMap: fm})
 }
 
-// NoContentResponse is sent to the commit phase when it's initiated from
-// ResponseWriter.NoContent.
+// NoContentResponse is used to write a "No Content" response.
 type NoContentResponse struct{}
 
 // RedirectResponse is used to generate redirect responses.
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 1dc0f71..9a7ef22 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -26,11 +26,6 @@ type ResponseWriter interface {
 	// Write writes a safe response.
 	Write(resp Response) Result
 
-	// NoContent responds with a 204 No Content response.
-	//
-	// If the ResponseWriter has already been written to, then this method panics.
-	NoContent() Result
-
 	// WriteError writes an error response (400-599).
 	//
 	// If the ResponseWriter has already been written to, then this method panics.


Test patch:
diff --git a/safehttp/default_dispatcher_test.go b/safehttp/default_dispatcher_test.go
index b474e07..14885fd 100644
--- a/safehttp/default_dispatcher_test.go
+++ b/safehttp/default_dispatcher_test.go
@@ -126,8 +126,18 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 				return d.Write(w, safehttp.RedirectResponse{Request: r, Location: "/anotherpath", Code: safehttp.StatusFound})
 			},
 			wantHeaders: map[string][]string{"Location": {"/anotherpath"}},
+			wantStatus:  safehttp.StatusFound,
 			wantBody:    "<a href=\"/anotherpath\">Found</a>.\n\n",
 		},
+		{
+			name: "No Content Response",
+			write: func(w http.ResponseWriter) error {
+				d := &safehttp.DefaultDispatcher{}
+				return d.Write(w, safehttp.NoContentResponse{})
+			},
+			wantBody:   "",
+			wantStatus: safehttp.StatusNoContent,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -148,6 +158,15 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 					t.Errorf("response header %q: -want +got %s", k, diff)
 				}
 			}
+
+			wantStatus := tt.wantStatus
+			if wantStatus == 0 {
+				wantStatus = 200
+			}
+
+			if got := rw.Code; got != int(wantStatus) {
+				t.Errorf("Status: got %d, want %d", got, wantStatus)
+			}
 		})
 	}
 }
diff --git a/safehttp/flight_test.go b/safehttp/flight_test.go
index 1a3b467..8bcc6d8 100644
--- a/safehttp/flight_test.go
+++ b/safehttp/flight_test.go
@@ -124,7 +124,6 @@ func TestFlightDoubleWritePanics(t *testing.T) {
 		"Write": func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 			return w.Write(safehtml.HTMLEscaped("Hello"))
 		},
-		"NoContent": func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result { return w.NoContent() },
 		"WriteError": func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 			return w.WriteError(safehttp.StatusPreconditionFailed)
 		},