Repository: google/go-safeweb (go)

Issue:
HSTS Plugin documentation
There should be documentation for the HSTS plugin in `/docs/`. There should also be some document listing all plugins that are provided by the framework.

Pull request description:
Added HSTS plugin documentation and structure for adding documentation about other plugins in the future.
Fixes #51 

Documentation for the HSTS plugin has been added. Also added a `plugins.md` file which can be used to describe plugins more in the future. `plugins.md` also contains a list of all the plugins we have currently implemented.

I also updated a godoc comment in the `hsts.Plugin` struct.

The links will point to the right place once this is merged.

Patch:
diff --git a/docs/high-level-design.md b/docs/high-level-design.md
index 48f1413..e1a3338 100644
--- a/docs/high-level-design.md
+++ b/docs/high-level-design.md
@@ -74,7 +74,7 @@ address, following issues (not an exhaustive list):
 - [**CSP (content security policy)**](https://csp.withgoogle.com/docs/index.html) -
   e.g. by automatically adding script nonces to HTML responses, adding relevant
   security headers
-- **Transport Security** - e.g. by enforcing HSTS support
+- **Transport Security** - e.g. by [enforcing HSTS support](plugins/hsts.md)
 - **IFraming** - e.g. by setting relevant HTTP headers to restrict framing or
   providing server-side support for origin selection
 - **Auth (access control)** - e.g. by providing infrastructure for plugging in
diff --git a/docs/plugins.md b/docs/plugins.md
new file mode 100644
index 0000000..16a77d1
--- /dev/null
+++ b/docs/plugins.md
@@ -0,0 +1,10 @@
+# Plugins
+
+The framework supports plugins. They provide a flexible way to address security 
+issues by setting non-overridable security headers or interrupting and
+responding to incoming requests before they reach the handler.
+
+## Provided plugins
+
+- [HSTS](plugins/hsts.md): Automatically redirects HTTP traffic to HTTPS and
+sets the `Strict-Transport-Security` header.
diff --git a/docs/plugins/hsts.md b/docs/plugins/hsts.md
new file mode 100644
index 0000000..b1e558f
--- /dev/null
+++ b/docs/plugins/hsts.md
@@ -0,0 +1,34 @@
+_Authors: grenfeldt@google.com_
+
+# HSTS Plugin
+
+HSTS<sup>1</sup> (HTTP Strict Transport Security) informs browsers that a
+website should only be accessed using HTTPS and not HTTP. This plugin enforces
+HSTS by redirecting all HTTP traffic to HTTPS and by setting the
+`Strict-Transport-Security` header on all outgoing HTTPS traffic.
+
+1) HSTS: [RFC](https://tools.ietf.org/html/rfc6797),
+[MDN](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Strict-Transport-Security),
+[Wikipedia](https://en.wikipedia.org/wiki/HTTP_Strict_Transport_Security)
+
+## Usage
+
+To construct the plugin with safe default settings, use: `hsts.NewPlugin()`.
+
+## Options
+
+**Option (Default value)**: Description.
+
+- **MaxAge (`2 years`)**: The amount of time that the browser should remember to
+only access this site using HTTPS. The default value is 2 years as recommended
+by https://hstspreload.org/.
+- **IncludeSubDomains (`enabled`)**: When this is enabled, the browser knows
+that all subdomains should also only be accessed via HTTPS.
+- **Preload (`disabled`)**: If enabled, the domain is eligible to be added as
+part of the browser HSTS preload list. This is disabled by default to prevent
+adding domains to the preload list unintentionally. For more info, see:
+https://hstspreload.org/.
+- **BehindProxy (`disabled`)**: If this server is placed behind a proxy that
+terminates HTTPS traffic then only HTTP traffic will be received by this server.
+This option disables redirection of HTTP to HTTPS and always sends the
+`Strict-Transport-Security` header, even over HTTP.
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/incoming_request.go b/safehttp/incoming_request.go
index 41ffba3..eeabb5e 100644
--- a/safehttp/incoming_request.go
+++ b/safehttp/incoming_request.go
@@ -32,7 +32,9 @@ type IncomingRequest struct {
 	URL       *url.URL
 }
 
-func newIncomingRequest(req *http.Request) *IncomingRequest {
+// NewIncomingRequest creates an IncomingRequest
+// from an http.Request.
+func NewIncomingRequest(req *http.Request) *IncomingRequest {
 	return &IncomingRequest{
 		req:    req,
 		Header: newHeader(req.Header),
diff --git a/safehttp/mux.go b/safehttp/mux.go
index 4dbea25..76ddc98 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -141,5 +141,5 @@ func (m methodHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	h.ServeHTTP(newResponseWriter(m.disp, w), newIncomingRequest(r))
+	h.ServeHTTP(NewResponseWriter(m.disp, w), NewIncomingRequest(r))
 }
diff --git a/safehttp/plugins/hsts/hsts.go b/safehttp/plugins/hsts/hsts.go
index 0382c52..e2bfc55 100644
--- a/safehttp/plugins/hsts/hsts.go
+++ b/safehttp/plugins/hsts/hsts.go
@@ -28,7 +28,7 @@ import (
 type Plugin struct {
 	// MaxAge is the duration that the browser should remember
 	// that a site is only to be accessed using HTTPS. MaxAge
-	// must be positive.
+	// must be positive. It will be rounded to seconds before use.
 	MaxAge time.Duration
 
 	// DisableIncludeSubDomains disables the includeSubDomains directive.
@@ -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/prototype.go b/safehttp/prototype.go
index fc22c19..5729f94 100644
--- a/safehttp/prototype.go
+++ b/safehttp/prototype.go
@@ -31,7 +31,7 @@ func NewMachinery(h HandleFunc, d Dispatcher) *Machinery {
 
 // HandleRequest TODO
 func (m *Machinery) HandleRequest(w http.ResponseWriter, req *http.Request) {
-	rw := newResponseWriter(m.d, w)
-	ir := newIncomingRequest(req)
+	rw := NewResponseWriter(m.d, w)
+	ir := NewIncomingRequest(req)
 	m.h(rw, ir)
 }
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 57857f8..1467141 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -29,7 +29,9 @@ type ResponseWriter struct {
 	header Header
 }
 
-func newResponseWriter(d Dispatcher, rw http.ResponseWriter) ResponseWriter {
+// NewResponseWriter creates a ResponseWriter from
+// a safehttp.Dispatcher and a http.ResponseWriter.
+func NewResponseWriter(d Dispatcher, rw http.ResponseWriter) ResponseWriter {
 	header := newHeader(rw.Header())
 	return ResponseWriter{d: d, rw: rw, header: header}
 }
@@ -53,8 +55,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..0c62afd 100644
--- a/safehttp/plugins/hsts/hsts_test.go
+++ b/safehttp/plugins/hsts/hsts_test.go
@@ -15,68 +15,20 @@
 package hsts_test
 
 import (
-	"html/template"
-	"io"
-	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/hsts"
-	"github.com/google/safehtml"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
 )
 
-type testDispatcher struct{}
-
-func (testDispatcher) Write(rw http.ResponseWriter, resp safehttp.Response) error {
-	switch x := resp.(type) {
-	case safehtml.HTML:
-		_, err := rw.Write([]byte(x.String()))
-		return err
-	default:
-		panic("not a safe response type")
-	}
-}
-
-func (testDispatcher) ExecuteTemplate(rw http.ResponseWriter, t safehttp.Template, data interface{}) error {
-	switch x := t.(type) {
-	case *template.Template:
-		return x.Execute(rw, data)
-	default:
-		panic("not a safe response type")
-	}
-}
-
-type responseRecorder struct {
-	header http.Header
-	writer io.Writer
-	status int
-}
-
-func newResponseRecorder(w io.Writer) *responseRecorder {
-	return &responseRecorder{header: http.Header{}, writer: w, status: http.StatusOK}
-}
-
-func (r *responseRecorder) Header() http.Header {
-	return r.header
-}
-
-func (r *responseRecorder) WriteHeader(statusCode int) {
-	r.status = statusCode
-}
-
-func (r *responseRecorder) Write(data []byte) (int, error) {
-	return r.writer.Write(data)
-}
-
 func TestHSTS(t *testing.T) {
 	var test = []struct {
 		name        string
 		plugin      hsts.Plugin
-		req         *http.Request
+		req         *safehttp.IncomingRequest
 		wantStatus  int
 		wantHeaders map[string][]string
 		wantBody    string
@@ -84,7 +36,7 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "HTTPS",
 			plugin:     hsts.NewPlugin(),
-			req:        httptest.NewRequest("GET", "https://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "https://localhost/", nil),
 			wantStatus: 200,
 			wantHeaders: map[string][]string{
 				"Strict-Transport-Security": {"max-age=63072000; includeSubDomains"}, // 63072000 seconds is two years
@@ -94,7 +46,7 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "HTTP",
 			plugin:     hsts.NewPlugin(),
-			req:        httptest.NewRequest("GET", "http://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "http://localhost/", nil),
 			wantStatus: 301,
 			wantHeaders: map[string][]string{
 				"Content-Type": {"text/html; charset=utf-8"},
@@ -105,7 +57,7 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "HTTP behind proxy",
 			plugin:     hsts.Plugin{BehindProxy: true},
-			req:        httptest.NewRequest("GET", "http://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "http://localhost/", nil),
 			wantStatus: 200,
 			wantHeaders: map[string][]string{
 				// max-age=0 tells the browser to expire the HSTS protection.
@@ -116,7 +68,7 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "Preload",
 			plugin:     hsts.Plugin{Preload: true, DisableIncludeSubDomains: true},
-			req:        httptest.NewRequest("GET", "https://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "https://localhost/", nil),
 			wantStatus: 200,
 			wantHeaders: map[string][]string{
 				// max-age=0 tells the browser to expire the HSTS protection.
@@ -127,7 +79,7 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "Preload and IncludeSubDomains",
 			plugin:     hsts.Plugin{Preload: true},
-			req:        httptest.NewRequest("GET", "https://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "https://localhost/", nil),
 			wantStatus: 200,
 			wantHeaders: map[string][]string{
 				// max-age=0 tells the browser to expire the HSTS protection.
@@ -138,7 +90,7 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "No preload and no includeSubDomains",
 			plugin:     hsts.Plugin{DisableIncludeSubDomains: true},
-			req:        httptest.NewRequest("GET", "https://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "https://localhost/", nil),
 			wantStatus: 200,
 			wantHeaders: map[string][]string{
 				// max-age=0 tells the browser to expire the HSTS protection.
@@ -149,33 +101,41 @@ func TestHSTS(t *testing.T) {
 		{
 			name:       "Custom maxage",
 			plugin:     hsts.Plugin{MaxAge: 3600 * time.Second},
-			req:        httptest.NewRequest("GET", "https://localhost/", nil),
+			req:        safehttptest.NewRequest("GET", "https://localhost/", nil),
 			wantStatus: 200,
 			wantHeaders: map[string][]string{
 				"Strict-Transport-Security": {"max-age=3600; includeSubDomains"}, // 3600 seconds is 1 hour
 			},
 			wantBody: "",
 		},
+		{
+			name:       "Negative maxage",
+			plugin:     hsts.Plugin{MaxAge: -1 * time.Second},
+			req:        safehttptest.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 {
 		t.Run(tt.name, func(t *testing.T) {
-			m := safehttp.NewMachinery(tt.plugin.Before, &testDispatcher{})
-
-			b := &strings.Builder{}
-			rr := newResponseRecorder(b)
+			rr := safehttptest.NewResponseRecorder()
 
-			m.HandleRequest(rr, tt.req)
+			tt.plugin.Before(rr.ResponseWriter, tt.req)
 
-			if rr.status != tt.wantStatus {
-				t.Errorf("status code got: %v want: %v", rr.status, tt.wantStatus)
+			if rr.Status() != tt.wantStatus {
+				t.Errorf("status code got: %v want: %v", rr.Status(), tt.wantStatus)
 			}
 
 			if diff := cmp.Diff(tt.wantHeaders, map[string][]string(rr.Header())); diff != "" {
 				t.Errorf("rr.Header() mismatch (-want +got):\n%s", diff)
 			}
 
-			if got := b.String(); got != tt.wantBody {
+			if got := rr.Body(); got != tt.wantBody {
 				t.Errorf("response body got: %q want: %q", got, tt.wantBody)
 			}
 		})
@@ -184,108 +144,27 @@ 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")
-		return p.Before(w, r)
-	}, &testDispatcher{})
 
-	req := httptest.NewRequest("GET", "https://localhost/", nil)
+	req := safehttptest.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)
+	rr := safehttptest.NewResponseRecorder()
+	rr.ResponseWriter.Header().Claim("Strict-Transport-Security")
 
-		m.HandleRequest(rr, req)
+	p.Before(rr.ResponseWriter, 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)
-		}
-
-		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 := rr.Body(), "Internal Server Error\n"; got != want {
+		t.Errorf("response body got: %q want: %q", got, want)
+	}
 }
diff --git a/safehttp/safehttptest/recorder.go b/safehttp/safehttptest/recorder.go
new file mode 100644
index 0000000..5222848
--- /dev/null
+++ b/safehttp/safehttptest/recorder.go
@@ -0,0 +1,121 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package safehttptest
+
+import (
+	"html/template"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/safehtml"
+)
+
+type testDispatcher struct{}
+
+func (testDispatcher) Write(rw http.ResponseWriter, resp safehttp.Response) error {
+	switch x := resp.(type) {
+	case safehtml.HTML:
+		_, err := rw.Write([]byte(x.String()))
+		return err
+	default:
+		panic("not a safe response type")
+	}
+}
+
+func (testDispatcher) ExecuteTemplate(rw http.ResponseWriter, t safehttp.Template, data interface{}) error {
+	switch x := t.(type) {
+	case *template.Template:
+		return x.Execute(rw, data)
+	default:
+		panic("not a safe response type")
+	}
+}
+
+// ResponseRecorder encapsulates a safehttp.ResponseWriter that records
+// mutations for later inspection in tests. The safehttp.ResponseWriter
+// should be passed as part of the handler function in tests.
+type ResponseRecorder struct {
+	safehttp.ResponseWriter
+	rw *responseWriter
+	b  *strings.Builder
+}
+
+// NewResponseRecorder creates a ResponseRecorder from the default testDispatcher.
+func NewResponseRecorder() *ResponseRecorder {
+	var b strings.Builder
+	rw := newResponseWriter(&b)
+	return &ResponseRecorder{
+		rw:             rw,
+		b:              &b,
+		ResponseWriter: safehttp.NewResponseWriter(testDispatcher{}, rw),
+	}
+}
+
+// NewResponseRecorderFromDispatcher creates a ResponseRecorder from a
+// provided safehttp.Dispatcher.
+func NewResponseRecorderFromDispatcher(d safehttp.Dispatcher) *ResponseRecorder {
+	var b strings.Builder
+	rw := newResponseWriter(&b)
+	return &ResponseRecorder{
+		rw:             rw,
+		b:              &b,
+		ResponseWriter: safehttp.NewResponseWriter(d, rw),
+	}
+}
+
+// Header returns the recorded response headers.
+func (r *ResponseRecorder) Header() http.Header {
+	return r.rw.Header()
+}
+
+// Status returns the recorded response status code.
+func (r *ResponseRecorder) Status() int {
+	return r.rw.status
+}
+
+// Body returns the recorded response body.
+func (r *ResponseRecorder) Body() string {
+	return r.b.String()
+}
+
+// responseWriter is an implementation of the http.ResponseWriter interface used
+// for constructing an HTTP response.
+type responseWriter struct {
+	header http.Header
+	writer io.Writer
+	status int
+}
+
+func newResponseWriter(w io.Writer) *responseWriter {
+	return &responseWriter{
+		header: http.Header{},
+		writer: w,
+		status: http.StatusOK,
+	}
+}
+
+func (r *responseWriter) Header() http.Header {
+	return r.header
+}
+
+func (r *responseWriter) WriteHeader(statusCode int) {
+	r.status = statusCode
+}
+
+func (r *responseWriter) Write(data []byte) (int, error) {
+	return r.writer.Write(data)
+}
diff --git a/safehttp/safehttptest/request.go b/safehttp/safehttptest/request.go
new file mode 100644
index 0000000..b2e494a
--- /dev/null
+++ b/safehttp/safehttptest/request.go
@@ -0,0 +1,52 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package safehttptest
+
+import (
+	"io"
+	"net/http/httptest"
+
+	"github.com/google/go-safeweb/safehttp"
+)
+
+// NewRequest returns a new incoming server Request,
+// suitable for passing to an http.Handler for testing.
+//
+// The target is the RFC 7230 "request-target": it may
+// be either a path or an absolute URL. If target is an
+// absolute URL, the host name from the URL is used.
+// Otherwise, "example.com" is used.
+//
+// The TLS field is set to a non-nil dummy value if
+// target has scheme "https".
+//
+// The Request.Proto is always HTTP/1.1.
+//
+// An empty method means "GET".
+//
+// The provided body may be nil. If the body is of type
+// *bytes.Reader, *strings.Reader, or *bytes.Buffer, the
+// Request.ContentLength is set.
+//
+// NewRequest panics on error for ease of use in testing,
+// where a panic is acceptable.
+//
+// To generate a client HTTP request instead of a server
+// request, see the NewRequest function in the net/http
+// package.
+func NewRequest(method, target string, body io.Reader) *safehttp.IncomingRequest {
+	req := httptest.NewRequest(method, target, body)
+	return safehttp.NewIncomingRequest(req)
+}