Refactor the websocket message handler

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")
@@ -1085,26 +1086,8 @@ class MessaggingSocketManager:
         # Initialize subscriptions for this connection
         self.ws_subscriptions[websocket] = set()
         
-        ws_path = getattr(getattr(websocket, "url", None), "path", None)
+        # Import here to avoid circular import
-        if not ws_path and isinstance(getattr(websocket, "scope", None), dict):
+        from websocket.handlers import handler_registry
-            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,
-            )
 
         while True:
             try:
@@ -1113,698 +1096,33 @@ class MessaggingSocketManager:
                 logger.error(f"Error receiving WebSocket message: {e}")
                 break
             
-            type = data["type"]
+            message_type = data["type"]
+            handler_info = handler_registry.get_handler(message_type)
             
-            def get_current_user_inner() -> User | None:
+            if handler_info:
+                handler, authRequired = handler_info
                 try:
-                    # Ensure session is in a usable state before querying
+                    # Authenticate user before calling handler
-                    try:
+                    user = authenticate_user(data, db, authRequired)
-                        db.rollback()
+                    # Set user association for authenticated connections
-                    except Exception:
+                    if user:
-                        pass
+                        self.user_by_ws[websocket] = user.id
                     
-                    if data.get("credentials"):
+                    # Extract inner data to pass to handler
-                        dummy_request = SimpleNamespace()
+                    handler_data = data.get("data", {})
-                        dummy_request.state = SimpleNamespace()
+                    result = await handler(self, websocket, db, user, handler_data)
-                        return get_current_user(
+                    # If handler returns a value, send it as a WebSocket message
-                            dummy_request,
+                    if result is not None:
-                            HTTPAuthorizationCredentials(
+                        await websocket.send_json({"type": message_type, "data": result})
-                                scheme=data["credentials"]["scheme"],
+                except HTTPException as e:
-                                credentials=data["credentials"]["credentials"]
+                    await self.send_error(websocket, message_type, e)
-                            ),
+                except WebSocketDisconnect:
-                            db
+                    raise  # Re-raise to close connection
-                        )
-                    else:
-                        return None
                 except Exception as e:
-                    logger.error(f"Error getting current user: {e}")
+                    logger.error(f"Error in handler for {message_type}: {e}")
-                    try:
+                    await self.send_error(websocket, message_type, HTTPException(500, "Internal server error"))
-                        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)
             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:

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"]
+

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
+

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())
+

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
+