Repository: google/go-safeweb
Pull request #34 (2020-07-30 08:26:17)

Skeleton for safe multiplexer.
`safehttp.ServeMux` is a wrapper around the `http.ServeMux`. Currently, it supports registering based on methods and specifying the list of allowed domains.
Co-authored-by: mattiasgrenfeldt <grenfeldt@google.com>
Co-authored-by: mihalimara22 <maramihali@google.com>

mattiasgrenfeldt: Rebased on top of master, ready to merge

Patch:
diff --git a/safehttp/handler.go b/safehttp/handler.go
index 7fc641f..4a30a24 100644
--- a/safehttp/handler.go
+++ b/safehttp/handler.go
@@ -16,3 +16,16 @@ package safehttp
 
 // HandleFunc TODO
 type HandleFunc func(ResponseWriter, *IncomingRequest) Result
+
+// Handler TODO
+type Handler interface {
+	ServeHTTP(ResponseWriter, *IncomingRequest) Result
+}
+
+// HandlerFunc TODO
+type HandlerFunc func(ResponseWriter, *IncomingRequest) Result
+
+// ServeHTTP calls f(w, r).
+func (f HandlerFunc) ServeHTTP(w ResponseWriter, r *IncomingRequest) Result {
+	return f(w, r)
+}
diff --git a/safehttp/incoming_request.go b/safehttp/incoming_request.go
index 91ea432..5d4dd6a 100644
--- a/safehttp/incoming_request.go
+++ b/safehttp/incoming_request.go
@@ -28,8 +28,8 @@ type IncomingRequest struct {
 	parseOnce sync.Once
 }
 
-func newIncomingRequest(req *http.Request) IncomingRequest {
-	return IncomingRequest{req: req, Header: newHeader(req.Header)}
+func newIncomingRequest(req *http.Request) *IncomingRequest {
+	return &IncomingRequest{req: req, Header: newHeader(req.Header)}
 }
 
 // PostForm parses the form parameters provided in the body of a POST, PATCH or
diff --git a/safehttp/mux.go b/safehttp/mux.go
new file mode 100644
index 0000000..4dbea25
--- /dev/null
+++ b/safehttp/mux.go
@@ -0,0 +1,145 @@
+// 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"
+)
+
+// TODO: add the missing methods
+const (
+	// HTTP GET request
+	MethodGet = "GET"
+	// HTTP Post request
+	MethodPost = "POST"
+)
+
+// ServeMux is an HTTP request multiplexer. It matches the URL of each incoming
+// 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.
+// TODO(@mihalimara22, @mattiasgrenfeldt): add a link to docs/ explaining
+// why this is done.
+//
+// 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
+// both "/images/" and "/images/thumbnails/", the latter handler will be called
+// for paths beginning "/images/thumbnails/" and the former will receive
+// requests for any other paths in the "/images/" subtree.
+//
+// Note that since a pattern ending in a slash names a rooted subtree, the
+// pattern "/" matches all paths not matched by other registered patterns,
+// not just the URL with Path == "/".
+//
+// If a subtree has been registered and a request is received naming the subtree
+// root without its trailing slash, ServeMux redirects that request to
+// the subtree root (adding the trailing slash). This behavior can be overridden
+// with a separate registration for the path without the trailing slash. For
+// example, registering "/images/" causes ServeMux to redirect a request for
+// "/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
+// 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/".
+//
+// ServeMux also takes care of sanitizing the URL request path and the Host
+// header, stripping the port number and redirecting any request containing . or
+// .. elements or repeated slashes to an equivalent, cleaner URL.
+//
+// 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
+
+	// Maps patterns to handlers supporting multiple HTTP methods.
+	handlers map[string]methodHandler
+}
+
+// NewServeMux allocates and returns a new ServeMux
+func NewServeMux(d Dispatcher, domains ...string) *ServeMux {
+	// 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{},
+	}
+}
+
+// Handle registers a handler for the given pattern and method. If another
+// handler is already registered for the same pattern and method, Handle panics.
+func (m *ServeMux) Handle(pattern string, method string, h Handler) {
+	mh, ok := m.handlers[pattern]
+	if !ok {
+		mh := methodHandler{
+			handlers: map[string]Handler{method: h},
+			domains:  m.domains,
+			disp:     m.disp,
+		}
+
+		m.handlers[pattern] = mh
+		m.mux.Handle(pattern, mh)
+		return
+	}
+
+	if _, ok := mh.handlers[method]; ok {
+		panic("method already registered")
+	}
+	mh.handlers[method] = h
+}
+
+// ServeHTTP dispatches the request to the handler whose method matches the
+// incoming request and whose pattern most closely matches the request URL.
+func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	m.mux.ServeHTTP(w, r)
+}
+
+// methodHandler is collection of handlers based on the request method.
+type methodHandler struct {
+	// Maps an HTTP method to its handler
+	handlers map[string]Handler
+	domains  map[string]bool
+	disp     Dispatcher
+}
+
+// ServeHTTP dispatches the request to the handler associated with
+// the incoming request's 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)
+		return
+	}
+
+	h.ServeHTTP(newResponseWriter(m.disp, w), newIncomingRequest(r))
+}
diff --git a/safehttp/prototype.go b/safehttp/prototype.go
index 9542a6e..fc22c19 100644
--- a/safehttp/prototype.go
+++ b/safehttp/prototype.go
@@ -33,5 +33,5 @@ func NewMachinery(h HandleFunc, d Dispatcher) *Machinery {
 func (m *Machinery) HandleRequest(w http.ResponseWriter, req *http.Request) {
 	rw := newResponseWriter(m.d, w)
 	ir := newIncomingRequest(req)
-	m.h(rw, &ir)
+	m.h(rw, ir)
 }


Test patch:
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
new file mode 100644
index 0000000..d80dffc
--- /dev/null
+++ b/safehttp/mux_test.go
@@ -0,0 +1,208 @@
+// 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 (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/safehtml"
+	"github.com/google/safehtml/template"
+)
+
+type testDispatcher struct{}
+
+func (testDispatcher) Write(rw http.ResponseWriter, resp Response) error {
+	switch x := resp.(type) {
+	case safehtml.HTML:
+		_, err := rw.Write([]byte(x.String()))
+		return err
+	default:
+		panic("not a safe response type")
+	}
+}
+
+func (testDispatcher) ExecuteTemplate(rw http.ResponseWriter, t Template, data interface{}) error {
+	switch x := t.(type) {
+	case *template.Template:
+		return x.Execute(rw, data)
+	default:
+		panic("not a safe response type")
+	}
+}
+
+type responseRecorder struct {
+	header http.Header
+	writer io.Writer
+	status int
+}
+
+func newResponseRecorder(w io.Writer) *responseRecorder {
+	return &responseRecorder{header: http.Header{}, writer: w, status: http.StatusOK}
+}
+
+func (r *responseRecorder) Header() http.Header {
+	return r.header
+}
+
+func (r *responseRecorder) WriteHeader(statusCode int) {
+	r.status = statusCode
+}
+
+func (r *responseRecorder) Write(data []byte) (int, error) {
+	return r.writer.Write(data)
+}
+
+func TestMuxOneHandlerOneRequest(t *testing.T) {
+	var test = []struct {
+		name       string
+		req        *http.Request
+		wantStatus int
+		wantHeader map[string][]string
+		wantBody   string
+	}{
+		{
+			name:       "Valid Request",
+			req:        httptest.NewRequest("GET", "http://foo.com/", nil),
+			wantStatus: 200,
+			wantHeader: map[string][]string{},
+			wantBody:   "&lt;h1&gt;Hello World!&lt;/h1&gt;",
+		},
+		{
+			name:       "Invalid Host",
+			req:        httptest.NewRequest("GET", "http://bar.com/", nil),
+			wantStatus: 404,
+			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("POST", "http://foo.com/", nil),
+			wantStatus: 405,
+			wantHeader: map[string][]string{
+				"Content-Type":           {"text/plain; charset=utf-8"},
+				"X-Content-Type-Options": {"nosniff"},
+			},
+			wantBody: "Method Not Allowed\n",
+		},
+	}
+
+	for _, tt := range test {
+		t.Run(tt.name, func(t *testing.T) {
+			mux := NewServeMux(testDispatcher{}, "foo.com")
+
+			h := HandlerFunc(func(w ResponseWriter, r *IncomingRequest) Result {
+				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
+			})
+			mux.Handle("/", MethodGet, h)
+
+			b := &strings.Builder{}
+			rw := newResponseRecorder(b)
+
+			mux.ServeHTTP(rw, tt.req)
+
+			if rw.status != tt.wantStatus {
+				t.Errorf("rw.status: got %v want %v", rw.status, tt.wantStatus)
+			}
+
+			if diff := cmp.Diff(tt.wantHeader, map[string][]string(rw.header)); diff != "" {
+				t.Errorf("rw.header mismatch (-want +got):\n%s", diff)
+			}
+
+			if got := b.String(); got != tt.wantBody {
+				t.Errorf("response body: got %q want %q", got, tt.wantBody)
+			}
+		})
+	}
+}
+
+func TestMuxServeTwoHandlers(t *testing.T) {
+	var tests = []struct {
+		name        string
+		req         *http.Request
+		hf          HandlerFunc
+		wantStatus  int
+		wantHeaders map[string][]string
+		wantBody    string
+	}{
+		{
+			name: "GET Handler",
+			req:  httptest.NewRequest("GET", "http://foo.com/bar", nil),
+			hf: HandlerFunc(func(w ResponseWriter, r *IncomingRequest) Result {
+				return w.Write(safehtml.HTMLEscaped("<h1>Hello World! GET</h1>"))
+			}),
+			wantStatus:  200,
+			wantHeaders: map[string][]string{},
+			wantBody:    "&lt;h1&gt;Hello World! GET&lt;/h1&gt;",
+		},
+		{
+			name: "POST Handler",
+			req:  httptest.NewRequest("POST", "http://foo.com/bar", nil),
+			hf: HandlerFunc(func(w ResponseWriter, r *IncomingRequest) Result {
+				return w.Write(safehtml.HTMLEscaped("<h1>Hello World! POST</h1>"))
+			}),
+			wantStatus:  200,
+			wantHeaders: map[string][]string{},
+			wantBody:    "&lt;h1&gt;Hello World! POST&lt;/h1&gt;",
+		},
+	}
+
+	d := testDispatcher{}
+	mux := NewServeMux(d, "foo.com")
+	mux.Handle("/bar", MethodGet, tests[0].hf)
+	mux.Handle("/bar", MethodPost, tests[1].hf)
+
+	for _, test := range tests {
+		b := &strings.Builder{}
+		rw := newResponseRecorder(b)
+		mux.ServeHTTP(rw, test.req)
+		if want := test.wantStatus; rw.status != want {
+			t.Errorf("rw.status: got %v want %v", rw.status, want)
+		}
+
+		if diff := cmp.Diff(test.wantHeaders, map[string][]string(rw.header)); diff != "" {
+			t.Errorf("rw.header mismatch (-want +got):\n%s", diff)
+		}
+
+		if got, want := b.String(), test.wantBody; got != want {
+			t.Errorf("response body: got %q want %q", got, want)
+		}
+	}
+}
+
+func TestMuxHandleSameMethodTwice(t *testing.T) {
+	d := testDispatcher{}
+	mux := NewServeMux(d, "foo.com")
+
+	registeredHandler := HandlerFunc(func(w ResponseWriter, r *IncomingRequest) Result {
+		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
+	})
+	mux.Handle("/bar", MethodGet, registeredHandler)
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf(`mux.Handle("/bar", MethodGet, registeredHandler) expected panic`)
+		}
+	}()
+
+	mux.Handle("/bar", MethodGet, registeredHandler)
+}