From 731b36b456cdb001ac1dab89989d6cb7ea88aa7f Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Thu, 29 Jan 2026 20:15:01 +0000 Subject: [PATCH] Cleanup auth WS code, refactor to remove duplication and pass proper context. --- paskia/db/operations.py | 4 +++- paskia/fastapi/admin.py | 2 +- paskia/fastapi/remote.py | 44 ++++++++---------------------------- paskia/fastapi/ws.py | 45 ++++++++++--------------------------- paskia/fastapi/wschat.py | 48 ++++++++++++++++++++++++++++++++++++++-- 5 files changed, 71 insertions(+), 72 deletions(-) diff --git a/paskia/db/operations.py b/paskia/db/operations.py index 07f235f..ef940e1 100644 --- a/paskia/db/operations.py +++ b/paskia/db/operations.py @@ -554,6 +554,7 @@ def create_reset_token( token_type: str, *, ctx: SessionContext | None = None, + user: str | None = None, ) -> None: """Create a reset token from a passphrase. @@ -561,13 +562,14 @@ def create_reset_token( For self-service (user creating own recovery link), pass user's ctx. For admin operations, pass admin's ctx. For system operations (bootstrap), pass neither to log no user. + For API operations where ctx is not available but user is known, pass user. """ key = _reset_key(passphrase) if key in _db.reset_tokens: raise ValueError("Reset token already exists") if user_uuid not in _db.users: raise ValueError(f"User {user_uuid} not found") - with _db.transaction("create_reset_token", ctx): + with _db.transaction("create_reset_token", ctx, user=user): _db.reset_tokens[key] = ResetToken( user_uuid=user_uuid, expiry=expiry, token_type=token_type ) diff --git a/paskia/fastapi/admin.py b/paskia/fastapi/admin.py index 738f217..6e54fe6 100644 --- a/paskia/fastapi/admin.py +++ b/paskia/fastapi/admin.py @@ -127,7 +127,7 @@ async def admin_create_org( db.create_org(org, ctx=ctx) # Grant requested permissions to the new org for perm in permissions: - db.add_permission_to_org(str(org.uuid), perm) + db.add_permission_to_org(str(org.uuid), perm, ctx=ctx) return {"uuid": str(org.uuid)} diff --git a/paskia/fastapi/remote.py b/paskia/fastapi/remote.py index 05a50fe..bd19ac6 100644 --- a/paskia/fastapi/remote.py +++ b/paskia/fastapi/remote.py @@ -18,9 +18,9 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect from paskia import db, remoteauth from paskia.authsession import expires from paskia.fastapi.session import infodict -from paskia.fastapi.wschat import authenticate_chat +from paskia.fastapi.wschat import authenticate_and_login from paskia.fastapi.wsutil import validate_origin, websocket_error_handler -from paskia.util import hostutil, passphrase, pow, useragent +from paskia.util import passphrase, pow, useragent # Create a FastAPI subapp for remote auth WebSocket endpoints app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) @@ -270,7 +270,7 @@ async def websocket_remote_auth_permit(ws: WebSocket): 7. Server sends {status: "success", message: "..."} """ - origin = validate_origin(ws) + validate_origin(ws) if remoteauth.instance is None: raise ValueError("Remote authentication is not available") @@ -310,56 +310,30 @@ async def websocket_remote_auth_permit(ws: WebSocket): # Handle authenticate request (no PoW needed - already validated during lookup) if msg.get("authenticate") and request is not None: - cred, new_sign_count = await authenticate_chat(ws, origin) + ctx = await authenticate_and_login(ws) - # Create a session for the REQUESTING device - assert cred.uuid is not None - - session_token = None + session_token = ctx.session.key reset_token = None if request.action == "register": # For registration, create a reset token for device addition - token_str = passphrase.generate() expiry = expires() db.create_reset_token( - user_uuid=cred.user_uuid, + user_uuid=ctx.user.uuid, passphrase=token_str, expiry=expiry, token_type="device addition", + user=str(ctx.user.uuid), ) reset_token = token_str - # Also create a session so the device is logged in - normalized_host = hostutil.normalize_host(request.host) - session_token = db.login( - user_uuid=cred.user_uuid, - credential_uuid=cred.uuid, - sign_count=new_sign_count, - host=normalized_host, - ip=request.ip, - user_agent=request.user_agent, - expiry=expires(), - ) - else: - # Default login action - - normalized_host = hostutil.normalize_host(request.host) - session_token = db.login( - user_uuid=cred.user_uuid, - credential_uuid=cred.uuid, - sign_count=new_sign_count, - host=normalized_host, - ip=request.ip, - user_agent=request.user_agent, - expiry=expires(), - ) # Complete the remote auth request (notifies the waiting device) + cred = db.data().credentials[ctx.session.credential_uuid] completed = await remoteauth.instance.complete_request( token=request.key, session_token=session_token, - user_uuid=cred.user_uuid, + user_uuid=ctx.user.uuid, credential_uuid=cred.uuid, reset_token=reset_token, ) diff --git a/paskia/fastapi/ws.py b/paskia/fastapi/ws.py index 8355f08..97264c0 100644 --- a/paskia/fastapi/ws.py +++ b/paskia/fastapi/ws.py @@ -1,13 +1,13 @@ from fastapi import FastAPI, WebSocket from paskia import db -from paskia.authsession import expires, get_reset +from paskia.authsession import get_reset from paskia.fastapi import authz, remote from paskia.fastapi.session import AUTH_COOKIE, infodict -from paskia.fastapi.wschat import authenticate_chat, register_chat +from paskia.fastapi.wschat import authenticate_and_login, register_chat from paskia.fastapi.wsutil import validate_origin, websocket_error_handler from paskia.globals import passkey -from paskia.util import hostutil, passphrase +from paskia.util import passphrase # Create a FastAPI subapp for WebSocket endpoints app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) @@ -46,7 +46,7 @@ async def websocket_register_add( s = ctx.session # Get user information and determine effective user_name for this registration - user = db.data().users.get(user_uuid) + user = db.data().users[user_uuid] user_name = user.display_name if name is not None: stripped = name.strip() @@ -59,7 +59,7 @@ async def websocket_register_add( # Create a new session and store everything in database metadata = infodict(ws, "authenticated") - token = db.create_credential_session( # type: ignore[attr-defined] + token = db.create_credential_session( user_uuid=user_uuid, credential=credential, reset_key=(s.key if reset is not None else None), @@ -91,41 +91,20 @@ async def websocket_authenticate(ws: WebSocket, auth=AUTH_COOKIE): session_user_uuid = None credential_ids = None if auth: - ctx = db.data().session_ctx(auth, host) - if ctx: - session_user_uuid = ctx.user.uuid + existing_ctx = db.data().session_ctx(auth, host) + if existing_ctx: + session_user_uuid = existing_ctx.user.uuid credential_ids = db.get_user_credential_ids(session_user_uuid) or None - cred, new_sign_count = await authenticate_chat(ws, origin, credential_ids) + ctx = await authenticate_and_login(ws, credential_ids) # If reauth mode, verify the credential belongs to the session's user - if session_user_uuid and cred.user_uuid != session_user_uuid: + if session_user_uuid and ctx.user.uuid != session_user_uuid: raise ValueError("This passkey belongs to a different account") - # Create session and update user/credential in a single transaction - assert cred.uuid is not None - metadata = infodict(ws, "auth") - normalized_host = hostutil.normalize_host(host) - if not normalized_host: - raise ValueError("Host required for session creation") - hostname = normalized_host.split(":")[0] - rp_id = passkey.instance.rp_id - if not (hostname == rp_id or hostname.endswith(f".{rp_id}")): - raise ValueError(f"Host must be the same as or a subdomain of {rp_id}") - - token = db.login( - user_uuid=cred.user_uuid, - credential_uuid=cred.uuid, - sign_count=new_sign_count, - host=normalized_host, - ip=metadata["ip"], - user_agent=metadata["user_agent"], - expiry=expires(), - ) - await ws.send_json( { - "user": str(cred.user_uuid), - "session_token": token, + "user": str(ctx.user.uuid), + "session_token": ctx.session.key, } ) diff --git a/paskia/fastapi/wschat.py b/paskia/fastapi/wschat.py index 316c0ad..0f1c680 100644 --- a/paskia/fastapi/wschat.py +++ b/paskia/fastapi/wschat.py @@ -7,8 +7,12 @@ from uuid import UUID from fastapi import WebSocket from paskia import db -from paskia.db import Credential +from paskia.authsession import expires +from paskia.db import Credential, SessionContext +from paskia.fastapi.session import infodict +from paskia.fastapi.wsutil import validate_origin from paskia.globals import passkey +from paskia.util import hostutil async def register_chat( @@ -31,7 +35,6 @@ async def register_chat( async def authenticate_chat( ws: WebSocket, - origin: str, credential_ids: list[bytes] | None = None, ) -> tuple[Credential, int]: """Run WebAuthn authentication flow and return the credential and new sign count. @@ -39,6 +42,7 @@ async def authenticate_chat( Returns: tuple of (credential, new_sign_count) where new_sign_count comes from WebAuthn verification """ + origin = validate_origin(ws) options, challenge = passkey.instance.auth_generate_options( credential_ids=credential_ids ) @@ -60,3 +64,43 @@ async def authenticate_chat( verification = passkey.instance.auth_verify(authcred, challenge, cred, origin) return cred, verification.new_sign_count + + +async def authenticate_and_login( + ws: WebSocket, + credential_ids: list[bytes] | None = None, +) -> SessionContext: + """Run WebAuthn authentication flow, create session, and return the session context. + + Returns: + SessionContext for the authenticated session + """ + origin = validate_origin(ws) + host = origin.split("://", 1)[1] + normalized_host = hostutil.normalize_host(host) + if not normalized_host: + raise ValueError("Host required for session creation") + hostname = normalized_host.split(":")[0] + rp_id = passkey.instance.rp_id + if not (hostname == rp_id or hostname.endswith(f".{rp_id}")): + raise ValueError(f"Host must be the same as or a subdomain of {rp_id}") + metadata = infodict(ws, "auth") + + cred, new_sign_count = await authenticate_chat(ws, credential_ids) + + # Create session and update user/credential + token = db.login( + user_uuid=cred.user_uuid, + credential_uuid=cred.uuid, + sign_count=new_sign_count, + host=normalized_host, + ip=metadata["ip"], + user_agent=metadata["user_agent"], + expiry=expires(), + ) + + # Fetch and return the full session context + ctx = db.data().session_ctx(token, normalized_host) + if not ctx: + raise ValueError("Failed to create session context") + return ctx