Refactoring done, bugs gone.

passkeyauth/main.py

@@ -5,19 +5,19 @@ This module provides a simple WebAuthn implementation that:
 - Uses WebSocket for real-time communication
 - Supports Resident Keys (discoverable credentials) for passwordless authentication
 - Maintains challenges locally per connection
-- Uses SQLite database for persistent storage of users and credentials
-- Enables true passwordless authentication where users don't need to enter a username
+- Uses async SQLite database for persistent storage of users and credentials
+- Enables true passwordless authentication where users don't need to enter a user_name
 """
 
 from pathlib import Path
 
-import db
 import uuid7
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
-from passkeyauth.passkey import Passkey
+from .db import Credential, db
+from .passkey import Passkey
 
 STATIC_DIR = Path(__file__).parent.parent / "static"
 
@@ -30,6 +30,12 @@ passkey = Passkey(
 app = FastAPI(title="Passkey Auth")
 
 
+@app.on_event("startup")
+async def startup_event():
+    """Initialize database on startup."""
+    await db.init_database()
+
+
 @app.websocket("/ws/new_user_registration")
 async def websocket_register_new(ws: WebSocket):
     """Register a new user and with a new passkey credential."""
@@ -38,40 +44,53 @@ async def websocket_register_new(ws: WebSocket):
         form = await ws.receive_json()
         user_id = uuid7.create().bytes
         user_name = form["user_name"]
-        await register_chat(ws, user_id, username)
+
+        # Generate registration options and handle registration
+        credential, verified = await register_chat(ws, user_id, user_name)
+
         # Store the user in the database
-        await db.create_user(user_name, user_id)
+        await db.create_user(user_id, user_name)
+        await db.store_credential(
+            Credential(
+                credential_id=credential.raw_id,
+                user_id=user_id,
+                aaguid=b"",  # verified.aaguid,
+                public_key=verified.credential_public_key,
+                sign_count=verified.sign_count,
+            )
+        )
         await ws.send_json({"status": "success", "user_id": user_id.hex()})
     except WebSocketDisconnect:
         pass
 
 
-async def register_chat(ws: WebSocket, user_id: bytes, username: str):
+async def register_chat(ws: WebSocket, user_id: bytes, user_name: str):
     """Generate registration options and send them to the client."""
     options, challenge = passkey.reg_generate_options(
         user_id=user_id,
-        username=username,
+        user_name=user_name,
     )
-    await ws.send_text(options)
+    await ws.send_json(options)
     # Wait for the client to use his authenticator to register
     credential = passkey.reg_credential(await ws.receive_json())
-    passkey.reg_verify(credential, challenge)
+    verified_registration = passkey.reg_verify(credential, challenge)
+    return credential, verified_registration
 
 
 @app.websocket("/ws/authenticate")
 async def websocket_authenticate(ws: WebSocket):
     await ws.accept()
     try:
-        options = passkey.auth_generate_options()
+        options, challenge = await passkey.auth_generate_options()
         await ws.send_json(options)
         # Wait for the client to use his authenticator to authenticate
         credential = passkey.auth_credential(await ws.receive_json())
         # Fetch from the database by credential ID
-        stored_cred = await db.fetch_credential(credential.raw_id)
-        # Verify the credential matches the stored data, that is also updated
-        passkey.auth_verify(credential, stored_cred)
-        # Update the credential in the database
+        stored_cred = await db.get_credential_by_id(credential.raw_id)
+        # Verify the credential matches the stored data
+        _ = await passkey.auth_verify(credential, challenge, stored_cred)
         await db.update_credential(stored_cred)
+        await ws.send_json({"status": "success"})
     except WebSocketDisconnect:
         pass
 
@@ -99,8 +118,5 @@ def main():
     )
 
 
-# Initialize database on startup
-db.init_database()
-
 if __name__ == "__main__":
     main()