From 3a658e990deaba75d3561b43308fc230774059f0 Mon Sep 17 00:00:00 2001 From: Tanvir Ahmed Date: Sun, 11 Oct 2026 20:35:01 +0000 Subject: [PATCH] feat: propagate hook context and add completion stage Make reqCtx.Request the source of truth so hook-injected context reaches route handlers. Add a HookOnComplete stage that runs exactly once on every exit path, including early short-circuits, for finalization that must always run. Expose the written HTTP status via models.ResponseStatusProvider and DeferredResponseWriter.StatusCode. --- internal/router/deferred_response_writer.go | 6 + .../router/deferred_response_writer_test.go | 47 ++++ models/context.go | 6 + models/plugin.go | 5 + router.go | 34 ++- router_hooks_test.go | 203 ++++++++++++++++++ 6 files changed, 300 insertions(+), 1 deletion(-) diff --git a/internal/router/deferred_response_writer.go b/internal/router/deferred_response_writer.go index c9a9e5e4..275b398b 100644 --- a/internal/router/deferred_response_writer.go +++ b/internal/router/deferred_response_writer.go @@ -26,6 +26,12 @@ func (w *DeferredResponseWriter) SetRequestContext(ctx *models.RequestContext) { w.ctx = ctx } +// StatusCode returns the HTTP status code written by the handler. +// It returns 0 if no status has been written yet. +func (w *DeferredResponseWriter) StatusCode() int { + return w.statusCode +} + func (w *DeferredResponseWriter) GetRequestContext() *models.RequestContext { return w.ctx } diff --git a/internal/router/deferred_response_writer_test.go b/internal/router/deferred_response_writer_test.go index 4149252b..078e114e 100644 --- a/internal/router/deferred_response_writer_test.go +++ b/internal/router/deferred_response_writer_test.go @@ -71,6 +71,53 @@ func TestFlushPreservesCookies(t *testing.T) { } } +func TestDeferredResponseWriterStatusCode(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + write func(w *DeferredResponseWriter) + wantStatus int + }{ + { + name: "reports status after WriteHeader", + write: func(w *DeferredResponseWriter) { + w.WriteHeader(http.StatusCreated) + }, + wantStatus: http.StatusCreated, + }, + { + name: "reports implicit 200 after Write", + write: func(w *DeferredResponseWriter) { + _, _ = w.Write([]byte("ok")) + }, + wantStatus: http.StatusOK, + }, + { + name: "reports zero before any write", + write: func(w *DeferredResponseWriter) {}, + wantStatus: 0, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + drw := &DeferredResponseWriter{ + Wrapped: httptest.NewRecorder(), + Logger: &internaltests.MockLogger{}, + } + + tc.write(drw) + + if got := drw.StatusCode(); got != tc.wantStatus { + t.Fatalf("expected status %d, got %d", tc.wantStatus, got) + } + }) + } +} + func TestFlushPreservesMultipleCookies(t *testing.T) { // Create a mock ResponseWriter originalWriter := httptest.NewRecorder() diff --git a/models/context.go b/models/context.go index f0db97c7..24f00060 100644 --- a/models/context.go +++ b/models/context.go @@ -66,6 +66,12 @@ type RequestContext struct { RedirectURL string // URL to redirect to (empty = no redirect) } +// ResponseStatusProvider is implemented by response writers that expose the +// HTTP status code written by handlers. +type ResponseStatusProvider interface { + StatusCode() int +} + // NewContextWithRequestContext returns a new context with the RequestContext attached func NewContextWithRequestContext(ctx context.Context, rc *RequestContext) context.Context { return context.WithValue(ctx, ContextRequestContext, rc) diff --git a/models/plugin.go b/models/plugin.go index 5cb44427..a52dc9c8 100644 --- a/models/plugin.go +++ b/models/plugin.go @@ -108,6 +108,11 @@ const ( HookAfter // HookOnResponse is executed after the response has been written HookOnResponse + // HookOnComplete is executed exactly once for every request, after all other + // stages (or after an early short-circuit), immediately before the response is + // flushed. Handlers do not see this stage. Use it for finalization that must + // run even when an earlier hook set Handled. + HookOnComplete ) // HookMatcher is a function that determines whether a hook should execute diff --git a/router.go b/router.go index 9850b10b..a2d29daa 100644 --- a/router.go +++ b/router.go @@ -399,6 +399,26 @@ func (r *Router) runHooks(stage models.HookStage, ctx *models.RequestContext) { } } +// runCompletionHooks runs HookOnComplete hooks regardless of ctx.Handled. +// Completion hooks must always run, even after an early short-circuit, so they +// are handled separately from the normal runHooks loop. +func (r *Router) runCompletionHooks(ctx *models.RequestContext) { + for _, hook := range r.hooks { + if hook.Stage != models.HookOnComplete { + continue + } + if hook.Matcher != nil && !hook.Matcher(ctx) { + continue + } + func() { + defer r.recoverFromPanic(string(HookTypeSync), hook.PluginID, models.HookOnComplete) + if err := hook.Handler(ctx); err != nil { + r.handleHookError(hook.PluginID, models.HookOnComplete, err, false) + } + }() + } +} + func (r *Router) recoverFromPanic(hookType, pluginID string, stage models.HookStage) { if err := recover(); err != nil { // Capture stack trace @@ -599,8 +619,10 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { } } - // Store context in request + // Store context in request and make it the single source of truth so that + // hook-injected context reaches handlers. reqWithCtx := req.WithContext(models.NewContextWithRequestContext(req.Context(), reqCtx)) + reqCtx.Request = reqWithCtx if req.Method == http.MethodOptions { r.applyCORS(req, wrappedWriter) @@ -615,6 +637,7 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { r.runHooks(models.HookOnRequest, reqCtx) if reqCtx.Handled { r.applyRedirectIfSet(reqCtx) + r.runCompletionHooks(reqCtx) r.finalizeResponse(reqCtx, wrappedWriter) return } @@ -623,17 +646,23 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { r.runHooks(models.HookBefore, reqCtx) if reqCtx.Handled { r.applyRedirectIfSet(reqCtx) + r.runCompletionHooks(reqCtx) r.finalizeResponse(reqCtx, wrappedWriter) return } // Stage 3: Route handler (via Chi) + // Re-derive the request so mutations made by hooks are honored by handlers. + if reqCtx.Request != nil { + reqWithCtx = reqCtx.Request + } r.router.ServeHTTP(wrappedWriter, reqWithCtx) // Stage 4: After hooks r.runHooks(models.HookAfter, reqCtx) if reqCtx.Handled { r.applyRedirectIfSet(reqCtx) + r.runCompletionHooks(reqCtx) r.finalizeResponse(reqCtx, wrappedWriter) return } @@ -645,6 +674,9 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { // Stage 5: OnResponse hooks r.runHooks(models.HookOnResponse, reqCtx) + // Completion hooks run on every exit path, after all other stages. + r.runCompletionHooks(reqCtx) + // Flush deferred writes or captured response r.finalizeResponse(reqCtx, wrappedWriter) } diff --git a/router_hooks_test.go b/router_hooks_test.go index 9bfff4e8..3d0ec211 100644 --- a/router_hooks_test.go +++ b/router_hooks_test.go @@ -1,6 +1,7 @@ package authula import ( + "context" "fmt" "net/http" "net/http/httptest" @@ -696,3 +697,205 @@ func TestRouterRedirectHandledRequest(t *testing.T) { }) } } + +func TestRouterHookContextPropagationToHandler(t *testing.T) { + t.Parallel() + + type ctxKey string + + const actorKey ctxKey = "test.actor" + + router := newTestRouter("/api/auth", nil) + + router.RegisterHook(models.Hook{ + Stage: models.HookOnRequest, + Handler: func(ctx *models.RequestContext) error { + ctx.Request = ctx.Request.WithContext( + context.WithValue(ctx.Request.Context(), actorKey, "hook-actor"), + ) + return nil + }, + }) + + var handlerValue any + registerTestRoute(router, http.MethodGet, "/test", func(w http.ResponseWriter, r *http.Request) { + handlerValue = r.Context().Value(actorKey) + w.WriteHeader(http.StatusOK) + }) + + w := performRequest(router, http.MethodGet, "/api/auth/test") + if w.Code != http.StatusOK { + t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code) + } + + if handlerValue != "hook-actor" { + t.Fatalf("expected handler to read hook-injected value %q, got %v", "hook-actor", handlerValue) + } +} + +func TestRouterHookActorReachesDownstreamHandler(t *testing.T) { + t.Parallel() + + router := newTestRouter("/api/auth", nil) + + router.RegisterHook(models.Hook{ + Stage: models.HookOnRequest, + Handler: func(ctx *models.RequestContext) error { + ctx.SetActorInContext(&models.Actor{ID: "actor-1", Type: models.ActorUser}) + return nil + }, + }) + + registerTestRoute(router, http.MethodGet, "/test", func(w http.ResponseWriter, r *http.Request) { + reqCtx, ok := models.GetRequestContext(r.Context()) + if !ok { + t.Fatal("expected request context to be available to handler") + } + if reqCtx.Actor == nil || reqCtx.Actor.ID != "actor-1" { + t.Fatalf("expected request context actor %q, got %+v", "actor-1", reqCtx.Actor) + } + + actor, ok := models.GetActorFromContext(r.Context()) + if !ok || actor.ID != "actor-1" { + t.Fatalf("expected actor %q from request context, got %+v", "actor-1", actor) + } + + w.WriteHeader(http.StatusOK) + }) + + w := performRequest(router, http.MethodGet, "/api/auth/test") + if w.Code != http.StatusOK { + t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code) + } +} + +func TestRouterCompletionHookRunsAfterShortCircuit(t *testing.T) { + t.Parallel() + + router := newTestRouter("/api/auth", nil) + + completionRuns := 0 + var seenStatus int + + router.RegisterHook(models.Hook{ + Stage: models.HookOnRequest, + Handler: func(ctx *models.RequestContext) error { + ctx.Handled = true + ctx.ResponseWriter.WriteHeader(http.StatusUnauthorized) + return nil + }, + }) + + router.RegisterHook(models.Hook{ + Stage: models.HookOnComplete, + Handler: func(ctx *models.RequestContext) error { + completionRuns++ + seenStatus = ctx.ResponseStatus + if seenStatus == 0 { + if provider, ok := ctx.ResponseWriter.(models.ResponseStatusProvider); ok { + seenStatus = provider.StatusCode() + } + } + return nil + }, + }) + + registerTestRoute(router, http.MethodGet, "/test", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + w := performRequest(router, http.MethodGet, "/api/auth/test") + if w.Code != http.StatusUnauthorized { + t.Fatalf("expected status %d, got %d", http.StatusUnauthorized, w.Code) + } + + if completionRuns != 1 { + t.Fatalf("expected completion hook to run exactly once, got %d", completionRuns) + } + + if seenStatus != http.StatusUnauthorized { + t.Fatalf("expected completion hook to see status %d, got %d", http.StatusUnauthorized, seenStatus) + } +} + +func TestRouterCompletionHookRunsExactlyOnce(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + run func(t *testing.T) + }{ + { + name: "normal path", + run: func(t *testing.T) { + router := newTestRouter("/api/auth", nil) + completionRuns := 0 + + router.RegisterHook(models.Hook{ + Stage: models.HookOnComplete, + Handler: func(ctx *models.RequestContext) error { + completionRuns++ + return nil + }, + }) + + registerTestRoute(router, http.MethodGet, "/test", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + w := performRequest(router, http.MethodGet, "/api/auth/test") + if w.Code != http.StatusOK { + t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code) + } + + if completionRuns != 1 { + t.Fatalf("expected completion hook to run exactly once, got %d", completionRuns) + } + }, + }, + { + name: "after HookAfter short-circuit", + run: func(t *testing.T) { + router := newTestRouter("/api/auth", nil) + completionRuns := 0 + + router.RegisterHook(models.Hook{ + Stage: models.HookAfter, + Handler: func(ctx *models.RequestContext) error { + ctx.Handled = true + ctx.SetResponse(http.StatusForbidden, nil, nil) + return nil + }, + }) + + router.RegisterHook(models.Hook{ + Stage: models.HookOnComplete, + Handler: func(ctx *models.RequestContext) error { + completionRuns++ + return nil + }, + }) + + registerTestRoute(router, http.MethodGet, "/test", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + w := performRequest(router, http.MethodGet, "/api/auth/test") + if w.Code != http.StatusForbidden { + t.Fatalf("expected status %d, got %d", http.StatusForbidden, w.Code) + } + + if completionRuns != 1 { + t.Fatalf("expected completion hook to run exactly once, got %d", completionRuns) + } + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + tc.run(t) + }) + } +}