diff_app_middleware_api_key_security.py

diff --git a/app/middleware/api_key_security.py b/app/middleware/api_key_security.py
index 3e342fd7..38ce1de7 100644
--- a/app/middleware/api_key_security.py
+++ b/app/middleware/api_key_security.py
@@ -3,17 +3,18 @@ Enhanced API Key Security Middleware
 Enforces rotation policies, IP whitelisting, and usage tracking
 """
 import logging
-import asyncio
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Optional
 from fastapi import Request, Response, HTTPException, status
 from starlette.middleware.base import BaseHTTPMiddleware
-from sqlalchemy.orm import Session
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.db import get_db
+from app.db import get_async_session
 from app.models import ApiKey, Tenant
 from app.services.api_key_security import api_key_security_service
 from app.utils.api_key_utils import hash_api_key
+from app.utils.auth import derive_api_key_role
 
 logger = logging.getLogger(__name__)
 
@@ -49,19 +50,20 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
         # Extract API key from request
         api_key = self._extract_api_key(request)
         if not api_key:
-            # No API key provided - let the request continue for public endpoints
-            return await call_next(request)
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="API key required"
+            )
 
-        try:
-            # Get database session
-            db = next(get_db())
+        endpoint = request.url.path
+        client_ip = self._get_client_ip(request)
 
-            try:
-                # Validate API key and get associated records
+        try:
+            async with get_async_session() as db:
                 api_key_record, tenant = await self._validate_api_key(db, api_key)
 
                 if not api_key_record:
-                    logger.warning(f"Invalid API key attempt from {self._get_client_ip(request)}")
+                    logger.warning(f"Invalid API key attempt from {client_ip}")
                     raise HTTPException(
                         status_code=status.HTTP_401_UNAUTHORIZED,
                         detail="Invalid or inactive API key"
@@ -87,7 +89,6 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
                     )
 
                 # Check IP whitelist
-                client_ip = self._get_client_ip(request)
                 if not self.security_service.is_ip_allowed(client_ip, api_key_record):
                     logger.warning(f"IP not whitelisted: {client_ip} for key {api_key_record.prefix}")
                     raise HTTPException(
@@ -98,21 +99,18 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
                         }
                     )
 
-                # Track usage (async, non-blocking)
-                endpoint = request.url.path
-                asyncio.create_task(
-                    self.security_service.track_usage(
-                        db, api_key_record.id, tenant.id, endpoint
-                    )
-                )
+                # Update last used timestamp and track usage
+                api_key_record.last_used_at = datetime.now(timezone.utc)
+                await self.security_service.track_usage(db, api_key_record.id, tenant.id, endpoint)
+                await db.commit()
 
-                # Update last used timestamp
-                api_key_record.last_used_at = datetime.now()
-                db.commit()
+                role = derive_api_key_role(api_key_record, tenant)
 
                 # Inject security context into request
                 request.state.api_key_id = api_key_record.id
                 request.state.tenant_id = tenant.id
+                request.state.tenant_role = role
+                request.state.tenant_active = tenant.active
                 request.state.tenant = tenant
                 request.state.api_key_prefix = api_key_record.prefix
 
@@ -123,15 +121,12 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
                         "tenant_id": tenant.id,
                         "client_ip": client_ip,
                         "endpoint": endpoint,
-                        "method": request.method
+                        "method": request.method,
+                        "role": role
                     }
                 )
 
-            finally:
-                db.close()
-
         except HTTPException:
-            # Re-raise HTTP exceptions
             raise
         except Exception as e:
             logger.error(f"API key security check failed: {str(e)}", exc_info=True)
@@ -150,29 +145,26 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
 
     def _should_skip_security(self, path: str) -> bool:
         """Check if security should be skipped for this path"""
-        skip_paths = [
+        normalized = path.rstrip("/") or "/"
+        skip_paths = {
             "/health",
             "/docs",
             "/openapi.json",
             "/redoc",
-            "/signup",  # Public signup endpoint
-            "/metrics",  # Monitoring endpoints
+            "/signup",
+            "/metrics",
             "/favicon.ico",
-        ]
-        return any(path.startswith(skip_path) for skip_path in skip_paths)
+        }
+        return normalized in skip_paths
 
     def _extract_api_key(self, request: Request) -> Optional[str]:
         """Extract API key from Authorization header"""
         auth_header = request.headers.get("Authorization")
-        if not auth_header:
-            return None
-
-        if not auth_header.startswith("Bearer "):
+        if not auth_header or not auth_header.startswith("Bearer "):
             return None
+        return auth_header.replace("Bearer ", "", 1)
 
-        return auth_header.replace("Bearer ", "")
-
-    async def _validate_api_key(self, db: Session, api_key: str) -> tuple:
+    async def _validate_api_key(self, db: AsyncSession, api_key: str) -> tuple:
         """
         Validate API key and return associated records
 
@@ -180,27 +172,28 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
             tuple: (api_key_record, tenant) or (None, None) if invalid
         """
         try:
-            # Hash the provided key for lookup
             key_hash = hash_api_key(api_key)
 
-            # Look up API key record
-            api_key_record = db.query(ApiKey).filter(
-                ApiKey.key_hash == key_hash,
-                ApiKey.active == True
-            ).first()
+            api_key_result = await db.execute(
+                select(ApiKey).where(
+                    ApiKey.key_hash == key_hash,
+                    ApiKey.active == True
+                )
+            )
+            api_key_record = api_key_result.scalar_one_or_none()
 
             if not api_key_record:
                 return None, None
 
-            # Check if API key has expired (expires_at field)
-            if api_key_record.expires_at and api_key_record.expires_at < datetime.now():
+            now = datetime.now(timezone.utc)
+            if api_key_record.expires_at and api_key_record.expires_at <= now:
                 logger.warning(f"Expired API key used: {api_key_record.prefix}")
                 return None, None
 
-            # Get associated tenant
-            tenant = db.query(Tenant).filter(
-                Tenant.id == api_key_record.tenant_id
-            ).first()
+            tenant_result = await db.execute(
+                select(Tenant).where(Tenant.id == api_key_record.tenant_id)
+            )
+            tenant = tenant_result.scalar_one_or_none()
 
             return api_key_record, tenant
 
@@ -214,18 +207,14 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
 
         Handles X-Forwarded-For headers for proxy setups
         """
-        # Check X-Forwarded-For header first (for reverse proxies)
         forwarded_for = request.headers.get("X-Forwarded-For")
         if forwarded_for:
-            # Take the first IP in the chain
             return forwarded_for.split(",")[0].strip()
 
-        # Check X-Real-IP header
         real_ip = request.headers.get("X-Real-IP")
         if real_ip:
             return real_ip.strip()
 
-        # Fallback to direct client IP
         if hasattr(request, "client") and request.client:
             return request.client.host
 
@@ -237,4 +226,4 @@ class ApiKeySecurityMiddleware(BaseHTTPMiddleware):
         response.headers["X-Tenant-ID"] = getattr(request.state, "tenant_id", "unknown")
         response.headers["X-Rate-Limit-Remaining"] = "1000"  # Placeholder for rate limiting
         response.headers["X-Content-Type-Options"] = "nosniff"
-        response.headers["X-Frame-Options"] = "DENY"
\ No newline at end of file
+        response.headers["X-Frame-Options"] = "DENY"