From 4213d4a55f0b2489737de5825c93103c1f1fb143 Mon Sep 17 00:00:00 2001
From: Raphael Simon <simon.raphael@gmail.com>
Date: Sun, 1 Oct 2023 13:44:16 -0700
Subject: [PATCH] Ensure Vars can be used in middleware

---
 http/mux.go | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/http/mux.go b/http/mux.go
index 11013ecb3a..146027172d 100644
--- a/http/mux.go
+++ b/http/mux.go
@@ -109,15 +109,18 @@ func (m *mux) Handle(method, pattern string, handler http.HandlerFunc) {
 
 // Vars extracts the path variables from the request context.
 func (m *mux) Vars(r *http.Request) map[string]string {
-	rctx := chi.RouteContext(r.Context())
-	params := rctx.URLParams
+	ctx := m.ensureContext(r)
+	if ctx == nil {
+		return nil
+	}
+	params := ctx.URLParams
 	if len(params.Keys) == 0 {
 		return nil
 	}
 	vars := make(map[string]string, len(params.Keys))
 	for i, k := range params.Keys {
 		if k == "*" {
-			wildcard := m.wildcards[r.Method+"::"+rctx.RoutePattern()]
+			wildcard := m.wildcards[r.Method+"::"+ctx.RoutePattern()]
 			vars[wildcard] = unescape(params.Values[i])
 			continue
 		}
@@ -143,11 +146,8 @@ func (m *mux) Use(f func(http.Handler) http.Handler) {
 // ResolvePattern returns the route pattern used to register the handler for the
 // given method and path.
 func (m *mux) ResolvePattern(r *http.Request) string {
-	ctx := chi.RouteContext(r.Context())
-	if ctx.RoutePattern() != "" {
-		return m.resolveWildcard(r.Method, ctx.RoutePattern())
-	}
-	if !m.Router.Match(ctx, r.Method, r.URL.Path) {
+	ctx := m.ensureContext(r)
+	if ctx == nil {
 		return ""
 	}
 	return m.resolveWildcard(r.Method, ctx.RoutePattern())
@@ -161,3 +161,19 @@ func (m *mux) resolveWildcard(method, pattern string) string {
 	}
 	return pattern
 }
+
+// ensureContext makes sure chi has initialized the request context if it
+// handles it, otherwise it returns nil.
+func (m *mux) ensureContext(r *http.Request) *chi.Context {
+	ctx := chi.RouteContext(r.Context())
+	if ctx == nil {
+		return nil // request not handled by chi
+	}
+	if ctx.RoutePattern() != "" {
+		return ctx // already initialized
+	}
+	if !m.Router.Match(ctx, r.Method, r.URL.Path) {
+		return nil // route not handled by chi
+	}
+	return ctx
+}