97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
import time
|
|
from functools import wraps
|
|
|
|
import msgspec
|
|
from sanic import errorpages
|
|
from sanic.exceptions import SanicException
|
|
from sanic.log import logger
|
|
from sanic.response import raw, redirect
|
|
|
|
from cista import auth
|
|
from cista.protocol import ErrorMsg
|
|
from cista.sanic_logging import log_ws_close, log_ws_open
|
|
|
|
|
|
def asend(ws, msg):
|
|
"""Send JSON message or bytes to a websocket"""
|
|
return ws.send(msg if isinstance(msg, bytes) else msgspec.json.encode(msg).decode())
|
|
|
|
|
|
def jres(data, **kwargs):
|
|
"""JSON Sanic response, using msgspec encoding"""
|
|
return raw(msgspec.json.encode(data), content_type="application/json", **kwargs)
|
|
|
|
|
|
async def handle_sanic_exception(request, e):
|
|
context, code = {}, 500
|
|
headers = None
|
|
message = str(e)
|
|
if isinstance(e, SanicException):
|
|
context = e.context or {}
|
|
code = e.status_code
|
|
headers = getattr(e, "headers", None)
|
|
if not message or (not request.app.debug and code == 500):
|
|
message = "Internal Server Error"
|
|
message = f"⚠️ {message}" if code < 500 else f"🛑 {message}"
|
|
if code == 500:
|
|
logger.exception(e)
|
|
# Non-browsers get JSON errors
|
|
if "text/html" not in request.headers.accept:
|
|
# Include auth context if present (for SSO auth required responses)
|
|
# Auth must be at top level for paskia library to detect it
|
|
response_data = {"code": code, "message": message, "detail": message, **context}
|
|
return jres(
|
|
response_data,
|
|
status=code,
|
|
headers=headers,
|
|
)
|
|
# Redirections flash the error message via cookies
|
|
if "redirect" in context:
|
|
res = redirect(context["redirect"])
|
|
res.cookies.add_cookie("message", message, max_age=5)
|
|
return res
|
|
# Otherwise use Sanic's default error page
|
|
return errorpages.HTMLRenderer(request, e, debug=request.app.debug).render()
|
|
|
|
|
|
def websocket_wrapper(handler):
|
|
"""Decorator for websocket handlers that catches exceptions and sends them back to the client"""
|
|
|
|
@wraps(handler)
|
|
async def wrapper(request, ws, *args, **kwargs):
|
|
username = getattr(request.ctx, "username", None)
|
|
extra = username or None
|
|
start = time.perf_counter()
|
|
ws_id = log_ws_open(request, extra=extra)
|
|
close_extra = None
|
|
try:
|
|
await auth.verify(request)
|
|
await handler(request, ws, *args, **kwargs)
|
|
except Exception as e:
|
|
context, code, message = {}, 500, str(e) or "Internal Server Error"
|
|
if isinstance(e, SanicException):
|
|
context = e.context or {}
|
|
code = e.status_code
|
|
message = f"⚠️ {message}" if code < 500 else f"🛑 {message}"
|
|
await asend(ws, ErrorMsg({"code": code, "message": message, **context}))
|
|
if not getattr(e, "quiet", False) or code == 500:
|
|
logger.exception(f"{code} {e!r}")
|
|
close_extra = f"{code} {message}"
|
|
raise
|
|
finally:
|
|
duration = time.perf_counter() - start
|
|
close_code = None
|
|
try:
|
|
p = ws.ws_proto
|
|
if p.close_rcvd is not None:
|
|
close_code = p.close_rcvd.code
|
|
elif p.close_sent is not None:
|
|
close_code = p.close_sent.code
|
|
elif getattr(p, "close_code", None) is not None:
|
|
close_code = p.close_code
|
|
except AttributeError:
|
|
pass
|
|
log_ws_close(ws_id, close_code, duration, extra=close_extra)
|
|
|
|
return wrapper
|