Allow scope access in handler (#40)

* feat: allow route scope access in handler

* feat: improve test coverage

* feat: improve test coverage

context.go

@@ -19,6 +19,7 @@ import (
 // to release resources after use.
 type ContextCloser interface {
 	Context
+	// Close releases the context to be reused later.
 	Close()
 }
 
@@ -82,12 +83,13 @@ type Context interface {
 	// This functionality is particularly beneficial for middlewares that need to wrap
 	// their custom ResponseWriter while preserving the state of the original Context.
 	CloneWith(w ResponseWriter, r *http.Request) ContextCloser
+	// Scope returns the HandlerScope associated with the current Context.
+	// This indicates the scope in which the handler is being executed, such as RouteHandler, NoRouteHandler, etc.
+	Scope() HandlerScope
 	// Tree is a local copy of the Tree in use to serve the request.
 	Tree() *Tree
 	// Fox returns the Router instance.
 	Fox() *Router
-	// Reset resets the Context to its initial state, attaching the provided ResponseWriter and http.Request.
-	Reset(w ResponseWriter, r *http.Request)
 }
 
 // cTx holds request-related information and allows interaction with the ResponseWriter.
@@ -105,6 +107,7 @@ type cTx struct {
 	fox         *Router
 	cachedQuery url.Values
 	rec         recorder
+	scope       HandlerScope
 	tsr         bool
 }
 
@@ -115,18 +118,20 @@ func (c *cTx) Reset(w ResponseWriter, r *http.Request) {
 	c.tsr = false
 	c.cachedQuery = nil
 	c.route = nil
+	c.scope = RouteHandler
 	*c.params = (*c.params)[:0]
 }
 
 // reset resets the Context to its initial state, attaching the provided http.ResponseWriter and http.Request.
-// Caution: You should always pass the original http.ResponseWriter to this method, not the ResponseWriter itself, to
+// Caution: always pass the original http.ResponseWriter to this method, not the ResponseWriter itself, to
 // avoid wrapping the ResponseWriter within itself. Use wisely!
 func (c *cTx) reset(w http.ResponseWriter, r *http.Request) {
 	c.rec.reset(w)
 	c.req = r
 	c.w = &c.rec
 	c.cachedQuery = nil
 	c.route = nil
+	c.scope = RouteHandler
 	*c.params = (*c.params)[:0]
 }
 
@@ -299,6 +304,7 @@ func (c *cTx) Clone() Context {
 		fox:   c.fox,
 		tree:  c.tree,
 		route: c.route,
+		scope: c.scope,
 	}
 
 	cp.rec.ResponseWriter = noopWriter{c.rec.Header().Clone()}
@@ -320,6 +326,7 @@ func (c *cTx) CloneWith(w ResponseWriter, r *http.Request) ContextCloser {
 	cp.req = r
 	cp.w = w
 	cp.route = c.route
+	cp.scope = c.scope
 	cp.cachedQuery = nil
 	if cap(*c.params) > cap(*cp.params) {
 		// Grow cp.params to a least cap(c.params)
@@ -332,6 +339,12 @@ func (c *cTx) CloneWith(w ResponseWriter, r *http.Request) ContextCloser {
 	return cp
 }
 
+// Scope returns the HandlerScope associated with the current Context.
+// This indicates the scope in which the handler is being executed, such as RouteHandler, NoRouteHandler, etc.
+func (c *cTx) Scope() HandlerScope {
+	return c.scope
+}
+
 // Close releases the context to be reused later.
 func (c *cTx) Close() {
 	// Put back the context, if not extended more than max params or max depth, allowing

context_test.go

@@ -265,6 +265,75 @@ func TestContext_Tree(t *testing.T) {
 	f.ServeHTTP(w, req)
 }
 
+func TestContext_Scope(t *testing.T) {
+	t.Parallel()
+
+	f := New(
+		WithRedirectTrailingSlash(true),
+		WithMiddlewareFor(RedirectHandler, func(next HandlerFunc) HandlerFunc {
+			return func(c Context) {
+				assert.Equal(t, RedirectHandler, c.Scope())
+				next(c)
+			}
+		}),
+		WithNoRouteHandler(func(c Context) {
+			assert.Equal(t, NoRouteHandler, c.Scope())
+		}),
+		WithNoMethodHandler(func(c Context) {
+			assert.Equal(t, NoMethodHandler, c.Scope())
+		}),
+		WithOptionsHandler(func(c Context) {
+			assert.Equal(t, OptionsHandler, c.Scope())
+		}),
+	)
+	require.NoError(t, f.Handle(http.MethodGet, "/foo", func(c Context) {
+		assert.Equal(t, RouteHandler, c.Scope())
+	}))
+
+	cases := []struct {
+		name string
+		req  *http.Request
+		w    http.ResponseWriter
+	}{
+		{
+			name: "route handler scope",
+			req:  httptest.NewRequest(http.MethodGet, "/foo", nil),
+			w:    httptest.NewRecorder(),
+		},
+		{
+			name: "redirect handler scope",
+			req:  httptest.NewRequest(http.MethodGet, "/foo/", nil),
+			w:    httptest.NewRecorder(),
+		},
+		{
+			name: "no method handler scope",
+			req:  httptest.NewRequest(http.MethodPost, "/foo", nil),
+			w:    httptest.NewRecorder(),
+		},
+		{
+			name: "options handler scope",
+			req:  httptest.NewRequest(http.MethodOptions, "/foo", nil),
+			w:    httptest.NewRecorder(),
+		},
+		{
+			name: "options handler scope",
+			req:  httptest.NewRequest(http.MethodOptions, "/foo", nil),
+			w:    httptest.NewRecorder(),
+		},
+		{
+			name: "no route handler scope",
+			req:  httptest.NewRequest(http.MethodOptions, "/bar", nil),
+			w:    httptest.NewRecorder(),
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			f.ServeHTTP(tc.w, tc.req)
+		})
+	}
+}
+
 func TestWrapF(t *testing.T) {
 	t.Parallel()
 

fox.go

@@ -66,21 +66,48 @@ func (f ClientIPStrategyFunc) ClientIP(c Context) (*net.IPAddr, error) {
 	return f(c)
 }
 
+// HandlerScope represents different scopes where a handler may be called. It also allows for fine-grained control
+// over where middleware is applied.
+type HandlerScope uint8
+
+const (
+	// RouteHandler scope applies to regular routes registered in the router.
+	RouteHandler HandlerScope = 1 << (8 - 1 - iota)
+	// NoRouteHandler scope applies to the NoRoute handler, which is invoked when no route matches the request.
+	NoRouteHandler
+	// NoMethodHandler scope applies to the NoMethod handler, which is invoked when a route exists, but the method is not allowed.
+	NoMethodHandler
+	// RedirectHandler scope applies to the internal redirect handler, used for handling requests with trailing slashes.
+	RedirectHandler
+	// OptionsHandler scope applies to the automatic OPTIONS handler, which handles pre-flight or cross-origin requests.
+	OptionsHandler
+	// AllHandlers is a combination of all the above scopes, which can be used to apply middlewares to all types of handlers.
+	AllHandlers = RouteHandler | NoRouteHandler | NoMethodHandler | RedirectHandler | OptionsHandler
+)
+
 // Route represent a registered route in the route tree.
 // Most of the Route API is EXPERIMENTAL and is likely to change in future release.
 type Route struct {
 	ipStrategy            ClientIPStrategy
-	base                  HandlerFunc
-	handler               HandlerFunc
+	hbase                 HandlerFunc
+	hself                 HandlerFunc
+	hall                  HandlerFunc
 	path                  string
 	mws                   []middleware
 	redirectTrailingSlash bool
 	ignoreTrailingSlash   bool
 }
 
-// Handle calls the base handler with the provided Context.
+// Handle calls the handler with the provided Context. See also HandleMiddleware.
 func (r *Route) Handle(c Context) {
-	r.base(c)
+	r.hbase(c)
+}
+
+// HandleMiddleware calls the handler with route-specific middleware applied, using the provided Context.
+func (r *Route) HandleMiddleware(c Context, _ ...struct{}) {
+	// The variadic parameter is intentionally added to prevent this method from having the same signature as HandlerFunc.
+	// This avoids accidental use of HandleMiddleware where a HandlerFunc is required.
+	r.hself(c)
 }
 
 // Path returns the route path.
@@ -127,7 +154,8 @@ type Router struct {
 
 type middleware struct {
 	m     MiddlewareFunc
-	scope MiddlewareScope
+	scope HandlerScope
+	g     bool
 }
 
 var _ http.Handler = (*Router)(nil)
@@ -363,7 +391,7 @@ func (fox *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if !tsr && n != nil {
 		c.route = n.route
 		c.tsr = tsr
-		n.route.handler(c)
+		n.route.hall(c)
 		// Put back the context, if not extended more than max params or max depth, allowing
 		// the slice to naturally grow within the constraint.
 		if cap(*c.params) <= int(tree.maxParams.Load()) && cap(*c.skipNds) <= int(tree.maxDepth.Load()) {
@@ -376,7 +404,7 @@ func (fox *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		if n.route.ignoreTrailingSlash {
 			c.route = n.route
 			c.tsr = tsr
-			n.route.handler(c)
+			n.route.hall(c)
 			c.Close()
 			return
 		}
@@ -385,6 +413,7 @@ func (fox *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			// Reset params as it may have recorded wildcard segment (the context may still be used in a middleware)
 			*c.params = (*c.params)[:0]
 			c.tsr = false
+			c.scope = RedirectHandler
 			fox.tsrRedirect(c)
 			c.Close()
 			return
@@ -425,6 +454,7 @@ NoMethodFallback:
 			sb.WriteString(", ")
 			sb.WriteString(http.MethodOptions)
 			w.Header().Set(HeaderAllow, sb.String())
+			c.scope = OptionsHandler
 			fox.autoOptions(c)
 			c.Close()
 			return
@@ -443,12 +473,14 @@ NoMethodFallback:
 		}
 		if sb.Len() > 0 {
 			w.Header().Set(HeaderAllow, sb.String())
+			c.scope = NoMethodHandler
 			fox.noMethod(c)
 			c.Close()
 			return
 		}
 	}
 
+	c.scope = NoRouteHandler
 	fox.noRoute(c)
 	c.Close()
 }
@@ -645,7 +677,7 @@ func isRemovable(method string) bool {
 	return true
 }
 
-func applyMiddleware(scope MiddlewareScope, mws []middleware, h HandlerFunc) HandlerFunc {
+func applyMiddleware(scope HandlerScope, mws []middleware, h HandlerFunc) HandlerFunc {
 	m := h
 	for i := len(mws) - 1; i >= 0; i-- {
 		if mws[i].scope&scope != 0 {
@@ -655,6 +687,21 @@ func applyMiddleware(scope MiddlewareScope, mws []middleware, h HandlerFunc) Han
 	return m
 }
 
+func applyRouteMiddleware(mws []middleware, base HandlerFunc) (HandlerFunc, HandlerFunc) {
+	rte := base
+	all := base
+	for i := len(mws) - 1; i >= 0; i-- {
+		if mws[i].scope&RouteHandler != 0 {
+			all = mws[i].m(all)
+			// route specific only
+			if !mws[i].g {
+				rte = mws[i].m(rte)
+			}
+		}
+	}
+	return rte, all
+}
+
 // localRedirect redirect the client to the new path, but it does not convert relative paths to absolute paths
 // like Redirect does. If the Content-Type header has not been set, localRedirect sets it to "text/html; charset=utf-8"
 // and writes a small HTML body. Setting the Content-Type header to any value, including nil, disables that behavior.