Repository: google/go-safeweb
Pull request #220 (2020-11-13 10:31:14)

Consolidated request handling and the ResponseWriter implementation to one object
The state needed to end-to-end handle a single request has been
scattered among many objects with unclear responsibilities. This is an
attempt to simplify the implementation in order to reason about
recursive calls (e.g. done by interceptors writing) and error handling.

Summary of changes:
- the `responseWriter` has been removed and most of it's responsibilities
have been moved to a new struct - `task` (the name is not ideal, TODO);
task implements `safehttp.ResponseWriter`
- the `safehttp/mux.go:handler` type has been removed in favor of
`HandlerConfig`, task and a `processRequest` function; `HandlerConfig` is just
data (no functions) and `processRequest` is responsible for bootstrapping
the task
- the `task` (a replacement for the `responseWriter`) does not hold a reference to
the `Dispatcher` anymore; previously both the `handler` (now removed) and the
`responseWriter` needed a reference to the `Dispatcher`


Patch:
diff --git a/safehttp/flight.go b/safehttp/flight.go
new file mode 100644
index 0000000..99090d8
--- /dev/null
+++ b/safehttp/flight.go
@@ -0,0 +1,264 @@
+// 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 safehttp
+
+import (
+	"fmt"
+	"net/http"
+)
+
+// A single request "flight".
+type flight struct {
+	rw  http.ResponseWriter
+	req *IncomingRequest
+
+	cfg HandlerConfig
+
+	code   StatusCode
+	header Header
+
+	// TODO(kele): we need to distinguish between calling Write and actually
+	// writing to the net/http.ResponseWriter.
+	written      bool
+	writtenError bool
+}
+
+// DeprecatedNewResponseWriter creates a ResponseWriter implementation that has
+// been historically used for testing. DO NOT USE this function, it's going
+// to be removed soon.
+//
+// The main problem with this is the implementation of the returned
+// ResponseWriter is incomplete, i.e. lacks the Handler and Interceptors fields
+// of the HandlerConfig needed by the flight. For the existing tests that use it,
+// it's fine, but they should be migrated.
+//
+// TODO(kele): remove this once we have a better option to provide a
+// ResponseWriter implementation for interceptor testing.
+func DeprecatedNewResponseWriter(rw http.ResponseWriter, dispatcher Dispatcher) ResponseWriter {
+	return &flight{
+		cfg:    HandlerConfig{Dispatcher: dispatcher},
+		rw:     rw,
+		header: newHeader(rw.Header()),
+	}
+}
+
+// HandlerConfig is the safe HTTP handler configuration, including the
+// dispatcher and interceptors.
+type HandlerConfig struct {
+	Handler      Handler
+	Dispatcher   Dispatcher
+	Interceptors []ConfiguredInterceptor
+}
+
+func processRequest(cfg HandlerConfig, rw http.ResponseWriter, req *http.Request) {
+	f := &flight{
+		cfg:    cfg,
+		rw:     rw,
+		header: newHeader(rw.Header()),
+		req:    NewIncomingRequest(req),
+	}
+
+	// The net/http package recovers handler panics, but we cannot rely on that
+	// behavior here. The reason is, we might need to run some interceptor
+	// stages interceptors before we respond with a 500 Internal Server Error.
+	// Therefore we're calling WriteError.
+	defer func() {
+		if r := recover(); r != nil {
+			f.WriteError(StatusInternalServerError)
+		}
+	}()
+
+	for _, it := range f.cfg.Interceptors {
+		it.Before(f, f.req)
+		if f.written {
+			return
+		}
+	}
+
+	f.cfg.Handler.ServeHTTP(f, f.req)
+	if !f.written {
+		f.NoContent()
+	}
+}
+
+// Write dispatches the response to the Dispatcher. This will be written to the
+// underlying http.ResponseWriter if the Dispatcher decides it's safe to do so.
+//
+// TODO: replace panics with proper error handling when writing the response fails.
+func (f *flight) Write(resp Response) Result {
+	if f.written {
+		panic("ResponseWriter was already written to")
+	}
+	f.written = true
+	f.commitPhase(resp)
+
+	ct, err := f.cfg.Dispatcher.ContentType(resp)
+	if err != nil {
+		panic(err)
+	}
+	f.rw.Header().Set("Content-Type", ct)
+
+	if f.code == 0 {
+		f.code = StatusOK
+	}
+	f.rw.WriteHeader(int(f.code))
+
+	if err := f.cfg.Dispatcher.Write(f.rw, resp); err != nil {
+		panic(err)
+	}
+	return Result{}
+}
+
+// NoContent responds with a 204 No Content response.
+//
+// If the ResponseWriter has already been written to, then this method will panic.
+func (f *flight) NoContent() Result {
+	// TODO: should NoContent call Write under the hood? Should the dispatcher
+	// handle this too?
+	if f.written {
+		panic("ResponseWriter was already written to")
+	}
+	f.written = true
+	f.commitPhase(NoContentResponse{})
+	f.rw.WriteHeader(int(StatusNoContent))
+	return Result{}
+}
+
+// WriteError writes an error response (400-599) according to the provided
+// status code.
+//
+// If the ResponseWriter has already been written to, then this method will panic.
+func (f *flight) WriteError(code StatusCode) Result {
+	// TODO: accept custom error responses that need to go through the dispatcher.
+	if f.writtenError {
+		panic("ResponseWriter.WriteError called twice")
+	}
+	f.written = true
+	f.writtenError = true
+	// TODO: we cannot really write if the Dispatcher already started writing
+	// but panicked, resulting in a WriteError call.
+	resp := &ErrorResponse{Code: code}
+	f.errorPhase(resp)
+	http.Error(f.rw, http.StatusText(int(resp.Code)), int(resp.Code))
+	return Result{}
+}
+
+// Redirect responds with a redirect to the given url, using code as the status code.
+//
+// If the ResponseWriter has already been written to, then this method will panic.
+func (f *flight) Redirect(r *IncomingRequest, url string, code StatusCode) Result {
+	if code < 300 || code >= 400 {
+		panic(fmt.Sprintf("wrong method called: redirect with status %d", code))
+	}
+	if f.written {
+		panic("ResponseWriter was already written to")
+	}
+	f.written = true
+	http.Redirect(f.rw, r.req, url, int(code))
+	return Result{}
+}
+
+// Header returns the collection of headers that will be set on the response.
+// Headers must be set before writing a response.
+func (f *flight) Header() Header {
+	return f.header
+}
+
+// SetCookie adds a Set-Cookie header to the provided ResponseWriter's headers.
+// The provided cookie must have a valid Name, otherwise an error will be
+// returned.
+//
+// TODO: should this be named AddCookie?
+func (f *flight) SetCookie(c *Cookie) error {
+	return f.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 (f *flight) SetCode(code StatusCode) {
+	if f.written {
+		return
+	}
+	if code < 100 || code >= 600 {
+		panic("invalid status code")
+	}
+	f.code = code
+}
+
+// commitPhase calls the Commit phases of all the interceptors. This stage will
+// run before a response is written to the ResponseWriter. If a response is
+// written to the ResponseWriter in a Commit phase then the Commit phases of the
+// remaining interceptors won'f execute.
+func (f *flight) commitPhase(resp Response) {
+	for i := len(f.cfg.Interceptors) - 1; i >= 0; i-- {
+		f.cfg.Interceptors[i].Commit(f, f.req, resp)
+	}
+}
+
+// errorPhase calls the OnError phases of all the interceptors associated with
+// a handler. This stage runs before an error response is written to the
+// ResponseWriter.
+func (f *flight) errorPhase(resp Response) {
+	for i := len(f.cfg.Interceptors) - 1; i >= 0; i-- {
+		f.cfg.Interceptors[i].OnError(f, f.req, resp)
+	}
+}
+
+// Dispatcher is responsible for writing a response received from the
+// ResponseWriter to the underlying http.ResponseWriter.
+//
+// The implementation of a custom Dispatcher should be thoroughly reviewed by
+// the security team to avoid introducing vulnerabilities.
+type Dispatcher interface {
+	// Content-Type returns the Content-Type of the provided response if it is
+	// of a safe type, supported by the Dispatcher, and should return an error
+	// otherwise.
+	//
+	// Sending a response to the http.ResponseWriter without properly setting
+	// CT is error-prone and could introduce vulnerabilities. Therefore, this
+	// method should be used to set the Content-Type header before calling
+	// Dispatcher.Write. Writing should not proceed if ContentType returns an
+	// error.
+	ContentType(resp Response) (string, error)
+
+	// Write writes a Response to the underlying http.ResponseWriter.
+	//
+	// It should return an error if the writing operation fails or if the
+	// provided Response should not be written to the http.ResponseWriter
+	// because it's unsafe.
+	Write(rw http.ResponseWriter, resp Response) error
+}
+
+// Result is the result of writing an HTTP response.
+//
+// Use ResponseWriter methods to obtain it.
+type Result struct{}
+
+// NotWritten returns a Result which indicates that nothing has been written yet. It
+// can be used in all functions that return a Result, such as in the ServeHTTP method
+// of a Handler or in the Before method of an Interceptor. When returned, NotWritten
+// indicates that the writing of the response should take place later. When this
+// is returned by the Before method in Interceptors the next Interceptor in line
+// is run. When this is returned by a Handler, a 204 No Content response is written.
+func NotWritten() Result {
+	return Result{}
+}
diff --git a/safehttp/incoming_request.go b/safehttp/incoming_request.go
index 7f49a13..791f931 100644
--- a/safehttp/incoming_request.go
+++ b/safehttp/incoming_request.go
@@ -44,6 +44,9 @@ type IncomingRequest struct {
 // NewIncomingRequest creates an IncomingRequest
 // from the underlying http.Request.
 func NewIncomingRequest(req *http.Request) *IncomingRequest {
+	if req == nil {
+		return nil
+	}
 	return &IncomingRequest{
 		req:    req,
 		Header: newHeader(req.Header),
diff --git a/safehttp/mux.go b/safehttp/mux.go
index f8570be..6dc72ed 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -70,17 +70,15 @@ type ServeMux struct {
 	mux *http.ServeMux
 }
 
-func registerHandlers(mux *http.ServeMux, handlers map[string]map[string]handler) {
+func registerHandlers(mux *http.ServeMux, handlers map[string]map[string]HandlerConfig) {
 	for pattern, handlersPerMethod := range handlers {
-		pattern := pattern
-		handlersPerMethod := handlersPerMethod
 		mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
-			h, ok := handlersPerMethod[r.Method]
+			cfg, ok := handlersPerMethod[r.Method]
 			if !ok {
 				http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
 				return
 			}
-			h.ServeHTTP(w, r)
+			processRequest(cfg, w, r)
 		})
 	}
 }
@@ -153,20 +151,22 @@ func (s *ServeMuxConfig) Mux() *ServeMux {
 	if dispatcher == nil {
 		dispatcher = DefaultDispatcher{}
 	}
-	// pattern -> method -> handler
-	handlers := map[string]map[string]handler{}
+	// pattern -> method -> HandlerConfig
+	handlers := map[string]map[string]HandlerConfig{}
 	for _, hr := range s.handlers {
 		if handlers[hr.pattern] == nil {
-			handlers[hr.pattern] = map[string]handler{}
+			handlers[hr.pattern] = map[string]HandlerConfig{}
 		}
 		if _, ok := handlers[hr.pattern][hr.method]; ok {
+			// TODO: this should be done in ServeMuxConfig.Handle, not here.
 			panic(fmt.Sprintf("double registration of (pattern = %q, method = %q)", hr.pattern, hr.method))
 		}
-		handlers[hr.pattern][hr.method] = handler{
-			handler:   hr.handler,
-			interceps: configureInterceptors(s.interceptors, hr.cfgs),
-			disp:      dispatcher,
-		}
+		handlers[hr.pattern][hr.method] =
+			HandlerConfig{
+				Dispatcher:   dispatcher,
+				Handler:      hr.handler,
+				Interceptors: configureInterceptors(s.interceptors, hr.cfgs),
+			}
 	}
 	m := http.NewServeMux()
 	registerHandlers(m, handlers)
@@ -179,6 +179,8 @@ func configureInterceptors(interceptors []Interceptor, cfgs []InterceptorConfig)
 		var cfg InterceptorConfig
 		for _, c := range cfgs {
 			if c.Match(it) {
+				// TODO: there should be a validation check that there is at
+				// most one config per interceptor.
 				cfg = c
 				break
 			}
@@ -201,74 +203,3 @@ func (s *ServeMuxConfig) Clone() *ServeMuxConfig {
 	copy(c.interceptors, s.interceptors)
 	return c
 }
-
-// handler encapsulates a handler and its corresponding
-// interceptors.
-type handler struct {
-	handler   Handler
-	interceps []ConfiguredInterceptor
-	disp      Dispatcher
-}
-
-// ServeHTTP calls the Before method of all the interceptors and then calls the
-// underlying handler. Any panics that occur during the Before phase, during handler
-// execution or during the Commit phase are recovered here and a 500 Internal Server
-// Error response is sent.
-func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	ir := NewIncomingRequest(r)
-	rw := newResponseWriter(h.disp, w, ir)
-	rw.handler = h
-
-	// The `net/http` package recovers handler panics, but we cannot rely on that behavior here.
-	// The reason is, we might need to run After/Commit stages of the interceptors before we
-	// respond with a 500 Internal Server Error.
-	defer func() {
-		if r := recover(); r != nil {
-			rw.WriteError(StatusInternalServerError)
-		}
-	}()
-
-	for _, it := range h.interceps {
-		it.Before(rw, ir)
-		if rw.written {
-			return
-		}
-	}
-
-	h.handler.ServeHTTP(rw, ir)
-	if !rw.written {
-		rw.NoContent()
-	}
-}
-
-// commitPhase calls the Commit phases of all the interceptors. This stage will
-// run before a response is written to the ResponseWriter. If a response is
-// written to the ResponseWriter in a Commit phase then the Commit phases of the
-// remaining interceptors won't execute.
-//
-// TODO: BIG WARNING, if ResponseWriter.Write and ResponseWriter.WriteTemplate
-// are called in Commit then this will recurse. CommitResponseWriter was an
-// attempt to prevent this by not giving access to Write and WriteTemplate in
-// the Commit phase.
-func (h handler) commitPhase(w *responseWriter, resp Response) {
-	for i := len(h.interceps) - 1; i >= 0; i-- {
-		it := h.interceps[i]
-		it.Commit(w, w.req, resp)
-		if w.written {
-			return
-		}
-	}
-}
-
-// errorPhase calls the OnError phases of all the interceptors associated with
-// a handler. This stage runs before an error response is written to the
-// ResponseWriter.
-//
-// TODO: BIG WARNING, if an interceptor attempts to write to the ResponseWriter
-// in the OnError phase will result in an irrecoverable error.
-func (h handler) errorPhase(w *responseWriter, resp Response) {
-	for i := len(h.interceps) - 1; i >= 0; i-- {
-		it := h.interceps[i]
-		it.OnError(w, w.req, resp)
-	}
-}
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 529ebc4..41e32d6 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -14,10 +14,6 @@
 
 package safehttp
 
-import (
-	"net/http"
-)
-
 // ResponseWriter is used to construct an HTTP response. When a Response is
 // passed to the ResponseWriter, it will invoke the Dispatcher with the
 // Response. An attempt to write to the ResponseWriter twice will
@@ -61,190 +57,3 @@ type ResponseWriter interface {
 	// code set by the latter will be the actual response status.
 	SetCode(code StatusCode)
 }
-
-type responseWriter struct {
-	d  Dispatcher
-	rw http.ResponseWriter
-
-	// 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
-	written bool
-}
-
-// NewResponseWriter creates a ResponseWriter from a safehttp.Dispatcher, an
-// http.ResponseWriter and a reference to the current IncomingRequest being served.
-// The IncomingRequest will only be used by the commit phase, which only runs when
-// the ServeMux is used, and can be passed as nil in tests.
-// TODO: remove the IncomingRequest parameter.
-func NewResponseWriter(d Dispatcher, rw http.ResponseWriter, req *IncomingRequest) ResponseWriter {
-	return newResponseWriter(d, rw, req)
-}
-
-func newResponseWriter(d Dispatcher, rw http.ResponseWriter, req *IncomingRequest) *responseWriter {
-	header := newHeader(rw.Header())
-	return &responseWriter{
-		d:      d,
-		rw:     rw,
-		header: header,
-		req:    req,
-	}
-}
-
-// Result is the result of writing an HTTP response.
-//
-// Use ResponseWriter methods to obtain it.
-type Result struct{}
-
-// NotWritten returns a Result which indicates that nothing has been written yet. It
-// can be used in all functions that return a Result, such as in the ServeHTTP method
-// of a Handler or in the Before method of an Interceptor. When returned, NotWritten
-// indicates that the writing of the response should take place later. When this
-// is returned by the Before method in Interceptors the next Interceptor in line
-// is run. When this is returned by a Handler, a 204 No Content response is written.
-func NotWritten() Result {
-	return Result{}
-}
-
-// Write dispatches the response to the Dispatcher. This will be written to the
-// underlying http.ResponseWriter if the Dispatcher decides it's safe to do so.
-//
-// TODO: replace panics with proper error handling when writing the response fails.
-func (w *responseWriter) Write(resp Response) Result {
-	if w.written {
-		panic("ResponseWriter was already written to")
-	}
-	w.handler.commitPhase(w, resp)
-	if w.written {
-		return Result{}
-	}
-
-	w.markWritten()
-
-	ct, err := w.d.ContentType(resp)
-	if err != nil {
-		panic(err)
-	}
-	w.rw.Header().Set("Content-Type", ct)
-
-	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)
-	}
-	return Result{}
-}
-
-// NoContent responds with a 204 No Content response.
-//
-// If the ResponseWriter has already been written to, then this method will panic.
-func (w *responseWriter) NoContent() Result {
-	if w.written {
-		panic("ResponseWriter was already written to")
-	}
-	w.handler.commitPhase(w, NoContentResponse{})
-	if w.written {
-		return Result{}
-	}
-	w.markWritten()
-	w.rw.WriteHeader(int(StatusNoContent))
-	return Result{}
-}
-
-// WriteError writes an error response (400-599) according to the provided
-// status code.
-//
-// If the ResponseWriter has already been written to, then this method will panic.
-func (w *responseWriter) WriteError(code StatusCode) Result {
-	w.markWritten()
-	resp := &ErrorResponse{Code: code}
-	w.handler.errorPhase(w, resp)
-	http.Error(w.rw, http.StatusText(int(resp.Code)), int(resp.Code))
-	return Result{}
-}
-
-// Redirect responds with a redirect to the given url, using code as the status code.
-//
-// If the ResponseWriter has already been written to, then this method will panic.
-func (w *responseWriter) Redirect(r *IncomingRequest, url string, code StatusCode) Result {
-	if code < 300 || code >= 400 {
-		panic("wrong method called")
-	}
-	w.markWritten()
-	http.Redirect(w.rw, r.req, url, int(code))
-	return Result{}
-}
-
-// markWritten ensures that the ResponseWriter is only written to once by panicking
-// if it is written more than once.
-func (w *responseWriter) markWritten() {
-	if w.written {
-		panic("ResponseWriter was already written to")
-	}
-	w.written = true
-}
-
-// Header returns the collection of headers that will be set
-// on the response. Headers must be set before writing a
-// response (e.g. Write, WriteTemplate).
-func (w responseWriter) Header() Header {
-	return w.header
-}
-
-// SetCookie adds a Set-Cookie header to the provided ResponseWriter's headers.
-// The provided cookie must have a valid Name, otherwise an error will be
-// returned.
-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.
-//
-// The implementation of a custom Dispatcher should be thoroughly reviewed by
-// the security team to avoid introducing vulnerabilities.
-type Dispatcher interface {
-	// Content-Type returns the Content-Type of the provided response if it is
-	// of a safe type, supported by the Dispatcher, and should return an error
-	// otherwise.
-	//
-	// Sending a response to the http.ResponseWriter without properly setting
-	// CT is error-prone and could introduce vulnerabilities. Therefore, this
-	// method should be used to set the Content-Type header before calling
-	// Dispatcher.Write. Writing should not proceed if ContentType returns an
-	// error.
-	ContentType(resp Response) (string, error)
-
-	// Write writes a Response to the underlying http.ResponseWriter.
-	//
-	// It should return an error if the writing operation fails or if the
-	// provided Response should not be written to the http.ResponseWriter
-	// because it's unsafe.
-	Write(rw http.ResponseWriter, resp Response) error
-}


Test patch:
diff --git a/safehttp/safehttptest/recorder.go b/safehttp/safehttptest/recorder.go
index 8cf2fb1..39d6051 100644
--- a/safehttp/safehttptest/recorder.go
+++ b/safehttp/safehttptest/recorder.go
@@ -31,26 +31,25 @@ type ResponseRecorder struct {
 	b  *strings.Builder
 }
 
-// NewResponseRecorder creates a ResponseRecorder from the safehttp.DefaultDispatcher.
+// NewResponseRecorder creates a ResponseRecorder.
 func NewResponseRecorder() *ResponseRecorder {
 	var b strings.Builder
 	rw := NewTestResponseWriter(&b)
 	return &ResponseRecorder{
 		rw:             rw,
 		b:              &b,
-		ResponseWriter: safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, rw, nil),
+		ResponseWriter: safehttp.DeprecatedNewResponseWriter(rw, nil),
 	}
 }
 
-// NewResponseRecorderFromDispatcher creates a ResponseRecorder from a
-// provided safehttp.Dispatcher.
+// NewResponseRecorderFromDispatcher creates a ResponseRecorder.
 func NewResponseRecorderFromDispatcher(d safehttp.Dispatcher) *ResponseRecorder {
 	var b strings.Builder
 	rw := NewTestResponseWriter(&b)
 	return &ResponseRecorder{
 		rw:             rw,
 		b:              &b,
-		ResponseWriter: safehttp.NewResponseWriter(d, rw, nil),
+		ResponseWriter: safehttp.DeprecatedNewResponseWriter(rw, d),
 	}
 }
 
diff --git a/safehttp/response_writer_test.go b/safehttp/task_test.go
similarity index 85%
rename from safehttp/response_writer_test.go
rename to safehttp/task_test.go
index 76ea225..ea0010a 100644
--- a/safehttp/response_writer_test.go
+++ b/safehttp/task_test.go
@@ -16,6 +16,7 @@ package safehttp_test
 
 import (
 	"math"
+	"net/http"
 	"strings"
 	"testing"
 	"text/template"
@@ -26,9 +27,16 @@ import (
 	"github.com/google/safehtml"
 )
 
-func TestResponseWriterSetCookie(t *testing.T) {
+func stubResponseWriter(dispatcher safehttp.Dispatcher, rw http.ResponseWriter) safehttp.ResponseWriter {
+	if dispatcher == nil {
+		dispatcher = safehttp.DefaultDispatcher{}
+	}
+	return safehttp.DeprecatedNewResponseWriter(rw, dispatcher)
+}
+
+func TestTaskSetCookie(t *testing.T) {
 	testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
-	rw := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, testRw, nil)
+	rw := stubResponseWriter(safehttp.DefaultDispatcher{}, testRw)
 
 	c := safehttp.NewCookie("foo", "bar")
 	err := rw.SetCookie(c)
@@ -44,9 +52,9 @@ func TestResponseWriterSetCookie(t *testing.T) {
 	}
 }
 
-func TestResponseWriterSetInvalidCookie(t *testing.T) {
+func TestTaskSetInvalidCookie(t *testing.T) {
 	testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
-	rw := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, testRw, nil)
+	rw := stubResponseWriter(safehttp.DefaultDispatcher{}, testRw)
 
 	c := safehttp.NewCookie("f=oo", "bar")
 	err := rw.SetCookie(c)
@@ -55,7 +63,7 @@ func TestResponseWriterSetInvalidCookie(t *testing.T) {
 	}
 }
 
-func TestResponseWriterWriteTwicePanic(t *testing.T) {
+func TestTaskWriteTwicePanic(t *testing.T) {
 	tests := []struct {
 		name  string
 		write func(w safehttp.ResponseWriter)
@@ -93,7 +101,7 @@ func TestResponseWriterWriteTwicePanic(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}), nil)
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r != nil {
 					return
@@ -105,7 +113,7 @@ func TestResponseWriterWriteTwicePanic(t *testing.T) {
 	}
 }
 
-func TestResponseWriterUnsafeResponse(t *testing.T) {
+func TestTaskUnsafeResponse(t *testing.T) {
 	tests := []struct {
 		name  string
 		write func(w safehttp.ResponseWriter)
@@ -133,7 +141,7 @@ func TestResponseWriterUnsafeResponse(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}), nil)
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r != nil {
 					return
@@ -145,7 +153,7 @@ func TestResponseWriterUnsafeResponse(t *testing.T) {
 	}
 }
 
-func TestResponseWriterStatusCode(t *testing.T) {
+func TestTaskStatusCode(t *testing.T) {
 	tests := []struct {
 		name       string
 		write      func(w safehttp.ResponseWriter)
@@ -202,7 +210,7 @@ func TestResponseWriterStatusCode(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			testRw := safehttptest.NewTestResponseWriter(&strings.Builder{})
-			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, testRw, nil)
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, testRw)
 			tt.write(w)
 			if gotStatus := testRw.Status(); gotStatus != tt.wantStatus {
 				t.Errorf("tt.write(w): got %v, want %v", gotStatus, tt.wantStatus)
@@ -211,7 +219,7 @@ func TestResponseWriterStatusCode(t *testing.T) {
 	}
 }
 
-func TestResponseWriterInvalidStatusCode(t *testing.T) {
+func TestTaskInvalidStatusCode(t *testing.T) {
 	tests := []struct {
 		name  string
 		write func(w safehttp.ResponseWriter)
@@ -233,7 +241,7 @@ func TestResponseWriterInvalidStatusCode(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}), nil)
+			w := stubResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r != nil {
 					return