Repository: google/go-safeweb (go)

Issue:
Add support to set 2XX statuses that are not 200 and 204
There currently is no way to set any of the 2XX statuses, so Partial Content is not doable.

Pull request description:
Allow the user to set a custom status code in the safehttp.ResponseWriter
Fixes #151

Previously, there was no way of setting a status code other than `safehttp.NoContent` as the write method in the `safehttp.ResponseWriter` would always set the status code to `safehttp.StatusOK`. Implemented initial functionality and tests to make this possible.

Patch:
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 6ac4a62..a306061 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -30,6 +30,7 @@ type ResponseWriter struct {
 
 	// Having this field unexported is essential for security. Otherwise one can
 	// easily overwrite the struct bypassing all our safety guarantees.
+	code    StatusCode
 	header  Header
 	handler handler
 	req     *IncomingRequest
@@ -86,7 +87,12 @@ func (w *ResponseWriter) Write(resp Response) Result {
 		panic(err)
 	}
 	w.rw.Header().Set("Content-Type", ct)
-	w.rw.WriteHeader(int(StatusOK))
+
+	if w.code != 0 {
+		w.rw.WriteHeader(int(w.code))
+	} else {
+		w.rw.WriteHeader(int(StatusOK))
+	}
 
 	if err := w.d.Write(w.rw, resp); err != nil {
 		panic(err)
@@ -157,6 +163,25 @@ func (w *ResponseWriter) SetCookie(c *Cookie) error {
 	return w.header.addCookie(c)
 }
 
+// SetCode allows setting a response status. If the response was already
+// written, trying to set the status code will have no effect. This method will
+// panic if an invalid status code is passed (i.e. not in the range 1XX-5XX).
+//
+// If SetCode was called before NoContent, Redirect or WriteError, the status
+// code set by the latter will be the actual response status.
+//
+// TODO(empijei@, kele@, maramihali@): decide what should be done if the
+// code passed is either 3XX (redirect) or 4XX-5XX (client/server error).
+func (w *ResponseWriter) SetCode(code StatusCode) {
+	if w.written {
+		return
+	}
+	if code < 100 || code >= 600 {
+		panic("invalid status code")
+	}
+	w.code = code
+}
+
 // Dispatcher is responsible for writing a response received from the
 // ResponseWriter to the underlying http.ResponseWriter.
 //


Test patch:
diff --git a/safehttp/response_writer_test.go b/safehttp/response_writer_test.go
index cf13eae..7c781a5 100644
--- a/safehttp/response_writer_test.go
+++ b/safehttp/response_writer_test.go
@@ -144,3 +144,103 @@ func TestResponseWriterUnsafeResponse(t *testing.T) {
 		})
 	}
 }
+
+func TestResponseWriterStatusCode(t *testing.T) {
+	tests := []struct {
+		name       string
+		write      func(w *safehttp.ResponseWriter)
+		wantStatus safehttp.StatusCode
+	}{
+		{
+			name: "Default status code",
+			write: func(w *safehttp.ResponseWriter) {
+				w.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
+			},
+			wantStatus: safehttp.StatusOK,
+		},
+		{
+			name: "Custom status code",
+			write: func(w *safehttp.ResponseWriter) {
+				w.SetCode(safehttp.StatusPartialContent)
+				w.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
+			},
+			wantStatus: safehttp.StatusPartialContent,
+		},
+		{
+			name: "Status code set after Write",
+			write: func(w *safehttp.ResponseWriter) {
+				w.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
+				w.SetCode(safehttp.StatusPartialContent)
+			},
+			wantStatus: safehttp.StatusOK,
+		},
+		{
+			name: "SetCode then NoContent",
+			write: func(w *safehttp.ResponseWriter) {
+				w.SetCode(safehttp.StatusPartialContent)
+				w.NoContent()
+			},
+			wantStatus: safehttp.StatusNoContent,
+		},
+		{
+			name: "SetCode then Redirect",
+			write: func(w *safehttp.ResponseWriter) {
+				w.SetCode(safehttp.StatusPartialContent)
+				w.Redirect(safehttptest.NewRequest(safehttp.MethodGet, "/", nil), "", safehttp.StatusTemporaryRedirect)
+			},
+			wantStatus: safehttp.StatusTemporaryRedirect,
+		},
+		{
+			name: "SetCode then WriteError",
+			write: func(w *safehttp.ResponseWriter) {
+				w.SetCode(safehttp.StatusPartialContent)
+				w.WriteError(safehttp.StatusInsufficientStorage)
+			},
+			wantStatus: safehttp.StatusInsufficientStorage,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
+			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, testRw, nil)
+			tt.write(w)
+			if gotStatus := testRw.Status(); gotStatus != tt.wantStatus {
+				t.Errorf("tt.write(w): got %v, want %v", gotStatus, tt.wantStatus)
+			}
+		})
+	}
+}
+
+func TestResponseWriterInvalidStatusCode(t *testing.T) {
+	tests := []struct {
+		name  string
+		write func(w *safehttp.ResponseWriter)
+	}{
+		{
+			name: "Status code smaller than 100",
+			write: func(w *safehttp.ResponseWriter) {
+				w.SetCode(50)
+				w.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
+			},
+		},
+		{
+			name: "Status code bigger than 599",
+			write: func(w *safehttp.ResponseWriter) {
+				w.SetCode(1000)
+				w.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}), nil)
+			defer func() {
+				if r := recover(); r != nil {
+					return
+				}
+				t.Errorf("tt.write(w) expected panic")
+			}()
+			tt.write(w)
+		})
+	}
+}