155 lines
6.1 KiB
Python
155 lines
6.1 KiB
Python
import asyncio
|
|
import time
|
|
import uuid
|
|
from collections import defaultdict, deque
|
|
from datetime import datetime
|
|
|
|
from fastapi import Depends, HTTPException, Request, WebSocket, WebSocketDisconnect, status
|
|
from fastapi.security import OAuth2PasswordBearer
|
|
from slowapi import Limiter
|
|
from sqlalchemy.orm import Session
|
|
|
|
from core.auth import decode_access_token, decode_access_token_payload, token_issued_before
|
|
from core.database import get_db
|
|
from core.models import User as UserModel
|
|
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
|
|
|
|
|
|
def escape_like(value: str) -> str:
|
|
r"""Neutralise LIKE wildcards in user input used as a search term.
|
|
|
|
Without this, '%' matches every row and '_' matches any character — so a one-char
|
|
search returns the whole table. Use with .ilike(..., escape='\\').
|
|
"""
|
|
return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
|
|
|
|
def parse_uuid(value: str, field: str = "id") -> uuid.UUID:
|
|
"""Path and body params reach us as arbitrary strings; a bare uuid.UUID() on one
|
|
raises ValueError, which surfaces as a 500 and a stack trace instead of a 400."""
|
|
try:
|
|
return uuid.UUID(value)
|
|
except (ValueError, AttributeError, TypeError):
|
|
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid {field}")
|
|
|
|
|
|
def get_client_ip(conn: Request | WebSocket) -> str:
|
|
"""The real external client, as seen through Nginx Proxy Manager.
|
|
|
|
NPM sets X-Forwarded-For to $proxy_add_x_forwarded_for, which *appends* the peer
|
|
it saw to whatever the client sent. So the last entry is always written by NPM and
|
|
anything a client forges lands to its left. Take the rightmost, never the first.
|
|
|
|
This trusts the header because port 8000 is not published: the only things that can
|
|
reach it are containers sharing a network with this one: NPM and the frontend on
|
|
`proxy`, and the database on `tcg-internal`. All three are trusted by topology. If a
|
|
port is ever published, or an untrusted container joins either network, the header
|
|
becomes forgeable and this needs revisiting.
|
|
"""
|
|
forwarded = conn.headers.get("x-forwarded-for")
|
|
if forwarded:
|
|
client = forwarded.rpartition(",")[2].strip()
|
|
if client:
|
|
return client
|
|
# No header: a direct connection, so the peer address is the client.
|
|
return conn.client.host if conn.client else "unknown"
|
|
|
|
|
|
# Per-user key where we can, real client IP otherwise. Without this every request
|
|
# keys on NPM's container address, i.e. one global bucket for the whole internet —
|
|
# eleven logins a minute would lock every user out.
|
|
def get_user_id_from_request(request: Request) -> str:
|
|
auth = request.headers.get("Authorization", "")
|
|
if auth.startswith("Bearer "):
|
|
user_id = decode_access_token(auth[7:])
|
|
if user_id:
|
|
return f"user:{user_id}"
|
|
return get_client_ip(request)
|
|
|
|
|
|
# Default key_func for every @limiter.limit that doesn't override it.
|
|
limiter = Limiter(key_func=get_user_id_from_request)
|
|
|
|
|
|
def get_current_user(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)) -> UserModel:
|
|
payload = decode_access_token_payload(token)
|
|
if not payload or not payload.get("sub"):
|
|
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
|
|
try:
|
|
user_uuid = uuid.UUID(payload["sub"])
|
|
except ValueError:
|
|
# Signature-verified, so not reachable without the key, but a 500 would be wrong.
|
|
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
|
|
user = db.query(UserModel).filter(UserModel.id == user_uuid).first()
|
|
if not user:
|
|
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found")
|
|
# Single chokepoint for revocation: the user row is already loaded, so it's free.
|
|
if token_issued_before(payload, user.token_valid_after):
|
|
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Session expired, please log in again")
|
|
# Throttle to one write per 5 minutes so every authenticated request doesn't hammer the DB
|
|
now = datetime.now()
|
|
if not user.last_active_at or (now - user.last_active_at).total_seconds() > 300:
|
|
user.last_active_at = now
|
|
db.commit()
|
|
return user
|
|
|
|
|
|
## WebSocket connection limiting
|
|
|
|
# slowapi decorates HTTP routes only, and WS handlers can't authenticate before
|
|
# accept() anyway, so connection attempts are limited by client IP instead.
|
|
WS_MAX_CONNECTIONS = 30
|
|
WS_WINDOW_SECONDS = 60
|
|
WS_AUTH_TIMEOUT_SECONDS = 10
|
|
|
|
_ws_attempts: dict[str, deque[float]] = defaultdict(deque)
|
|
|
|
|
|
async def check_ws_rate_limit(websocket: WebSocket) -> bool:
|
|
"""Reject and close if this client is opening sockets too fast.
|
|
|
|
Call before accept(). Closing an un-accepted WebSocket sends an HTTP rejection,
|
|
so a flooding client never gets a live connection or a DB session.
|
|
"""
|
|
key = get_client_ip(websocket)
|
|
now = time.monotonic()
|
|
attempts = _ws_attempts[key]
|
|
while attempts and now - attempts[0] > WS_WINDOW_SECONDS:
|
|
attempts.popleft()
|
|
if len(attempts) >= WS_MAX_CONNECTIONS:
|
|
await websocket.close(code=1008)
|
|
return False
|
|
attempts.append(now)
|
|
# Keys are unbounded otherwise: drop idle ones whenever the table gets large.
|
|
if len(_ws_attempts) > 10000:
|
|
for k in [k for k, v in _ws_attempts.items() if not v or now - v[-1] > WS_WINDOW_SECONDS]:
|
|
del _ws_attempts[k]
|
|
return True
|
|
|
|
|
|
async def accept_and_authenticate_ws(websocket: WebSocket) -> str | None:
|
|
"""Rate-limit, accept, then take the auth token from the first frame.
|
|
|
|
Returns the user id, or None having already closed the socket.
|
|
|
|
The token can't be read before accept() — ASGI delivers no frames until then — so
|
|
the flood limit above is what guards the pre-auth window, and the receive timeout
|
|
stops an accepted-but-never-authenticated socket from pinning a DB session.
|
|
"""
|
|
if not await check_ws_rate_limit(websocket):
|
|
return None
|
|
await websocket.accept()
|
|
try:
|
|
token = await asyncio.wait_for(websocket.receive_text(), timeout=WS_AUTH_TIMEOUT_SECONDS)
|
|
except asyncio.TimeoutError:
|
|
await websocket.close(code=1008)
|
|
return None
|
|
except WebSocketDisconnect:
|
|
return None
|
|
user_id = decode_access_token(token)
|
|
if not user_id:
|
|
await websocket.close(code=1008)
|
|
return None
|
|
return user_id
|