Repository: google/go-safeweb (go)

Issue:
Consider making host checks in a plugin rather than in the core mux implementation
Currently, host checking is done here: https://github.com/google/go-safeweb/blob/0c930b079f42206df056e531d88cf9009caa6f97/safehttp/mux.go#L173

It might make sense to do it in an interceptor. Especially since we might need to provide convenient ways of configuring the check (e.g. for local development against `localhost`).

Pull request:
Host checks are now inside a hostcheck plugin.
Host checks were performed inside the mux to protect from DNS rebinding
and HTTP request smuggling attacks. There is no real reason why these
should be inside the mux implementation instead of a plugin.

Fixes #168.

Patch:
diff --git a/safehttp/mux.go b/safehttp/mux.go
index 6073840..8070cd1 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -35,11 +35,6 @@ const (
 // request against a list of registered patterns and calls the handler for
 // the pattern that most closely matches the URL.
 //
-// When creating the multiplexer, the user needs to specify a list of allowed
-// domains. The server will only serve requests target to those domains and
-// otherwise will reply with HTTP 404 Not Found in order to protect agains DNS
-// rebinding and HTTP request smuggling.
-//
 // Patterns names are fixed, rooted paths, like "/favicon.ico", or rooted
 // subtrees like "/images/" (note the trailing slash). Longer patterns take
 // precedence over shorter ones, so that if there are handlers registered for
@@ -59,8 +54,7 @@ const (
 // "/images" to "/images/", unless "/images" has been registered separately.
 //
 // Patterns may optionally begin with a host name, restricting matches to URLs
-// on that host only. This host name must be in the list of allowed domains passed
-// when creating the ServeMux. Host-specific patterns take precedence over general
+// on that host only.  Host-specific patterns take precedence over general
 // patterns, so that a handler might register for the two patterns "/codesearch"
 // and "codesearch.google.com/" without also taking over requests for
 // "http://www.google.com/".
@@ -72,30 +66,22 @@ const (
 // Multiple handlers can be registered for a single pattern, as long as they
 // handle different HTTP methods.
 type ServeMux struct {
-	mux     *http.ServeMux
-	domains map[string]bool
-	disp    Dispatcher
+	mux  *http.ServeMux
+	disp Dispatcher
 
 	// Maps patterns to handlers supporting multiple HTTP methods.
 	handlers  map[string]methodHandler
 	interceps []Interceptor
 }
 
-// NewServeMux allocates and returns a new ServeMux. Passing a nil Dispatcher
-// will cause the method to panic.
-func NewServeMux(d Dispatcher, domains ...string) *ServeMux {
+// NewServeMux allocates and returns a new ServeMux. If the provided dispatcher
+// is nil, the DefaultDispatcher is used.
+func NewServeMux(d Dispatcher) *ServeMux {
 	if d == nil {
 		d = DefaultDispatcher{}
 	}
-
-	// TODO(@mattiasgrenfeldt, @mihalimara22): make domains a variadic of string **literals**.
-	dm := map[string]bool{}
-	for _, host := range domains {
-		dm[host] = true
-	}
 	return &ServeMux{
 		mux:      http.NewServeMux(),
-		domains:  dm,
 		disp:     d,
 		handlers: map[string]methodHandler{},
 	}
@@ -136,7 +122,6 @@ func (m *ServeMux) Handle(pattern string, method string, h Handler, cfgs ...Inte
 	if !ok {
 		mh := methodHandler{
 			handlers: map[string]handlerWithInterceptors{method: hi},
-			domains:  m.domains,
 		}
 
 		m.handlers[pattern] = mh
@@ -185,17 +170,11 @@ func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 type methodHandler struct {
 	// Maps an HTTP method to its handlerWithInterceptors
 	handlers map[string]handlerWithInterceptors
-	domains  map[string]bool
 }
 
 // ServeHTTP dispatches the request to the handlerWithInterceptors associated
 // with the IncomingRequest method.
 func (m methodHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	if !m.domains[r.Host] {
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
-	}
-
 	h, ok := m.handlers[r.Method]
 	if !ok {
 		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
diff --git a/safehttp/plugins/hostcheck/hostcheck.go b/safehttp/plugins/hostcheck/hostcheck.go
new file mode 100644
index 0000000..f82ab2e
--- /dev/null
+++ b/safehttp/plugins/hostcheck/hostcheck.go
@@ -0,0 +1,56 @@
+// 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 hostcheck provides a plugin that checks whether the request is
+// intended to be sent to a given host.
+//
+// This is a protection mechanism against
+// DNS rebinding attacks (https://en.wikipedia.org/wiki/DNS_rebinding) and HTTP
+// request smuggling (https://portswigger.net/web-security/request-smuggling).
+package hostcheck
+
+import (
+	"github.com/google/go-safeweb/safehttp"
+)
+
+// Interceptor checks whether the Host header of the incoming request is in an
+// allowlist.
+type Interceptor struct {
+	hosts map[string]bool
+}
+
+var _ safehttp.Interceptor = Interceptor{}
+
+// New creates an Interceptor.
+func New(hosts ...string) Interceptor {
+	it := Interceptor{hosts: map[string]bool{}}
+	for _, h := range hosts {
+		it.hosts[h] = true
+	}
+	return it
+}
+
+// Before checks whether the request's Host header is in the list of allowed
+// hosts. If it's not, it responds with 404 Not Found.
+func (it Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, _ safehttp.InterceptorConfig) safehttp.Result {
+	if !it.hosts[r.Host()] {
+		return w.WriteError(safehttp.StatusNotFound)
+	}
+	return safehttp.NotWritten()
+}
+
+// Commit is a no-op, required to satisfy the safehttp.Interceptor interface.
+func (Interceptor) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, resp safehttp.Response, _ safehttp.InterceptorConfig) safehttp.Result {
+	return safehttp.NotWritten()
+}


Test patch:
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index db62576..b9dbdb4 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -45,16 +45,6 @@ func TestMuxOneHandlerOneRequest(t *testing.T) {
 			},
 			wantBody: "&lt;h1&gt;Hello World!&lt;/h1&gt;",
 		},
-		{
-			name:       "Invalid Host",
-			req:        httptest.NewRequest(safehttp.MethodGet, "http://bar.com/", nil),
-			wantStatus: safehttp.StatusNotFound,
-			wantHeader: map[string][]string{
-				"Content-Type":           {"text/plain; charset=utf-8"},
-				"X-Content-Type-Options": {"nosniff"},
-			},
-			wantBody: "Not Found\n",
-		},
 		{
 			name:       "Invalid Method",
 			req:        httptest.NewRequest(safehttp.MethodPost, "http://foo.com/", nil),
@@ -69,7 +59,7 @@ func TestMuxOneHandlerOneRequest(t *testing.T) {
 
 	for _, tt := range test {
 		t.Run(tt.name, func(t *testing.T) {
-			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 			h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
@@ -132,7 +122,7 @@ func TestMuxServeTwoHandlers(t *testing.T) {
 	}
 
 	d := safehttp.DefaultDispatcher{}
-	mux := safehttp.NewServeMux(d, "foo.com")
+	mux := safehttp.NewServeMux(d)
 	mux.Handle("/bar", safehttp.MethodGet, tests[0].hf)
 	mux.Handle("/bar", safehttp.MethodPost, tests[1].hf)
 
@@ -156,7 +146,7 @@ func TestMuxServeTwoHandlers(t *testing.T) {
 
 func TestMuxHandleSameMethodTwice(t *testing.T) {
 	d := safehttp.DefaultDispatcher{}
-	mux := safehttp.NewServeMux(d, "foo.com")
+	mux := safehttp.NewServeMux(d)
 
 	registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
@@ -261,7 +251,7 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Install ServeMux Interceptor before handler registration",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 				mux.Install(setHeaderInterceptor{
 					name:  "Foo",
 					value: "bar",
@@ -283,7 +273,7 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Install Interrupting Interceptor",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 				mux.Install(internalErrorInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
@@ -303,7 +293,7 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Handler Communication With ServeMux Interceptor",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 				mux.Install(&claimHeaderInterceptor{headerToClaim: "Foo"})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
@@ -324,7 +314,7 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Panicking interceptor recovered by ServeMux",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 				mux.Install(panickingInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
@@ -344,7 +334,7 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Commit phase sets header",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 				mux.Install(committerInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
@@ -461,7 +451,7 @@ func TestMuxInterceptorConfigs(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 			mux.Install(setHeaderConfigInterceptor{})
 
 			registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
@@ -543,7 +533,7 @@ func (interceptorThree) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingR
 }
 
 func TestMuxDeterministicInterceptorOrder(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	mux.Install(interceptorOne{})
 	mux.Install(interceptorTwo{})
 	mux.Install(interceptorThree{})
@@ -581,7 +571,7 @@ func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 }
 
 func TestMuxHandlerReturnsNotWritten(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return safehttp.NotWritten()
 	})
diff --git a/safehttp/plugins/hostcheck/hostcheck_test.go b/safehttp/plugins/hostcheck/hostcheck_test.go
new file mode 100644
index 0000000..ef6657c
--- /dev/null
+++ b/safehttp/plugins/hostcheck/hostcheck_test.go
@@ -0,0 +1,66 @@
+// 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 hostcheck_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/plugins/hostcheck"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+	"github.com/google/safehtml"
+)
+
+func TestInterceptor(t *testing.T) {
+	var test = []struct {
+		name       string
+		req        *http.Request
+		wantStatus safehttp.StatusCode
+	}{
+		{
+			name:       "Valid Host",
+			req:        httptest.NewRequest(safehttp.MethodGet, "http://foo.com/", nil),
+			wantStatus: safehttp.StatusOK,
+		},
+		{
+			name:       "Invalid Host",
+			req:        httptest.NewRequest(safehttp.MethodGet, "http://bar.com/", nil),
+			wantStatus: safehttp.StatusNotFound,
+		},
+	}
+
+	for _, tt := range test {
+		t.Run(tt.name, func(t *testing.T) {
+			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+			mux.Install(hostcheck.New("foo.com"))
+
+			h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
+			})
+			mux.Handle("/", safehttp.MethodGet, h)
+
+			b := &strings.Builder{}
+			rw := safehttptest.NewTestResponseWriter(b)
+			mux.ServeHTTP(rw, tt.req)
+
+			if rw.Status() != tt.wantStatus {
+				t.Errorf("rw.Status(): got %v want %v", rw.Status(), tt.wantStatus)
+			}
+		})
+	}
+}
diff --git a/tests/integration/csp/csp_test.go b/tests/integration/csp/csp_test.go
index 6e17497..1f82e8d 100644
--- a/tests/integration/csp/csp_test.go
+++ b/tests/integration/csp/csp_test.go
@@ -15,18 +15,19 @@
 package csp_test
 
 import (
+	"net/http/httptest"
+	"strings"
+	"testing"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/csp"
 	"github.com/google/go-safeweb/safehttp/safehttptest"
-	safetemplate "github.com/google/safehtml/template"
-	"net/http/httptest"
-	"strings"
-	"testing"
+	"github.com/google/safehtml/template"
 )
 
 func TestServeMuxInstallCSP(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	it := csp.Default("")
 	mux.Install(&it)
 
@@ -41,7 +42,7 @@ func TestServeMuxInstallCSP(t *testing.T) {
 		// installed, but auto-injection isn't yet supported and has to be done
 		// manually by the handler).
 		nonce, err = csp.Nonce(r.Context())
-		t := safetemplate.Must(safetemplate.New("name").Funcs(fns).Parse(`<script nonce="{{CSPNonce}}" type="application/javascript">alert("script")</script><h1>{{.}}</h1>`))
+		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")
 	})
diff --git a/tests/integration/header/header_test.go b/tests/integration/header/header_test.go
index 3a139a7..6842f1d 100644
--- a/tests/integration/header/header_test.go
+++ b/tests/integration/header/header_test.go
@@ -16,19 +16,20 @@ package header
 
 import (
 	"bufio"
-	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/safehtml"
 )
 
 func TestAccessIncomingHeaders(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
 		if got, want := ir.Header.Get("A"), "B"; got != want {
 			t.Errorf(`ir.Header.Get("A") got: %v want: %v`, got, want)
@@ -51,7 +52,7 @@ func TestAccessIncomingHeaders(t *testing.T) {
 }
 
 func TestChangingResponseHeaders(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
 		rw.Header().Set("pIZZA", "Pasta")
 		return rw.Write(safehtml.HTMLEscaped("hello"))
diff --git a/tests/integration/hsts/hsts_test.go b/tests/integration/hsts/hsts_test.go
index 74965eb..1931b3b 100644
--- a/tests/integration/hsts/hsts_test.go
+++ b/tests/integration/hsts/hsts_test.go
@@ -15,12 +15,13 @@
 package hsts_test
 
 import (
-	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/hsts"
@@ -28,7 +29,7 @@ import (
 )
 
 func TestHSTSServeMuxInstall(t *testing.T) {
-	mux := safehttp.NewServeMux(&safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(&safehttp.DefaultDispatcher{})
 
 	mux.Install(hsts.Default())
 	handler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
diff --git a/tests/integration/mux/mux_test.go b/tests/integration/mux/mux_test.go
index 23a19c6..13aa986 100644
--- a/tests/integration/mux/mux_test.go
+++ b/tests/integration/mux/mux_test.go
@@ -15,16 +15,17 @@
 package mux_test
 
 import (
-	"github.com/google/go-cmp/cmp"
-	"github.com/google/go-safeweb/safehttp"
-	"github.com/google/go-safeweb/safehttp/safehttptest"
-	"github.com/google/safehtml"
-	safetemplate "github.com/google/safehtml/template"
 	"html/template"
 	"math"
 	"net/http/httptest"
 	"strings"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+	"github.com/google/safehtml"
+	safetemplate "github.com/google/safehtml/template"
 )
 
 func TestMuxDefaultDispatcher(t *testing.T) {
@@ -37,7 +38,7 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Safe HTML Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
@@ -53,7 +54,7 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Safe HTML Template Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.WriteTemplate(safetemplate.
@@ -71,7 +72,7 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Valid JSON Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					data := struct {
@@ -119,7 +120,7 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 		{
 			name: "Unsafe HTML Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write("<h1>Hello World!</h1>")
@@ -131,7 +132,7 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 		{
 			name: "Unsafe Template Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.WriteTemplate(template.
@@ -145,7 +146,7 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 		{
 			name: "Invalid JSON Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.WriteJSON(math.Inf(1))
diff --git a/tests/integration/safehtml/safehtml_test.go b/tests/integration/safehtml/safehtml_test.go
index f43dd60..8a185c0 100644
--- a/tests/integration/safehtml/safehtml_test.go
+++ b/tests/integration/safehtml/safehtml_test.go
@@ -26,7 +26,7 @@ import (
 )
 
 func TestHandleRequestWrite(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
 		return rw.Write(safehtml.HTMLEscaped("<h1>Escaped, so not really a heading</h1>"))
 	}))
@@ -46,7 +46,7 @@ func TestHandleRequestWrite(t *testing.T) {
 }
 
 func TestHandleRequestWriteTemplate(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	mux.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.")
 	}))
diff --git a/tests/integration/staticheaders/staticheaders_test.go b/tests/integration/staticheaders/staticheaders_test.go
index 0fa6b64..e68a1df 100644
--- a/tests/integration/staticheaders/staticheaders_test.go
+++ b/tests/integration/staticheaders/staticheaders_test.go
@@ -15,12 +15,13 @@
 package staticheaders_test
 
 import (
-	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/go-safeweb/safehttp/plugins/staticheaders"
@@ -28,7 +29,7 @@ import (
 )
 
 func TestServeMuxInstallStaticHeaders(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 
 	mux.Install(staticheaders.Interceptor{})
 	handler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
diff --git a/tests/integration/xsrf/xsrf_test.go b/tests/integration/xsrf/xsrf_test.go
index 6403050..83053ab 100644
--- a/tests/integration/xsrf/xsrf_test.go
+++ b/tests/integration/xsrf/xsrf_test.go
@@ -15,17 +15,18 @@
 package xsrf_test
 
 import (
-	"github.com/google/go-safeweb/safehttp"
-	"github.com/google/go-safeweb/safehttp/plugins/xsrf"
-	"github.com/google/go-safeweb/safehttp/safehttptest"
-	safetemplate "github.com/google/safehtml/template"
 	"net/http/httptest"
 	"strings"
 	"testing"
+
+	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/plugins/xsrf"
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+	"github.com/google/safehtml/template"
 )
 
 func TestServeMuxInstallXSRF(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{}, "foo.com")
+	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
 	it := xsrf.Interceptor{SecretAppKey: "testSecretAppKey"}
 	mux.Install(&it)
 
@@ -40,7 +41,7 @@ func TestServeMuxInstallXSRF(t *testing.T) {
 		// installed, but auto-injection isn't yet supported and has to be done
 		// manually by the handler).
 		token, err = xsrf.Token(r)
-		t := safetemplate.Must(safetemplate.New("name").Funcs(fns).Parse(`<form><input type="hidden" name="token" value="{{XSRFToken}}">{{.}}</form>`))
+		t := template.Must(template.New("name").Funcs(fns).Parse(`<form><input type="hidden" name="token" value="{{XSRFToken}}">{{.}}</form>`))
 
 		return w.WriteTemplate(t, "Content")
 	})