Files
multica/server/internal/middleware/auth_test.go
Bohan Jiang 86e7de3e41 feat(server/auth): cache auth token lookups in Redis with 10m TTL
* feat(server/auth): cache PAT lookups in Redis with 60s TTL

Personal access tokens used to hit Postgres on every request: a SELECT
to resolve token_hash → user_id, plus a fire-and-forget UPDATE of
last_used_at. For a CLI / daemon making many requests per second this
is wasted DB load — the token is the same and the answer hasn't changed.

Add a Redis-backed cache (auth.PATCache) keyed by token hash, TTL 60s:

- On cache hit, the auth middleware skips both the SELECT and the
  last_used_at UPDATE. last_used_at is now refreshed at most once per
  TTL window per token, not per request.
- On cache miss the middleware falls back to today's behavior: query
  Postgres, populate the cache, async-update last_used_at.
- On revoke, the handler invalidates the cache entry so revocation
  takes effect immediately rather than waiting for the TTL to expire.
  This required changing RevokePersonalAccessToken from :exec to :one
  RETURNING token_hash.

The cache is nil-safe: when REDIS_URL isn't configured, NewPATCache
returns nil and the middleware degrades to today's always-hit-DB
behavior. JWT validation is untouched (already DB-free).

Tested with REDIS_TEST_URL — same gating pattern the rest of the
suite uses for Redis-backed tests. New tests cover nil-safety, set/
get/invalidate, TTL, and the middleware short-circuit on cache hit.

* fix(server/auth): clamp PAT cache TTL to token's remaining lifetime

GPT-Boy review caught: a PAT expiring in <60s would still be cached
for the full PATCacheTTL window, so the token could continue passing
auth on cache hit for up to ~60s after its expires_at. The DB query
filters expired tokens (revoked = FALSE AND expires_at > now()), but
that filter never ran on a cache hit.

Make Set take an explicit ttl, and add TTLForExpiry to compute it:
  - no expires_at      → full PATCacheTTL
  - expires_at far     → full PATCacheTTL
  - expires_at <60s    → time until expiry
  - already expired    → 0, Set skips caching (TOCTOU defense between
                         the SELECT and the Set, since the SELECT
                         already filters expired rows)

Regression test pins the clamp behavior end-to-end against Redis.

* feat(server/auth): cache daemon-token + PAT lookups in DaemonAuth, bump TTL to 10m

Daemon /api/daemon/* requests (heartbeat, claim task) hit DaemonAuth
which previously did its own GetDaemonTokenByHash on every request and
*also* duplicated the PAT lookup on the mul_ fallback — bypassing the
cache added in 1cdd674c. Today's daemons authenticate via mul_ PATs
(mdt_ minting isn't wired up yet), so the duplicate PAT path is the one
that actually matters for hot-path DB load.

Three changes:

1. New auth.DaemonTokenCache mirrors PATCache for the mdt_ path
   (key = mul:auth:daemon:<sha256>, JSON value = {workspace_id, daemon_id}).
   Forward-looking infrastructure for when daemon tokens get minted; the
   middleware short-circuits the DB SELECT on cache hit. TTL clamped to
   the token's expires_at via the shared TTLForExpiry helper.

2. DaemonAuth now also consults PATCache on its mul_ fallback, sharing
   the same cache as the regular Auth middleware. A daemon making 4 hb/min
   collapses from 4 GetPersonalAccessTokenByHash + 4 last_used_at writes
   per minute to ~1 of each per AuthCacheTTL window (~10 minutes).

3. Rename PATCacheTTL → AuthCacheTTL and bump from 60s to 10 minutes.
   The constant is now shared between PAT and daemon caches; 10m matches
   the user-requested longer TTL for further DB write reduction. Revoke
   latency on the happy path is still instant via active invalidation;
   the worst-case (Redis Del miss / direct-DB revoke) grows from ~60s to
   ~10m.

Tests cover nil-safety, set/get/invalidate, TTL, clamped TTL on near-
expiry tokens, and the middleware short-circuit for both cache paths
(mdt_ via DaemonTokenCache, mul_ fallback via PATCache).

* feat(server/auth): cache PAT lookups on the WebSocket auth path

The third place a PAT is resolved — patResolver.ResolveToken used by
realtime.HandleWebSocket — was still hitting Postgres on every /ws
auth and firing an unconditional last_used_at UPDATE, bypassing the
cache added in 1cdd674c. Wire it through the same shared PATCache so
revoking a token through any path (Auth middleware, DaemonAuth PAT
fallback, or WS auth) hits all three caches with one Invalidate.

Also leaves a comment on DeleteDaemonTokensByWorkspaceAndDaemon —
the query has no caller today, but a future deregister/rotate flow
must remember to call DaemonTokenCache.Invalidate(hash) for each
deleted row, otherwise deleted daemon tokens stay valid until TTL.
2026-04-29 17:07:54 +08:00

268 lines
7.8 KiB
Go

package middleware
import (
"context"
"net/http"
"net/http/httptest"
"os"
"testing"
"time"
"github.com/golang-jwt/jwt/v5"
"github.com/multica-ai/multica/server/internal/auth"
"github.com/redis/go-redis/v9"
)
// newRedisTestClient connects to REDIS_TEST_URL, flushes, and skips when
// unset — same gating pattern the rest of the suite uses for Redis-backed
// tests, so `go test ./...` works on a stock laptop without a Redis.
func newRedisTestClient(t *testing.T) *redis.Client {
t.Helper()
url := os.Getenv("REDIS_TEST_URL")
if url == "" {
t.Skip("REDIS_TEST_URL not set")
}
opts, err := redis.ParseURL(url)
if err != nil {
t.Fatalf("parse REDIS_TEST_URL: %v", err)
}
rdb := redis.NewClient(opts)
ctx := context.Background()
if err := rdb.Ping(ctx).Err(); err != nil {
t.Skipf("REDIS_TEST_URL unreachable: %v", err)
}
if err := rdb.FlushDB(ctx).Err(); err != nil {
t.Fatalf("flushdb: %v", err)
}
t.Cleanup(func() {
rdb.FlushDB(context.Background())
rdb.Close()
})
return rdb
}
func generateToken(claims jwt.MapClaims, secret []byte) string {
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
s, _ := token.SignedString(secret)
return s
}
func validClaims() jwt.MapClaims {
return jwt.MapClaims{
"sub": "test-user-id",
"email": "test@multica.ai",
"exp": time.Now().Add(time.Hour).Unix(),
}
}
// authMiddleware returns the Auth middleware with nil queries (JWT-only tests).
func authMiddleware(next http.Handler) http.Handler {
return Auth(nil, nil)(next)
}
func TestAuth_MissingHeader(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
req := httptest.NewRequest("GET", "/api/me", nil)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
if body := w.Body.String(); body != `{"error":"missing authorization"}`+"\n" {
t.Fatalf("unexpected body: %s", body)
}
}
func TestAuth_NoBearerPrefix(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Token some-token")
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
// Non-Bearer Authorization header with no cookie falls through to "missing authorization".
if body := w.Body.String(); body != `{"error":"missing authorization"}`+"\n" {
t.Fatalf("unexpected body: %s", body)
}
}
func TestAuth_InvalidToken(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer not-a-valid-jwt")
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
}
func TestAuth_ExpiredToken(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
claims := validClaims()
claims["exp"] = time.Now().Add(-time.Hour).Unix()
token := generateToken(claims, auth.JWTSecret())
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer "+token)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
}
func TestAuth_WrongSecret(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
token := generateToken(validClaims(), []byte("wrong-secret"))
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer "+token)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
}
func TestAuth_WrongSigningMethod(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
// Use "none" signing method
token := jwt.NewWithClaims(jwt.SigningMethodNone, validClaims())
s, _ := token.SignedString(jwt.UnsafeAllowNoneSignatureType)
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer "+s)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
}
func TestAuth_ValidToken(t *testing.T) {
var gotUserID, gotEmail string
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotUserID = r.Header.Get("X-User-ID")
gotEmail = r.Header.Get("X-User-Email")
w.WriteHeader(http.StatusOK)
}))
token := generateToken(validClaims(), auth.JWTSecret())
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer "+token)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d", w.Code)
}
if gotUserID != "test-user-id" {
t.Fatalf("expected X-User-ID 'test-user-id', got '%s'", gotUserID)
}
if gotEmail != "test@multica.ai" {
t.Fatalf("expected X-User-Email 'test@multica.ai', got '%s'", gotEmail)
}
}
func TestAuth_MissingClaims(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
// Token with no sub or email claims, only exp
claims := jwt.MapClaims{
"exp": time.Now().Add(time.Hour).Unix(),
}
token := generateToken(claims, auth.JWTSecret())
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer "+token)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
}
func TestAuth_InvalidPAT(t *testing.T) {
handler := authMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer mul_invalid_token_here")
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d", w.Code)
}
}
// TestAuth_PATCacheHit pins the optimization: when the PAT cache already
// holds an entry for this token, the middleware MUST NOT call into queries
// — it short-circuits before the DB lookup and the last_used_at update.
//
// We exploit that by passing nil queries: a cache miss would dereference
// the nil and panic; a cache hit must not. Reaching the next handler with
// the cached user_id therefore proves the short-circuit fired.
func TestAuth_PATCacheHit(t *testing.T) {
rdb := newRedisTestClient(t)
cache := auth.NewPATCache(rdb)
if cache == nil {
t.Fatal("expected non-nil cache")
}
const rawToken = "mul_cache_hit_test_token"
hash := auth.HashToken(rawToken)
cache.Set(context.Background(), hash, "cached-user-id", auth.AuthCacheTTL)
var gotUserID string
mw := Auth(nil, cache) // nil queries — only safe on cache hit
handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotUserID = r.Header.Get("X-User-ID")
w.WriteHeader(http.StatusOK)
}))
req := httptest.NewRequest("GET", "/api/me", nil)
req.Header.Set("Authorization", "Bearer "+rawToken)
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 on cache hit, got %d", w.Code)
}
if gotUserID != "cached-user-id" {
t.Fatalf("expected cached X-User-ID, got %q", gotUserID)
}
}