passkey/fastapi/api.py

@@ -30,15 +30,20 @@ def register_api_routes(app: FastAPI):
     @app.post("/auth/validate")
     async def validate_token(response: Response, auth=Cookie(None)):
         """Lightweight token validation endpoint."""
+        try:
             s = await get_session(auth)
             return {
                 "valid": True,
                 "user_uuid": str(s.user_uuid),
             }
+        except ValueError:
+            response.status_code = 401
+            return {"valid": False}
 
     @app.post("/auth/user-info")
     async def api_user_info(response: Response, auth=Cookie(None)):
         """Get full user information for the authenticated user."""
+        try:
             reset = passphrase.is_well_formed(auth)
             s = await (get_reset if reset else get_session)(auth)
             u = await db.instance.get_user_by_uuid(s.user_uuid)
@@ -91,6 +96,12 @@ def register_api_routes(app: FastAPI):
                 "credentials": credentials,
                 "aaguid_info": aaguid_info,
             }
+        except ValueError as e:
+            response.status_code = 400
+            return {"detail": f"Failed to get user info: {e}"}
+        except Exception:
+            response.status_code = 500
+            return {"detail": "Failed to get user info"}
 
     @app.post("/auth/logout")
     async def api_logout(response: Response, auth=Cookie(None)):
@@ -108,6 +119,7 @@ def register_api_routes(app: FastAPI):
     @app.post("/auth/set-session")
     async def api_set_session(response: Response, auth=Depends(bearer_auth)):
         """Set session cookie from Authorization header. Fetched after login by WebSocket."""
+        try:
             user = await get_session(auth.credentials)
             if not user:
                 raise ValueError("Invalid Authorization header.")
@@ -118,10 +130,25 @@ def register_api_routes(app: FastAPI):
                 "user_uuid": str(user.user_uuid),
             }
 
+        except ValueError as e:
+            response.status_code = 400
+            return {"detail": str(e)}
+        except Exception:
+            response.status_code = 500
+            return {"detail": "Failed to set session"}
+
     @app.delete("/auth/credential/{uuid}")
     async def api_delete_credential(
         response: Response, uuid: UUID, auth: str = Cookie(None)
     ):
         """Delete a specific credential for the current user."""
+        try:
             await delete_credential(uuid, auth)
             return {"message": "Credential deleted successfully"}
+
+        except ValueError as e:
+            response.status_code = 400
+            return {"detail": str(e)}
+        except Exception:
+            response.status_code = 500
+            return {"detail": "Failed to delete credential"}

passkey/fastapi/mainapp.py

@@ -1,10 +1,9 @@
 import contextlib
-import logging
 from contextlib import asynccontextmanager
 from pathlib import Path
 
 from fastapi import Cookie, FastAPI, Request, Response
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
 from ..authsession import get_session
@@ -31,21 +30,6 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
-
-# Global exception handlers
-@app.exception_handler(ValueError)
-async def value_error_handler(request: Request, exc: ValueError):
-    """Handle ValueError exceptions globally with 400 status code."""
-    return JSONResponse(status_code=400, content={"detail": str(exc)})
-
-
-@app.exception_handler(Exception)
-async def general_exception_handler(request: Request, exc: Exception):
-    """Handle all other exceptions globally with 500 status code."""
-    logging.exception("Internal Server Error")
-    return JSONResponse(status_code=500, content={"detail": "Internal server error"})
-
-
 # Mount the WebSocket subapp
 app.mount("/auth/ws", ws.app)
 

passkey/fastapi/reset.py

@@ -15,6 +15,7 @@ def register_reset_routes(app):
     @app.post("/auth/create-link")
     async def api_create_link(request: Request, response: Response, auth=Cookie(None)):
         """Create a device addition link for the authenticated user."""
+        try:
             # Require authentication
             s = await get_session(auth)
 
@@ -37,6 +38,13 @@ def register_reset_routes(app):
                 "expires": expires().isoformat(),
             }
 
+        except ValueError:
+            response.status_code = 401
+            return {"detail": "Authentication required"}
+        except Exception as e:
+            response.status_code = 500
+            return {"detail": f"Failed to create registration link: {str(e)}"}
+
     @app.get("/auth/{reset_token}")
     async def reset_authentication(
         request: Request,

passkey/fastapi/ws.py

@@ -1,6 +1,15 @@
+"""
+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
@@ -14,25 +23,6 @@ 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()
 
@@ -63,12 +53,13 @@ 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."""
-    origin = ws.headers["origin"]
+    await ws.accept()
+    origin = ws.headers.get("origin")
+    try:
         user_uuid = uuid7.create()
         # WebAuthn registration
         credential = await register_chat(ws, user_uuid, user_name, origin=origin)
@@ -94,13 +85,21 @@ 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)
@@ -112,7 +111,9 @@ 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)
@@ -133,12 +134,20 @@ 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
@@ -164,3 +173,11 @@ 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"})