Modify SetCode to not panic if it's called after the response was alr…

…eady written

safehttp/response_writer.go

@@ -163,9 +163,9 @@ func (w *ResponseWriter) SetCookie(c *Cookie) error {
 	return w.header.addCookie(c)
 }
 
-// SetCode allows setting a response status. This method will panic if an
-// invalid status code is passed (i.e. not in the range 1XX-5XX) or if the
-// response has already been written.
+// 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.
@@ -174,7 +174,7 @@ func (w *ResponseWriter) SetCookie(c *Cookie) error {
 // code passed is either 3XX (redirect) or 4XX-5XX (client/server error).
 func (w *ResponseWriter) SetCode(code StatusCode) {
 	if w.written {
-		panic("ResponseWriter was already written to")
+		return
 	}
 	if code < 100 || code >= 600 {
 		panic("invalid status code")

safehttp/response_writer_test.go

@@ -166,6 +166,14 @@ func TestResponseWriterStatusCode(t *testing.T) {
 			},
 			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) {
@@ -222,13 +230,6 @@ func TestResponseWriterInvalidStatusCode(t *testing.T) {
 				w.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
 			},
 		},
-		{
-			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)
-			},
-		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {