Repository: google/go-safeweb (go)

Issue:
Consider removing WriteTemplate and WriteJSON, implement helpers to just rely on Write
### Rationale

Currently we have a quite low-level API (`Write(Response)`) and two very high-level one, which are tailored to a specific function.

My proposal is to just stick to the general purpose API (Write) which is the one responsible for the wiring with the Dispatcher.

This will keep the interface for Dispatcher and ResponseWriter simpler, and will make it clear that the responsibility for wrapping and handling different types is on the Dispatcher implementers.

This will also remove the duplicated code that we have that asks for content-type in every call to WriteTemplate, WriteJSON etc.

### Template

The specific type we require is a `Template` interface but we actually ask more of it: for example when we will need to nonce-inject we will need Clone and Funcs methods, and we cannot express the Clonable interface, so we will be asking users to implement something, but will **implicitly ask for more**.

Moreover this API does not work if we have to call ExecuteTemplate instead, which will require an additional string parameter.

To then render templates we would offer helpers that would look like these: `w.Write(mytpl.WithData(data))` `w.Write(mytpl.NamedWithData("foo", data))`. If we don't want to change the packages that export the templates we can also put those helpers and wrappers in safehttp directly (which we are already doing anyways, even if in another form). With the extra added benefit that the signature will match the exact types we expect.

### JSON

Currently WriteJSON just wraps things in a JsonResponse. We could just create an helper for that and pass JSONResponse to Write. Example usage would look like `w.Write(safehttp.JsonResponse{data})` instead of `w.WriteJSON(data)` so it will be slightly more verbose but I think the overall API will benefit from the simplification.

### An optional extra step
We could even make the so-created `Response` implementations have an (unexported) method that takes care of properly setting the content-type and one for serializing the body, which would allow us to actually limit the `safehttp.Response` implementers to be checked at compile-time.

At that point the dispatcher responsibility would just be to check that something extraneous to safehttp is an acceptably safe response when put in a generic wrapper.

Overall we would reduce the typeswitching a lot and make the API cleaner.

Pull request description:
ResponseWriter and Dispatcher now have a single Write method
Fixes #150 

Unified the functionality of `WriteJSON`, `ExecuteTemplate` and `Write` inside the new `safehttp.DefaultDispatcher.Write` method, also responsible for setting the `Content-Type` header and writing the status code.

Patch:
diff --git a/safehttp/default_dispatcher.go b/safehttp/default_dispatcher.go
index a2366b3..099f4c5 100644
--- a/safehttp/default_dispatcher.go
+++ b/safehttp/default_dispatcher.go
@@ -26,11 +26,17 @@ import (
 // DefaultDispatcher is responsible for writing safe responses.
 type DefaultDispatcher struct{}
 
-// ContentType returns the Content-Type of a safe response if it's of a
-// safe-type and an error otherwise.
+// ContentType returns the Content-Type of a response if it's deemed safe and
+// an error otherwise.
 func (DefaultDispatcher) ContentType(resp Response) (string, error) {
-	switch resp.(type) {
-	case safehtml.HTML, *template.Template:
+	switch x := resp.(type) {
+	case safehtml.HTML:
+		return "text/html; charset=utf-8", nil
+	case TemplateResponse:
+		_, ok := (x.Template).(*template.Template)
+		if !ok {
+			return "", fmt.Errorf("%T is not a safe response type, a Content-Type cannot be provided", resp)
+		}
 		return "text/html; charset=utf-8", nil
 	case JSONResponse:
 		return "application/json; charset=utf-8", nil
@@ -39,49 +45,39 @@ func (DefaultDispatcher) ContentType(resp Response) (string, error) {
 	}
 }
 
-// Write writes the response to the Response Writer if it's safe HTML. It
-// returns a non-nil error if the response is not safe HTML or if writing the
-// response fails.
-func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
-	x, ok := resp.(safehtml.HTML)
-	if !ok {
-		return fmt.Errorf("%T is not a safe response type and it cannot be written", resp)
-	}
-	_, err := io.WriteString(rw, x.String())
-	return err
-
-}
-
-// WriteJSON serialises and writes a JSON to the underlying http.ResponseWriter.
-// If the provided response is not valid JSON or writing the response fails, the
-// method will return an error.
-func (DefaultDispatcher) WriteJSON(rw http.ResponseWriter, resp JSONResponse) error {
-	io.WriteString(rw, ")]}',\n") // Break parsing of JavaScript in order to prevent XSSI.
-	return json.NewEncoder(rw).Encode(resp.Data)
-}
-
-// ExecuteTemplate applies the parsed template to the provided data object, if
-// the template is a safe HTML template, writing the output to the  http.
-// ResponseWriter. If the funcMap is non-nil, its elements override the
-// existing names to functions mappings in the template. An attempt to define a
-// new name to function mapping that is not already in the template will result
-// in a panic.
+// Write writes the response to the http.ResponseWriter if it's deemed safe. It
+// returns a non-nil error if the response is deemed unsafe or if the writing
+// operation fails.
 //
-// If an error occurs executing the template or writing its output,
-// execution stops, but partial results may already have been written to
-// the Response Writer.
-func (DefaultDispatcher) ExecuteTemplate(rw http.ResponseWriter, resp TemplateResponse) error {
-	t := *resp.Template
-	x, ok := t.(*template.Template)
-	if !ok {
-		return fmt.Errorf("%T is not a safe template and it cannot be parsed and written", resp.Template)
-	}
-	if len(resp.FuncMap) == 0 {
-		return x.Execute(rw, resp.Data)
-	}
-	cloned, err := x.Clone()
-	if err != nil {
+// For JSONResponses, the underlying object is serialised and written if it's a
+// valid JSON.
+//
+// For TemplateResponses, the parsed template is applied to the provided data
+// object. If the funcMap is non-nil, its elements override the  existing names
+// to functions mappings in the template. An attempt to define a new name to
+// function mapping that is not already in the template will result  in a panic.
+func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
+	switch x := resp.(type) {
+	case JSONResponse:
+		io.WriteString(rw, ")]}',\n") // Break parsing of JavaScript in order to prevent XSSI.
+		return json.NewEncoder(rw).Encode(x.Data)
+	case TemplateResponse:
+		t, ok := (x.Template).(*template.Template)
+		if !ok {
+			return fmt.Errorf("%T is not a safe template and it cannot be parsed and written", t)
+		}
+		if len(x.FuncMap) == 0 {
+			return t.Execute(rw, x.Data)
+		}
+		cloned, err := t.Clone()
+		if err != nil {
+			return err
+		}
+		return cloned.Funcs(x.FuncMap).Execute(rw, x.Data)
+	case safehtml.HTML:
+		_, err := io.WriteString(rw, x.String())
 		return err
+	default:
+		return fmt.Errorf("%T is not a safe response type and it cannot be written", resp)
 	}
-	return cloned.Funcs(resp.FuncMap).Execute(rw, resp.Data)
 }
diff --git a/safehttp/response.go b/safehttp/response.go
index 7ebb2d2..2c7719b 100644
--- a/safehttp/response.go
+++ b/safehttp/response.go
@@ -29,6 +29,13 @@ type JSONResponse struct {
 	Data interface{}
 }
 
+// WriteJSON creates a JSONResponse from the data object and calls the Write
+// function of the ResponseWriter, passing the response. The data object should
+// be valid JSON, otherwise an error will occur.
+func WriteJSON(w *ResponseWriter, data interface{}) Result {
+	return w.Write(JSONResponse{data})
+}
+
 // Template implements a template.
 type Template interface {
 	// Execute applies data to the template and then writes the result to
@@ -43,11 +50,25 @@ type Template interface {
 // TemplateResponse bundles a Template with its data and names to function
 // mappings to be passed together to the commit phase.
 type TemplateResponse struct {
-	Template *Template
-	Data     *interface{}
+	Template Template
+	Data     interface{}
 	FuncMap  map[string]interface{}
 }
 
+// ExecuteTemplate creates a TemplateResponse from the provided Template and its
+// data and calls the Write function of the ResponseWriter, passing the
+// response.
+func ExecuteTemplate(w *ResponseWriter, t Template, data interface{}) Result {
+	return w.Write(TemplateResponse{t, data, nil})
+}
+
+// ExecuteTemplateWithFuncs creates a TemplateResponse from the provided
+// Template, its data and the name to function mappings and calls the Write
+// function of the ResponseWriter, passing the response.
+func ExecuteTemplateWithFuncs(w *ResponseWriter, t Template, data interface{}, fm map[string]interface{}) Result {
+	return w.Write(TemplateResponse{t, data, fm})
+}
+
 // NoContentResponse is sent to the commit phase when it's initiated from
 // ResponseWriter.NoContent.
 type NoContentResponse struct{}
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 73324ce..6ac4a62 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -66,12 +66,10 @@ func NotWritten() Result {
 	return Result{}
 }
 
-// Write dispatches the response to the Dispatcher, setting the Content-Type and
-// response status if the provided response is a safe response. The
-// Dispatcher will then write the response to the underlying Response Writer.
+// 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 getting the response
-// Content-Type or writing the response fails.
+// 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")
@@ -80,71 +78,17 @@ func (w *ResponseWriter) Write(resp Response) Result {
 	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)
-	w.rw.WriteHeader(int(StatusOK))
-	if err := w.d.Write(w.rw, resp); err != nil {
-		panic(err)
-	}
-	return Result{}
-}
 
-// WriteJSON encapsulates data into a JSON response and dispatches it to the
-// Dispatcher, to be serialised and written to the ResponseWriter. It will also
-// set the Content-Type and response status will be set.
-//
-// TODO: replace panics with proper error handling when getting the response
-// Content-Type or writing the response fails.
-func (w *ResponseWriter) WriteJSON(data interface{}) Result {
-	if w.written {
-		panic("ResponseWriter was already written to")
-	}
-	w.handler.commitPhase(w, JSONResponse{Data: data})
-	if w.written {
-		return Result{}
-	}
 	w.markWritten()
-	resp := JSONResponse{Data: data}
+
 	ct, err := w.d.ContentType(resp)
 	if err != nil {
 		panic(err)
 	}
 	w.rw.Header().Set("Content-Type", ct)
 	w.rw.WriteHeader(int(StatusOK))
-	if err := w.d.WriteJSON(w.rw, resp); err != nil {
-		panic(err)
-	}
-	return Result{}
-}
 
-// WriteTemplate dispatches a parsed template and a data object to the
-// Dispatcher to be executed and written to the underlying Response Writer, in
-// case the template is a safe HTML template. If it is a safe HTML Template, the
-// Content-Type  and response status will also be set.
-//
-// TODO: replace panics with proper error handling when getting the response
-// Content-Type or writing the response fails.
-func (w *ResponseWriter) WriteTemplate(t Template, data interface{}) Result {
-	if w.written {
-		panic("ResponseWriter was already written to")
-	}
-	resp := TemplateResponse{Template: &t, Data: &data, FuncMap: map[string]interface{}{}}
-	w.handler.commitPhase(w, resp)
-	if w.written {
-		return Result{}
-	}
-	w.markWritten()
-	ct, err := w.d.ContentType(t)
-	if err != nil {
-		panic(err)
-	}
-	w.rw.Header().Set("Content-Type", ct)
-	w.rw.WriteHeader(int(StatusOK))
-	if err := w.d.ExecuteTemplate(w.rw, resp); err != nil {
+	if err := w.d.Write(w.rw, resp); err != nil {
 		panic(err)
 	}
 	return Result{}
@@ -219,36 +163,21 @@ func (w *ResponseWriter) SetCookie(c *Cookie) error {
 // The implementation of a custom Dispatcher should be thoroughly reviewed by
 // the security team to avoid introducing vulnerabilities.
 type Dispatcher interface {
-	// Write writes a Response to the 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
-	Write(rw http.ResponseWriter, resp Response) error
-
-	// WriteJSON serialises and writes a JSON object to the http.ResponseWriter.
-	//
-	// This should return an error if the provided response is not a safe
-	// JSONResponse  or if writing the object to the http.ResponseWriter fails.
-	WriteJSON(rw http.ResponseWriter, resp JSONResponse) error
-
-	// ExecuteTemplate applies the response parsed template to its data object
-	// and writes the output to the http.ResponseWriter. If the response funcMap
-	// is non-nil, the Dispatcher will try to update the names to functions
-	// mappings of the template.
-	//
-	// This should return an error if the template contained by the
-	// TemplateResponse is not a safe template or if an error occurs executing
-	// or writing the template.
-	ExecuteTemplate(rw http.ResponseWriter, resp TemplateResponse) error
-
 	// Content-Type returns the Content-Type of the provided response if it is
-	// of a type supported by the Dispatcher and should return an error
+	// 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 any
-	// of the writing methods in the Dispatcher. Writing should not proceed
-	// if ContentType returns an error.
+	// 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/default_dispatcher_test.go b/safehttp/default_dispatcher_test.go
index ed58b41..ed4dda5 100644
--- a/safehttp/default_dispatcher_test.go
+++ b/safehttp/default_dispatcher_test.go
@@ -51,11 +51,7 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 						Parse("<h1>{{ . }}</h1>")))
 				var data interface{}
 				data = "This is an actual heading, though."
-				resp := safehttp.TemplateResponse{
-					Template: &t,
-					Data:     &data,
-				}
-				return d.ExecuteTemplate(w, resp)
+				return d.Write(w, safehttp.TemplateResponse{t, data, nil})
 			},
 			wantBody: "<h1>This is an actual heading, though.</h1>",
 		},
@@ -73,13 +69,7 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 				fm := map[string]interface{}{
 					"Token": func() string { return "Token-secret" },
 				}
-				resp := safehttp.TemplateResponse{
-					Template: &t,
-					Data:     &data,
-					FuncMap:  fm,
-				}
-
-				return d.ExecuteTemplate(w, resp)
+				return d.Write(w, safehttp.TemplateResponse{t, data, fm})
 			},
 			wantBody: `<form><input type="hidden" name="token" value="Token-secret">Content</form>`,
 		},
@@ -97,13 +87,7 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 				fm := map[string]interface{}{
 					"Nonce": func() string { return "Nonce-secret" },
 				}
-				resp := safehttp.TemplateResponse{
-					Template: &t,
-					Data:     &data,
-					FuncMap:  fm,
-				}
-
-				return d.ExecuteTemplate(w, resp)
+				return d.Write(w, safehttp.TemplateResponse{t, data, fm})
 			},
 			wantBody: `<script nonce="Nonce-secret" type="application/javascript">alert("script")</script><h1>Content</h1>`,
 		},
@@ -114,7 +98,7 @@ func TestDefaultDispatcherValidResponse(t *testing.T) {
 				data := struct {
 					Field string `json:"field"`
 				}{Field: "myField"}
-				return d.WriteJSON(w, safehttp.JSONResponse{Data: data})
+				return d.Write(w, safehttp.JSONResponse{data})
 			},
 			wantBody: ")]}',\n{\"field\":\"myField\"}\n",
 		},
@@ -160,11 +144,7 @@ func TestDefaultDispatcherInvalidResponse(t *testing.T) {
 						Parse("<h1>{{ . }}</h1>")))
 				var data interface{}
 				data = "This is an actual heading, though."
-				resp := safehttp.TemplateResponse{
-					Template: &t,
-					Data:     &data,
-				}
-				return d.ExecuteTemplate(w, resp)
+				return d.Write(w, safehttp.TemplateResponse{t, data, nil})
 			},
 			want: "",
 		},
@@ -172,7 +152,7 @@ func TestDefaultDispatcherInvalidResponse(t *testing.T) {
 			name: "Invalid JSON Response",
 			write: func(w http.ResponseWriter) error {
 				d := &safehttp.DefaultDispatcher{}
-				return d.WriteJSON(w, safehttp.JSONResponse{Data: math.Inf(1)})
+				return d.Write(w, safehttp.JSONResponse{math.Inf(1)})
 			},
 			want: ")]}',\n",
 		},
diff --git a/safehttp/response_writer_test.go b/safehttp/response_writer_test.go
index bf8d159..cf13eae 100644
--- a/safehttp/response_writer_test.go
+++ b/safehttp/response_writer_test.go
@@ -15,6 +15,7 @@
 package safehttp_test
 
 import (
+	"math"
 	"strings"
 	"testing"
 	"text/template"
@@ -67,52 +68,69 @@ func TestResponseWriterWriteTwicePanic(t *testing.T) {
 			},
 		},
 		{
-			name: "Call WriteJSON twice",
+			name: "Call NoContent twice",
 			write: func(w *safehttp.ResponseWriter) {
-				obj := struct{ Field string }{Field: "myField"}
-				w.WriteJSON(obj)
-				w.WriteJSON(obj)
+				w.NoContent()
+				w.NoContent()
 			},
 		},
 		{
-			name: "Call Write then WriteJSON",
+			name: "Call WriteError twice",
 			write: func(w *safehttp.ResponseWriter) {
-				obj := struct{ Field string }{Field: "myField"}
-				w.Write(obj)
-				w.WriteJSON(obj)
+				w.WriteError(safehttp.StatusInternalServerError)
+				w.WriteError(safehttp.StatusInternalServerError)
 			},
 		},
 		{
-			name: "Call WriteTemplate twice",
+			name: "Call Redirect twice",
 			write: func(w *safehttp.ResponseWriter) {
-				w.WriteTemplate(template.Must(template.New("name").Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
-				w.WriteTemplate(template.Must(template.New("name").Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
+				ir := safehttptest.NewRequest(safehttp.MethodGet, "/", nil)
+				w.Redirect(ir, "/asdf", safehttp.StatusMovedPermanently)
+				w.Redirect(ir, "/asdf", safehttp.StatusMovedPermanently)
 			},
 		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}), nil)
+			defer func() {
+				if r := recover(); r != nil {
+					return
+				}
+				t.Errorf("tt.write(w) expected panic")
+			}()
+			tt.write(w)
+		})
+	}
+}
+
+func TestResponseWriterUnsafeResponse(t *testing.T) {
+	tests := []struct {
+		name  string
+		write func(w *safehttp.ResponseWriter)
+	}{
 		{
-			name: "Call NoContent twice",
+			name: "Unsafe HTML Response",
 			write: func(w *safehttp.ResponseWriter) {
-				w.NoContent()
-				w.NoContent()
+				w.Write("<h1>Hello World!</h1>")
 			},
 		},
 		{
-			name: "Call WriteError twice",
+			name: "Invalid JSON Response",
 			write: func(w *safehttp.ResponseWriter) {
-				w.WriteError(safehttp.StatusInternalServerError)
-				w.WriteError(safehttp.StatusInternalServerError)
+				safehttp.WriteJSON(w, math.Inf(1))
 			},
 		},
 		{
-			name: "Call Redirect twice",
+			name: "Unsafe Template Response",
 			write: func(w *safehttp.ResponseWriter) {
-				ir := safehttptest.NewRequest(safehttp.MethodGet, "/", nil)
-				w.Redirect(ir, "/asdf", safehttp.StatusMovedPermanently)
-				w.Redirect(ir, "/asdf", safehttp.StatusMovedPermanently)
+				safehttp.ExecuteTemplate(w, template.
+					Must(template.New("name").
+						Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 			},
 		},
 	}
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			w := safehttp.NewResponseWriter(safehttp.DefaultDispatcher{}, safehttptest.NewTestResponseWriter(&strings.Builder{}), nil)
diff --git a/tests/integration/csp/csp_test.go b/tests/integration/csp/csp_test.go
index 827b0f8..fbe409a 100644
--- a/tests/integration/csp/csp_test.go
+++ b/tests/integration/csp/csp_test.go
@@ -44,7 +44,7 @@ func TestServeMuxInstallCSP(t *testing.T) {
 		nonce, err = csp.Nonce(r.Context())
 		t := template.Must(template.New("name").Funcs(fns).Parse(`<script nonce="{{CSPNonce}}" type="application/javascript">alert("script")</script><h1>{{.}}</h1>`))
 
-		return w.WriteTemplate(t, "Content")
+		return safehttp.ExecuteTemplateWithFuncs(w, t, "Content", fns)
 	})
 	mb.Handle("/bar", safehttp.MethodGet, handler)
 
diff --git a/tests/integration/mux/mux_test.go b/tests/integration/mux/mux_test.go
index f4c874b..13eed03 100644
--- a/tests/integration/mux/mux_test.go
+++ b/tests/integration/mux/mux_test.go
@@ -48,7 +48,7 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Safe HTML Template Response",
 			handler: safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-				return w.WriteTemplate(safetemplate.
+				return safehttp.ExecuteTemplate(w, safetemplate.
 					Must(safetemplate.New("name").
 						Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 			}),
@@ -63,7 +63,7 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 				data := struct {
 					Field string `json:"field"`
 				}{Field: "myField"}
-				return w.WriteJSON(data)
+				return safehttp.WriteJSON(w, data)
 			}),
 			wantHeaders: map[string][]string{
 				"Content-Type": {"application/json; charset=utf-8"},
@@ -110,7 +110,7 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 		{
 			name: "Unsafe Template Response",
 			handler: safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-				return w.WriteTemplate(template.
+				return safehttp.ExecuteTemplate(w, template.
 					Must(template.New("name").
 						Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 			}),
@@ -118,7 +118,7 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 		{
 			name: "Invalid JSON Response",
 			handler: safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-				return w.WriteJSON(math.Inf(1))
+				return safehttp.WriteJSON(w, math.Inf(1))
 			}),
 		},
 	}
diff --git a/tests/integration/safehtml/safehtml_test.go b/tests/integration/safehtml/safehtml_test.go
index a8303a3..1a0f79b 100644
--- a/tests/integration/safehtml/safehtml_test.go
+++ b/tests/integration/safehtml/safehtml_test.go
@@ -49,7 +49,7 @@ func TestHandleRequestWrite(t *testing.T) {
 func TestHandleRequestWriteTemplate(t *testing.T) {
 	mb := &safehttp.ServeMuxConfig{}
 	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
-		return rw.WriteTemplate(template.Must(template.New("name").Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
+		return safehttp.ExecuteTemplate(rw, template.Must(template.New("name").Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 	}))
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/", nil)
diff --git a/tests/integration/xsrf/xsrf_test.go b/tests/integration/xsrf/xsrf_test.go
index 85f92d8..5c60f27 100644
--- a/tests/integration/xsrf/xsrf_test.go
+++ b/tests/integration/xsrf/xsrf_test.go
@@ -43,7 +43,7 @@ func TestServeMuxInstallXSRF(t *testing.T) {
 		token, err = xsrf.Token(r)
 		t := template.Must(template.New("name").Funcs(fns).Parse(`<form><input type="hidden" name="token" value="{{XSRFToken}}">{{.}}</form>`))
 
-		return w.WriteTemplate(t, "Content")
+		return safehttp.ExecuteTemplateWithFuncs(w, t, "Content", fns)
 	})
 	mb.Handle("/bar", safehttp.MethodGet, handler)