Repository: google/go-safeweb
Pull request #191 (2020-10-06 13:34:28)

Introducing ServeMuxBuilder
Related issues: #117, #88.

This is the first draft of a `ServeMuxBuilder`. Right now it's just a thin wrapper that contains virtually no logic.

In the nearest future, `ServeMux` will be simplified and some of the logic will be moved into `ServeMuxBuilder` instead.

kele: > Unrelated question but why rename to `standardHandler`?

My reasoning here was to indicate that this will actually be registered on the `standard` mux.

Patch:
diff --git a/safehttp/mux.go b/safehttp/mux.go
index ff4910f..4b4345a 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -74,9 +74,9 @@ type ServeMux struct {
 	interceps []Interceptor
 }
 
-// NewServeMux allocates and returns a new ServeMux. If the provided dispatcher
+// newServeMux allocates and returns a new ServeMux. If the provided dispatcher
 // is nil, the DefaultDispatcher is used.
-func NewServeMux(d Dispatcher) *ServeMux {
+func newServeMux(d Dispatcher) *ServeMux {
 	if d == nil {
 		d = DefaultDispatcher{}
 	}
@@ -87,15 +87,7 @@ func NewServeMux(d Dispatcher) *ServeMux {
 	}
 }
 
-// Handle registers a handler for the given pattern and method. If another
-// handler is already registered for the same pattern and method, Handle panics.
-//
-// InterceptorConfigs can be passed in order to modify the behavior of the
-// interceptors on a registered handler. Passing an InterceptorConfig whose
-// corresponding Interceptor was not installed will produce no effect. If
-// multiple configurations are passed for the same Interceptor, only the first
-// one will take effect.
-func (m *ServeMux) Handle(pattern string, method string, h Handler, cfgs ...InterceptorConfig) {
+func (m *ServeMux) handle(pattern string, method string, h Handler, cfgs ...InterceptorConfig) {
 	var interceps []ConfiguredInterceptor
 	for _, it := range m.interceps {
 		var cfg InterceptorConfig
@@ -130,10 +122,7 @@ func (m *ServeMux) Handle(pattern string, method string, h Handler, cfgs ...Inte
 	mh.handlers[method] = hi
 }
 
-// Install installs an Interceptor.
-//
-// Interceptors order is undetermined and should not be relied on.
-func (m *ServeMux) Install(i Interceptor) {
+func (m *ServeMux) install(i Interceptor) {
 	m.interceps = append(m.interceps, i)
 }
 
@@ -161,6 +150,58 @@ func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	m.mux.ServeHTTP(w, r)
 }
 
+// ServeMuxBuilder is a builder for ServeMux.
+type ServeMuxBuilder struct {
+	dispatcher   Dispatcher
+	handlers     []handlerRegistration
+	interceptors []Interceptor
+}
+
+type handlerRegistration struct {
+	pattern string
+	method  string
+	handler Handler
+	cfgs    []InterceptorConfig
+}
+
+// Handle registers a handler for the given pattern and method. If a handler is
+// registered twice for the same pattern and method, Build will panic.
+//
+// InterceptorConfigs can be passed in order to modify the behavior of the
+// interceptors on a registered handler. Passing an InterceptorConfig whose
+// corresponding Interceptor was not installed will produce no effect. If
+// multiple configurations are passed for the same Interceptor, only the first
+// one will take effect.
+func (b *ServeMuxBuilder) Handle(pattern string, method string, h Handler, cfgs ...InterceptorConfig) *ServeMuxBuilder {
+	b.handlers = append(b.handlers, handlerRegistration{
+		pattern: pattern,
+		method:  method,
+		handler: h,
+		cfgs:    cfgs,
+	})
+	return b
+}
+
+// Install installs an Interceptor.
+//
+// Interceptors order is undetermined and should not be relied on.
+func (b *ServeMuxBuilder) Install(i Interceptor) *ServeMuxBuilder {
+	b.interceptors = append(b.interceptors, i)
+	return b
+}
+
+// Build builds the ServeMux.
+func (b *ServeMuxBuilder) Build() *ServeMux {
+	m := newServeMux(b.dispatcher)
+	for _, it := range b.interceptors {
+		m.install(it)
+	}
+	for _, hr := range b.handlers {
+		m.handle(hr.pattern, hr.method, hr.handler, hr.cfgs...)
+	}
+	return m
+}
+
 // standardHandler is a collection of handler based on the request method.
 type standardHandler struct {
 	// Maps an HTTP method to its handler


Test patch:
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index b9dbdb4..274ccf0 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -59,16 +59,17 @@ func TestMuxOneHandlerOneRequest(t *testing.T) {
 
 	for _, tt := range test {
 		t.Run(tt.name, func(t *testing.T) {
-			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+			mb := &safehttp.ServeMuxBuilder{}
 
 			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)
+			mb.Handle("/", safehttp.MethodGet, h)
 
 			b := &strings.Builder{}
 			rw := safehttptest.NewTestResponseWriter(b)
 
+			mux := mb.Build()
 			mux.ServeHTTP(rw, tt.req)
 
 			if rw.Status() != tt.wantStatus {
@@ -121,10 +122,10 @@ func TestMuxServeTwoHandlers(t *testing.T) {
 		},
 	}
 
-	d := safehttp.DefaultDispatcher{}
-	mux := safehttp.NewServeMux(d)
-	mux.Handle("/bar", safehttp.MethodGet, tests[0].hf)
-	mux.Handle("/bar", safehttp.MethodPost, tests[1].hf)
+	mb := &safehttp.ServeMuxBuilder{}
+	mb.Handle("/bar", safehttp.MethodGet, tests[0].hf)
+	mb.Handle("/bar", safehttp.MethodPost, tests[1].hf)
+	mux := mb.Build()
 
 	for _, test := range tests {
 		b := &strings.Builder{}
@@ -145,13 +146,12 @@ func TestMuxServeTwoHandlers(t *testing.T) {
 }
 
 func TestMuxHandleSameMethodTwice(t *testing.T) {
-	d := safehttp.DefaultDispatcher{}
-	mux := safehttp.NewServeMux(d)
+	mb := &safehttp.ServeMuxBuilder{}
 
 	registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
 	defer func() {
 		if r := recover(); r == nil {
@@ -159,7 +159,8 @@ func TestMuxHandleSameMethodTwice(t *testing.T) {
 		}
 	}()
 
-	mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+	mb.Build()
 }
 
 type setHeaderInterceptor struct {
@@ -251,8 +252,8 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Install ServeMux Interceptor before handler registration",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-				mux.Install(setHeaderInterceptor{
+				mb := &safehttp.ServeMuxBuilder{}
+				mb.Install(setHeaderInterceptor{
 					name:  "Foo",
 					value: "bar",
 				})
@@ -260,8 +261,8 @@ func TestMuxInterceptors(t *testing.T) {
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
-				return mux
+				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				return mb.Build()
 			}(),
 			wantStatus: safehttp.StatusOK,
 			wantHeaders: map[string][]string{
@@ -273,15 +274,15 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Install Interrupting Interceptor",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-				mux.Install(internalErrorInterceptor{})
+				mb := &safehttp.ServeMuxBuilder{}
+				mb.Install(internalErrorInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mux
+				return mb.Build()
 			}(),
 			wantStatus: safehttp.StatusInternalServerError,
 			wantHeaders: map[string][]string{
@@ -293,16 +294,16 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Handler Communication With ServeMux Interceptor",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-				mux.Install(&claimHeaderInterceptor{headerToClaim: "Foo"})
+				mb := &safehttp.ServeMuxBuilder{}
+				mb.Install(&claimHeaderInterceptor{headerToClaim: "Foo"})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					claimInterceptorSetHeader(w, r, "bar")
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mux
+				return mb.Build()
 			}(),
 			wantStatus: safehttp.StatusOK,
 			wantHeaders: map[string][]string{
@@ -314,15 +315,15 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Panicking interceptor recovered by ServeMux",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-				mux.Install(panickingInterceptor{})
+				mb := &safehttp.ServeMuxBuilder{}
+				mb.Install(panickingInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mux
+				return mb.Build()
 			}(),
 			wantStatus: safehttp.StatusInternalServerError,
 			wantHeaders: map[string][]string{
@@ -334,15 +335,15 @@ func TestMuxInterceptors(t *testing.T) {
 		{
 			name: "Commit phase sets header",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-				mux.Install(committerInterceptor{})
+				mb := &safehttp.ServeMuxBuilder{}
+				mb.Install(committerInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mux
+				return mb.Build()
 			}(),
 			wantStatus: safehttp.StatusOK,
 			wantHeaders: map[string][]string{
@@ -451,19 +452,20 @@ func TestMuxInterceptorConfigs(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-			mux.Install(setHeaderConfigInterceptor{})
+			mb := &safehttp.ServeMuxBuilder{}
+			mb.Install(setHeaderConfigInterceptor{})
 
 			registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 			})
-			mux.Handle("/bar", safehttp.MethodGet, registeredHandler, tt.config)
+			mb.Handle("/bar", safehttp.MethodGet, registeredHandler, tt.config)
 
 			b := &strings.Builder{}
 			rw := safehttptest.NewTestResponseWriter(b)
 
 			req := httptest.NewRequest("GET", "http://foo.com/bar", nil)
 
+			mux := mb.Build()
 			mux.ServeHTTP(rw, req)
 
 			if rw.Status() != tt.wantStatus {
@@ -533,21 +535,22 @@ func (interceptorThree) Commit(w *safehttp.ResponseWriter, r *safehttp.IncomingR
 }
 
 func TestMuxDeterministicInterceptorOrder(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-	mux.Install(interceptorOne{})
-	mux.Install(interceptorTwo{})
-	mux.Install(interceptorThree{})
+	mb := &safehttp.ServeMuxBuilder{}
+	mb.Install(interceptorOne{})
+	mb.Install(interceptorTwo{})
+	mb.Install(interceptorThree{})
 
 	registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
 	b := &strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(b)
 
 	req := httptest.NewRequest("GET", "http://foo.com/bar", nil)
 
+	mux := mb.Build()
 	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusOK; rw.Status() != want {
@@ -571,16 +574,17 @@ func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 }
 
 func TestMuxHandlerReturnsNotWritten(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+	mb := &safehttp.ServeMuxBuilder{}
 	h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return safehttp.NotWritten()
 	})
-	mux.Handle("/bar", safehttp.MethodGet, h)
+	mb.Handle("/bar", safehttp.MethodGet, h)
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
 
 	b := &strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(b)
 
+	mux := mb.Build()
 	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusNoContent; rw.Status() != want {
diff --git a/safehttp/plugins/hostcheck/hostcheck_test.go b/safehttp/plugins/hostcheck/hostcheck_test.go
index ef6657c..c7e7d6d 100644
--- a/safehttp/plugins/hostcheck/hostcheck_test.go
+++ b/safehttp/plugins/hostcheck/hostcheck_test.go
@@ -46,16 +46,17 @@ func TestInterceptor(t *testing.T) {
 
 	for _, tt := range test {
 		t.Run(tt.name, func(t *testing.T) {
-			mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-			mux.Install(hostcheck.New("foo.com"))
+			mb := &safehttp.ServeMuxBuilder{}
+			mb.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)
+			mb.Handle("/", safehttp.MethodGet, h)
 
 			b := &strings.Builder{}
 			rw := safehttptest.NewTestResponseWriter(b)
+			mux := mb.Build()
 			mux.ServeHTTP(rw, tt.req)
 
 			if rw.Status() != tt.wantStatus {
diff --git a/tests/integration/csp/csp_test.go b/tests/integration/csp/csp_test.go
index 1f82e8d..b3a2f29 100644
--- a/tests/integration/csp/csp_test.go
+++ b/tests/integration/csp/csp_test.go
@@ -27,9 +27,9 @@ import (
 )
 
 func TestServeMuxInstallCSP(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+	mb := &safehttp.ServeMuxBuilder{}
 	it := csp.Default("")
-	mux.Install(&it)
+	mb.Install(&it)
 
 	var nonce string
 	var err error
@@ -46,13 +46,14 @@ func TestServeMuxInstallCSP(t *testing.T) {
 
 		return w.WriteTemplate(t, "Content")
 	})
-	mux.Handle("/bar", safehttp.MethodGet, handler)
+	mb.Handle("/bar", safehttp.MethodGet, handler)
 
 	b := strings.Builder{}
 	rr := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/bar", nil)
 
+	mux := mb.Build()
 	mux.ServeHTTP(rr, req)
 
 	if err != nil {
diff --git a/tests/integration/header/header_test.go b/tests/integration/header/header_test.go
index 6842f1d..a3fd571 100644
--- a/tests/integration/header/header_test.go
+++ b/tests/integration/header/header_test.go
@@ -29,8 +29,8 @@ import (
 )
 
 func TestAccessIncomingHeaders(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mb := &safehttp.ServeMuxBuilder{}
+	mb.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)
 		}
@@ -48,12 +48,13 @@ func TestAccessIncomingHeaders(t *testing.T) {
 	b := &strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(b)
 
+	mux := mb.Build()
 	mux.ServeHTTP(rw, req)
 }
 
 func TestChangingResponseHeaders(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mb := &safehttp.ServeMuxBuilder{}
+	mb.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"))
 	}))
@@ -63,6 +64,7 @@ func TestChangingResponseHeaders(t *testing.T) {
 	b := &strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(b)
 
+	mux := mb.Build()
 	mux.ServeHTTP(rw, req)
 
 	want := []string{"Pasta"}
diff --git a/tests/integration/hsts/hsts_test.go b/tests/integration/hsts/hsts_test.go
index 1931b3b..3627f81 100644
--- a/tests/integration/hsts/hsts_test.go
+++ b/tests/integration/hsts/hsts_test.go
@@ -29,20 +29,20 @@ import (
 )
 
 func TestHSTSServeMuxInstall(t *testing.T) {
-	mux := safehttp.NewServeMux(&safehttp.DefaultDispatcher{})
+	mb := &safehttp.ServeMuxBuilder{}
 
-	mux.Install(hsts.Default())
+	mb.Install(hsts.Default())
 	handler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mux.Handle("/asdf", safehttp.MethodGet, handler)
+	mb.Handle("/asdf", safehttp.MethodGet, handler)
 
 	b := strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
-	mux.ServeHTTP(rw, req)
+	mb.Build().ServeHTTP(rw, req)
 
 	if want := safehttp.StatusOK; rw.Status() != want {
 		t.Errorf("rw.Status() got: %v want: %v", rw.Status(), want)
diff --git a/tests/integration/mux/mux_test.go b/tests/integration/mux/mux_test.go
index 13aa986..5686547 100644
--- a/tests/integration/mux/mux_test.go
+++ b/tests/integration/mux/mux_test.go
@@ -38,13 +38,13 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Safe HTML Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+				mb := &safehttp.ServeMuxBuilder{}
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mux.Handle("/pizza", safehttp.MethodGet, h)
-				return mux
+				mb.Handle("/pizza", safehttp.MethodGet, h)
+				return mb.Build()
 			}(),
 			wantHeaders: map[string][]string{
 				"Content-Type": {"text/html; charset=utf-8"},
@@ -54,15 +54,15 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Safe HTML Template Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+				mb := &safehttp.ServeMuxBuilder{}
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.WriteTemplate(safetemplate.
 						Must(safetemplate.New("name").
 							Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 				})
-				mux.Handle("/pizza", safehttp.MethodGet, h)
-				return mux
+				mb.Handle("/pizza", safehttp.MethodGet, h)
+				return mb.Build()
 			}(),
 			wantHeaders: map[string][]string{
 				"Content-Type": {"text/html; charset=utf-8"},
@@ -72,7 +72,7 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 		{
 			name: "Valid JSON Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+				mb := &safehttp.ServeMuxBuilder{}
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					data := struct {
@@ -80,8 +80,8 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 					}{Field: "myField"}
 					return w.WriteJSON(data)
 				})
-				mux.Handle("/pizza", safehttp.MethodGet, h)
-				return mux
+				mb.Handle("/pizza", safehttp.MethodGet, h)
+				return mb.Build()
 			}(),
 			wantHeaders: map[string][]string{
 				"Content-Type": {"application/json; charset=utf-8"},
@@ -120,39 +120,39 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 		{
 			name: "Unsafe HTML Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+				mb := &safehttp.ServeMuxBuilder{}
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write("<h1>Hello World!</h1>")
 				})
-				mux.Handle("/pizza", safehttp.MethodGet, h)
-				return mux
+				mb.Handle("/pizza", safehttp.MethodGet, h)
+				return mb.Build()
 			}(),
 		},
 		{
 			name: "Unsafe Template Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+				mb := &safehttp.ServeMuxBuilder{}
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.WriteTemplate(template.
 						Must(template.New("name").
 							Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 				})
-				mux.Handle("/pizza", safehttp.MethodGet, h)
-				return mux
+				mb.Handle("/pizza", safehttp.MethodGet, h)
+				return mb.Build()
 			}(),
 		},
 		{
 			name: "Invalid JSON Response",
 			mux: func() *safehttp.ServeMux {
-				mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+				mb := &safehttp.ServeMuxBuilder{}
 
 				h := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.WriteJSON(math.Inf(1))
 				})
-				mux.Handle("/pizza", safehttp.MethodGet, h)
-				return mux
+				mb.Handle("/pizza", safehttp.MethodGet, h)
+				return mb.Build()
 			}(),
 		},
 	}
diff --git a/tests/integration/safehtml/safehtml_test.go b/tests/integration/safehtml/safehtml_test.go
index 8a185c0..efa74a2 100644
--- a/tests/integration/safehtml/safehtml_test.go
+++ b/tests/integration/safehtml/safehtml_test.go
@@ -15,19 +15,20 @@
 package safehtml_test
 
 import (
-	"github.com/google/go-safeweb/safehttp/safehttptest"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"github.com/google/go-safeweb/safehttp/safehttptest"
+
 	"github.com/google/go-safeweb/safehttp"
 	"github.com/google/safehtml"
 	"github.com/google/safehtml/template"
 )
 
 func TestHandleRequestWrite(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mb := &safehttp.ServeMuxBuilder{}
+	mb.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>"))
 	}))
 
@@ -36,7 +37,7 @@ func TestHandleRequestWrite(t *testing.T) {
 	b := &strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(b)
 
-	mux.ServeHTTP(rw, req)
+	mb.Build().ServeHTTP(rw, req)
 
 	body := b.String()
 
@@ -46,8 +47,8 @@ func TestHandleRequestWrite(t *testing.T) {
 }
 
 func TestHandleRequestWriteTemplate(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
-	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw *safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mb := &safehttp.ServeMuxBuilder{}
+	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.")
 	}))
 
@@ -56,7 +57,7 @@ func TestHandleRequestWriteTemplate(t *testing.T) {
 	b := &strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(b)
 
-	mux.ServeHTTP(rw, req)
+	mb.Build().ServeHTTP(rw, req)
 
 	body := b.String()
 
diff --git a/tests/integration/staticheaders/staticheaders_test.go b/tests/integration/staticheaders/staticheaders_test.go
index e68a1df..fc083e0 100644
--- a/tests/integration/staticheaders/staticheaders_test.go
+++ b/tests/integration/staticheaders/staticheaders_test.go
@@ -29,20 +29,20 @@ import (
 )
 
 func TestServeMuxInstallStaticHeaders(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+	mb := &safehttp.ServeMuxBuilder{}
 
-	mux.Install(staticheaders.Interceptor{})
+	mb.Install(staticheaders.Interceptor{})
 	handler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mux.Handle("/asdf", safehttp.MethodGet, handler)
+	mb.Handle("/asdf", safehttp.MethodGet, handler)
 
 	b := strings.Builder{}
 	rw := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
-	mux.ServeHTTP(rw, req)
+	mb.Build().ServeHTTP(rw, req)
 
 	if want := safehttp.StatusOK; rw.Status() != want {
 		t.Errorf("rw.Status() got: %v want: %v", rw.Status(), want)
diff --git a/tests/integration/xsrf/xsrf_test.go b/tests/integration/xsrf/xsrf_test.go
index 83053ab..cad2547 100644
--- a/tests/integration/xsrf/xsrf_test.go
+++ b/tests/integration/xsrf/xsrf_test.go
@@ -26,9 +26,9 @@ import (
 )
 
 func TestServeMuxInstallXSRF(t *testing.T) {
-	mux := safehttp.NewServeMux(safehttp.DefaultDispatcher{})
+	mb := &safehttp.ServeMuxBuilder{}
 	it := xsrf.Interceptor{SecretAppKey: "testSecretAppKey"}
-	mux.Install(&it)
+	mb.Install(&it)
 
 	var token string
 	var err error
@@ -45,14 +45,14 @@ func TestServeMuxInstallXSRF(t *testing.T) {
 
 		return w.WriteTemplate(t, "Content")
 	})
-	mux.Handle("/bar", safehttp.MethodGet, handler)
+	mb.Handle("/bar", safehttp.MethodGet, handler)
 
 	b := strings.Builder{}
 	rr := safehttptest.NewTestResponseWriter(&b)
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/bar", nil)
 
-	mux.ServeHTTP(rr, req)
+	mb.Build().ServeHTTP(rr, req)
 
 	if err != nil {
 		t.Fatalf("xsrf.Token: got error %v", err)