Cleanup auth WS code, refactor to remove duplication and pass proper context.

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
         )

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)}
 

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
+                session_token = ctx.session.key
-                assert cred.uuid is not None
-
-                session_token = None
                 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,
                 )

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)
+        existing_ctx = db.data().session_ctx(auth, host)
-        if ctx:
+        if existing_ctx:
-            session_user_uuid = ctx.user.uuid
+            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),
+            "user": str(ctx.user.uuid),
-            "session_token": token,
+            "session_token": ctx.session.key,
         }
     )

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