diff --git a/mux.go b/mux.go index 3009e97..bfe528b 100644 --- a/mux.go +++ b/mux.go @@ -2,6 +2,7 @@ package pat import ( + "context" "net/http" "net/url" "strings" @@ -97,15 +98,52 @@ type PatternServeMux struct { // NotFound, if set, is used whenever the request doesn't match any // pattern for its method. NotFound should be set before serving any // requests. - NotFound http.Handler - handlers map[string][]*patHandler + NotFound http.Handler + handlers map[string][]*patHandler + middlewares []MiddlewareFunc } +// MiddlewareFunc is a function which receives an http.Handler and returns +// another http.Handler. +// Typically, the returned handler is a closure which does something with the +// http.ResponseWriter and http.Request passed +// to it, and then calls the handler passed as parameter to the MiddlewareFunc. +type MiddlewareFunc func(http.Handler) http.Handler + +// middleware interface is anything which implements a MiddlewareFunc named +// Middleware. +type middleware interface { + Middleware(handler http.Handler) http.Handler +} + +// Middleware allows MiddlewareFunc to implement the middleware interface. +func (mw MiddlewareFunc) Middleware(handler http.Handler) http.Handler { + return mw(handler) +} + +type contextKey int + +const ( + // RouteKey is inspired by mux and other routers to preserve the matched + // pattern that can be referenced in the lifetime of a handler. + // This is useful for telemetry and instrumentation to use the pattern + // AND not the whole URL, which can result in cardinality explosion. + // For compatibility with most other mux like gorilla, mux, use count=1 + RouteKey contextKey = iota + 1 +) + // New returns a new PatternServeMux. func New() *PatternServeMux { return &PatternServeMux{handlers: make(map[string][]*patHandler)} } +// Use appends a MiddlewareFunc to the chain. Middleware can be used to intercept or otherwise modify requests and/or responses, and are executed in the order that they are applied to the Router. +func (p *PatternServeMux) Use(mwf ...MiddlewareFunc) { + for _, fn := range mwf { + p.middlewares = append(p.middlewares, fn) + } +} + // ServeHTTP matches r.URL.Path against its routing table using the rules // described above. func (p *PatternServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -114,7 +152,17 @@ func (p *PatternServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) { if len(params) > 0 && !ph.redirect { r.URL.RawQuery = url.Values(params).Encode() + "&" + r.URL.RawQuery } - ph.ServeHTTP(w, r) + + // Set the routeKey in context to the current pattern. + ctx := context.WithValue(r.Context(), RouteKey, ph.pat) + + h := ph.Handler + // Build middleware chain if no error was found + for i := len(p.middlewares) - 1; i >= 0; i-- { + h = p.middlewares[i].Middleware(h) + } + + h.ServeHTTP(w, r.WithContext(ctx)) return } } diff --git a/mux_test.go b/mux_test.go index 5ad1f1d..4e65c65 100644 --- a/mux_test.go +++ b/mux_test.go @@ -72,6 +72,15 @@ func TestPatRoutingHit(t *testing.T) { if got, want := r.URL.Query().Get(":name"), "keith"; got != want { t.Errorf("got %q, want %q", got, want) } + + if rk := r.Context().Value(RouteKey); rk != nil { + if rk.(string) != "/foo/:name" { + t.Errorf("routeKey %v does not match /foo/:name", rk) + } + } else { + t.Error("Should've found routeKey /foo/:name") + } + })) p.ServeHTTP(nil, newRequest("GET", "/foo/keith?a=b", nil)) @@ -121,6 +130,14 @@ func TestPatNoParams(t *testing.T) { if r.URL.RawQuery != "" { t.Errorf("RawQuery was %q; should be empty", r.URL.RawQuery) } + + if rk := r.Context().Value(RouteKey); rk != nil { + if rk.(string) != "/foo/" { + t.Errorf("routeKey %v does not match /foo/:name", rk) + } + } else { + t.Error("Should've found routeKey /foo/:name") + } })) p.ServeHTTP(nil, newRequest("GET", "/foo/", nil)) @@ -262,6 +279,40 @@ func TestEscapedUrl(t *testing.T) { } } +func TestMiddleware(t *testing.T) { + mdf := func(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + q.Set("middleware", "passed") + + r.URL.RawQuery = q.Encode() + + h.ServeHTTP(w, r) + }) + } + + p := New() + + var middlewareCalled bool + p.Get("/foo/:name", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Logf("%#v", r.URL.Query()) + if got, want := r.URL.Query().Get("middleware"), "passed"; got != want { + t.Errorf("got %q, want %q", got, want) + } else { + middlewareCalled = true + } + })) + + // use middleware + p.Use(mdf) + + p.ServeHTTP(nil, newRequest("GET", "/foo/bad", nil)) + + if !middlewareCalled { + t.Error("middleware not called") + } +} + func newRequest(method, urlStr string, body io.Reader) *http.Request { req, err := http.NewRequest(method, urlStr, body) if err != nil {