Repository: google/go-safeweb (go)

Issue:
ResponseWriter.ClientError support (4xx) similar to ResponseWriter.ServerError (5xx)
While a `safehttp.ResponseWriter.ServerError` method is currently being implemented, there is no support for for returning client errors when a request fails.

`ServerError` is supposed to handle `5xx` errors, this is about `4xx`.

Pull request:
ResponeWriter.ClientError support
Fixes #52 .

Adds functionality to the `safehttp.ResponseWriter` so that a response with the `4xx` status code can be returned when the request fails because of a client error.

Patch:
diff --git a/safehttp/header.go b/safehttp/header.go
index 5327192..4531085 100644
--- a/safehttp/header.go
+++ b/safehttp/header.go
@@ -24,31 +24,40 @@ import (
 // The keys will be in canonical form, as returned by
 // textproto.CanonicalMIMEHeaderKey.
 type Header struct {
-	wrapped   http.Header
-	immutable map[string]bool
+	wrapped http.Header
+	claimed map[string]bool
 }
 
 func newHeader(h http.Header) Header {
-	return Header{wrapped: h, immutable: map[string]bool{}}
+	return Header{
+		wrapped: h,
+		claimed: map[string]bool{},
+	}
 }
 
-// MarkImmutable marks the header with the given name as immutable.
-// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
-// This header is now read-only. If the header previously was
-// undefined, then it will forever be undefined. Other methods in
+// Claim claims the header with the given name and returns a function
+// which can be used to set the header. The name is first canonicalized
+// using textproto.CanonicalMIMEHeaderKey. Other methods in
 // the struct can't write to, change or delete the header with this
-// name. These methods will instead fail when applied on an immutable
-// header.
-func (h Header) MarkImmutable(name string) {
+// name. These methods will instead fail when applied on an claimed
+// header. The only way to modify the header is to use the returned
+// function. The Set-Cookie header can't be claimed.
+func (h Header) Claim(name string) (set func([]string), err error) {
 	name = textproto.CanonicalMIMEHeaderKey(name)
-	h.immutable[name] = true
+	if err := h.writableHeader(name); err != nil {
+		return nil, err
+	}
+	h.claimed[name] = true
+	return func(v []string) {
+		h.wrapped[name] = v
+	}, nil
 }
 
 // Set sets the header with the given name to the given value.
 // The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
 // This method first removes all other values associated with this
 // header before setting the new value. Returns an error when
-// applied on immutable headers or on the Set-Cookie header.
+// applied on claimed headers or on the Set-Cookie header.
 func (h Header) Set(name, value string) error {
 	name = textproto.CanonicalMIMEHeaderKey(name)
 	if err := h.writableHeader(name); err != nil {
@@ -61,7 +70,7 @@ func (h Header) Set(name, value string) error {
 // Add adds a new header with the given name and the given value to
 // the collection of headers. The name is first canonicalized using
 // textproto.CanonicalMIMEHeaderKey. Returns an error when applied
-// on immutable headers or on the Set-Cookie header.
+// on claimed headers or on the Set-Cookie header.
 func (h Header) Add(name, value string) error {
 	name = textproto.CanonicalMIMEHeaderKey(name)
 	if err := h.writableHeader(name); err != nil {
@@ -73,7 +82,7 @@ func (h Header) Add(name, value string) error {
 
 // Del deletes all headers with the given name. The name is first
 // canonicalized using textproto.CanonicalMIMEHeaderKey. Returns an
-// error when applied on immutable headers or on the Set-Cookie
+// error when applied on claimed headers or on the Set-Cookie
 // header.
 func (h Header) Del(name string) error {
 	name = textproto.CanonicalMIMEHeaderKey(name)
@@ -126,8 +135,8 @@ func (h Header) writableHeader(name string) error {
 	if name == "Set-Cookie" {
 		return errors.New("can't write to Set-Cookie header")
 	}
-	if h.immutable[name] {
-		return errors.New("immutable header")
+	if h.claimed[name] {
+		return errors.New("claimed header")
 	}
 	return nil
 }
diff --git a/safehttp/plugins/hsts/hsts.go b/safehttp/plugins/hsts/hsts.go
index 0382c52..dd2ceac 100644
--- a/safehttp/plugins/hsts/hsts.go
+++ b/safehttp/plugins/hsts/hsts.go
@@ -67,13 +67,13 @@ func NewPlugin() Plugin {
 // response.
 func (p *Plugin) Before(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 	if p.MaxAge < 0 {
-		return w.ServerError(safehttp.StatusInternalServerError, "Internal Server Error")
+		return w.ServerError(safehttp.StatusInternalServerError)
 	}
 
 	if !p.BehindProxy && r.TLS == nil {
 		u, err := url.Parse(r.URL.String())
 		if err != nil {
-			return w.ServerError(safehttp.StatusInternalServerError, "Internal Server Error")
+			return w.ServerError(safehttp.StatusInternalServerError)
 		}
 		u.Scheme = "https"
 		return w.Redirect(r, u.String(), safehttp.StatusMovedPermanently)
@@ -90,10 +90,10 @@ func (p *Plugin) Before(w safehttp.ResponseWriter, r *safehttp.IncomingRequest)
 	}
 	h := w.Header()
 	if err := h.Set("Strict-Transport-Security", value.String()); err != nil {
-		// TODO(@mattiasgrenfeldt): Replace the response with an actual saferesponse somehow.
-		return w.ServerError(safehttp.StatusInternalServerError, "Internal Server Error")
+		return w.ServerError(safehttp.StatusInternalServerError)
+	}
+	if _, err := h.Claim("Strict-Transport-Security"); err != nil {
+		return w.ServerError(safehttp.StatusInternalServerError)
 	}
-	// TODO: Implement header claiming.
-	h.MarkImmutable("Strict-Transport-Security")
 	return safehttp.Result{}
 }
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 57857f8..a532d89 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -53,8 +53,25 @@ func (w *ResponseWriter) WriteTemplate(t Template, data interface{}) Result {
 	return Result{}
 }
 
+// ClientError TODO
+func (w *ResponseWriter) ClientError(code StatusCode) Result {
+	if code < 400 || code >= 500 {
+		// TODO(@mihalimara22): Replace this when we decide how to handle this case
+		panic("wrong method called")
+	}
+	http.Error(w.rw, http.StatusText(int(code)), int(code))
+	return Result{}
+}
+
 // ServerError TODO
-func (w *ResponseWriter) ServerError(code StatusCode, resp Response) Result {
+func (w *ResponseWriter) ServerError(code StatusCode) Result {
+	if code < 500 || code >= 600 {
+		// TODO(@mattiasgrenfeldt, @mihalimara22, @kele, @empijei): Decide how it should
+		// be communicated to the user of the framework that they've called the wrong
+		// method.
+		return Result{}
+	}
+	http.Error(w.rw, http.StatusText(int(code)), int(code))
 	return Result{}
 }
 


Test patch:
diff --git a/safehttp/header_test.go b/safehttp/header_test.go
index 7b25dc0..42fbc5e 100644
--- a/safehttp/header_test.go
+++ b/safehttp/header_test.go
@@ -67,12 +67,14 @@ func TestSetEmptySetCookie(t *testing.T) {
 	}
 }
 
-func TestSetImmutable(t *testing.T) {
+func TestSetClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
 	if err := h.Set("Foo-Key", "Pizza-Value"); err != nil {
 		t.Errorf(`h.Set("Foo-Key", "Pizza-Value") got: %v want: nil`, err)
 	}
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	if err := h.Set("Foo-Key", "Bar-Value"); err == nil {
 		t.Error(`h.Set("Foo-Key", "Bar-Value") got: nil want: error`)
 	}
@@ -81,9 +83,11 @@ func TestSetImmutable(t *testing.T) {
 	}
 }
 
-func TestSetEmptyImmutable(t *testing.T) {
+func TestSetEmptyClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	if err := h.Set("Foo-Key", "Bar-Value"); err == nil {
 		t.Error(`h.Set("Foo-Key", "Bar-Value") got: nil want: error`)
 	}
@@ -144,12 +148,14 @@ func TestAddEmptySetCookie(t *testing.T) {
 	}
 }
 
-func TestAddImmutable(t *testing.T) {
+func TestAddClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
 	if err := h.Add("Foo-Key", "Bar-Value"); err != nil {
 		t.Errorf(`h.Add("Foo-Key", "Bar-Value") got err: %v want: nil`, err)
 	}
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	if err := h.Add("Foo-Key", "Pizza-Value"); err == nil {
 		t.Error(`h.Add("Foo-Key", "Pizza-Value") got: nil want: error`)
 	}
@@ -158,9 +164,11 @@ func TestAddImmutable(t *testing.T) {
 	}
 }
 
-func TestAddEmptyImmutable(t *testing.T) {
+func TestAddEmptyClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	if err := h.Add("Foo-Key", "Pizza-Value"); err == nil {
 		t.Error(`h.Add("Foo-Key", "Pizza-Value") got: nil want: error`)
 	}
@@ -222,12 +230,14 @@ func TestDelEmptySetCookie(t *testing.T) {
 	}
 }
 
-func TestDelImmutable(t *testing.T) {
+func TestDelClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
 	if err := h.Set("Foo-Key", "Bar-Value"); err != nil {
 		t.Errorf(`h.Set("Foo-Key", "Bar-Value") got err: %v want: nil`, err)
 	}
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	if err := h.Del("Foo-Key"); err == nil {
 		t.Error(`h.Del("Foo-Key") got: nil want: error`)
 	}
@@ -236,9 +246,11 @@ func TestDelImmutable(t *testing.T) {
 	}
 }
 
-func TestDelEmptyImmutable(t *testing.T) {
+func TestDelEmptyClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	if err := h.Del("Foo-Key"); err == nil {
 		t.Error(`h.Del("Foo-Key") got: nil want: error`)
 	}
@@ -265,16 +277,18 @@ func TestSetCookieInvalidName(t *testing.T) {
 	}
 }
 
-// TestValuesModifyImmutable verifies that modifying the
+// TestValuesModifyClaimed verifies that modifying the
 // slice returned by Values() doesn't modify the underlying
 // slice. The test ensures that Values() returns a copy
 // of the underlying slice.
-func TestValuesModifyImmutable(t *testing.T) {
+func TestValuesModifyClaimed(t *testing.T) {
 	h := newHeader(http.Header{})
 	if err := h.Set("Foo-Key", "Bar-Value"); err != nil {
 		t.Errorf(`h.Set("Foo-Key", "Bar-Value") got err: %v want: nil`, err)
 	}
-	h.MarkImmutable("Foo-Key")
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
 	v := h.Values("Foo-Key")
 	if diff := cmp.Diff([]string{"Bar-Value"}, v); diff != "" {
 		t.Errorf("h.Values(\"Foo-Key\") mismatch (-want +got):\n%s", diff)
@@ -354,3 +368,49 @@ func TestValuesEmptyHeader(t *testing.T) {
 		t.Errorf("h.Values(\"Foo-Key\") mismatch (-want +got):\n%s", diff)
 	}
 }
+
+func TestClaim(t *testing.T) {
+	h := newHeader(http.Header{})
+	set, err := h.Claim("Foo-Key")
+	if err != nil {
+		t.Fatalf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
+	set([]string{"Bar-Value", "Pizza-Value"})
+	if diff := cmp.Diff([]string{"Bar-Value", "Pizza-Value"}, h.Values("Foo-Key")); diff != "" {
+		t.Errorf("h.Values(\"Foo-Key\") mismatch (-want +got):\n%s", diff)
+	}
+}
+
+// TestClaimCanonicalization verifies that names of headers
+// are canonicalized before being interpreted as header
+// names.
+// Note that the casing of the header name is different
+// when accessing and modifying the same header.
+func TestClaimCanonicalization(t *testing.T) {
+	h := newHeader(http.Header{})
+	set, err := h.Claim("fOO-kEY")
+	if err != nil {
+		t.Fatalf(`h.Claim("fOO-kEY") got: %v want: nil`, err)
+	}
+	set([]string{"Bar-Value", "Pizza-Value"})
+	if diff := cmp.Diff([]string{"Bar-Value", "Pizza-Value"}, h.Values("fOo-kEy")); diff != "" {
+		t.Errorf("h.Values(\"fOo-kEy\") mismatch (-want +got):\n%s", diff)
+	}
+}
+
+func TestClaimSetCookie(t *testing.T) {
+	h := newHeader(http.Header{})
+	if _, err := h.Claim("Set-Cookie"); err == nil {
+		t.Error(`h.Claim("Set-Cookie") got: nil want: error`)
+	}
+}
+
+func TestClaimClaimed(t *testing.T) {
+	h := newHeader(http.Header{})
+	if _, err := h.Claim("Foo-Key"); err != nil {
+		t.Errorf(`h.Claim("Foo-Key") got: %v want: nil`, err)
+	}
+	if _, err := h.Claim("Foo-Key"); err == nil {
+		t.Error(`h.Claim("Foo-Key") got: nil want: error`)
+	}
+}
diff --git a/safehttp/plugins/hsts/hsts_test.go b/safehttp/plugins/hsts/hsts_test.go
index dbc63b8..eb1e846 100644
--- a/safehttp/plugins/hsts/hsts_test.go
+++ b/safehttp/plugins/hsts/hsts_test.go
@@ -156,6 +156,17 @@ func TestHSTS(t *testing.T) {
 			},
 			wantBody: "",
 		},
+		{
+			name:       "Negative maxage",
+			plugin:     hsts.Plugin{MaxAge: -1 * time.Second},
+			req:        httptest.NewRequest("GET", "https://localhost/", nil),
+			wantStatus: 500,
+			wantHeaders: map[string][]string{
+				"Content-Type":           {"text/plain; charset=utf-8"},
+				"X-Content-Type-Options": {"nosniff"},
+			},
+			wantBody: "Internal Server Error\n",
+		},
 	}
 
 	for _, tt := range test {
@@ -185,107 +196,30 @@ func TestHSTS(t *testing.T) {
 func TestStrictTransportSecurityAlreadyImmutable(t *testing.T) {
 	p := hsts.NewPlugin()
 	m := safehttp.NewMachinery(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-		w.Header().MarkImmutable("Strict-Transport-Security")
+		w.Header().Claim("Strict-Transport-Security")
 		return p.Before(w, r)
 	}, &testDispatcher{})
 
 	req := httptest.NewRequest("GET", "https://localhost/", nil)
 
-	// TODO(@mattiasgrenfeldt): Change from current to desired behavior once
-	// safehttp.ResponseWriter.ServerError has changed signature and been
-	// implemented.
-	t.Run("Current Behavior", func(t *testing.T) {
-		b := &strings.Builder{}
-		rr := newResponseRecorder(b)
-
-		m.HandleRequest(rr, req)
-
-		if want := 200; rr.status != want {
-			t.Errorf("status code got: %v want: %v", rr.status, want)
-		}
-
-		wantHeaders := map[string][]string{}
-		if diff := cmp.Diff(wantHeaders, map[string][]string(rr.Header())); diff != "" {
-			t.Errorf("rr.Header() mismatch (-want +got):\n%s", diff)
-		}
+	b := &strings.Builder{}
+	rr := newResponseRecorder(b)
 
-		if got, want := b.String(), ""; got != want {
-			t.Errorf("response body got: %q want: %q", got, want)
-		}
-	})
+	m.HandleRequest(rr, req)
 
-	t.Run("Desired Behavior", func(t *testing.T) {
-		t.Skip()
-		b := &strings.Builder{}
-		rr := newResponseRecorder(b)
-
-		m.HandleRequest(rr, req)
-
-		if want := 500; rr.status != want {
-			t.Errorf("status code got: %v want: %v", rr.status, want)
-		}
-
-		wantHeaders := map[string][]string{}
-		if diff := cmp.Diff(wantHeaders, map[string][]string(rr.Header())); diff != "" {
-			t.Errorf("rr.Header() mismatch (-want +got):\n%s", diff)
-		}
-
-		if got, want := b.String(), "Internal Server Error"; got != want {
-			t.Errorf("response body got: %q want: %q", got, want)
-		}
-	})
-}
-
-func TestNegativeMaxAge(t *testing.T) {
-	p := hsts.Plugin{MaxAge: -1 * time.Second}
-	m := safehttp.NewMachinery(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-		w.Header().MarkImmutable("Strict-Transport-Security")
-		return p.Before(w, r)
-	}, &testDispatcher{})
-
-	req := httptest.NewRequest("GET", "https://localhost/", nil)
-
-	// TODO(@mattiasgrenfeldt): Change from current to desired behavior once
-	// safehttp.ResponseWriter.ServerError has changed signature and been
-	// implemented.
-	t.Run("Current Behavior", func(t *testing.T) {
-		b := &strings.Builder{}
-		rr := newResponseRecorder(b)
-
-		m.HandleRequest(rr, req)
-
-		if want := 200; rr.status != want {
-			t.Errorf("status code got: %v want: %v", rr.status, want)
-		}
-
-		wantHeaders := map[string][]string{}
-		if diff := cmp.Diff(wantHeaders, map[string][]string(rr.Header())); diff != "" {
-			t.Errorf("rr.Header() mismatch (-want +got):\n%s", diff)
-		}
-
-		if got, want := b.String(), ""; got != want {
-			t.Errorf("response body got: %q want: %q", got, want)
-		}
-	})
-
-	t.Run("Desired Behavior", func(t *testing.T) {
-		t.Skip()
-		b := &strings.Builder{}
-		rr := newResponseRecorder(b)
-
-		m.HandleRequest(rr, req)
-
-		if want := 500; rr.status != want {
-			t.Errorf("status code got: %v want: %v", rr.status, want)
-		}
+	if want := 500; rr.status != want {
+		t.Errorf("status code got: %v want: %v", rr.status, want)
+	}
 
-		wantHeaders := map[string][]string{}
-		if diff := cmp.Diff(wantHeaders, map[string][]string(rr.Header())); diff != "" {
-			t.Errorf("rr.Header() mismatch (-want +got):\n%s", diff)
-		}
+	wantHeaders := map[string][]string{
+		"Content-Type":           {"text/plain; charset=utf-8"},
+		"X-Content-Type-Options": {"nosniff"},
+	}
+	if diff := cmp.Diff(wantHeaders, map[string][]string(rr.Header())); diff != "" {
+		t.Errorf("rr.Header() mismatch (-want +got):\n%s", diff)
+	}
 
-		if got, want := b.String(), "Internal Server Error"; got != want {
-			t.Errorf("response body got: %q want: %q", got, want)
-		}
-	})
+	if got, want := b.String(), "Internal Server Error\n"; got != want {
+		t.Errorf("response body got: %q want: %q", got, want)
+	}
 }