diff --git a/backend/routes/messaging.py b/backend/routes/messaging.py index b9ff6f8..7f2a2cd 100644 --- a/backend/routes/messaging.py +++ b/backend/routes/messaging.py @@ -28,6 +28,7 @@ from better_profanity import profanity as _bp from security.audit import log_access, log_dm, log_public_chat, log_security from security.profanity import censor_text from security.rate_limit import rate_limit_per_ip +from websocket.utils import authenticate_user router = APIRouter() logger = logging.getLogger("uvicorn.error") @@ -1084,27 +1085,9 @@ class MessaggingSocketManager: async def handle_connection(self, websocket: WebSocket, db: Session): # Initialize subscriptions for this connection self.ws_subscriptions[websocket] = set() - - ws_path = getattr(getattr(websocket, "url", None), "path", None) - if not ws_path and isinstance(getattr(websocket, "scope", None), dict): - ws_path = websocket.scope.get("path") - ws_path = ws_path or "unknown" - headers = {} - if isinstance(getattr(websocket, "scope", None), dict): - headers = {k.decode("latin1"): v.decode("latin1") for k, v in websocket.scope.get("headers", [])} - xff = headers.get("x-forwarded-for") - client_ip = xff.split(",")[0].strip() if xff else (websocket.client.host if websocket.client else None) - - def _log_ws(event: str, user: User | None, **extra: Any) -> None: - log_access( - "ws_event", - path=ws_path, - event=event, - user=user.username if user else None, - user_id=user.id if user else None, - ip=client_ip, - **extra, - ) + + # Import here to avoid circular import + from websocket.handlers import handler_registry while True: try: @@ -1113,698 +1096,33 @@ class MessaggingSocketManager: logger.error(f"Error receiving WebSocket message: {e}") break - type = data["type"] - - def get_current_user_inner() -> User | None: + message_type = data["type"] + handler_info = handler_registry.get_handler(message_type) + + if handler_info: + handler, authRequired = handler_info try: - # Ensure session is in a usable state before querying - try: - db.rollback() - except Exception: - pass + # Authenticate user before calling handler + user = authenticate_user(data, db, authRequired) + # Set user association for authenticated connections + if user: + self.user_by_ws[websocket] = user.id - if data.get("credentials"): - dummy_request = SimpleNamespace() - dummy_request.state = SimpleNamespace() - return get_current_user( - dummy_request, - HTTPAuthorizationCredentials( - scheme=data["credentials"]["scheme"], - credentials=data["credentials"]["credentials"] - ), - db - ) - else: - return None + # Extract inner data to pass to handler + handler_data = data.get("data", {}) + result = await handler(self, websocket, db, user, handler_data) + # If handler returns a value, send it as a WebSocket message + if result is not None: + await websocket.send_json({"type": message_type, "data": result}) + except HTTPException as e: + await self.send_error(websocket, message_type, e) + except WebSocketDisconnect: + raise # Re-raise to close connection except Exception as e: - logger.error(f"Error getting current user: {e}") - try: - db.rollback() - except Exception: - pass - return None - if data.get("credentials"): - dummy_request = SimpleNamespace() - dummy_request.state = SimpleNamespace() - return get_current_user( - dummy_request, - HTTPAuthorizationCredentials( - scheme=data["credentials"]["scheme"], - credentials=data["credentials"]["credentials"] - ), - db - ) - else: - return None - - if type == "getUpdates": - # Handle gap detection - client requests updates from a specific sequence number - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - last_seq = data.get("data", {}).get("lastSeq", 0) - self.last_seq_by_ws[websocket] = last_seq - current_seq = self.sequence_numbers.get(current_user.id, 0) - - # Query database for missed updates - missed_updates = [] - if last_seq > 0 and last_seq < current_seq: - try: - import json - # Get all updates between last_seq and current_seq - update_logs = db.query(UpdateLog).filter( - UpdateLog.user_id == current_user.id, - UpdateLog.sequence > last_seq, - UpdateLog.sequence <= current_seq - ).order_by(UpdateLog.sequence.asc()).all() - - # Each log entry contains a batch of updates with the same sequence number - for log in update_logs: - updates = json.loads(log.updates) - missed_updates.append({ - "seq": log.sequence, - "updates": updates - }) - except Exception as e: - logger.error(f"Failed to retrieve missed updates: {e}") - - # Send missed updates - for batch in missed_updates: - await websocket.send_json({ - "type": "updates", - "seq": batch["seq"], - "updates": batch["updates"] - }) - - await websocket.send_json({ - "type": "getUpdates", - "data": { - "status": "ok", - "lastSeq": current_seq, - "missedCount": len(missed_updates) - } - }) - # Update the websocket's last sequence tracking - self.last_seq_by_ws[websocket] = current_seq - _log_ws("getUpdates", current_user, last_seq=last_seq, current_seq=current_seq, missed_count=len(missed_updates)) - except HTTPException as e: - _log_ws("getUpdates_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "ping": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if current_user: - self.user_by_ws[websocket] = current_user.id - # Set user online in DB - current_user.online = True - current_user.last_seen = datetime.now() - db.commit() - # Add to online users - self.online_users.add(current_user.id) - # Broadcast status change - await self.broadcast_status_change(current_user.id, True, current_user.last_seen.isoformat(), db) - else: - await websocket.send_json({ - "type": "ping", - "data": { - "status": "error", - "error": { - "detail": "Failed to authorize", - "code": 401 - } - } - }) - _log_ws("ping_error", current_user) - except HTTPException: - await websocket.send_json({ - "type": "ping", - "data": { - "status": "error", - "error": { - "detail": "Failed to authorize", - "code": 401 - } - } - }) - _log_ws("ping_error", current_user) - await websocket.send_json({"type": "ping", "data": {"status": "success"}}) - _log_ws("ping", current_user) - elif type == "getMessages": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - self.user_by_ws[websocket] = current_user.id - - await websocket.send_json({"type": type, "data": await get_messages(current_user, db)}) - _log_ws("getMessages", current_user) - except HTTPException as e: - _log_ws("getMessages_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "sendMessage": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - self.user_by_ws[websocket] = current_user.id - - message_request: SendMessageRequest = SendMessageRequest.model_validate(data["data"]) - - # Call internal function directly (rate limiting is handled at infrastructure level via Caddy) - response = await _send_message_internal(message_request, current_user, db, []) - await self.broadcast({ - "type": "newMessage", - "data": response["message"] - }, db) - - await websocket.send_json({"type": type, "data": response}) - _log_ws("sendMessage", current_user, message_id=response["message"]["id"]) - except HTTPException as e: - _log_ws("sendMessage_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "dmSend": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - self.user_by_ws[websocket] = current_user.id - payload = data["data"] - required = ["recipientId", "iv", "ciphertext", "salt", "iv2", "wrappedMk"] - for key in required: - if key not in payload: - raise HTTPException(status_code=400, detail=f"Missing {key}") - env = DMEnvelope( - sender_id=current_user.id, - recipient_id=int(payload["recipientId"]), - iv_b64=payload["iv"], - ciphertext_b64=payload["ciphertext"], - salt_b64=payload["salt"], - iv2_b64=payload["iv2"], - wrapped_mk_b64=payload["wrappedMk"], - reply_to_id=payload.get("replyToId") if isinstance(payload.get("replyToId"), int) else None, - ) - db.add(env) - db.commit() - db.refresh(env) - - payload = { - "type": "dmNew", - "data": { - "id": env.id, - "senderId": env.sender_id, - "recipientId": env.recipient_id, - "iv": env.iv_b64, - "ciphertext": env.ciphertext_b64, - "salt": env.salt_b64, - "iv2": env.iv2_b64, - "wrappedMk": env.wrapped_mk_b64, - "timestamp": env.timestamp.isoformat(), - "replyToId": env.reply_to_id, - } - } - - # Send push notification for DM - try: - await push_service.send_dm_notification(db, env, current_user) - except Exception as e: - logger.error(f"Failed to send push notification for DM {env.id}: {e}") - - await self.send_update_to_user(env.recipient_id, "dmNew", payload["data"], db); - await websocket.send_json({"type": type, "data": {"status": "ok", "id": env.id}}); - await self.send_update_to_user(env.sender_id, "dmNew", payload["data"], db); - - _log_ws("dmSend", current_user, dm_envelope_id=env.id, recipient_id=env.recipient_id) - log_dm( - "message_sent_ws", - dm_envelope_id=env.id, - sender_id=current_user.id, - sender_username=current_user.username, - recipient_id=env.recipient_id, - reply_to=env.reply_to_id, - ) - except HTTPException as e: - _log_ws("dmSend_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "editMessage": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - message_id = data["data"]["message_id"] - request: EditMessageRequest = EditMessageRequest.model_validate(data["data"]) - - response = await edit_message(message_id, request, current_user, db) - await self.broadcast({ - "type": "messageEdited", - "data": response["message"] - }, db) - - await websocket.send_json({"type": type, "data": response}) - _log_ws("editMessage", current_user, message_id=message_id) - except HTTPException as e: - _log_ws("editMessage_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "dmEdit": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - payload = data["data"] - env_id = int(payload["id"]) - env: DMEnvelope | None = db.query(DMEnvelope).filter(DMEnvelope.id == env_id).first() - if not env: - raise HTTPException(status_code=404, detail="DM not found") - if env.sender_id != current_user.id: - raise HTTPException(status_code=403, detail="You can only edit your own messages") - - # Replace ciphertext and iv - env.iv_b64 = payload["iv"] - env.ciphertext_b64 = payload["ciphertext"] - env.iv2_b64 = payload["iv2"] - env.wrapped_mk_b64 = payload["wrappedMk"] - env.salt_b64 = payload["salt"] - db.commit() - db.refresh(env) - - payload_ws = { - "type": "dmEdited", - "data": { - "id": env.id, - "senderId": env.sender_id, - "recipientId": env.recipient_id, - "iv": env.iv_b64, - "ciphertext": env.ciphertext_b64, - "iv2": env.iv2_b64, - "wrappedMk": env.wrapped_mk_b64, - "salt": env.salt_b64, - "timestamp": env.timestamp.isoformat(), - } - } - await self.send_update_to_user(env.recipient_id, "dmEdited", payload_ws["data"], db) - await self.send_update_to_user(env.sender_id, "dmEdited", payload_ws["data"], db) - await websocket.send_json({"type": type, "data": {"status": "ok", "id": env.id}}) - - _log_ws("dmEdit", current_user, dm_envelope_id=env.id) - log_dm( - "message_edited", - dm_envelope_id=env.id, - user_id=current_user.id, - username=current_user.username, - ) - except HTTPException as e: - _log_ws("dmEdit_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "dmDelete": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - payload = data["data"] - env_id = int(payload["id"]) - env: DMEnvelope | None = db.query(DMEnvelope).filter(DMEnvelope.id == env_id).first() - if not env: - raise HTTPException(status_code=404, detail="DM not found") - if env.sender_id != current_user.id: - raise HTTPException(status_code=403, detail="You can only delete your own messages") - - db.delete(env) - db.commit() - - payload_ws = { - "type": "dmDeleted", - "data": { - "id": env_id, - "senderId": current_user.id, - "recipientId": payload.get("recipientId") - } - } - await self.send_update_to_user(env.recipient_id, "dmDeleted", payload_ws["data"], db) - await websocket.send_json({"type": type, "data": {"status": "ok", "id": env_id}}) - await self.send_update_to_user(env.sender_id, "dmDeleted", payload_ws["data"], db) - - _log_ws("dmDelete", current_user, dm_envelope_id=env_id) - log_dm( - "message_deleted", - dm_envelope_id=env_id, - user_id=current_user.id, - username=current_user.username, - recipient_id=env.recipient_id, - ) - except HTTPException as e: - _log_ws("dmDelete_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "deleteMessage": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - message_id = data["data"]["message_id"] - response = await delete_message(message_id, current_user, db) - await self.broadcast({ - "type": "messageDeleted", - "data": {"message_id": message_id} - }, db) - - await websocket.send_json({"type": type, "data": response}) - _log_ws("deleteMessage", current_user, message_id=message_id) - except HTTPException as e: - _log_ws("deleteMessage_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "addReaction": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - request_data = data["data"] - reaction_request = ReactionRequest( - message_id=request_data["message_id"], - emoji=request_data["emoji"] - ) - - response = await add_reaction(reaction_request, current_user, db) - - # Broadcast reaction update - await self.broadcast({ - "type": "reactionUpdate", - "data": { - "message_id": request_data["message_id"], - "emoji": request_data["emoji"], - "action": response["action"], - "user_id": current_user.id, - "username": current_user.username, - "reactions": response["reactions"] - } - }, db) - - await websocket.send_json({"type": type, "data": response}) - _log_ws("addReaction", current_user, message_id=request_data["message_id"], emoji=request_data["emoji"], action=response["action"]) - except HTTPException as e: - _log_ws("addReaction_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "addDmReaction": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - request_data = data["data"] - reaction_request = DMReactionRequest( - dm_envelope_id=request_data["dm_envelope_id"], - emoji=request_data["emoji"] - ) - - response = await add_dm_reaction(reaction_request, current_user, db) - - # Broadcast reaction update - await self.broadcast({ - "type": "dmReactionUpdate", - "data": { - "dm_envelope_id": request_data["dm_envelope_id"], - "emoji": request_data["emoji"], - "action": response["action"], - "user_id": current_user.id, - "username": current_user.username, - "reactions": response["reactions"] - } - }, db) - - await websocket.send_json({"type": type, "data": response}) - _log_ws("addDmReaction", current_user, dm_envelope_id=request_data["dm_envelope_id"], emoji=request_data["emoji"], action=response["action"]) - except HTTPException as e: - _log_ws("addDmReaction_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "call_signaling": - # Forward WebRTC signaling between peers - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - self.user_by_ws[websocket] = current_user.id - - payload = data.get("data") or {} - to_user_id = int(payload.get("toUserId") or 0) - if not to_user_id: - raise HTTPException(status_code=400, detail="Missing toUserId") - - # Ensure sender is set by the server - payload["fromUserId"] = current_user.id - payload["fromUsername"] = current_user.username - - await self.send_to_user(to_user_id, { - "type": "call_signaling", - "data": payload - }) - - # Optional ack - await websocket.send_json({"type": "call_signaling", "data": {"status": "ok"}}) - _log_ws("call_signaling", current_user, to_user_id=to_user_id) - except HTTPException as e: - _log_ws("call_signaling_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - elif type == "call_video_toggle": - # Forward video toggle state between peers - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - self.user_by_ws[websocket] = current_user.id - - payload = data.get("data") or {} - to_user_id = int(payload.get("toUserId") or 0) - if not to_user_id: - raise HTTPException(status_code=400, detail="Missing toUserId") - - # Ensure sender is set by the server - payload["fromUserId"] = current_user.id - - await self.send_update_to_user(to_user_id, "call_signaling", { - "type": "call_video_toggle", - "fromUserId": current_user.id, - "toUserId": to_user_id, - "data": {"enabled": payload.get("enabled", False)} - }, db) - - await websocket.send_json({"type": "call_video_toggle", "data": {"status": "ok"}}) - except HTTPException as e: - _log_ws("call_video_toggle_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - else: - _log_ws("call_video_toggle", current_user, to_user_id=to_user_id, enabled=payload.get("enabled", False)) - elif type == "call_screen_share_toggle": - # Forward screen share toggle state between peers - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - self.user_by_ws[websocket] = current_user.id - - payload = data.get("data") or {} - to_user_id = int(payload.get("toUserId") or 0) - if not to_user_id: - raise HTTPException(status_code=400, detail="Missing toUserId") - - # Ensure sender is set by the server - payload["fromUserId"] = current_user.id - - await self.send_update_to_user(to_user_id, "call_signaling", { - "type": "call_screen_share_toggle", - "fromUserId": current_user.id, - "toUserId": to_user_id, - "data": {"enabled": payload.get("enabled", False)} - }, db) - - await websocket.send_json({"type": "call_screen_share_toggle", "data": {"status": "ok"}}) - except HTTPException as e: - _log_ws("call_screen_share_toggle_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - else: - _log_ws("call_screen_share_toggle", current_user, to_user_id=to_user_id, enabled=payload.get("enabled", False)) - elif type == "subscribeStatus": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - user_id_to_subscribe = int(data["data"]["userId"]) - self.ws_subscriptions[websocket].add(user_id_to_subscribe) - - # Get current status of the user - target_user = db.query(User).filter(User.id == user_id_to_subscribe).first() - if target_user: - await websocket.send_json({ - "type": "statusUpdate", - "data": { - "userId": user_id_to_subscribe, - "online": target_user.online, - "lastSeen": target_user.last_seen.isoformat() if target_user.last_seen else None - } - }) - else: - await websocket.send_json({ - "type": "subscribeStatus", - "data": {"status": "error", "error": "User not found"} - }) - except HTTPException as e: - _log_ws("subscribeStatus_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - else: - _log_ws("subscribeStatus", current_user, target_user_id=user_id_to_subscribe) - elif type == "unsubscribeStatus": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - user_id_to_unsubscribe = int(data["data"]["userId"]) - self.ws_subscriptions[websocket].discard(user_id_to_unsubscribe) - - await websocket.send_json({"type": "unsubscribeStatus", "data": {"status": "ok"}}) - except HTTPException as e: - _log_ws("unsubscribeStatus_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - else: - _log_ws("unsubscribeStatus", current_user, target_user_id=user_id_to_unsubscribe) - elif type == "typing": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - was_typing = self.typing_state.get(current_user.id, False) - self.typing_users[current_user.id] = time.time() - is_now_typing = True - - # Only send update if state changed (started typing) - if not was_typing: - self.typing_state[current_user.id] = True - # Broadcast to all connected users - await self.broadcast({ - "type": "typing", - "data": { - "userId": current_user.id, - "username": current_user.username - } - }, db) - - # No confirmation response - privacy protection - except HTTPException as e: - _log_ws("typing_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - else: - _log_ws("typing", current_user) - elif type == "stopTyping": - current_user: User | None = None - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - was_typing = self.typing_state.get(current_user.id, False) - if current_user.id in self.typing_users: - del self.typing_users[current_user.id] - - # Only send update if state changed (stopped typing) - if was_typing: - self.typing_state[current_user.id] = False - # Broadcast to all connected users - await self.broadcast({ - "type": "stopTyping", - "data": { - "userId": current_user.id, - "username": current_user.username - } - }, db) - - # No confirmation response - privacy protection - except HTTPException as e: - _log_ws("stopTyping_error", current_user, detail=str(getattr(e, "detail", e))) - await self.send_error(websocket, type, e) - else: - _log_ws("stopTyping", current_user) - elif type == "dmTyping": - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - recipient_id = int(data["data"]["recipientId"]) - - if current_user.id not in self.dm_typing_users: - self.dm_typing_users[current_user.id] = {} - if current_user.id not in self.dm_typing_state: - self.dm_typing_state[current_user.id] = {} - - was_typing = self.dm_typing_state[current_user.id].get(recipient_id, False) - self.dm_typing_users[current_user.id][recipient_id] = time.time() - - # Only send update if state changed (started typing) - if not was_typing: - self.dm_typing_state[current_user.id][recipient_id] = True - # Send only to recipient - await self.send_update_to_user(recipient_id, "dmTyping", { - "userId": current_user.id, - "username": current_user.username - }, db) - - # No confirmation response - privacy protection - except HTTPException as e: - await self.send_error(websocket, type, e) - elif type == "stopDmTyping": - try: - current_user = get_current_user_inner() - if not current_user: - raise HTTPException(401) - - recipient_id = int(data["data"]["recipientId"]) - - was_typing = False - if current_user.id in self.dm_typing_state: - was_typing = self.dm_typing_state[current_user.id].get(recipient_id, False) - - if current_user.id in self.dm_typing_users and recipient_id in self.dm_typing_users[current_user.id]: - del self.dm_typing_users[current_user.id][recipient_id] - if not self.dm_typing_users[current_user.id]: - del self.dm_typing_users[current_user.id] - - # Only send update if state changed (stopped typing) - if was_typing: - if current_user.id in self.dm_typing_state: - self.dm_typing_state[current_user.id][recipient_id] = False - # Send only to recipient - await self.send_update_to_user(recipient_id, "stopDmTyping", { - "userId": current_user.id, - "username": current_user.username - }, db) - - # No confirmation response - privacy protection - except HTTPException as e: - await self.send_error(websocket, type, e) + logger.error(f"Error in handler for {message_type}: {e}") + await self.send_error(websocket, message_type, HTTPException(500, "Internal server error")) else: - await websocket.send_json({"type": type, "error": {"code": 400, "detail": "Invalid type"}}) + await websocket.send_json({"type": message_type, "error": {"code": 400, "detail": "Invalid type"}}) async def disconnect(self, websocket: WebSocket, code: int = 1000, message: str | None = None): try: diff --git a/backend/websocket/__init__.py b/backend/websocket/__init__.py new file mode 100644 index 0000000..efe848a --- /dev/null +++ b/backend/websocket/__init__.py @@ -0,0 +1,7 @@ +from websocket.registry import WebSocketHandlerRegistry + +# Note: handler_registry and websocket_handler are not imported here to avoid circular dependency +# Import them directly from websocket.handlers when needed + +__all__ = ["WebSocketHandlerRegistry"] + diff --git a/backend/websocket/handlers.py b/backend/websocket/handlers.py new file mode 100644 index 0000000..6b949e6 --- /dev/null +++ b/backend/websocket/handlers.py @@ -0,0 +1,576 @@ +from datetime import datetime +import json +import logging +import time +from typing import Any +from fastapi import HTTPException, WebSocket +from sqlalchemy.orm import Session + +from websocket.registry import WebSocketHandlerRegistry +from websocket.utils import authenticate_user +from routes.messaging import ( + MessaggingSocketManager, + convert_message, + convert_dm_envelope, + _send_message_internal, + get_messages, + edit_message, + delete_message, + add_reaction, + add_dm_reaction, +) +from models import ( + User, + SendMessageRequest, + EditMessageRequest, + DMEnvelope, + ReactionRequest, + DMReactionRequest, + UpdateLog, +) +from security.audit import log_access, log_dm, log_public_chat +from routes.account import convert_user + +logger = logging.getLogger("uvicorn.error") + +# Create global registry instance +handler_registry = WebSocketHandlerRegistry() + +# Create decorator alias +websocket_handler = handler_registry.register + + +def log(manager: MessaggingSocketManager, websocket: WebSocket, user: User | None, event: str, **extra: Any) -> None: + """Log WebSocket event.""" + ws_path = getattr(getattr(websocket, "url", None), "path", None) + if not ws_path and isinstance(getattr(websocket, "scope", None), dict): + ws_path = websocket.scope.get("path") + ws_path = ws_path or "unknown" + headers = {} + if isinstance(getattr(websocket, "scope", None), dict): + headers = {k.decode("latin1"): v.decode("latin1") for k, v in websocket.scope.get("headers", [])} + xff = headers.get("x-forwarded-for") + client_ip = xff.split(",")[0].strip() if xff else (websocket.client.host if websocket.client else None) + + log_access( + "ws_event", + path=ws_path, + event=event, + user=user.username if user else None, + user_id=user.id if user else None, + ip=client_ip, + **extra, + ) + + +@websocket_handler("getUpdates", authRequired=True) +async def getUpdates(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Handle gap detection - client requests updates from a specific sequence number.""" + last_seq = data.get("lastSeq", 0) + manager.last_seq_by_ws[websocket] = last_seq + current_seq = manager.sequence_numbers.get(user.id, 0) + + # Query database for missed updates + missed_updates = [] + if last_seq > 0 and last_seq < current_seq: + try: + # Get all updates between last_seq and current_seq + update_logs = db.query(UpdateLog).filter( + UpdateLog.user_id == user.id, + UpdateLog.sequence > last_seq, + UpdateLog.sequence <= current_seq + ).order_by(UpdateLog.sequence.asc()).all() + + # Each log entry contains a batch of updates with the same sequence number + for log_entry in update_logs: + updates = json.loads(log_entry.updates) + missed_updates.append({ + "seq": log_entry.sequence, + "updates": updates + }) + except Exception as e: + logger.error(f"Failed to retrieve missed updates: {e}") + + # Send missed updates directly (not through return value) + for batch in missed_updates: + await websocket.send_json({ + "type": "updates", + "seq": batch["seq"], + "updates": batch["updates"] + }) + + # Update the websocket's last sequence tracking + manager.last_seq_by_ws[websocket] = current_seq + log(manager, websocket, user, "getUpdates", last_seq=last_seq, current_seq=current_seq, missed_count=len(missed_updates)) + + return { + "status": "ok", + "lastSeq": current_seq, + "missedCount": len(missed_updates) + } + + +@websocket_handler("ping", authRequired=True) +async def ping(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Handle ping - authenticate and set user online.""" + # Set user online in DB + user.online = True + user.last_seen = datetime.now() + db.commit() + # Add to online users + manager.online_users.add(user.id) + # Broadcast status change + await manager.broadcast_status_change(user.id, True, user.last_seen.isoformat(), db) + + log(manager, websocket, user, "ping") + return {"status": "success"} + + +@websocket_handler("getMessages", authRequired=True) +async def getMessages(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Get all public chat messages.""" + result = await get_messages(user, db) + log(manager, websocket, user, "getMessages") + return result + + +@websocket_handler("sendMessage", authRequired=True) +async def sendMessage(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Send a public chat message.""" + message_request: SendMessageRequest = SendMessageRequest.model_validate(data) + + # Call internal function directly (rate limiting is handled at infrastructure level via Caddy) + response = await _send_message_internal(message_request, user, db, []) + await manager.broadcast({ + "type": "newMessage", + "data": response["message"] + }, db) + + log(manager, websocket, user, "sendMessage", message_id=response["message"]["id"]) + return response + + +@websocket_handler("dmSend", authRequired=True) +async def dmSend(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Send a direct message.""" + payload = data + required = ["recipientId", "iv", "ciphertext", "salt", "iv2", "wrappedMk"] + for key in required: + if key not in payload: + raise HTTPException(status_code=400, detail=f"Missing {key}") + + env = DMEnvelope( + sender_id=user.id, + recipient_id=int(payload["recipientId"]), + iv_b64=payload["iv"], + ciphertext_b64=payload["ciphertext"], + salt_b64=payload["salt"], + iv2_b64=payload["iv2"], + wrapped_mk_b64=payload["wrappedMk"], + reply_to_id=payload.get("replyToId") if isinstance(payload.get("replyToId"), int) else None, + ) + db.add(env) + db.commit() + db.refresh(env) + + payload_ws = { + "type": "dmNew", + "data": { + "id": env.id, + "senderId": env.sender_id, + "recipientId": env.recipient_id, + "iv": env.iv_b64, + "ciphertext": env.ciphertext_b64, + "salt": env.salt_b64, + "iv2": env.iv2_b64, + "wrappedMk": env.wrapped_mk_b64, + "timestamp": env.timestamp.isoformat(), + "replyToId": env.reply_to_id, + } + } + + # Send push notification for DM + try: + from push_service import push_service + await push_service.send_dm_notification(db, env, user) + except Exception as e: + logger.error(f"Failed to send push notification for DM {env.id}: {e}") + + await manager.send_update_to_user(env.recipient_id, "dmNew", payload_ws["data"], db) + await manager.send_update_to_user(env.sender_id, "dmNew", payload_ws["data"], db) + + log(manager, websocket, user, "dmSend", dm_envelope_id=env.id, recipient_id=env.recipient_id) + log_dm( + "message_sent_ws", + dm_envelope_id=env.id, + sender_id=user.id, + sender_username=user.username, + recipient_id=env.recipient_id, + reply_to=env.reply_to_id, + ) + + return {"status": "ok", "id": env.id} + + +@websocket_handler("editMessage", authRequired=True) +async def editMessage(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Edit a public chat message.""" + from types import SimpleNamespace + + message_id = data["message_id"] + request: EditMessageRequest = EditMessageRequest.model_validate(data) + + # Create a dummy request object for the HTTP endpoint function + dummy_request = SimpleNamespace() + response = await edit_message(dummy_request, message_id, request, user, db) + await manager.broadcast({ + "type": "messageEdited", + "data": response["message"] + }, db) + + log(manager, websocket, user, "editMessage", message_id=message_id) + return response + + +@websocket_handler("dmEdit", authRequired=True) +async def dmEdit(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Edit a direct message.""" + payload = data + env_id = int(payload["id"]) + env: DMEnvelope | None = db.query(DMEnvelope).filter(DMEnvelope.id == env_id).first() + if not env: + raise HTTPException(status_code=404, detail="DM not found") + if env.sender_id != user.id: + raise HTTPException(status_code=403, detail="You can only edit your own messages") + + # Replace ciphertext and iv + env.iv_b64 = payload["iv"] + env.ciphertext_b64 = payload["ciphertext"] + env.iv2_b64 = payload["iv2"] + env.wrapped_mk_b64 = payload["wrappedMk"] + env.salt_b64 = payload["salt"] + db.commit() + db.refresh(env) + + payload_ws = { + "type": "dmEdited", + "data": { + "id": env.id, + "senderId": env.sender_id, + "recipientId": env.recipient_id, + "iv": env.iv_b64, + "ciphertext": env.ciphertext_b64, + "iv2": env.iv2_b64, + "wrappedMk": env.wrapped_mk_b64, + "salt": env.salt_b64, + "timestamp": env.timestamp.isoformat(), + } + } + await manager.send_update_to_user(env.recipient_id, "dmEdited", payload_ws["data"], db) + await manager.send_update_to_user(env.sender_id, "dmEdited", payload_ws["data"], db) + + log(manager, websocket, user, "dmEdit", dm_envelope_id=env.id) + log_dm( + "message_edited", + dm_envelope_id=env.id, + user_id=user.id, + username=user.username, + ) + + return {"status": "ok", "id": env.id} + + +@websocket_handler("dmDelete", authRequired=True) +async def dmDelete(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Delete a direct message.""" + payload = data + env_id = int(payload["id"]) + env: DMEnvelope | None = db.query(DMEnvelope).filter(DMEnvelope.id == env_id).first() + if not env: + raise HTTPException(status_code=404, detail="DM not found") + if env.sender_id != user.id: + raise HTTPException(status_code=403, detail="You can only delete your own messages") + + db.delete(env) + db.commit() + + payload_ws = { + "type": "dmDeleted", + "data": { + "id": env_id, + "senderId": user.id, + "recipientId": payload.get("recipientId") + } + } + await manager.send_update_to_user(env.recipient_id, "dmDeleted", payload_ws["data"], db) + await manager.send_update_to_user(env.sender_id, "dmDeleted", payload_ws["data"], db) + + log(manager, websocket, user, "dmDelete", dm_envelope_id=env_id) + log_dm( + "message_deleted", + dm_envelope_id=env_id, + user_id=user.id, + username=user.username, + recipient_id=env.recipient_id, + ) + + return {"status": "ok", "id": env_id} + + +@websocket_handler("deleteMessage", authRequired=True) +async def deleteMessage(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Delete a public chat message.""" + message_id = data["message_id"] + response = await delete_message(message_id, user, db) + await manager.broadcast({ + "type": "messageDeleted", + "data": {"message_id": message_id} + }, db) + + log(manager, websocket, user, "deleteMessage", message_id=message_id) + return response + + +@websocket_handler("addReaction", authRequired=True) +async def addReaction(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Add or remove a reaction to a public chat message.""" + reaction_request = ReactionRequest( + message_id=data["message_id"], + emoji=data["emoji"] + ) + + response = await add_reaction(reaction_request, user, db) + + # Broadcast reaction update + await manager.broadcast({ + "type": "reactionUpdate", + "data": { + "message_id": data["message_id"], + "emoji": data["emoji"], + "action": response["action"], + "user_id": user.id, + "username": user.username, + "reactions": response["reactions"] + } + }, db) + + log(manager, websocket, user, "addReaction", message_id=data["message_id"], emoji=data["emoji"], action=response["action"]) + return response + + +@websocket_handler("addDmReaction", authRequired=True) +async def addDmReaction(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Add or remove a reaction to a direct message.""" + reaction_request = DMReactionRequest( + dm_envelope_id=data["dm_envelope_id"], + emoji=data["emoji"] + ) + + response = await add_dm_reaction(reaction_request, user, db) + + # Broadcast reaction update + await manager.broadcast({ + "type": "dmReactionUpdate", + "data": { + "dm_envelope_id": data["dm_envelope_id"], + "emoji": data["emoji"], + "action": response["action"], + "user_id": user.id, + "username": user.username, + "reactions": response["reactions"] + } + }, db) + + log(manager, websocket, user, "addDmReaction", dm_envelope_id=data["dm_envelope_id"], emoji=data["emoji"], action=response["action"]) + return response + + +@websocket_handler("call_signaling", authRequired=True) +async def call_signaling(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Forward WebRTC signaling between peers.""" + payload = data or {} + to_user_id = int(payload.get("toUserId") or 0) + if not to_user_id: + raise HTTPException(status_code=400, detail="Missing toUserId") + + # Ensure sender is set by the server + payload["fromUserId"] = user.id + payload["fromUsername"] = user.username + + await manager.send_to_user(to_user_id, { + "type": "call_signaling", + "data": payload + }) + + log(manager, websocket, user, "call_signaling", to_user_id=to_user_id) + return {"status": "ok"} + + +@websocket_handler("call_video_toggle", authRequired=True) +async def call_video_toggle(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Forward video toggle state between peers.""" + payload = data or {} + to_user_id = int(payload.get("toUserId") or 0) + if not to_user_id: + raise HTTPException(status_code=400, detail="Missing toUserId") + + await manager.send_update_to_user(to_user_id, "call_signaling", { + "type": "call_video_toggle", + "fromUserId": user.id, + "toUserId": to_user_id, + "data": {"enabled": payload.get("enabled", False)} + }, db) + + log(manager, websocket, user, "call_video_toggle", to_user_id=to_user_id, enabled=payload.get("enabled", False)) + return {"status": "ok"} + + +@websocket_handler("call_screen_share_toggle", authRequired=True) +async def call_screen_share_toggle(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Forward screen share toggle state between peers.""" + payload = data or {} + to_user_id = int(payload.get("toUserId") or 0) + if not to_user_id: + raise HTTPException(status_code=400, detail="Missing toUserId") + + await manager.send_update_to_user(to_user_id, "call_signaling", { + "type": "call_screen_share_toggle", + "fromUserId": user.id, + "toUserId": to_user_id, + "data": {"enabled": payload.get("enabled", False)} + }, db) + + log(manager, websocket, user, "call_screen_share_toggle", to_user_id=to_user_id, enabled=payload.get("enabled", False)) + return {"status": "ok"} + + +@websocket_handler("subscribeStatus", authRequired=True) +async def subscribeStatus(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Subscribe to status updates for a user.""" + user_id_to_subscribe = int(data["userId"]) + manager.ws_subscriptions[websocket].add(user_id_to_subscribe) + + # Get current status of the user + target_user = db.query(User).filter(User.id == user_id_to_subscribe).first() + if target_user: + # Send current status directly (not through return value) + await websocket.send_json({ + "type": "statusUpdate", + "data": { + "userId": user_id_to_subscribe, + "online": target_user.online, + "lastSeen": target_user.last_seen.isoformat() if target_user.last_seen else None + } + }) + log(manager, websocket, user, "subscribeStatus", target_user_id=user_id_to_subscribe) + return {"status": "ok"} + else: + log(manager, websocket, user, "subscribeStatus_error", target_user_id=user_id_to_subscribe, error="User not found") + raise HTTPException(status_code=404, detail="User not found") + + +@websocket_handler("unsubscribeStatus", authRequired=True) +async def unsubscribeStatus(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> dict | None: + """Unsubscribe from status updates for a user.""" + user_id_to_unsubscribe = int(data["userId"]) + manager.ws_subscriptions[websocket].discard(user_id_to_unsubscribe) + + log(manager, websocket, user, "unsubscribeStatus", target_user_id=user_id_to_unsubscribe) + return {"status": "ok"} + + +@websocket_handler("typing", authRequired=True) +async def typing(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> None: + """Handle typing indicator start for public chat.""" + was_typing = manager.typing_state.get(user.id, False) + manager.typing_users[user.id] = time.time() + + # Only send update if state changed (started typing) + if not was_typing: + manager.typing_state[user.id] = True + # Broadcast to all connected users + await manager.broadcast({ + "type": "typing", + "data": { + "userId": user.id, + "username": user.username + } + }, db) + + # No confirmation response - privacy protection + + +@websocket_handler("stopTyping", authRequired=True) +async def stopTyping(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> None: + """Handle typing indicator stop for public chat.""" + was_typing = manager.typing_state.get(user.id, False) + if user.id in manager.typing_users: + del manager.typing_users[user.id] + + # Only send update if state changed (stopped typing) + if was_typing: + manager.typing_state[user.id] = False + # Broadcast to all connected users + await manager.broadcast({ + "type": "stopTyping", + "data": { + "userId": user.id, + "username": user.username + } + }, db) + + # No confirmation response - privacy protection + log(manager, websocket, user, "stopTyping") + + +@websocket_handler("dmTyping", authRequired=True) +async def dmTyping(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> None: + """Handle typing indicator start for DM.""" + recipient_id = int(data["recipientId"]) + + if user.id not in manager.dm_typing_users: + manager.dm_typing_users[user.id] = {} + if user.id not in manager.dm_typing_state: + manager.dm_typing_state[user.id] = {} + + was_typing = manager.dm_typing_state[user.id].get(recipient_id, False) + manager.dm_typing_users[user.id][recipient_id] = time.time() + + # Only send update if state changed (started typing) + if not was_typing: + manager.dm_typing_state[user.id][recipient_id] = True + # Send only to recipient + await manager.send_update_to_user(recipient_id, "dmTyping", { + "userId": user.id, + "username": user.username + }, db) + + # No confirmation response - privacy protection + + +@websocket_handler("stopDmTyping", authRequired=True) +async def stopDmTyping(manager: MessaggingSocketManager, websocket: WebSocket, db: Session, user: User, data: dict) -> None: + """Handle typing indicator stop for DM.""" + recipient_id = int(data["recipientId"]) + + was_typing = False + if user.id in manager.dm_typing_state: + was_typing = manager.dm_typing_state[user.id].get(recipient_id, False) + + if user.id in manager.dm_typing_users and recipient_id in manager.dm_typing_users[user.id]: + del manager.dm_typing_users[user.id][recipient_id] + if not manager.dm_typing_users[user.id]: + del manager.dm_typing_users[user.id] + + # Only send update if state changed (stopped typing) + if was_typing: + if user.id in manager.dm_typing_state: + manager.dm_typing_state[user.id][recipient_id] = False + # Send only to recipient + await manager.send_update_to_user(recipient_id, "stopDmTyping", { + "userId": user.id, + "username": user.username + }, db) + + # No confirmation response - privacy protection + diff --git a/backend/websocket/registry.py b/backend/websocket/registry.py new file mode 100644 index 0000000..d9a6271 --- /dev/null +++ b/backend/websocket/registry.py @@ -0,0 +1,33 @@ +from typing import Callable + + +class WebSocketHandlerRegistry: + """Registry for WebSocket message handlers with authentication support.""" + + def __init__(self): + self._handlers: dict[str, tuple[Callable, bool]] = {} + + def register(self, message_type: str, authRequired: bool = True): + """Register a handler for a message type. + + Args: + message_type: The WebSocket message type to handle + authRequired: If True, handler will receive authenticated User (not None) or raise 401 + """ + def decorator(func: Callable): + self._handlers[message_type] = (func, authRequired) + return func + return decorator + + def get_handler(self, message_type: str) -> tuple[Callable, bool] | None: + """Get handler and authRequired flag for a message type. + + Returns: + Tuple of (handler function, authRequired flag) or None if not found + """ + return self._handlers.get(message_type) + + def get_all_types(self) -> list[str]: + """Get all registered message types for debugging/logging.""" + return list(self._handlers.keys()) + diff --git a/backend/websocket/utils.py b/backend/websocket/utils.py new file mode 100644 index 0000000..1688706 --- /dev/null +++ b/backend/websocket/utils.py @@ -0,0 +1,92 @@ +from fastapi import HTTPException +from fastapi.security import HTTPAuthorizationCredentials +from sqlalchemy.orm import Session +from types import SimpleNamespace +from dependencies import get_current_user +from models import User + + +def extract_token_from_data(data: dict) -> str | None: + """Extract authentication token from WebSocket message data. + + Args: + data: WebSocket message data dictionary + + Returns: + Token string or None if not present + """ + credentials = data.get("credentials") + if credentials and isinstance(credentials, dict): + return credentials.get("credentials") + return None + + +def get_current_user_from_token(token: str, db: Session) -> User | None: + """Get user from authentication token. + + Args: + token: JWT token string + db: Database session + + Returns: + User object or None if token is invalid + """ + try: + # Ensure session is in a usable state before querying + try: + db.rollback() + except Exception: + pass + + dummy_request = SimpleNamespace() + dummy_request.state = SimpleNamespace() + + try: + from fastapi.security import HTTPBearer + security = HTTPBearer() + # We need to create credentials manually + credentials = HTTPAuthorizationCredentials( + scheme="Bearer", + credentials=token + ) + return get_current_user(dummy_request, credentials, db) + except HTTPException: + return None + except Exception: + try: + db.rollback() + except Exception: + pass + return None + + +def authenticate_user(data: dict, db: Session, authRequired: bool) -> User | None: + """Authenticate user from WebSocket message data. + + Args: + data: WebSocket message data dictionary + db: Database session + authRequired: If True, raises 401 on missing/invalid token + + Returns: + User object (guaranteed not None if authRequired=True) or None + + Raises: + HTTPException: 401 if authRequired=True and token is missing/invalid + """ + token = extract_token_from_data(data) + + if authRequired: + if not token: + raise HTTPException(status_code=401, detail="Missing credentials") + + user = get_current_user_from_token(token, db) + if not user: + raise HTTPException(status_code=401, detail="Invalid credentials") + + return user + else: + if token: + return get_current_user_from_token(token, db) + return None +