// src/index.ts
import { Router } from 'itty-router';
import { json, error } from 'itty-router-extras';

const router = Router();

interface SessionPayload {
  user_id: string;
  tenant_id: string;
  permissions: string[];
  email: string;
  expires_at: number;
}

// Validate JWT and fetch/cache session
async function getSessionFromKV(
  token: string,
  kv: KVNamespace,
  env: any
): Promise<SessionPayload | null> {
  // 1. Check KV first (immediate response)
  const cached = await kv.get(`session:${token}`, 'json');
  if (cached) {
    // Stale-while-revalidate: serve cached, revalidate in background
    if (cached.expires_at - Date.now() < 5 * 60 * 1000) {
      // Less than 5 min left, revalidate in background
      kv.put(
        `session:${token}`,
        JSON.stringify(cached),
        {
          expirationTtl: cached.ttl,
        }
      ).catch(() => {}); // Fire and forget
    }
    return cached;
  }

  // 2. Cache miss, fetch from origin
  const originResponse = await fetch(
    `${env.ORIGIN_URL}/api/auth/validate-session`,
    {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${token}`,
      },
    }
  );

  if (!originResponse.ok) return null;

  const session = (await originResponse.json()) as SessionPayload;

  // 3. Write to KV with TTL (default 1 hour)
  const ttl = Math.floor((session.expires_at - Date.now()) / 1000);
  if (ttl > 0) {
    await kv.put(`session:${token}`, JSON.stringify(session), {
      expirationTtl: Math.min(ttl, 3600), // Cap at 1 hour
    });
  }

  return session;
}

// Middleware: require valid session
router.all('*', async (req, env, ctx) => {
  const token = req.headers
    .get('authorization')
    ?.replace('Bearer ', '')
    ?.trim();

  if (!token) {
    return error(401, { error: 'Missing authorization token' });
  }

  const session = await getSessionFromKV(token, env.KV_CACHE, env);

  if (!session || session.expires_at < Date.now()) {
    return error(401, { error: 'Invalid or expired session' });
  }

  // Attach to request context
  req.session = session;
});

// Example: API route behind auth
router.post('/api/tenant/:tenant_id/data', (req) => {
  const { tenant_id } = req.params;
  const { session } = req;

  // RBAC check
  if (session.tenant_id !== tenant_id) {
    return error(403, { error: 'Tenant mismatch' });
  }

  if (!session.permissions.includes('data:write')) {
    return error(403, { error: 'Permission denied' });
  }

  // Proxy to origin
  return fetch(`${req.env.ORIGIN_URL}${req.url}`, {
    method: req.method,
    headers: req.headers,
    body: req.body,
  });
});

export default router;

python
# app/routers/auth.py
from fastapi import APIRouter, Depends, HTTPException, Header
from sqlalchemy.orm import Session
from datetime import datetime, timedelta
import jwt

from app.db import get_db
from app.models import User, Tenant, Permission
from app.schemas import SessionResponse
from app.config import settings

router = APIRouter(prefix="/api/auth", tags=["auth"])

def get_session_payload(
    token: str,
    db: Session
) -> dict:
    """
    Validate JWT and return full session payload.
    This runs ONLY on KV miss, so we optimize for correctness over speed.
    """
    try:
        payload = jwt.decode(
            token,
            settings.SECRET_KEY,
            algorithms=["HS256"]
        )
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token expired")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid token")

    user_id = payload.get("sub")
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid token")

    # Fetch user and tenant (this hits DB, but infrequently due to KV)
    user = db.query(User).filter(User.id == user_id).first()
    if not user or not user.is_active:
        raise HTTPException(status_code=401, detail="User inactive")

    tenant = db.query(Tenant).filter(Tenant.id == user.tenant_id).first()
    if not tenant or not tenant.is_active:
        raise HTTPException(status_code=401, detail="Tenant inactive")

    # Fetch permissions (this is the expensive query KV saves)
    permissions = (
        db.query(Permission.code)
        .join(User.roles)
        .join(Permission.roles)
        .filter(User.id == user_id)
        .distinct()
        .all()
    )

    return {
        "user_id": str(user.id),
        "tenant_id": str(user.tenant_id),
        "email": user.email,
        "permissions": [p[0] for p in permissions],
        "expires_at": int((datetime.utcnow() + timedelta(hours=1)).timestamp() * 1000),
        "ttl": 3600,
    }

@router.post("/validate-session", response_model=dict)
async def validate_session(
    authorization: str = Header(None),
    db: Session = Depends(get_db),
):
    """
    Called by Cloudflare Worker on KV miss.
    Expensive query runs here, not on every request.
    """
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing token")

    token = authorization.replace("Bearer ", "").strip()
    return get_session_payload(token, db)

@router.post("/logout")
async def logout(
    token: str,
    db: Session = Depends(get_db),
):
    """
    Invalidate KV cache on logout.
    Send deletion request to Cloudflare.
    """
    # Delete from KV (optional, TTL handles it)
    # In production, call Cloudflare API to delete session:{token}

    # Also invalidate user's tokens in DB (refresh token revocation)
    db.query(RefreshToken).filter(
        RefreshToken.user_id == user_id
    ).delete()
    db.commit()

    return {"status": "logged out"}

@router.post("/update-permissions/:user_id")
async def update_permissions(
    user_id: str,
    db: Session = Depends(get_db),
):
    """
    When permissions change, invalidate all KV sessions for this user.
    This ensures cache coherency.
    """
    # Update permissions in DB
    # ...

    # Invalidate all active sessions for this user in KV
    # Query Redis-