feat: propagate hook context and add completion stage

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
 }

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()

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)

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

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)
 }

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)
+		})
+	}
+}