Middleware (handler wrapping)

mux.go

@@ -48,6 +48,8 @@ type Router struct {
 	strictSlash bool
 	// If true, do not clear the request context after handling the request
 	KeepContext bool
+	// Middleware layers to be wrapped around matched handler
+	layers []Middleware
 }
 
 // Match matches registered routes against the request.
@@ -87,6 +89,12 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	if !r.KeepContext {
 		defer context.Clear(req)
 	}
+
+	for i := len(r.layers) - 1; i >= 0; i-- {
+		layer := r.layers[i]
+		handler = layer(handler)
+	}
+
 	handler.ServeHTTP(w, req)
 }
 
@@ -114,6 +122,21 @@ func (r *Router) StrictSlash(value bool) *Router {
 	return r
 }
 
+// ----------------------------------------------------------------------------
+// Middleware
+// ----------------------------------------------------------------------------
+
+// Middleware wraps an http.Handler to provide additional functionality
+type Middleware func(http.Handler) http.Handler
+
+// Use registers middleware layers to be used
+func (r *Router) Use(layers ...Middleware) *Router {
+	for _, layer := range layers {
+		r.layers = append(r.layers, layer)
+	}
+	return r
+}
+
 // ----------------------------------------------------------------------------
 // parentRoute
 // ----------------------------------------------------------------------------

mux_test.go

@@ -718,6 +718,39 @@ func TestSubrouterHeader(t *testing.T) {
 	}
 }
 
+func TestUse(t *testing.T) {
+	expected := "foobarbaz"
+
+	mw1 := func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			fmt.Fprintf(w, "foo")
+			next.ServeHTTP(w, req)
+		})
+	}
+
+	mw2 := func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			fmt.Fprintf(w, "bar")
+			next.ServeHTTP(w, req)
+		})
+	}
+
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		fmt.Fprintf(w, "baz")
+	}
+
+	r := NewRouter().Use(mw1, mw2)
+	r.HandleFunc("/", handler)
+
+	req, _ := http.NewRequest("GET", "http://localhost/", nil)
+	resp := NewRecorder()
+
+	r.ServeHTTP(resp, req)
+	if resp.Body.String() != expected {
+		t.Errorf("Expecting %q got %q", expected, resp.Body.String())
+	}
+}
+
 // mapToPairs converts a string map to a slice of string pairs
 func mapToPairs(m map[string]string) []string {
 	var i int