100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
"""
|
|
Shared WebSocket utilities for FastAPI endpoints.
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from functools import wraps
|
|
|
|
import base64url
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
from webauthn.helpers.exceptions import InvalidAuthenticationResponse
|
|
|
|
from paskia.fastapi import authz
|
|
from paskia.fastapi.logging import log_ws_close, log_ws_open
|
|
from paskia.globals import passkey
|
|
from paskia.util import pow
|
|
|
|
|
|
def websocket_error_handler(func):
|
|
"""Decorator for WebSocket endpoints that handles common errors."""
|
|
|
|
@wraps(func)
|
|
async def wrapper(ws: WebSocket, *args, **kwargs):
|
|
start = time.perf_counter()
|
|
ws_id = log_ws_open(ws)
|
|
close_code = None
|
|
|
|
try:
|
|
await ws.accept()
|
|
return await func(ws, *args, **kwargs)
|
|
except WebSocketDisconnect as e:
|
|
close_code = e.code
|
|
except authz.AuthException as e:
|
|
await ws.send_json(
|
|
{
|
|
"status": e.status_code,
|
|
**(await authz.auth_error_content(e)),
|
|
}
|
|
)
|
|
except (ValueError, InvalidAuthenticationResponse) as e:
|
|
await ws.send_json({"status": 401, "detail": str(e)})
|
|
except Exception:
|
|
logging.exception("Internal Server Error")
|
|
await ws.send_json({"status": 500, "detail": "Internal Server Error"})
|
|
finally:
|
|
log_ws_close(ws_id, close_code, time.perf_counter() - start)
|
|
|
|
return wrapper
|
|
|
|
|
|
async def require_pow(ws: WebSocket, work: int | None = None) -> None:
|
|
"""Send a PoW challenge and verify the client's solution.
|
|
|
|
Sends: {"pow": {"challenge": "<base64>", "work": 10}}
|
|
Expects: {"pow": "<base64-solution>"}
|
|
|
|
Args:
|
|
ws: WebSocket connection
|
|
work: PoW difficulty level (default: pow.DEFAULT_WORK)
|
|
|
|
Raises:
|
|
ValueError: If the PoW solution is invalid
|
|
"""
|
|
challenge = pow.generate_challenge()
|
|
if work is None:
|
|
work = pow.DEFAULT_WORK
|
|
|
|
await ws.send_json(
|
|
{
|
|
"pow": {
|
|
"challenge": base64url.enc(challenge),
|
|
"work": work,
|
|
}
|
|
}
|
|
)
|
|
|
|
response = await ws.receive_json()
|
|
solution_b64 = response.get("pow")
|
|
if not solution_b64:
|
|
raise ValueError("PoW solution required")
|
|
|
|
try:
|
|
solution = base64url.dec(solution_b64)
|
|
except Exception:
|
|
raise ValueError("Invalid PoW solution encoding")
|
|
|
|
pow.verify_pow(challenge, solution, work)
|
|
|
|
|
|
def validate_origin(ws: WebSocket) -> str:
|
|
"""Extract and validate origin from WebSocket request headers.
|
|
|
|
Raises:
|
|
ValueError: If origin header is missing or not in allowed list
|
|
"""
|
|
origin = ws.headers.get("origin")
|
|
if not origin:
|
|
raise ValueError("Origin header is required for WebSocket connections")
|
|
return passkey.instance.validate_origin(origin)
|