Feature router in request context

mux.go

@@ -94,6 +94,9 @@ type routeConf struct {
 	// If true, the http.Request context will not contain the Route.
 	omitRouteFromContext bool
 
+	// if true, the the http.Request context will not contain the router
+	omitRouterFromContext bool
+
 	// Manager for the variables from host and path.
 	regexp routeRegexpGroup
 
@@ -207,6 +210,10 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 			} else {
 				req = requestWithRouteAndVars(req, match.Route, match.Vars)
 			}
+
+			if !r.omitRouterFromContext {
+				req = requestWithRouter(req, r)
+			}
 		}
 	}
 
@@ -279,6 +286,15 @@ func (r *Router) OmitRouteFromContext(value bool) *Router {
 	return r
 }
 
+// OmitRouterFromContext defines the behavior of omitting the Router from the
+// http.Request context.
+//
+// RouterFromRequest will yield nil with this option.
+func (r *Router) OmitRouterFromContext(value bool) *Router {
+	r.omitRouterFromContext = value
+	return r
+}
+
 // UseEncodedPath tells the router to match the encoded original path
 // to the routes.
 // For eg. "/path/foo%2Fbar/to" will match the path "/path/{var}/to".
@@ -443,6 +459,7 @@ type contextKey int
 const (
 	varsKey contextKey = iota
 	routeKey
+	routerKey
 )
 
 // Vars returns the route variables for the current request, if any.
@@ -464,6 +481,13 @@ func CurrentRoute(r *http.Request) *Route {
 	return nil
 }
 
+func CurrentRouter(r *http.Request) *Router {
+	if rv := r.Context().Value(routerKey); rv != nil {
+		return rv.(*Router)
+	}
+	return nil
+}
+
 // requestWithVars adds the matched vars to the request ctx.
 // It shortcuts the operation when the vars are empty.
 func requestWithVars(r *http.Request, vars map[string]string) *http.Request {
@@ -486,6 +510,11 @@ func requestWithRouteAndVars(r *http.Request, route *Route, vars map[string]stri
 	return r.WithContext(ctx)
 }
 
+func requestWithRouter(r *http.Request, router *Router) *http.Request {
+	ctx := context.WithValue(r.Context(), routerKey, router)
+	return r.WithContext(ctx)
+}
+
 // ----------------------------------------------------------------------------
 // Helpers
 // ----------------------------------------------------------------------------

mux_test.go

@@ -1768,6 +1768,73 @@ func TestPanicOnCapturingGroups(t *testing.T) {
 	NewRouter().NewRoute().Path("/{type:(promo|special)}/{promoId}.json")
 }
 
+func TestRouterInContext(t *testing.T) {
+	router := NewRouter()
+	router.HandleFunc("/r1", func(w http.ResponseWriter, r *http.Request) {
+		contextRouter := CurrentRouter(r)
+		if contextRouter == nil {
+			t.Fatal("Router not found in context")
+			return
+		}
+
+		route := contextRouter.Get("r2")
+		if route == nil {
+			t.Fatal("Route with name not found")
+			return
+		}
+
+		url, err := route.URL()
+		if err != nil {
+			t.Fatal("Error while getting url for r2: ", err)
+			return
+		}
+
+		_, err = w.Write([]byte(url.String()))
+		if err != nil {
+			t.Fatalf("Failed writing HTTP response: %v", err)
+		}
+	}).Name("r1")
+
+	noRouterMsg := []byte("no-router")
+	haveRouterMsg := []byte("have-router")
+	router.HandleFunc("/r2", func(w http.ResponseWriter, r *http.Request) {
+		var msg []byte
+
+		contextRouter := CurrentRouter(r)
+		if contextRouter == nil {
+			msg = noRouterMsg
+		} else {
+			msg = haveRouterMsg
+		}
+
+		_, err := w.Write(msg)
+		if err != nil {
+			t.Fatalf("Failed writing HTTP response: %v", err)
+		}
+	}).Name("r2")
+
+	t.Run("router in request context get route by name", func(t *testing.T) {
+		rw := NewRecorder()
+		req := newRequest("GET", "/r1")
+
+		router.ServeHTTP(rw, req)
+		if !bytes.Equal(rw.Body.Bytes(), []byte("/r2")) {
+			t.Fatalf("Expected output to be '/r1' but got '%s'", rw.Body.String())
+		}
+	})
+
+	t.Run("omit router from request context", func(t *testing.T) {
+		rw := NewRecorder()
+		req := newRequest("GET", "/r2")
+
+		router.OmitRouterFromContext(true)
+		router.ServeHTTP(rw, req)
+		if !bytes.Equal(rw.Body.Bytes(), noRouterMsg) {
+			t.Fatal("Router not omitted from context")
+		}
+	})
+}
+
 // ----------------------------------------------------------------------------
 // Helpers
 // ----------------------------------------------------------------------------