Repository: google/go-safeweb (go)

Issue:
Support header claiming as a means of expressing ownership of (security-sensitive) headers
Plugins might need to mark a header as immutable and then change it at a later stage. A method called `Claim` will therefore be implemented that sets the header as immutable and returns something that can later change the header. When `Claim` is implemented there will be no need for `MarkImmutable` (for now) and therefore it will be removed.

Pull request description:
Header Claiming
Fixes #44

Plugins might need to mark a header as immutable and then change it at a later stage. A method called Claim will therefore be implemented that sets the header as immutable and returns something that can later change the header. When Claim is implemented there will be no need for MarkImmutable (for now) and therefore it will be removed.

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..b1cdd8e 100644
--- a/safehttp/plugins/hsts/hsts.go
+++ b/safehttp/plugins/hsts/hsts.go
@@ -93,7 +93,8 @@ func (p *Plugin) Before(w safehttp.ResponseWriter, r *safehttp.IncomingRequest)
 		// TODO(@mattiasgrenfeldt): Replace the response with an actual saferesponse somehow.
 		return w.ServerError(safehttp.StatusInternalServerError, "Internal Server Error")
 	}
-	// TODO: Implement header claiming.
-	h.MarkImmutable("Strict-Transport-Security")
+	if _, err := h.Claim("Strict-Transport-Security"); err != nil {
+		return w.ServerError(safehttp.StatusInternalServerError, "Internal Server Error")
+	}
 	return safehttp.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..122fdf1 100644
--- a/safehttp/plugins/hsts/hsts_test.go
+++ b/safehttp/plugins/hsts/hsts_test.go
@@ -185,7 +185,7 @@ 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{})
 
@@ -238,10 +238,7 @@ func TestStrictTransportSecurityAlreadyImmutable(t *testing.T) {
 
 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{})
+	m := safehttp.NewMachinery(p.Before, &testDispatcher{})
 
 	req := httptest.NewRequest("GET", "https://localhost/", nil)