Error handling cleanup for WS too.

passkey/fastapi/ws.py

@@ -1,15 +1,6 @@
-"""
-WebSocket handlers for passkey authentication operations.
-
-This module contains all WebSocket endpoints for:
-- User registration
-- Adding credentials to existing users
-- Device credential addition via token
-- Authentication
-"""
-
 import logging
 from datetime import datetime
+from functools import wraps
 from uuid import UUID
 
 import uuid7
@@ -23,6 +14,25 @@ from ..util import passphrase
 from ..util.tokens import create_token, session_key
 from .session import infodict
 
+
+# WebSocket error handling decorator
+def websocket_error_handler(func):
+    @wraps(func)
+    async def wrapper(ws: WebSocket, *args, **kwargs):
+        try:
+            await ws.accept()
+            return await func(ws, *args, **kwargs)
+        except WebSocketDisconnect:
+            pass
+        except (ValueError, InvalidAuthenticationResponse) as e:
+            await ws.send_json({"detail": str(e)})
+        except Exception:
+            logging.exception("Internal Server Error")
+            await ws.send_json({"detail": "Internal Server Error"})
+
+    return wrapper
+
+
 # Create a FastAPI subapp for WebSocket endpoints
 app = FastAPI()
 
@@ -53,13 +63,12 @@ async def register_chat(
 
 
 @app.websocket("/register")
+@websocket_error_handler
 async def websocket_register_new(
     ws: WebSocket, user_name: str = Query(""), auth=Cookie(None)
 ):
     """Register a new user and with a new passkey credential."""
-    await ws.accept()
-    origin = ws.headers.get("origin")
-    try:
+    origin = ws.headers["origin"]
     user_uuid = uuid7.create()
     # WebAuthn registration
     credential = await register_chat(ws, user_uuid, user_name, origin=origin)
@@ -85,21 +94,13 @@ async def websocket_register_new(
             "session_token": token,
         }
     )
-    except ValueError as e:
-        await ws.send_json({"detail": str(e)})
-    except WebSocketDisconnect:
-        pass
-    except Exception:
-        logging.exception("Internal Server Error")
-        await ws.send_json({"detail": "Internal Server Error"})
 
 
 @app.websocket("/add_credential")
+@websocket_error_handler
 async def websocket_register_add(ws: WebSocket, auth=Cookie(None)):
     """Register a new credential for an existing user."""
-    await ws.accept()
     origin = ws.headers["origin"]
-    try:
     # Try to get either a regular session or a reset session
     reset = passphrase.is_well_formed(auth)
     s = await (get_reset if reset else get_session)(auth)
@@ -111,9 +112,7 @@ async def websocket_register_add(ws: WebSocket, auth=Cookie(None)):
     challenge_ids = await db.instance.get_credentials_by_user_uuid(user_uuid)
 
     # WebAuthn registration
-        credential = await register_chat(
-            ws, user_uuid, user_name, challenge_ids, origin
-        )
+    credential = await register_chat(ws, user_uuid, user_name, challenge_ids, origin)
     if reset:
         # Replace reset session with a new session
         await db.instance.delete_session(s.key)
@@ -134,20 +133,12 @@ async def websocket_register_add(ws: WebSocket, auth=Cookie(None)):
             "message": "New credential added successfully",
         }
     )
-    except ValueError as e:
-        await ws.send_json({"detail": str(e)})
-    except WebSocketDisconnect:
-        pass
-    except Exception:
-        logging.exception("Internal Server Error")
-        await ws.send_json({"detail": "Internal Server Error"})
 
 
 @app.websocket("/authenticate")
+@websocket_error_handler
 async def websocket_authenticate(ws: WebSocket):
-    await ws.accept()
     origin = ws.headers["origin"]
-    try:
     options, challenge = passkey.auth_generate_options()
     await ws.send_json(options)
     # Wait for the client to use his authenticator to authenticate
@@ -173,11 +164,3 @@ async def websocket_authenticate(ws: WebSocket):
             "session_token": token,
         }
     )
-    except (ValueError, InvalidAuthenticationResponse) as e:
-        logging.exception("ValueError")
-        await ws.send_json({"detail": str(e)})
-    except WebSocketDisconnect:
-        pass
-    except Exception:
-        logging.exception("Internal Server Error")
-        await ws.send_json({"detail": "Internal Server Error"})