New token arch

This commit is contained in:
Urtzi Alfaro
2026-01-10 21:45:37 +01:00
parent cc53037552
commit bf1db7cb9e
26 changed files with 1751 additions and 107 deletions

View File

@@ -5,7 +5,7 @@ FIXED VERSION - Proper JWT verification and token structure handling
"""
import structlog
from fastapi import Request, HTTPException
from fastapi import Request, HTTPException, status
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
@@ -60,6 +60,16 @@ class AuthMiddleware(BaseHTTPMiddleware):
if request.method == "OPTIONS":
return await call_next(request)
# SECURITY: Remove any incoming x-subscription-* headers
# These will be re-injected from verified JWT only
sanitized_headers = [
(k, v) for k, v in request.headers.raw
if not k.decode().lower().startswith('x-subscription-')
and not k.decode().lower().startswith('x-user-')
and not k.decode().lower().startswith('x-tenant-')
]
request.headers.__dict__["_list"] = sanitized_headers
# Skip authentication for public routes
if self._is_public_route(request.url.path):
return await call_next(request)
@@ -168,7 +178,14 @@ class AuthMiddleware(BaseHTTPMiddleware):
)
# Get tenant subscription tier and inject into user context
subscription_tier = await self._get_tenant_subscription_tier(tenant_id, request)
# NEW: Use JWT data if available, skip HTTP call
if user_context.get("subscription_from_jwt"):
subscription_tier = user_context.get("subscription_tier")
logger.debug("Using subscription tier from JWT", tier=subscription_tier)
else:
# Only for old tokens - remove after full rollout
subscription_tier = await self._get_tenant_subscription_tier(tenant_id, request)
if subscription_tier:
user_context["subscription_tier"] = subscription_tier
@@ -255,6 +272,22 @@ class AuthMiddleware(BaseHTTPMiddleware):
if payload and self._validate_token_payload(payload):
logger.debug("Token validated locally")
# NEW: Check token freshness for subscription changes (async)
if payload.get("tenant_id") and request:
try:
is_fresh = await self._verify_token_freshness(payload, payload["tenant_id"])
if not is_fresh:
logger.warning("Stale token detected - subscription changed since token was issued",
user_id=payload.get("user_id"),
tenant_id=payload.get("tenant_id"))
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Token is stale - subscription has changed"
)
except Exception as e:
logger.warning("Token freshness check failed, allowing token", error=str(e))
# Allow token if check fails (fail open for availability)
# Check if token is near expiry and set flag for response header
if request:
import time
@@ -321,6 +354,78 @@ class AuthMiddleware(BaseHTTPMiddleware):
if time_until_expiry < 300: # 5 minutes
logger.warning(f"Token expires in {int(time_until_expiry)} seconds for user {payload.get('email')}")
# NEW: Check token freshness for subscription changes
if payload.get("tenant_id"):
try:
# Note: We can't await here because this is a sync function
# Token freshness will be checked in the async dispatch method
# For now, just log that we would check freshness
logger.debug("Token freshness check would be performed in async context",
tenant_id=payload.get("tenant_id"))
except Exception as e:
logger.warning("Token freshness check setup failed", error=str(e))
return True
def _validate_jwt_integrity(self, payload: Dict[str, Any]) -> bool:
"""
Validate JWT payload integrity beyond signature verification.
Prevents edge cases where payload might be malformed.
"""
# Required fields must exist
required_fields = ["user_id", "email", "exp", "iat", "iss"]
if not all(field in payload for field in required_fields):
logger.warning("JWT missing required fields", missing=[f for f in required_fields if f not in payload])
return False
# Issuer must be our auth service
if payload.get("iss") != "bakery-auth":
logger.warning("JWT has invalid issuer", issuer=payload.get("iss"))
return False
# Token type must be valid
if payload.get("type") not in ["access", "service"]:
logger.warning("JWT has invalid type", token_type=payload.get("type"))
return False
# Subscription tier must be valid if present
valid_tiers = ["starter", "professional", "enterprise"]
if payload.get("subscription"):
tier = payload["subscription"].get("tier", "").lower()
if tier and tier not in valid_tiers:
logger.warning("JWT has invalid subscription tier", tier=tier)
return False
return True
async def _verify_token_freshness(self, payload: Dict[str, Any], tenant_id: str) -> bool:
"""
Verify token was issued after the last subscription change.
Prevents use of stale tokens with old subscription data.
"""
if not self.redis_client:
return True # Skip check if no Redis
try:
subscription_changed_at = await self.redis_client.get(
f"tenant:{tenant_id}:subscription_changed_at"
)
if subscription_changed_at:
changed_timestamp = float(subscription_changed_at)
token_issued_at = payload.get("iat", 0)
if token_issued_at < changed_timestamp:
logger.warning(
"Token issued before subscription change",
token_iat=token_issued_at,
subscription_changed=changed_timestamp,
tenant_id=tenant_id
)
return False # Token is stale
except Exception as e:
logger.warning("Failed to check token freshness", error=str(e))
return True
def _jwt_payload_to_user_context(self, payload: Dict[str, Any]) -> Dict[str, Any]:
@@ -328,6 +433,13 @@ class AuthMiddleware(BaseHTTPMiddleware):
Convert JWT payload to user context format
FIXED: Proper mapping between JWT structure and user context
"""
# NEW: Validate JWT integrity before processing
if not self._validate_jwt_integrity(payload):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid JWT payload"
)
base_context = {
"user_id": payload["user_id"],
"email": payload["email"],
@@ -336,6 +448,20 @@ class AuthMiddleware(BaseHTTPMiddleware):
"role": payload.get("role", "user"),
}
# NEW: Extract subscription from JWT
if payload.get("tenant_id"):
base_context["tenant_id"] = payload["tenant_id"]
base_context["tenant_role"] = payload.get("tenant_role", "member")
if payload.get("subscription"):
sub = payload["subscription"]
base_context["subscription_tier"] = sub.get("tier", "starter")
base_context["subscription_status"] = sub.get("status", "active")
base_context["subscription_from_jwt"] = True # Flag to skip HTTP
if payload.get("tenant_access"):
base_context["tenant_access"] = payload["tenant_access"]
if payload.get("service"):
service_name = payload["service"]
base_context["service"] = service_name
@@ -571,4 +697,4 @@ class AuthMiddleware(BaseHTTPMiddleware):
except Exception as e:
logger.error(f"Error getting tenant subscription tier: {e}")
return "starter" # Default to starter on error
return "starter" # Default to starter on error

View File

@@ -203,6 +203,9 @@ class DemoMiddleware(BaseHTTPMiddleware):
)
# This allows the request to pass through AuthMiddleware
# NEW: Extract subscription tier from demo account type
subscription_tier = "enterprise" if session_info.get("demo_account_type") == "enterprise" else "professional"
request.state.user = {
"user_id": demo_user_id, # Use actual demo user UUID
"email": f"demo-{session_id}@demo.local",
@@ -211,7 +214,11 @@ class DemoMiddleware(BaseHTTPMiddleware):
"is_demo": True,
"demo_session_id": session_id,
"demo_account_type": session_info.get("demo_account_type", "professional"),
"demo_session_status": current_status
"demo_session_status": current_status,
# NEW: Subscription context (no HTTP call needed!)
"subscription_tier": subscription_tier,
"subscription_status": "active",
"subscription_from_jwt": True # Flag to skip HTTP calls
}
# Update activity

View File

@@ -12,6 +12,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
import httpx
from typing import Dict, Any, Optional, List
import asyncio
from datetime import datetime, timezone
from app.core.config import settings
from app.utils.subscription_error_responses import create_upgrade_required_response
@@ -30,9 +31,10 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
- Operations routes (/api/v1/tenants/{tenant_id}/{service}/operations/*): ALL tiers (role-based)
"""
def __init__(self, app, tenant_service_url: str):
def __init__(self, app, tenant_service_url: str, redis_client=None):
super().__init__(app)
self.tenant_service_url = tenant_service_url.rstrip('/')
self.redis_client = redis_client # Optional Redis client for abuse detection
# Define route patterns that require subscription validation
# Using new standardized URL structure
@@ -236,20 +238,60 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
Dict with 'allowed' boolean and additional metadata
"""
try:
# Use the same authentication pattern as gateway routes
# Check if JWT already has subscription
if hasattr(request.state, 'user') and request.state.user:
user_context = request.state.user
user_id = user_context.get('user_id', 'unknown')
if user_context.get("subscription_from_jwt"):
# Use JWT data directly - NO HTTP CALL!
current_tier = user_context.get("subscription_tier", "starter")
logger.debug("Using subscription tier from JWT (no HTTP call)",
tenant_id=tenant_id,
current_tier=current_tier,
minimum_tier=minimum_tier,
allowed_tiers=allowed_tiers)
if current_tier not in [tier.lower() for tier in allowed_tiers]:
tier_names = ', '.join(allowed_tiers)
return {
'allowed': False,
'message': f'This feature requires a {tier_names} subscription plan',
'current_tier': current_tier
}
await self._log_subscription_access(
tenant_id,
user_id,
feature,
current_tier,
True,
"jwt"
)
return {
'allowed': True,
'message': 'Access granted (JWT subscription)',
'current_tier': current_tier
}
# Use the same authentication pattern as gateway routes for fallback
headers = dict(request.headers)
headers.pop("host", None)
# Extract user_id for logging (fallback path)
user_id = 'unknown'
# Add user context headers if available
if hasattr(request.state, 'user') and request.state.user:
user = request.state.user
headers["x-user-id"] = str(user.get('user_id', ''))
user_id = str(user.get('user_id', 'unknown'))
headers["x-user-id"] = user_id
headers["x-user-email"] = str(user.get('email', ''))
headers["x-user-role"] = str(user.get('role', 'user'))
headers["x-user-full-name"] = str(user.get('full_name', ''))
headers["x-tenant-id"] = str(user.get('tenant_id', ''))
# Call tenant service fast tier endpoint with caching
# Call tenant service fast tier endpoint with caching (fallback for old tokens)
timeout_config = httpx.Timeout(
connect=1.0, # Connection timeout - very short for cached endpoint
read=5.0, # Read timeout - short for cached lookup
@@ -291,6 +333,15 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
# Check if current tier is in allowed tiers
if current_tier not in [tier.lower() for tier in allowed_tiers]:
tier_names = ', '.join(allowed_tiers)
await self._log_subscription_access(
tenant_id,
user_id,
feature,
current_tier,
False,
"jwt"
)
return {
'allowed': False,
'message': f'This feature requires a {tier_names} subscription plan',
@@ -298,6 +349,15 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
}
# Tier check passed
await self._log_subscription_access(
tenant_id,
user_id,
feature,
current_tier,
True,
"database"
)
return {
'allowed': True,
'message': 'Access granted',
@@ -343,3 +403,64 @@ class SubscriptionMiddleware(BaseHTTPMiddleware):
'current_plan': 'unknown'
}
async def _log_subscription_access(
self,
tenant_id: str,
user_id: str,
requested_feature: str,
current_tier: str,
access_granted: bool,
source: str # "jwt" or "database"
):
"""
Log all subscription-gated access attempts for audit and anomaly detection.
"""
logger.info(
"Subscription access check",
tenant_id=tenant_id,
user_id=user_id,
feature=requested_feature,
tier=current_tier,
granted=access_granted,
source=source,
timestamp=datetime.now(timezone.utc).isoformat()
)
# For denied access, check for suspicious patterns
if not access_granted:
await self._check_for_abuse_patterns(tenant_id, user_id, requested_feature)
async def _check_for_abuse_patterns(
self,
tenant_id: str,
user_id: str,
feature: str
):
"""
Detect potential abuse patterns like repeated premium feature access attempts.
"""
if not self.redis_client:
return
# Track denied attempts in a sliding window
key = f"subscription_denied:{tenant_id}:{user_id}:{feature}"
try:
attempts = await self.redis_client.incr(key)
if attempts == 1:
await self.redis_client.expire(key, 3600) # 1 hour window
# Alert if too many denied attempts (potential bypass attempt)
if attempts > 10:
logger.warning(
"SECURITY: Excessive premium feature access attempts detected",
tenant_id=tenant_id,
user_id=user_id,
feature=feature,
attempts=attempts,
window="1 hour"
)
# Could trigger alert to security team here
except Exception as e:
logger.warning("Failed to track abuse patterns", error=str(e))