Repository: google/go-safeweb
Pull request #322 (2021-06-16 10:15:36)

Removed ServeMuxConfig.Handle in favor of ServeMux.Handle to be more consistent with net/http
Functional changes are inside safehttp/mux.go.


Patch:
diff --git a/examples/echo/echo.go b/examples/echo/echo.go
index 2b4c203..ebd69fc 100644
--- a/examples/echo/echo.go
+++ b/examples/echo/echo.go
@@ -20,7 +20,9 @@
 package main
 
 import (
+	"fmt"
 	"log"
+	"net/http"
 	"time"
 
 	"github.com/google/safehtml"
@@ -33,15 +35,17 @@ import (
 var start time.Time
 
 func main() {
-	mc := web.NewMuxConfig()
+	port := 8080
+	addr := fmt.Sprintf("localhost:%d", port)
+	m := web.NewMuxConfigDev(port).Mux()
 
-	mc.Handle("/echo", safehttp.MethodGet, safehttp.HandlerFunc(echo))
-	mc.Handle("/uptime", safehttp.MethodGet, safehttp.HandlerFunc(uptime))
+	m.Handle("/echo", safehttp.MethodGet, safehttp.HandlerFunc(echo))
+	m.Handle("/uptime", safehttp.MethodGet, safehttp.HandlerFunc(uptime))
 
 	start = time.Now()
-	log.Println("Visit http://localhost:8080")
-	log.Println("Listening on localhost:8080...")
-	log.Fatal(web.ListenAndServeDev(8080, mc))
+	log.Printf("Visit http://%s\n", addr)
+	log.Printf("Listening on %s...\n", addr)
+	log.Fatal(http.ListenAndServe(addr, m))
 }
 
 func echo(w safehttp.ResponseWriter, req *safehttp.IncomingRequest) safehttp.Result {
diff --git a/examples/echo/security/web/web.go b/examples/echo/security/web/web.go
index 2c34a80..4987d7b 100644
--- a/examples/echo/security/web/web.go
+++ b/examples/echo/security/web/web.go
@@ -21,7 +21,6 @@ package web
 
 import (
 	"fmt"
-	"net/http"
 
 	"github.com/google/go-safeweb/safehttp/plugins/coop"
 	"github.com/google/go-safeweb/safehttp/plugins/cors"
@@ -41,44 +40,51 @@ import (
 //  - Cross-Origin-Opener-Policy
 //  - Content-Security-Policy
 //  - Fetch Metadata
+//  - HSTS
+//  - CORS
+//  - Host checking (against DNS rebinding and request smuggling)
 //
 // Warning: XSRF protection is currently missing due to
 // https://github.com/google/go-safeweb/issues/171.
-func NewMuxConfig() *safehttp.ServeMuxConfig {
+func NewMuxConfig(addr string) *safehttp.ServeMuxConfig {
 	c := &safehttp.ServeMuxConfig{}
 
 	c.Intercept(coop.Default(""))
 	c.Intercept(csp.Default(""))
 	c.Intercept(&fetchmetadata.Interceptor{})
 	c.Intercept(staticheaders.Interceptor{})
+
+	c.Intercept(hsts.Default())
+	c.Intercept(cors.Default())
+	c.Intercept(hostcheck.New(addr))
 	return c
 }
 
-// ListenAndServe starts an HTTP server on the given address.
-// In addition to the security mechanisms applied by NewMuxConfig, it also adds:
+// NewMuxConfigDev returns a ServeMuxConfig with a set of interceptors already
+// installed for security reasons.
+// These include:
 //
-//  - HSTS
-//  - CORS
+//  - Cross-Origin-Opener-Policy
+//  - Content-Security-Policy
+//  - Fetch Metadata
 //  - Host checking (against DNS rebinding and request smuggling)
 //
-// If you need to start an HTTP on a localhost for development purposes, you'll
-// likely need to disable some of these protections. Use ListenAndServeDev
-// instead.
-func ListenAndServe(addr string, mc *safehttp.ServeMuxConfig) error {
-	mc.Intercept(hsts.Default())
-	mc.Intercept(cors.Default())
-	mc.Intercept(hostcheck.New(addr))
-	return http.ListenAndServe(addr, mc.Mux())
-}
-
-// ListenAndServeDev starts a HTTP server on localhost:port for development
-// purposes. Most notably, it doesn't include some of the security mechanisms
-// that ListenAndServe provides.
+// It DOES NOT include HSTS or CORS, as these are difficult to setup in a
+// development environment.
 //
 // Important: the host checking plugin will accept only requests coming to
 // localhost:port, not e.g. 127.0.0.1:port.
-func ListenAndServeDev(port int, mc *safehttp.ServeMuxConfig) error {
+func NewMuxConfigDev(port int) *safehttp.ServeMuxConfig {
+	c := &safehttp.ServeMuxConfig{}
+
+	c.Intercept(coop.Default(""))
+	c.Intercept(csp.Default(""))
+	c.Intercept(&fetchmetadata.Interceptor{})
+	c.Intercept(staticheaders.Interceptor{})
+
 	addr := fmt.Sprintf("localhost:%d", port)
-	mc.Intercept(hostcheck.New(addr))
-	return http.ListenAndServe(addr, mc.Mux())
+	c.Intercept(hostcheck.New(addr))
+	// No HSTS, no CORS
+
+	return c
 }
diff --git a/examples/sample-application/cmd/main.go b/examples/sample-application/cmd/main.go
index 608143d..3489dc1 100644
--- a/examples/sample-application/cmd/main.go
+++ b/examples/sample-application/cmd/main.go
@@ -45,9 +45,9 @@ func main() {
 	db := storage.NewDB()
 
 	addr := net.JoinHostPort("localhost", strconv.Itoa(*port))
-	cfg := secure.NewMuxConfig(db, addr)
-	server.Load(db, cfg)
+	mux := secure.NewMuxConfig(db, addr).Mux()
+	server.Load(db, mux)
 
 	log.Printf("Listening on %q", addr)
-	log.Fatal(http.ListenAndServe(addr, cfg.Mux()))
+	log.Fatal(http.ListenAndServe(addr, mux))
 }
diff --git a/examples/sample-application/server/server.go b/examples/sample-application/server/server.go
index 65de8d8..d247141 100644
--- a/examples/sample-application/server/server.go
+++ b/examples/sample-application/server/server.go
@@ -59,20 +59,20 @@ type serverDeps struct {
 	db *storage.DB
 }
 
-func Load(db *storage.DB, cfg *safehttp.ServeMuxConfig) {
+func Load(db *storage.DB, mux *safehttp.ServeMux) {
 	deps := &serverDeps{
 		db: db,
 	}
 
 	// Private endpoints, only accessible to authenticated users (default).
-	cfg.Handle("/notes/", "GET", getNotesHandler(deps))
-	cfg.Handle("/notes", "POST", postNotesHandler(deps))
-	cfg.Handle("/logout", "POST", logoutHandler(deps))
+	mux.Handle("/notes/", "GET", getNotesHandler(deps))
+	mux.Handle("/notes", "POST", postNotesHandler(deps))
+	mux.Handle("/logout", "POST", logoutHandler(deps))
 
 	// Public enpoints, no auth checks performed.
-	cfg.Handle("/login", "POST", postLoginHandler(deps), auth.Skip{})
-	cfg.Handle("/static/", "GET", safehttp.FileServerEmbed(staticFiles), auth.Skip{})
-	cfg.Handle("/", "GET", indexHandler(deps), auth.Skip{})
+	mux.Handle("/login", "POST", postLoginHandler(deps), auth.Skip{})
+	mux.Handle("/static/", "GET", safehttp.FileServerEmbed(staticFiles), auth.Skip{})
+	mux.Handle("/", "GET", indexHandler(deps), auth.Skip{})
 }
 
 func getNotesHandler(deps *serverDeps) safehttp.Handler {
diff --git a/examples/trustedtypes/server.go b/examples/trustedtypes/server.go
index 5e93c49..867e599 100644
--- a/examples/trustedtypes/server.go
+++ b/examples/trustedtypes/server.go
@@ -37,18 +37,19 @@ func main() {
 	addr := net.JoinHostPort(host, port)
 	var mc safehttp.ServeMuxConfig
 	mc.Intercept(csp.Default(""))
+	mux := mc.Mux()
 
 	safeTmplSrc, _ := template.TrustedSourceFromConstantDir("", template.TrustedSource{}, "safe.html")
 	safeTmpl := template.Must(htmlinject.LoadFiles(nil, htmlinject.LoadConfig{}, safeTmplSrc))
-	mc.Handle("/safe", safehttp.MethodGet, safehttp.HandlerFunc(handleTemplate(safeTmpl)))
+	mux.Handle("/safe", safehttp.MethodGet, safehttp.HandlerFunc(handleTemplate(safeTmpl)))
 
 	unsafeTmplSrc, _ := template.TrustedSourceFromConstantDir("", template.TrustedSource{}, "unsafe.html")
 	unsafeTmpl := template.Must(htmlinject.LoadFiles(nil, htmlinject.LoadConfig{}, unsafeTmplSrc))
-	mc.Handle("/unsafe", safehttp.MethodGet, safehttp.HandlerFunc(handleTemplate(unsafeTmpl)))
+	mux.Handle("/unsafe", safehttp.MethodGet, safehttp.HandlerFunc(handleTemplate(unsafeTmpl)))
 
 	log.Printf("Visit http://%s\n", addr)
 	log.Printf("Listening on %s...\n", addr)
-	log.Fatal(http.ListenAndServe(addr, mc.Mux()))
+	log.Fatal(http.ListenAndServe(addr, mux))
 }
 
 func handleTemplate(tmpl safehttp.Template) func(w safehttp.ResponseWriter, req *safehttp.IncomingRequest) safehttp.Result {
diff --git a/safehttp/migration.go b/safehttp/migration.go
index 6ccd3ca..add0ce0 100644
--- a/safehttp/migration.go
+++ b/safehttp/migration.go
@@ -32,5 +32,10 @@ import (
 // method will return the handler only when given "/foo/" as an argument, not
 // "/foo" nor "/foo/x".
 func RegisteredHandler(mux *ServeMux, pattern string) http.Handler {
-	return mux.handlerMap[pattern]
+	if h, ok := mux.handlers[pattern]; ok {
+		return h
+	}
+	// Keep this. Otherwise mux.handlerMap[pattern] returns a
+	// (*registeredHandler)(nil), which is not equal to an untyped nil.
+	return nil
 }
diff --git a/safehttp/mux.go b/safehttp/mux.go
index e2a4e33..1ba73d9 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -68,24 +68,12 @@ const (
 // Multiple handlers can be registered for a single pattern, as long as they
 // handle different HTTP methods.
 type ServeMux struct {
-	mux        *http.ServeMux
-	handlerMap map[string]http.Handler
-}
+	mux      *http.ServeMux
+	handlers map[string]*registeredHandler
 
-func handlerMap(handlers map[string]map[string]handlerConfig, methodNotAllowed handlerConfig) map[string]http.Handler {
-	m := make(map[string]http.Handler)
-	for pattern, handlersPerMethod := range handlers {
-		handlersPerMethod := handlersPerMethod
-		fn := func(w http.ResponseWriter, r *http.Request) {
-			cfg, ok := handlersPerMethod[r.Method]
-			if !ok {
-				cfg = methodNotAllowed
-			}
-			processRequest(cfg, w, r)
-		}
-		m[pattern] = http.HandlerFunc(fn)
-	}
-	return m
+	dispatcher       Dispatcher
+	interceptors     []Interceptor
+	methodNotAllowed handlerConfig
 }
 
 // ServeHTTP dispatches the request to the handler whose method matches the
@@ -112,12 +100,37 @@ func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	m.mux.ServeHTTP(w, r)
 }
 
+// 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, Mux will panic.
+func (m *ServeMux) Handle(pattern string, method string, h Handler, cfgs ...InterceptorConfig) {
+	if m.handlers[pattern] == nil {
+		m.handlers[pattern] = &registeredHandler{
+			pattern:          pattern,
+			methodNotAllowed: m.methodNotAllowed,
+			methods:          make(map[string]handlerConfig),
+		}
+		m.mux.Handle(pattern, m.handlers[pattern])
+	}
+	m.handlers[pattern].handleMethod(method,
+		handlerConfig{
+			Dispatcher:   m.dispatcher,
+			Handler:      h,
+			Interceptors: configureInterceptors(m.interceptors, cfgs),
+		})
+}
+
 // ServeMuxConfig is a builder for ServeMux.
 type ServeMuxConfig struct {
-	dispatcher              Dispatcher
-	handlers                []handlerRegistration
-	interceptors            []Interceptor
-	methodNotAllowedHandler handlerRegistration
+	dispatcher   Dispatcher
+	interceptors []Interceptor
+
+	methodNotAllowed     Handler
+	methodNotAllowedCfgs []InterceptorConfig
 }
 
 // NewServeMuxConfig crates a ServeMuxConfig with the provided Dispatcher. If
@@ -127,48 +140,21 @@ func NewServeMuxConfig(disp Dispatcher) *ServeMuxConfig {
 		disp = &DefaultDispatcher{}
 	}
 	return &ServeMuxConfig{
-		dispatcher: disp,
-		methodNotAllowedHandler: handlerRegistration{
-			handler: HandlerFunc(defaultMethotNotAllowed),
-		},
+		dispatcher:       disp,
+		methodNotAllowed: HandlerFunc(defaultMethotNotAllowed),
 	}
 }
 
-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, Mux will panic.
-func (s *ServeMuxConfig) Handle(pattern string, method string, h Handler, cfgs ...InterceptorConfig) {
-	s.handlers = append(s.handlers, handlerRegistration{
-		pattern: pattern,
-		method:  method,
-		handler: h,
-		cfgs:    cfgs,
-	})
-}
-
 // HandleMethodNotAllowed registers a handler that runs when a given method is
 // not allowed for a registered path.
 func (s *ServeMuxConfig) HandleMethodNotAllowed(h Handler, cfgs ...InterceptorConfig) {
-	s.methodNotAllowedHandler = handlerRegistration{
-		handler: h,
-		cfgs:    cfgs,
-	}
+	s.methodNotAllowed = h
+	s.methodNotAllowedCfgs = cfgs
 }
 
-func defaultMethotNotAllowed(w ResponseWriter, req *IncomingRequest) Result {
+var defaultMethotNotAllowed = HandlerFunc(func(w ResponseWriter, req *IncomingRequest) Result {
 	return w.WriteError(StatusMethodNotAllowed)
-}
+})
 
 // Intercept installs the given interceptors.
 //
@@ -191,35 +177,54 @@ func (s *ServeMuxConfig) Mux() *ServeMux {
 	if s.dispatcher == nil {
 		panic("Use NewServeMuxConfig instead of creating ServeMuxConfig using a composite literal.")
 	}
-	// pattern -> method -> handlerConfig
-	handlers := map[string]map[string]handlerConfig{}
-	for _, hr := range s.handlers {
-		if handlers[hr.pattern] == nil {
-			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] =
-			handlerConfig{
-				Dispatcher:   s.dispatcher,
-				Handler:      hr.handler,
-				Interceptors: configureInterceptors(s.interceptors, hr.cfgs),
-			}
-	}
+
 	methodNotAllowed := handlerConfig{
 		Dispatcher:   s.dispatcher,
-		Handler:      s.methodNotAllowedHandler.handler,
-		Interceptors: configureInterceptors(s.interceptors, s.methodNotAllowedHandler.cfgs),
+		Handler:      s.methodNotAllowed,
+		Interceptors: configureInterceptors(s.interceptors, s.methodNotAllowedCfgs),
+	}
+
+	m := &ServeMux{
+		mux:              http.NewServeMux(),
+		handlers:         make(map[string]*registeredHandler),
+		dispatcher:       s.dispatcher,
+		interceptors:     s.interceptors,
+		methodNotAllowed: methodNotAllowed,
+	}
+	return m
+}
+
+// Clone creates a copy of the current config.
+// This can be used to create several instances of Mux that share the same set of
+// plugins.
+func (s *ServeMuxConfig) Clone() *ServeMuxConfig {
+	return &ServeMuxConfig{
+		dispatcher:           s.dispatcher,
+		interceptors:         append([]Interceptor(nil), s.interceptors...),
+		methodNotAllowed:     s.methodNotAllowed,
+		methodNotAllowedCfgs: append([]InterceptorConfig(nil), s.methodNotAllowedCfgs...),
+	}
+}
+
+type registeredHandler struct {
+	pattern          string
+	methods          map[string]handlerConfig
+	methodNotAllowed handlerConfig
+}
+
+func (rh *registeredHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	cfg, ok := rh.methods[r.Method]
+	if !ok {
+		cfg = rh.methodNotAllowed
 	}
+	processRequest(cfg, w, r)
+}
 
-	mux := http.NewServeMux()
-	m := handlerMap(handlers, methodNotAllowed)
-	for pattern, handler := range m {
-		mux.Handle(pattern, handler)
+func (rh *registeredHandler) handleMethod(method string, cfg handlerConfig) {
+	if _, exists := rh.methods[method]; exists {
+		panic(fmt.Sprintf("double registration of (pattern = %q, method = %q)", rh.pattern, method))
 	}
-	return &ServeMux{mux: mux, handlerMap: m}
+	rh.methods[method] = cfg
 }
 
 func configureInterceptors(interceptors []Interceptor, cfgs []InterceptorConfig) []configuredInterceptor {
@@ -248,17 +253,3 @@ func configureInterceptors(interceptors []Interceptor, cfgs []InterceptorConfig)
 	}
 	return its
 }
-
-// Clone creates a copy of the current config.
-// This can be used to create several instances of Mux that share the same set of
-// plugins and some common handlers.
-func (s *ServeMuxConfig) Clone() *ServeMuxConfig {
-	c := &ServeMuxConfig{
-		dispatcher:   s.dispatcher,
-		handlers:     make([]handlerRegistration, len(s.handlers)),
-		interceptors: make([]Interceptor, len(s.interceptors)),
-	}
-	copy(c.handlers, s.handlers)
-	copy(c.interceptors, s.interceptors)
-	return c
-}


Test patch:
diff --git a/safehttp/defaults/defaults_test.go b/safehttp/defaults/defaults_test.go
index 5e33267..7589b25 100644
--- a/safehttp/defaults/defaults_test.go
+++ b/safehttp/defaults/defaults_test.go
@@ -28,7 +28,8 @@ func TestServeMuxConfig(t *testing.T) {
 	t.Run("can load in prod mode", func(t *testing.T) {
 		const resp = "response"
 		cfg, _ := ServeMuxConfig([]string{"test.host.example"}, "test-xsrf-key")
-		cfg.Handle("/test", "GET", safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		mux := cfg.Mux()
+		mux.Handle("/test", "GET", safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 			form, err := r.URL.Query()
 			if err != nil {
 				t.Errorf("Cannot parse GET form: %v", err)
@@ -39,7 +40,6 @@ func TestServeMuxConfig(t *testing.T) {
 			}
 			return w.Write(safehtml.HTMLEscaped(resp))
 		}))
-		mux := cfg.Mux()
 		w := httptest.NewRecorder()
 		r := httptest.NewRequest("GET", "https://test.host.example/test?test=true", nil)
 		mux.ServeHTTP(w, r)
diff --git a/safehttp/fileserver_1_16_test.go b/safehttp/fileserver_1_16_test.go
index 0eed006..eb1a10e 100644
--- a/safehttp/fileserver_1_16_test.go
+++ b/safehttp/fileserver_1_16_test.go
@@ -63,8 +63,8 @@ func TestFileServerEmbed(t *testing.T) {
 	}
 
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/", safehttp.MethodGet, safehttp.FileServerEmbed(testEmbeddedFS))
 	m := mb.Mux()
+	m.Handle("/", safehttp.MethodGet, safehttp.FileServerEmbed(testEmbeddedFS))
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
diff --git a/safehttp/fileserver_test.go b/safehttp/fileserver_test.go
index de19439..b3f901e 100644
--- a/safehttp/fileserver_test.go
+++ b/safehttp/fileserver_test.go
@@ -59,8 +59,8 @@ func TestFileServer(t *testing.T) {
 	}
 
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/", safehttp.MethodGet, safehttp.FileServer(tmpDir))
 	m := mb.Mux()
+	m.Handle("/", safehttp.MethodGet, safehttp.FileServer(tmpDir))
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
diff --git a/safehttp/flight_test.go b/safehttp/flight_test.go
index fb8af3f..1bdc354 100644
--- a/safehttp/flight_test.go
+++ b/safehttp/flight_test.go
@@ -65,13 +65,14 @@ func TestFlightInterceptorPanic(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			mb := safehttp.NewServeMuxConfig(nil)
 			mb.Intercept(tc.interceptor)
-			mb.Handle("/search", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+			mux := mb.Mux()
+
+			mux.Handle("/search", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				// IMPORTANT: We are setting the header here and expecting to be
 				// cleared if a panic occurs.
 				w.Header().Set("foo", "bar")
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 			}))
-			mux := mb.Mux()
 
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/search", nil)
 			rw := httptest.NewRecorder()
@@ -99,13 +100,14 @@ func TestFlightInterceptorPanic(t *testing.T) {
 
 func TestFlightHandlerPanic(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/search", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	mux := mb.Mux()
+
+	mux.Handle("/search", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		// IMPORTANT: We are setting the header here and expecting to be
 		// cleared if a panic occurs.
 		w.Header().Set("foo", "bar")
 		panic("handler")
 	}))
-	mux := mb.Mux()
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/search", nil)
 	rw := httptest.NewRecorder()
@@ -137,13 +139,13 @@ func TestFlightDoubleWritePanics(t *testing.T) {
 		for secondWriteName, secondWrite := range writeFuncs {
 			t.Run(fmt.Sprintf("%s->%s", firstWriteName, secondWriteName), func(t *testing.T) {
 				mb := safehttp.NewServeMuxConfig(nil)
-				mb.Handle("/search", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+				mux := mb.Mux()
+				mux.Handle("/search", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					firstWrite(w, r)
 					secondWrite(w, r) // this should panic
 					t.Fatal("should never reach this point")
 					return safehttp.Result{}
 				}))
-				mux := mb.Mux()
 
 				req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/search", nil)
 				rw := httptest.NewRecorder()
diff --git a/safehttp/flightvalues_test.go b/safehttp/flightvalues_test.go
index 8b468b7..3c7f01d 100644
--- a/safehttp/flightvalues_test.go
+++ b/safehttp/flightvalues_test.go
@@ -76,10 +76,10 @@ func handlerInteractingWithTheInterceptor(w safehttp.ResponseWriter, req *safeht
 func TestHandlerInteractingWithInterceptor(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
 	mb.Intercept(&safeHeadersInterceptor{})
+	m := mb.Mux()
 
-	mb.Handle("/safety", safehttp.MethodGet, safehttp.HandlerFunc(handlerInteractingWithTheInterceptor))
+	m.Handle("/safety", safehttp.MethodGet, safehttp.HandlerFunc(handlerInteractingWithTheInterceptor))
 
-	m := mb.Mux()
 	rr := httptest.NewRecorder()
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/safety?level=2", nil)
diff --git a/safehttp/migration_test.go b/safehttp/migration_test.go
index 901a5af..c50d8f9 100644
--- a/safehttp/migration_test.go
+++ b/safehttp/migration_test.go
@@ -27,11 +27,12 @@ import (
 
 func TestRegisteredHandler(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
+	safeMux := mb.Mux()
 
-	mb.Handle("/abc", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	safeMux.Handle("/abc", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("Welcome!"))
 	}))
-	mb.Handle("/abc", safehttp.MethodPost, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	safeMux.Handle("/abc", safehttp.MethodPost, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		f, err := r.PostForm()
 		if err != nil {
 			return w.WriteError(safehttp.StatusBadRequest)
@@ -42,12 +43,10 @@ func TestRegisteredHandler(t *testing.T) {
 		}
 		return w.Write(safehtml.HTMLEscaped(fmt.Sprintf("Added %s.", animal)))
 	}))
-	mb.Handle("/def", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	safeMux.Handle("/def", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("Bye!"))
 	}))
 
-	safeMux := mb.Mux()
-
 	mux := http.NewServeMux()
 	mux.Handle("/abc", safehttp.RegisteredHandler(safeMux, "/abc"))
 	mux.Handle("/def", safehttp.RegisteredHandler(safeMux, "/def"))
@@ -107,16 +106,16 @@ func TestRegisteredHandler(t *testing.T) {
 
 func TestRegisteredHandler_StrictPatterns(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
+	safeMux := mb.Mux()
 
-	mb.Handle("/foo/", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	safeMux.Handle("/foo/", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("Homepage!"))
 	}))
-	safeMux := mb.Mux()
 
 	if safehttp.RegisteredHandler(safeMux, "/foo/") == nil {
 		t.Error(`RegisteredHandler(_, "/foo/") got nil, want non-nil`)
 	}
-	if safehttp.RegisteredHandler(safeMux, "/foo/subpath") != nil {
-		t.Error(`RegisteredHandler(_, "/foo/subpath") got non-nil, want nil`)
+	if got := safehttp.RegisteredHandler(safeMux, "/foo/subpath"); got != nil {
+		t.Errorf(`RegisteredHandler(_, "/foo/subpath") got %v, want nil`, got)
 	}
 }
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index 9a99a0b..8f7d7c2 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -57,15 +57,15 @@ func TestMuxOneHandlerOneRequest(t *testing.T) {
 	for _, tt := range test {
 		t.Run(tt.name, func(t *testing.T) {
 			mb := safehttp.NewServeMuxConfig(nil)
+			mux := mb.Mux()
 
 			h := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 			})
-			mb.Handle("/", safehttp.MethodGet, h)
+			mux.Handle("/", safehttp.MethodGet, h)
 
 			rw := httptest.NewRecorder()
 
-			mux := mb.Mux()
 			mux.ServeHTTP(rw, tt.req)
 
 			if rw.Code != int(tt.wantStatus) {
@@ -119,10 +119,11 @@ func TestMuxServeTwoHandlers(t *testing.T) {
 	}
 
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/bar", safehttp.MethodGet, tests[0].hf)
-	mb.Handle("/bar", safehttp.MethodPost, tests[1].hf)
 	mux := mb.Mux()
 
+	mux.Handle("/bar", safehttp.MethodGet, tests[0].hf)
+	mux.Handle("/bar", safehttp.MethodPost, tests[1].hf)
+
 	for _, test := range tests {
 		rw := httptest.NewRecorder()
 		mux.ServeHTTP(rw, test.req)
@@ -166,9 +167,9 @@ func TestMuxRegisterCorrectHandlerAllPaths(t *testing.T) {
 	}
 
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/get", safehttp.MethodGet, tests[0].hf)
-	mb.Handle("/get2", safehttp.MethodGet, tests[1].hf)
 	mux := mb.Mux()
+	mux.Handle("/get", safehttp.MethodGet, tests[0].hf)
+	mux.Handle("/get2", safehttp.MethodGet, tests[1].hf)
 
 	for _, test := range tests {
 		rw := httptest.NewRecorder()
@@ -182,11 +183,12 @@ func TestMuxRegisterCorrectHandlerAllPaths(t *testing.T) {
 
 func TestMuxHandleSameMethodTwice(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
+	mux := mb.Mux()
 
 	registeredHandler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+	mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
 	defer func() {
 		if r := recover(); r != nil {
@@ -195,8 +197,7 @@ func TestMuxHandleSameMethodTwice(t *testing.T) {
 		t.Errorf(`mux.Handle("/bar", MethodGet, registeredHandler) expected panic`)
 	}()
 
-	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
-	mb.Mux()
+	mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
 }
 
 type setHeaderInterceptor struct {
@@ -293,12 +294,13 @@ func TestMuxInterceptors(t *testing.T) {
 					name:  "Foo",
 					value: "bar",
 				})
+				mux := mb.Mux()
 
 				registeredHandler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
-				return mb.Mux()
+				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				return mux
 			}(),
 			wantStatus: safehttp.StatusOK,
 			wantHeaders: map[string][]string{
@@ -312,13 +314,14 @@ func TestMuxInterceptors(t *testing.T) {
 			mux: func() *safehttp.ServeMux {
 				mb := safehttp.NewServeMuxConfig(nil)
 				mb.Intercept(internalErrorInterceptor{})
+				mux := mb.Mux()
 
 				registeredHandler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mb.Mux()
+				return mux
 			}(),
 			wantStatus: safehttp.StatusInternalServerError,
 			wantHeaders: map[string][]string{
@@ -332,14 +335,15 @@ func TestMuxInterceptors(t *testing.T) {
 			mux: func() *safehttp.ServeMux {
 				mb := safehttp.NewServeMuxConfig(nil)
 				mb.Intercept(&claimHeaderInterceptor{headerToClaim: "Foo"})
+				mux := mb.Mux()
 
 				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>"))
 				})
-				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mb.Mux()
+				return mux
 			}(),
 			wantStatus: safehttp.StatusOK,
 			wantHeaders: map[string][]string{
@@ -353,13 +357,14 @@ func TestMuxInterceptors(t *testing.T) {
 			mux: func() *safehttp.ServeMux {
 				mb := safehttp.NewServeMuxConfig(nil)
 				mb.Intercept(committerInterceptor{})
+				mux := mb.Mux()
 
 				registeredHandler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
-				mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
-				return mb.Mux()
+				return mux
 			}(),
 			wantStatus: safehttp.StatusOK,
 			wantHeaders: map[string][]string{
@@ -373,7 +378,6 @@ func TestMuxInterceptors(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			rw := httptest.NewRecorder()
-
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
 
 			tt.mux.ServeHTTP(rw, req)
@@ -499,17 +503,16 @@ func TestMuxInterceptorConfigs(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			mb := safehttp.NewServeMuxConfig(nil)
 			mb.Intercept(tt.interceptor)
+			mux := mb.Mux()
 
 			registeredHandler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 			})
-			mb.Handle("/bar", safehttp.MethodGet, registeredHandler, tt.config)
+			mux.Handle("/bar", safehttp.MethodGet, registeredHandler, tt.config)
 
 			rw := httptest.NewRecorder()
-
 			req := httptest.NewRequest("GET", "http://foo.com/bar", nil)
 
-			mux := mb.Mux()
 			mux.ServeHTTP(rw, req)
 
 			if rw.Code != int(tt.wantStatus) {
@@ -592,17 +595,16 @@ func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 	mb.Intercept(interceptorOne{})
 	mb.Intercept(interceptorTwo{})
 	mb.Intercept(interceptorThree{})
+	mux := mb.Mux()
 
 	registeredHandler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mb.Handle("/bar", safehttp.MethodGet, registeredHandler)
+	mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
 
 	rw := httptest.NewRecorder()
-
 	req := httptest.NewRequest("GET", "http://foo.com/bar", nil)
 
-	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusOK; rw.Code != int(want) {
@@ -630,12 +632,12 @@ func TestMuxHandlerReturnsNotWritten(t *testing.T) {
 	h := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return safehttp.NotWritten()
 	})
-	mb.Handle("/bar", safehttp.MethodGet, h)
-	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
+	mux := mb.Mux()
+	mux.Handle("/bar", safehttp.MethodGet, h)
 
+	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/bar", nil)
 	rw := httptest.NewRecorder()
 
-	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusNoContent; rw.Code != int(want) {
@@ -651,15 +653,15 @@ func TestMuxHandlerReturnsNotWritten(t *testing.T) {
 
 func TestMuxMethodNotAllowedDefaults(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
+	mux := mb.Mux()
 
 	h := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		panic("not tested")
 	})
-	mb.Handle("/", safehttp.MethodGet, h)
+	mux.Handle("/", safehttp.MethodGet, h)
 
 	rw := httptest.NewRecorder()
 
-	mux := mb.Mux()
 	mux.ServeHTTP(rw, httptest.NewRequest(safehttp.MethodPost, "http://foo.com/", nil))
 
 	if got, want := rw.Code, int(safehttp.StatusMethodNotAllowed); got != want {
@@ -728,17 +730,17 @@ type methodNotAllowedInterceptorConfig struct {
 func TestMuxMethodNotAllowedCustom(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(&methodNotAllowedDispatcher{})
 	mb.Intercept(&methodNotAllowedInterceptor{})
-
-	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-		panic("not tested")
-	}))
 	mb.HandleMethodNotAllowed(safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
 		return rw.WriteError(&methodNotAllowedError{"custom message"})
 	}), methodNotAllowedInterceptorConfig{before: "foo", commit: "bar"})
+	mux := mb.Mux()
+
+	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		panic("not tested")
+	}))
 
 	rw := httptest.NewRecorder()
 
-	mux := mb.Mux()
 	mux.ServeHTTP(rw, httptest.NewRequest(safehttp.MethodPost, "http://foo.com/", nil))
 
 	if got, want := rw.Code, int(safehttp.StatusMethodNotAllowed); got != want {
diff --git a/safehttp/plugins/hostcheck/hostcheck_test.go b/safehttp/plugins/hostcheck/hostcheck_test.go
index dcc753f..04302e4 100644
--- a/safehttp/plugins/hostcheck/hostcheck_test.go
+++ b/safehttp/plugins/hostcheck/hostcheck_test.go
@@ -46,14 +46,14 @@ func TestInterceptor(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			mb := safehttp.NewServeMuxConfig(nil)
 			mb.Intercept(hostcheck.New("foo.com"))
+			mux := mb.Mux()
 
 			h := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 			})
-			mb.Handle("/", safehttp.MethodGet, h)
+			mux.Handle("/", safehttp.MethodGet, h)
 
 			rw := httptest.NewRecorder()
-			mux := mb.Mux()
 			mux.ServeHTTP(rw, tt.req)
 
 			if rw.Code != int(tt.wantStatus) {
diff --git a/safehttp/server_test.go b/safehttp/server_test.go
index be6d9d7..3269d5c 100644
--- a/safehttp/server_test.go
+++ b/safehttp/server_test.go
@@ -30,13 +30,13 @@ func TestServer(t *testing.T) {
 	l := requesttesting.NewFakeListener()
 	defer l.Close()
 
-	cfg := NewServeMuxConfig(nil)
-	cfg.Handle("/", "GET", HandlerFunc(func(w ResponseWriter, r *IncomingRequest) Result {
+	mux := NewServeMuxConfig(nil).Mux()
+	mux.Handle("/", "GET", HandlerFunc(func(w ResponseWriter, r *IncomingRequest) Result {
 		return w.Write(safehtml.HTMLEscaped("response"))
 	}))
 	rtim := 10 * time.Second
 	s := Server{
-		Mux:         cfg.Mux(),
+		Mux:         mux,
 		ReadTimeout: rtim,
 	}
 	go s.Serve(l)
diff --git a/tests/integration/csp/csp_test.go b/tests/integration/csp/csp_test.go
index a7f6987..7d43d97 100644
--- a/tests/integration/csp/csp_test.go
+++ b/tests/integration/csp/csp_test.go
@@ -44,13 +44,13 @@ func TestServeMuxInstallCSP(t *testing.T) {
 
 		return safehttp.ExecuteTemplateWithFuncs(w, t, "Content", fns)
 	})
-	mb.Handle("/bar", safehttp.MethodGet, handler)
+	mux := mb.Mux()
+	mux.Handle("/bar", safehttp.MethodGet, handler)
 
 	rr := httptest.NewRecorder()
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/bar", nil)
 
-	mux := mb.Mux()
 	mux.ServeHTTP(rr, req)
 
 	if err != nil {
diff --git a/tests/integration/devmode/devmode_test.go b/tests/integration/devmode/devmode_test.go
index 9e57046..60b8998 100644
--- a/tests/integration/devmode/devmode_test.go
+++ b/tests/integration/devmode/devmode_test.go
@@ -33,7 +33,8 @@ func TestDevMode(t *testing.T) {
 		}
 		const resp = "response"
 		cfg, _ := defaults.ServeMuxConfig([]string{"test.host.example"}, "test-xsrf-key")
-		cfg.Handle("/test", "GET", safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		mux := cfg.Mux()
+		mux.Handle("/test", "GET", safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 			form, err := r.URL.Query()
 			if err != nil {
 				t.Errorf("Cannot parse GET form: %v", err)
@@ -45,7 +46,6 @@ func TestDevMode(t *testing.T) {
 			w.AddCookie(safehttp.NewCookie("test", "insecure"))
 			return w.Write(safehtml.HTMLEscaped(resp))
 		}))
-		mux := cfg.Mux()
 		w := httptest.NewRecorder()
 		r := httptest.NewRequest("GET", "https://test.host.example/test?test=true", nil)
 		mux.ServeHTTP(w, r)
diff --git a/tests/integration/devmodefreeze/devmodefreeze_test.go b/tests/integration/devmodefreeze/devmodefreeze_test.go
index b6ee15a..b79d8d8 100644
--- a/tests/integration/devmodefreeze/devmodefreeze_test.go
+++ b/tests/integration/devmodefreeze/devmodefreeze_test.go
@@ -29,7 +29,9 @@ func TestDevMode(t *testing.T) {
 	t.Run("can load in prod mode and can't change afterwards", func(t *testing.T) {
 		const resp = "response"
 		cfg, _ := defaults.ServeMuxConfig([]string{"test.host.example"}, "test-xsrf-key")
-		cfg.Handle("/test", "GET", safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		mux := cfg.Mux()
+
+		mux.Handle("/test", "GET", safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 			form, err := r.URL.Query()
 			if err != nil {
 				t.Errorf("Cannot parse GET form: %v", err)
@@ -41,7 +43,6 @@ func TestDevMode(t *testing.T) {
 			w.AddCookie(safehttp.NewCookie("test", "insecure"))
 			return w.Write(safehtml.HTMLEscaped(resp))
 		}))
-		mux := cfg.Mux()
 		w := httptest.NewRecorder()
 		r := httptest.NewRequest("GET", "https://test.host.example/test?test=true", nil)
 		mux.ServeHTTP(w, r)
diff --git a/tests/integration/errors/errors_test.go b/tests/integration/errors/errors_test.go
index cc38f04..96e54a4 100644
--- a/tests/integration/errors/errors_test.go
+++ b/tests/integration/errors/errors_test.go
@@ -55,8 +55,9 @@ func (m myDispatcher) Error(rw http.ResponseWriter, resp safehttp.ErrorResponse)
 // error responses go through the commit phase.
 func TestCustomErrors(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(myDispatcher{})
+	mux := mb.Mux()
 
-	mb.Handle("/compute", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	mux.Handle("/compute", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		qs, err := r.URL.Query()
 		if err != nil {
 			return w.WriteError(safehttp.StatusBadRequest)
@@ -71,13 +72,11 @@ func TestCustomErrors(t *testing.T) {
 		return w.Write(safehtml.HTMLEscaped(fmt.Sprintf("Result: %d", a*a)))
 	}))
 
-	m := mb.Mux()
-
 	t.Run("correct request", func(t *testing.T) {
 		rr := httptest.NewRecorder()
 
 		req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/compute?a=3", nil)
-		m.ServeHTTP(rr, req)
+		mux.ServeHTTP(rr, req)
 
 		if got, want := rr.Code, safehttp.StatusOK; got != int(want) {
 			t.Errorf("rr.Code got: %v want: %v", got, want)
@@ -92,7 +91,7 @@ func TestCustomErrors(t *testing.T) {
 		rr := httptest.NewRecorder()
 
 		req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/compute?foo=3", nil)
-		m.ServeHTTP(rr, req)
+		mux.ServeHTTP(rr, req)
 
 		if got, want := rr.Code, safehttp.StatusBadRequest; got != int(want) {
 			t.Errorf("rr.Code got: %v want: %v", got, want)
@@ -129,17 +128,17 @@ func (interceptor) Match(safehttp.InterceptorConfig) bool {
 func TestCustomErrorsInBefore(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(myDispatcher{})
 	mb.Intercept(interceptor{errBefore: true})
+	mux := mb.Mux()
 
-	mb.Handle("/compute", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+	mux.Handle("/compute", safehttp.MethodGet, safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("the handler code doesn't matter in this test case"))
 	}))
 
-	m := mb.Mux()
 	t.Run("error in Before", func(t *testing.T) {
 		rr := httptest.NewRecorder()
 
 		req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/compute?a=3", nil)
-		m.ServeHTTP(rr, req)
+		mux.ServeHTTP(rr, req)
 
 		if got, want := rr.Code, safehttp.StatusForbidden; got != int(want) {
 			t.Errorf("rr.Code got: %v want: %v", got, want)
diff --git a/tests/integration/header/header_test.go b/tests/integration/header/header_test.go
index 445e4a3..7d2ac46 100644
--- a/tests/integration/header/header_test.go
+++ b/tests/integration/header/header_test.go
@@ -28,7 +28,9 @@ import (
 
 func TestAccessIncomingHeaders(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mux := mb.Mux()
+
+	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)
 		}
@@ -44,22 +46,21 @@ func TestAccessIncomingHeaders(t *testing.T) {
 	}
 
 	rw := httptest.NewRecorder()
-	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 }
 
 func TestChangingResponseHeaders(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mux := mb.Mux()
+
+	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"))
 	}))
 
 	req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/", nil)
-
 	rw := httptest.NewRecorder()
 
-	mux := mb.Mux()
 	mux.ServeHTTP(rw, req)
 
 	want := []string{"Pasta"}
diff --git a/tests/integration/hsts/hsts_test.go b/tests/integration/hsts/hsts_test.go
index 144e26e..9725df5 100644
--- a/tests/integration/hsts/hsts_test.go
+++ b/tests/integration/hsts/hsts_test.go
@@ -27,18 +27,18 @@ import (
 
 func TestHSTSServeMuxInstall(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-
 	mb.Intercept(hsts.Default())
+	mux := mb.Mux()
+
 	handler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mb.Handle("/asdf", safehttp.MethodGet, handler)
+	mux.Handle("/asdf", safehttp.MethodGet, handler)
 
 	rw := httptest.NewRecorder()
-
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
-	mb.Mux().ServeHTTP(rw, req)
+	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusOK; rw.Code != int(want) {
 		t.Errorf("rw.Code got: %v want: %v", rw.Code, want)
@@ -59,18 +59,18 @@ func TestHSTSServeMuxInstall(t *testing.T) {
 
 func TestHSTSOnErrors(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-
 	mb.Intercept(hsts.Default())
+	mux := mb.Mux()
+
 	handler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.WriteError(safehttp.StatusNotFound)
 	})
-	mb.Handle("/asdf", safehttp.MethodGet, handler)
+	mux.Handle("/asdf", safehttp.MethodGet, handler)
 
 	rw := httptest.NewRecorder()
-
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
-	mb.Mux().ServeHTTP(rw, req)
+	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusNotFound; rw.Code != int(want) {
 		t.Errorf("rw.Code got: %v want: %v", rw.Code, want)
diff --git a/tests/integration/mux/mux_test.go b/tests/integration/mux/mux_test.go
index e968a95..c22da76 100644
--- a/tests/integration/mux/mux_test.go
+++ b/tests/integration/mux/mux_test.go
@@ -72,12 +72,14 @@ func TestMuxDefaultDispatcher(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			mb := safehttp.NewServeMuxConfig(nil)
-			mb.Handle("/pizza", safehttp.MethodGet, tt.handler)
+			mux := mb.Mux()
 
-			rw := httptest.NewRecorder()
+			mux.Handle("/pizza", safehttp.MethodGet, tt.handler)
 
+			rw := httptest.NewRecorder()
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/pizza", nil)
-			mb.Mux().ServeHTTP(rw, req)
+
+			mux.ServeHTTP(rw, req)
 
 			if wantStatus := safehttp.StatusOK; rw.Code != int(wantStatus) {
 				t.Errorf("rw.Code: got %v want %v", rw.Code, wantStatus)
@@ -128,11 +130,11 @@ func TestMuxDefaultDispatcherUnsafeResponses(t *testing.T) {
 			t.Skip()
 
 			mb := safehttp.NewServeMuxConfig(nil)
-			mb.Handle("/pizza", safehttp.MethodGet, tt.handler)
 			mux := mb.Mux()
 
-			rw := httptest.NewRecorder()
+			mux.Handle("/pizza", safehttp.MethodGet, tt.handler)
 
+			rw := httptest.NewRecorder()
 			req := httptest.NewRequest(safehttp.MethodGet, "http://foo.com/pizza", nil)
 
 			mux.ServeHTTP(rw, req)
diff --git a/tests/integration/safehtml/safehtml_test.go b/tests/integration/safehtml/safehtml_test.go
index 4a3fcca..1586abe 100644
--- a/tests/integration/safehtml/safehtml_test.go
+++ b/tests/integration/safehtml/safehtml_test.go
@@ -25,7 +25,9 @@ import (
 
 func TestHandleRequestWrite(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mux := mb.Mux()
+
+	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>"))
 	}))
 
@@ -33,7 +35,7 @@ func TestHandleRequestWrite(t *testing.T) {
 
 	rw := httptest.NewRecorder()
 
-	mb.Mux().ServeHTTP(rw, req)
+	mux.ServeHTTP(rw, req)
 
 	body := rw.Body.String()
 
@@ -44,7 +46,9 @@ func TestHandleRequestWrite(t *testing.T) {
 
 func TestHandleRequestWriteTemplate(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-	mb.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
+	mux := mb.Mux()
+
+	mux.Handle("/", safehttp.MethodGet, safehttp.HandlerFunc(func(rw safehttp.ResponseWriter, ir *safehttp.IncomingRequest) safehttp.Result {
 		return safehttp.ExecuteTemplate(rw, template.Must(template.New("name").Parse("<h1>{{ . }}</h1>")), "This is an actual heading, though.")
 	}))
 
@@ -52,7 +56,7 @@ func TestHandleRequestWriteTemplate(t *testing.T) {
 
 	rw := httptest.NewRecorder()
 
-	mb.Mux().ServeHTTP(rw, req)
+	mux.ServeHTTP(rw, req)
 
 	body := rw.Body.String()
 
diff --git a/tests/integration/staticheaders/staticheaders_test.go b/tests/integration/staticheaders/staticheaders_test.go
index 6e5fbd4..77e1095 100644
--- a/tests/integration/staticheaders/staticheaders_test.go
+++ b/tests/integration/staticheaders/staticheaders_test.go
@@ -27,18 +27,19 @@ import (
 
 func TestServeMuxInstallStaticHeaders(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-
 	mb.Intercept(staticheaders.Interceptor{})
+	mux := mb.Mux()
+
 	handler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
-	mb.Handle("/asdf", safehttp.MethodGet, handler)
+	mux.Handle("/asdf", safehttp.MethodGet, handler)
 
 	rw := httptest.NewRecorder()
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
-	mb.Mux().ServeHTTP(rw, req)
+	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusOK; rw.Code != int(want) {
 		t.Errorf("rw.Code got: %v want: %v", rw.Code, want)
@@ -60,18 +61,19 @@ func TestServeMuxInstallStaticHeaders(t *testing.T) {
 
 func TestStaticHeadersOnError(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
-
 	mb.Intercept(staticheaders.Interceptor{})
+	mux := mb.Mux()
+
 	handler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.WriteError(safehttp.StatusNotFound)
 	})
-	mb.Handle("/asdf", safehttp.MethodGet, handler)
+	mux.Handle("/asdf", safehttp.MethodGet, handler)
 
 	rw := httptest.NewRecorder()
 
 	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
 
-	mb.Mux().ServeHTTP(rw, req)
+	mux.ServeHTTP(rw, req)
 
 	if want := safehttp.StatusNotFound; rw.Code != int(want) {
 		t.Errorf("rw.Status() got: %v want: %v", rw.Code, want)
diff --git a/tests/integration/xsrf/xsrf_test.go b/tests/integration/xsrf/xsrf_test.go
index 02177c5..da68ae3 100644
--- a/tests/integration/xsrf/xsrf_test.go
+++ b/tests/integration/xsrf/xsrf_test.go
@@ -28,6 +28,7 @@ func TestServeMuxInstallXSRF(t *testing.T) {
 	mb := safehttp.NewServeMuxConfig(nil)
 	it := xsrfhtml.Interceptor{SecretAppKey: "testSecretAppKey"}
 	mb.Intercept(&it)
+	mux := mb.Mux()
 
 	handler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		fns := map[string]interface{}{
@@ -37,13 +38,13 @@ func TestServeMuxInstallXSRF(t *testing.T) {
 
 		return safehttp.ExecuteTemplateWithFuncs(w, t, "Content", fns)
 	})
-	mb.Handle("/bar", safehttp.MethodGet, handler)
+	mux.Handle("/bar", safehttp.MethodGet, handler)
 
 	rr := httptest.NewRecorder()
 
 	req := httptest.NewRequest(safehttp.MethodGet, "https://foo.com/bar", nil)
 
-	mb.Mux().ServeHTTP(rr, req)
+	mux.ServeHTTP(rr, req)
 
 	if got, want := rr.Code, safehttp.StatusOK; got != int(want) {
 		t.Errorf("rr.Code got: %v want: %v", got, int(want))