32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from app.core.middleware.audit_context import request_id_var, user_id_var, ip_address_var, user_agent_var
|
|
import ulid
|
|
|
|
class AuditMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
# 1. Set IP Address
|
|
x_forwarded_for = request.headers.get("x-forwarded-for")
|
|
ip = x_forwarded_for.split(",")[0].strip() if x_forwarded_for else (request.client.host if request.client else "127.0.0.1")
|
|
|
|
token_ip = ip_address_var.set(ip)
|
|
token_ua = user_agent_var.set(request.headers.get("user-agent", ""))
|
|
|
|
# 2. Set request ID
|
|
req_id = getattr(request.state, "request_id", None)
|
|
if not req_id:
|
|
req_id = str(ulid.ULID())
|
|
request.state.request_id = req_id
|
|
token_req = request_id_var.set(req_id)
|
|
|
|
# Default user_id to None
|
|
token_user = user_id_var.set(None)
|
|
|
|
try:
|
|
response = await call_next(request)
|
|
return response
|
|
finally:
|
|
ip_address_var.reset(token_ip)
|
|
user_agent_var.reset(token_ua)
|
|
request_id_var.reset(token_req)
|
|
user_id_var.reset(token_user)
|