Repository: google/go-safeweb (go)

Issue:
Allow using ServeMux as a factory for net/http.Handlers
Currently to use the `safehttp` package handlers together with a `net/http.ServeMux` (e.g. during migration) one needs to pass the whole `ServeMux` as a `net/http.Handler` (possible because `safehttp.ServeMux` satisfies the interface).

Instead, it'd make sense to add a `func (m *ServeMux) Handler(path string) http.Handler` method, mostly for clarity of intent and easy discoverability.

Note: additional benefit of doing this is avoiding double-routing that otherwise happens when `ServeMux` is used as a handler for `net/http.ServeMux`.

Pull request description:
A way to adopt Go Safe Web one handler at a time
Added a `RegisteredHandler` function that retrieves the registered HTTP handler.

This allows for using `safehttp.ServeMux` alongside `net/http.Handler`.

Fixes #245.

Patch:
diff --git a/go.mod b/go.mod
index e72a147..27e48d7 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,9 @@ go 1.16
 require (
 	github.com/google/go-cmp v0.5.0
 	github.com/google/safehtml v0.0.2
-	golang.org/x/net v0.0.0-20210410081132-afb366fc7cd1
-	golang.org/x/tools v0.1.0
+	github.com/yuin/goldmark v1.3.7 // indirect
+	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 // indirect
+	golang.org/x/net v0.0.0-20210510120150-4163338589ed
+	golang.org/x/sys v0.0.0-20210514084401-e8d321eab015 // indirect
+	golang.org/x/tools v0.1.1
 )
diff --git a/go.sum b/go.sum
index b7882fa..c2e3ccf 100644
--- a/go.sum
+++ b/go.sum
@@ -3,18 +3,29 @@ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/safehtml v0.0.2 h1:ZOt2VXg4x24bW0m2jtzAOkhoXV0iM8vNKc0paByCZqM=
 github.com/google/safehtml v0.0.2/go.mod h1:L4KWwDsUJdECRAEpZoBn3O64bQaywRscowZjJAzjHnU=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
+github.com/yuin/goldmark v1.3.7 h1:NSaHgaeJFCtWXCBkBKXw0rhgMuJ0VoE9FB5mWldcrQ4=
+github.com/yuin/goldmark v1.3.7/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 h1:psW17arqaxU48Z5kZ0CQnkZWQJsqcURM6tKiBApRjXI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/mod v0.3.0 h1:RM4zey1++hCTbCVQfnWeKs9/IEsaBLA8vTkd0WVtmH4=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo=
+golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/net v0.0.0-20210410081132-afb366fc7cd1 h1:4qWs8cYYH6PoEFy4dfhDFgoMGkwAcETd+MmPdCPMzUc=
 golang.org/x/net v0.0.0-20210410081132-afb366fc7cd1/go.mod h1:9tjilg8BloeKEkVJvy7fQ90B1CfIiPueXVOjqfkSzI8=
+golang.org/x/net v0.0.0-20210510120150-4163338589ed h1:p9UgmWI9wKpfYmgaV/IZKGdXc5qEK45tDwwwDyjS26I=
+golang.org/x/net v0.0.0-20210510120150-4163338589ed/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -22,6 +33,11 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44 h1:Bli41pIlzTzf3KEY06n+xnzK/BESIg2ze4Pgfh/aI8c=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210514084401-e8d321eab015 h1:hZR0X1kPW+nwyJ9xRxqZk1vx5RUObAPBdKVvXPDUH/E=
+golang.org/x/sys v0.0.0-20210514084401-e8d321eab015/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 h1:v+OssWQX+hTHEmOBgwxdZxK4zHq3yOs8F9J7mk0PY8E=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k=
@@ -32,6 +48,8 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.0 h1:po9/4sTYwZU9lPhi1tOrb4hCv3qrhiQ77LZfGa2OjwY=
 golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
+golang.org/x/tools v0.1.1 h1:wGiQel/hW0NnEkJUk8lbzkX2gFJU6PFxf1v5OlCfuOs=
+golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
diff --git a/safehttp/doc.go b/safehttp/doc.go
index e15294d..706e4b3 100644
--- a/safehttp/doc.go
+++ b/safehttp/doc.go
@@ -123,6 +123,26 @@
 //
 // TODO
 //
+// Incremental Adoption
+//
+// In order to migrate your service using http.Handlers to the safehttp package,
+// we recommend you start doing that one endpoint at a time. Use
+// RegisteredHandler to do this.
+//
+//  safeMuxConfig := /* configured ServeMuxConfig, including interceptors */
+//  safeMuxConfig.Handle("/bar", safehttp.MethodGET, barGETSafeHandler)
+//  safeMuxConfig.Handle("/bar", safehttp.MethodPOST, barPOSTSafeHandler)
+//  safeMuxConfig.Handle("/xyz", safehttp.MethodPOST, xyzSafeHandler)
+//  safeMux := safeMuxConfig.Mux()
+//
+//  // old, not yet migrated
+//  http.Handle("/foo", fooHandler)
+//
+//  // new, migrated
+//  http.Handle("/bar", safehttp.RegisteredHandler(safeMux, "/bar"))
+//  http.Handle("/xyz", safehttp.RegisteredHandler(safeMux, "/xyz"))
+//
+//
 // Restricting Risky APIs
 //
 // TODO
diff --git a/safehttp/migration.go b/safehttp/migration.go
new file mode 100644
index 0000000..6ccd3ca
--- /dev/null
+++ b/safehttp/migration.go
@@ -0,0 +1,36 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package safehttp
+
+import (
+	"net/http"
+)
+
+// RegisteredHandler returns the combined (all request methods) handler
+// registered for a given pattern. Returns nil if the exact pattern wasn't used
+// to register any handlers.
+//
+// This method is helpful for migrating services incrementally, endpoint by
+// endpoint. The handler runs all the installed interceptors and the dispatcher.
+//
+// Important
+//
+// This function does not attempt to do any kind of path matching. If the
+// handler was registered using the ServeMuxConfig for a pattern "/foo/", this
+// 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]
+}
diff --git a/safehttp/mux.go b/safehttp/mux.go
index 10d144d..1fde0fd 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -68,20 +68,24 @@ const (
 // Multiple handlers can be registered for a single pattern, as long as they
 // handle different HTTP methods.
 type ServeMux struct {
-	mux *http.ServeMux
+	mux        *http.ServeMux
+	handlerMap map[string]http.Handler
 }
 
-func registerHandlers(mux *http.ServeMux, handlers map[string]map[string]handlerConfig, methodNotAllowed handlerConfig) {
+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
-		mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
+		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
 }
 
 // ServeHTTP dispatches the request to the handler whose method matches the
@@ -206,9 +210,13 @@ func (s *ServeMuxConfig) Mux() *ServeMux {
 		Handler:      s.methodNotAllowedHandler.handler,
 		Interceptors: configureInterceptors(s.interceptors, s.methodNotAllowedHandler.cfgs),
 	}
-	m := http.NewServeMux()
-	registerHandlers(m, handlers, methodNotAllowed)
-	return &ServeMux{mux: m}
+
+	mux := http.NewServeMux()
+	m := handlerMap(handlers, methodNotAllowed)
+	for pattern, handler := range m {
+		mux.Handle(pattern, handler)
+	}
+	return &ServeMux{mux: mux, handlerMap: m}
 }
 
 func configureInterceptors(interceptors []Interceptor, cfgs []InterceptorConfig) []configuredInterceptor {


Test patch:
diff --git a/safehttp/migration_test.go b/safehttp/migration_test.go
new file mode 100644
index 0000000..901a5af
--- /dev/null
+++ b/safehttp/migration_test.go
@@ -0,0 +1,122 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package safehttp_test
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/safehtml"
+)
+
+func TestRegisteredHandler(t *testing.T) {
+	mb := safehttp.NewServeMuxConfig(nil)
+
+	mb.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 {
+		f, err := r.PostForm()
+		if err != nil {
+			return w.WriteError(safehttp.StatusBadRequest)
+		}
+		animal := f.String("animal", "")
+		if animal == "" {
+			return w.WriteError(safehttp.StatusBadRequest)
+		}
+		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 {
+		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"))
+
+	var tests = []struct {
+		name       string
+		req        *http.Request
+		wantStatus safehttp.StatusCode
+		wantBody   string
+	}{
+		{
+			name:       "Valid GET Request",
+			req:        httptest.NewRequest(safehttp.MethodGet, "http://foo.com/abc", nil),
+			wantStatus: safehttp.StatusOK,
+			wantBody:   "Welcome!",
+		},
+		{
+			name: "Valid POST Request",
+			req: func() *http.Request {
+				req := httptest.NewRequest(safehttp.MethodPost, "http://foo.com/abc", strings.NewReader("animal=cat"))
+				req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+				return req
+			}(),
+			wantStatus: safehttp.StatusOK,
+			wantBody:   "Added cat.",
+		},
+		{
+			name:       "Different handler",
+			req:        httptest.NewRequest(safehttp.MethodGet, "http://foo.com/def", nil),
+			wantStatus: safehttp.StatusOK,
+			wantBody:   "Bye!",
+		},
+		{
+			name:       "Invalid Method",
+			req:        httptest.NewRequest(safehttp.MethodHead, "http://foo.com/abc", nil),
+			wantStatus: safehttp.StatusMethodNotAllowed,
+			wantBody:   "Method Not Allowed\n",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rw := httptest.NewRecorder()
+
+			mux.ServeHTTP(rw, tt.req)
+
+			if rw.Code != int(tt.wantStatus) {
+				t.Errorf("rw.Code: got %v want %v", rw.Code, tt.wantStatus)
+			}
+
+			if got := rw.Body.String(); got != tt.wantBody {
+				t.Errorf("response body: got %q want %q", got, tt.wantBody)
+			}
+		})
+	}
+}
+
+func TestRegisteredHandler_StrictPatterns(t *testing.T) {
+	mb := safehttp.NewServeMuxConfig(nil)
+
+	mb.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`)
+	}
+}