Repository: google/go-safeweb (go)

Issue:
Get rid of Interceptor keys in ServeMux, since they aren't needed.
Currently `Interceptors` are installed and stored using unique keys in the `ServeMux`. We have decided these keys aren't needed and should be removed. If interaction is needed between a `Handler` and an `Interceptor` then the interceptors needs to implement package level functions that take `ResponseWriter`s and `IncomingRequest`s and modify the context when state needs to be passed.

Pull request:
Removed Interceptor keys from ServeMux, now interceptors and their configs are stored in a slice.
Now Interceptors are stored together with their configs in slices instead of maps. To interact with an interceptor from the handler that interceptor needs to implement package level functions which take ResponseWriters and IncomingRequests and do things. State can be stored in the context of the request.

Fixes #143 

Patch:
diff --git a/safehttp/mux.go b/safehttp/mux.go
index 95e0ffd..d91a855 100644
--- a/safehttp/mux.go
+++ b/safehttp/mux.go
@@ -77,11 +77,8 @@ type ServeMux struct {
 	disp    Dispatcher
 
 	// Maps patterns to handlers supporting multiple HTTP methods.
-	handlers map[string]methodHandler
-
-	// Maps interceptor key to interceptor.
-	interceps      map[string]Interceptor
-	intercepsOrder []string
+	handlers  map[string]methodHandler
+	interceps []Interceptor
 }
 
 // NewServeMux allocates and returns a new ServeMux
@@ -92,35 +89,41 @@ func NewServeMux(d Dispatcher, domains ...string) *ServeMux {
 		dm[host] = true
 	}
 	return &ServeMux{
-		mux:       http.NewServeMux(),
-		domains:   dm,
-		disp:      d,
-		handlers:  map[string]methodHandler{},
-		interceps: map[string]Interceptor{},
+		mux:      http.NewServeMux(),
+		domains:  dm,
+		disp:     d,
+		handlers: map[string]methodHandler{},
 	}
 }
 
+type appliedInterceptor struct {
+	it  Interceptor
+	cfg Config
+}
+
 // Handle registers a handler for the given pattern and method. If another
 // handler is already registered for the same pattern and method, Handle panics.
 //
 // Configs can be optionally passed in order to modify the behavior of the
 // interceptors on a registered handler. Passing a Config whose corresponding
-// Interceptor was not installed will produce no effect.
+// Interceptor was not installed will produce no effect. If multiple Configs are
+// passed for the same Interceptor, only the first one will take effect.
 func (m *ServeMux) Handle(pattern string, method string, h Handler, cfgs ...Config) {
-	cfgMap := map[string]Config{}
-	for _, cfg := range cfgs {
-		for k, it := range m.interceps {
-			if cfg.Match(it) {
-				cfgMap[k] = cfg
+	var interceps []appliedInterceptor
+	for _, it := range m.interceps {
+		var cfg Config
+		for _, c := range cfgs {
+			if c.Match(it) {
+				cfg = c
+				break
 			}
 		}
+		interceps = append(interceps, appliedInterceptor{it: it, cfg: cfg})
 	}
 	hi := handlerWithInterceptors{
-		handler:        h,
-		interceps:      m.interceps,
-		disp:           m.disp,
-		configs:        cfgMap,
-		intercepsOrder: m.intercepsOrder,
+		handler:   h,
+		interceps: interceps,
+		disp:      m.disp,
 	}
 
 	mh, ok := m.handlers[pattern]
@@ -141,18 +144,9 @@ func (m *ServeMux) Handle(pattern string, method string, h Handler, cfgs ...Conf
 	mh.handlers[method] = hi
 }
 
-// Install installs an Interceptor. Interceptor keys need to be unique. If an
-// Interceptor with the same key has already been installed, Install panics.
-//
-// TODO(@empijei, @grenfeldt, @kele, @mihalimara22): Right now you could install
-// the same interceptor twice with different keys, we need to figure out how
-// exactly we want to avoid that and how we define key uniqueness.
-func (m *ServeMux) Install(key string, i Interceptor) {
-	if _, exists := m.interceps[key]; exists {
-		panic("interceptor with same key already installed")
-	}
-	m.interceps[key] = i
-	m.intercepsOrder = append(m.intercepsOrder, key)
+// Install installs an Interceptor.
+func (m *ServeMux) Install(i Interceptor) {
+	m.interceps = append(m.interceps, i)
 }
 
 // ServeHTTP dispatches the request to the handler whose method matches the
@@ -188,17 +182,15 @@ func (m methodHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 // handlerWithInterceptors encapsulates a handler and its corresponding
 // interceptors.
 type handlerWithInterceptors struct {
-	handler        Handler
-	configs        map[string]Config
-	interceps      map[string]Interceptor
-	intercepsOrder []string
-	disp           Dispatcher
+	handler   Handler
+	interceps []appliedInterceptor
+	disp      Dispatcher
 }
 
 // ServeHTTP calls the Before method of all the interceptors and then calls the
 // underlying handler.
 func (h handlerWithInterceptors) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	rw := NewResponseWriter(h.disp, w, h.interceps)
+	rw := NewResponseWriter(h.disp, w)
 	ir := NewIncomingRequest(r)
 
 	// The `net/http` package recovers handler panics, but we cannot rely on that behavior here.
@@ -210,8 +202,8 @@ func (h handlerWithInterceptors) ServeHTTP(w http.ResponseWriter, r *http.Reques
 		}
 	}()
 
-	for _, key := range h.intercepsOrder {
-		h.interceps[key].Before(rw, ir, h.configs[key])
+	for _, interceptor := range h.interceps {
+		interceptor.it.Before(rw, ir, interceptor.cfg)
 		if rw.written {
 			return
 		}
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 21fc577..e8fc9ca 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -25,32 +25,21 @@ type ResponseWriter struct {
 
 	// Having this field unexported is essential for security. Otherwise one can
 	// easily overwrite the struct bypassing all our safety guarantees.
-	header       Header
-	muxInterceps map[string]Interceptor
-	written      bool
+	header  Header
+	written bool
 }
 
 // NewResponseWriter creates a ResponseWriter from a safehttp.Dispatcher, an
 // http.ResponseWriter and a list of interceptors associated with a ServeMux.
-func NewResponseWriter(d Dispatcher, rw http.ResponseWriter, muxInterceps map[string]Interceptor) *ResponseWriter {
+func NewResponseWriter(d Dispatcher, rw http.ResponseWriter) *ResponseWriter {
 	header := newHeader(rw.Header())
 	return &ResponseWriter{
-		d:            d,
-		rw:           rw,
-		header:       header,
-		muxInterceps: muxInterceps,
+		d:      d,
+		rw:     rw,
+		header: header,
 	}
 }
 
-// Interceptor returns the interceptor associated with the given key.
-func (w *ResponseWriter) Interceptor(key string) Interceptor {
-	mp, ok := w.muxInterceps[key]
-	if !ok {
-		return nil
-	}
-	return mp
-}
-
 // Result TODO
 type Result struct{}
 


Test patch:
diff --git a/safehttp/mux_test.go b/safehttp/mux_test.go
index baa82f9..2cf7159 100644
--- a/safehttp/mux_test.go
+++ b/safehttp/mux_test.go
@@ -15,6 +15,7 @@
 package safehttp_test
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -181,16 +182,19 @@ func (internalErrorInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.I
 
 type claimHeaderInterceptor struct {
 	headerToClaim string
-	setValue      func([]string)
 }
 
-func (p *claimHeaderInterceptor) Before(w *safehttp.ResponseWriter, _ *safehttp.IncomingRequest, cfg interface{}) safehttp.Result {
-	p.setValue = w.Header().Claim(p.headerToClaim)
+type claimCtxKey struct{}
+
+func (p *claimHeaderInterceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, cfg interface{}) safehttp.Result {
+	f := w.Header().Claim(p.headerToClaim)
+	r.SetContext(context.WithValue(r.Context(), claimCtxKey{}, f))
 	return safehttp.NotWritten()
 }
 
-func (p *claimHeaderInterceptor) SetHeader(value string) {
-	p.setValue([]string{value})
+func claimInterceptorSetHeader(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, value string) {
+	f := r.Context().Value(claimCtxKey{}).(func([]string))
+	f([]string{value})
 }
 
 type panickingInterceptor struct{}
@@ -211,7 +215,7 @@ func TestMuxInterceptors(t *testing.T) {
 			name: "Install ServeMux Interceptor before handler registration",
 			mux: func() *safehttp.ServeMux {
 				mux := safehttp.NewServeMux(testDispatcher{}, "foo.com")
-				mux.Install("test", setHeaderInterceptor{
+				mux.Install(setHeaderInterceptor{
 					name:  "Foo",
 					value: "bar",
 				})
@@ -230,7 +234,7 @@ func TestMuxInterceptors(t *testing.T) {
 			name: "Install Interrupting Interceptor",
 			mux: func() *safehttp.ServeMux {
 				mux := safehttp.NewServeMux(testDispatcher{}, "foo.com")
-				mux.Install("test", internalErrorInterceptor{})
+				mux.Install(internalErrorInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
@@ -250,11 +254,10 @@ func TestMuxInterceptors(t *testing.T) {
 			name: "Handler Communication With ServeMux Interceptor",
 			mux: func() *safehttp.ServeMux {
 				mux := safehttp.NewServeMux(testDispatcher{}, "foo.com")
-				mux.Install("claim", &claimHeaderInterceptor{headerToClaim: "Foo"})
+				mux.Install(&claimHeaderInterceptor{headerToClaim: "Foo"})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
-					p := w.Interceptor("claim").(*claimHeaderInterceptor)
-					p.SetHeader("bar")
+					claimInterceptorSetHeader(w, r, "bar")
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 				})
 				mux.Handle("/bar", safehttp.MethodGet, registeredHandler)
@@ -269,7 +272,7 @@ func TestMuxInterceptors(t *testing.T) {
 			name: "Panicking interceptor recovered by ServeMux",
 			mux: func() *safehttp.ServeMux {
 				mux := safehttp.NewServeMux(testDispatcher{}, "foo.com")
-				mux.Install("panic", panickingInterceptor{})
+				mux.Install(panickingInterceptor{})
 
 				registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 					return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
@@ -367,7 +370,7 @@ func TestMuxInterceptorConfigs(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			mux := safehttp.NewServeMux(testDispatcher{}, "foo.com")
-			mux.Install("setHeader", setHeaderConfigInterceptor{})
+			mux.Install(setHeaderConfigInterceptor{})
 
 			registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 				return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
@@ -427,9 +430,9 @@ func (interceptorThree) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingR
 
 func TestMuxDeterministicInterceptorOrder(t *testing.T) {
 	mux := safehttp.NewServeMux(testDispatcher{}, "foo.com")
-	mux.Install("one", interceptorOne{})
-	mux.Install("two", interceptorTwo{})
-	mux.Install("three", interceptorThree{})
+	mux.Install(interceptorOne{})
+	mux.Install(interceptorTwo{})
+	mux.Install(interceptorThree{})
 
 	registeredHandler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
diff --git a/safehttp/response_writer_test.go b/safehttp/response_writer_test.go
index d8ee1cc..9a0cc11 100644
--- a/safehttp/response_writer_test.go
+++ b/safehttp/response_writer_test.go
@@ -27,7 +27,7 @@ import (
 
 func TestResponseWriterSetCookie(t *testing.T) {
 	rr := newResponseRecorder(&strings.Builder{})
-	rw := safehttp.NewResponseWriter(testDispatcher{}, rr, nil)
+	rw := safehttp.NewResponseWriter(testDispatcher{}, rr)
 
 	c := safehttp.NewCookie("foo", "bar")
 	err := rw.SetCookie(c)
@@ -45,7 +45,7 @@ func TestResponseWriterSetCookie(t *testing.T) {
 
 func TestResponseWriterSetInvalidCookie(t *testing.T) {
 	rr := newResponseRecorder(&strings.Builder{})
-	rw := safehttp.NewResponseWriter(testDispatcher{}, rr, nil)
+	rw := safehttp.NewResponseWriter(testDispatcher{}, rr)
 
 	c := safehttp.NewCookie("f=oo", "bar")
 	err := rw.SetCookie(c)
@@ -99,7 +99,7 @@ func TestResponseWriterWriteTwicePanic(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			w := safehttp.NewResponseWriter(testDispatcher{}, newResponseRecorder(&strings.Builder{}), nil)
+			w := safehttp.NewResponseWriter(testDispatcher{}, newResponseRecorder(&strings.Builder{}))
 			defer func() {
 				if r := recover(); r == nil {
 					t.Errorf("tt.write(w) expected panic")
diff --git a/safehttp/safehttptest/recorder.go b/safehttp/safehttptest/recorder.go
index 894dc99..2941b95 100644
--- a/safehttp/safehttptest/recorder.go
+++ b/safehttp/safehttptest/recorder.go
@@ -61,7 +61,7 @@ func NewResponseRecorder() *ResponseRecorder {
 	return &ResponseRecorder{
 		rw:             rw,
 		b:              &b,
-		ResponseWriter: safehttp.NewResponseWriter(testDispatcher{}, rw, nil),
+		ResponseWriter: safehttp.NewResponseWriter(testDispatcher{}, rw),
 	}
 }
 
@@ -73,7 +73,7 @@ func NewResponseRecorderFromDispatcher(d safehttp.Dispatcher) *ResponseRecorder
 	return &ResponseRecorder{
 		rw:             rw,
 		b:              &b,
-		ResponseWriter: safehttp.NewResponseWriter(d, rw, nil),
+		ResponseWriter: safehttp.NewResponseWriter(d, rw),
 	}
 }
 
diff --git a/tests/integration/hsts/hsts_test.go b/tests/integration/hsts/hsts_test.go
index 0acc65a..1eee741 100644
--- a/tests/integration/hsts/hsts_test.go
+++ b/tests/integration/hsts/hsts_test.go
@@ -68,7 +68,7 @@ func (r *responseRecorder) Write(data []byte) (int, error) {
 func TestHSTSServeMuxInstall(t *testing.T) {
 	mux := safehttp.NewServeMux(&dispatcher{}, "foo.com")
 
-	mux.Install("hsts", hsts.Default())
+	mux.Install(hsts.Default())
 	handler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})
diff --git a/tests/integration/staticheaders/staticheaders_test.go b/tests/integration/staticheaders/staticheaders_test.go
index 03a299d..5b83a9b 100644
--- a/tests/integration/staticheaders/staticheaders_test.go
+++ b/tests/integration/staticheaders/staticheaders_test.go
@@ -74,7 +74,7 @@ func (r *responseRecorder) Write(data []byte) (int, error) {
 func TestServeMuxInstallStaticHeaders(t *testing.T) {
 	mux := safehttp.NewServeMux(dispatcher{}, "foo.com")
 
-	mux.Install("staticheaders", staticheaders.Plugin{})
+	mux.Install(staticheaders.Plugin{})
 	handler := safehttp.HandlerFunc(func(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
 	})