from datetime import datetime import html import logging from pathlib import Path import os import re import uuid import asyncio import time import unicodedata from collections import defaultdict, deque from difflib import SequenceMatcher from types import SimpleNamespace from typing import Any from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, UploadFile, File, Form, Request from fastapi.responses import FileResponse from fastapi.security import HTTPAuthorizationCredentials from sqlalchemy.orm import Session from dependencies import get_current_user, get_db from .account import convert_user from constants import OWNER_USERNAME from models import Message, SendMessageRequest, EditMessageRequest, User, DMEnvelope, MessageFile, DMFile, Reaction, ReactionRequest, ReactionResponse, DMReaction, DMReactionRequest, DMReactionResponse, UpdateLog from push_service import push_service from PIL import Image import io import json 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 router = APIRouter() logger = logging.getLogger("uvicorn.error") MAX_TOTAL_SIZE = 4 * 1024 * 1024 * 1024 # 4 GB FILES_BASE_DIR = Path("data/uploads/files") FILES_NORMAL_DIR = FILES_BASE_DIR / "normal" FILES_ENCRYPTED_DIR = FILES_BASE_DIR / "encrypted" os.makedirs(FILES_NORMAL_DIR, exist_ok=True) os.makedirs(FILES_ENCRYPTED_DIR, exist_ok=True) _SPAM_WINDOW_SECONDS = 45 _SPAM_SIMILARITY_THRESHOLD = 0.88 _SPAM_MESSAGE_LIMIT = 5 _BURST_WINDOW_SECONDS = 30 _BURST_COUNT_THRESHOLD = 20 _SHORT_MESSAGE_LENGTH = 8 _SHORT_MESSAGE_REPEAT_LIMIT = 4 _recent_message_cache: dict[int, deque[tuple[str, str, float]]] = defaultdict(deque) _message_rate_cache: dict[int, deque[float]] = defaultdict(deque) _burst_last_logged: dict[int, float] = {} def _normalize_for_spam(text: str) -> str: normalized = unicodedata.normalize("NFKC", text or "").casefold() # Remove whitespace and punctuation while keeping alphanumerics cleaned = re.sub(r"[^0-9a-zа-яё]+", "", normalized, flags=re.IGNORECASE) return cleaned def _monitor_public_message_activity(user: User, content: str, db: Session) -> None: now = time.time() def suspend(reason: str, event: str, **extra: Any) -> None: if user.suspended or user.id == 1: return user.suspended = True user.suspension_reason = reason db.commit() log_security( event, severity="warning", user_id=user.id, username=user.username, reason=reason, **extra, ) try: asyncio.create_task(messagingManager.send_suspension_to_user(user.id, reason)) except Exception: pass # Rate tracking for burst detection rate_bucket = _message_rate_cache[user.id] rate_bucket.append(now) while rate_bucket and now - rate_bucket[0] > _BURST_WINDOW_SECONDS: rate_bucket.popleft() burst_count = len(rate_bucket) if burst_count >= _BURST_COUNT_THRESHOLD: last_logged = _burst_last_logged.get(user.id) if not last_logged or now - last_logged > _BURST_WINDOW_SECONDS: log_security( "public_message_burst", severity="warning", user_id=user.id, username=user.username, count=burst_count, window_seconds=_BURST_WINDOW_SECONDS, ) _burst_last_logged[user.id] = now suspend( "Automatic suspension: excessive message rate", "auto_suspension_public_burst", count=burst_count, window_seconds=_BURST_WINDOW_SECONDS, ) # Similarity-based spam detection normalized = _normalize_for_spam(content) history = _recent_message_cache[user.id] while history and now - history[0][2] > _SPAM_WINDOW_SECONDS: history.popleft() prior_same = sum(1 for prev_norm, _, _ in history if prev_norm == normalized) prior_similar = sum( 1 for prev_norm, _, _ in history if prev_norm and normalized and prev_norm != normalized and SequenceMatcher(None, normalized, prev_norm).ratio() >= _SPAM_SIMILARITY_THRESHOLD ) history.append((normalized, content, now)) total_matches = prior_same + prior_similar + 1 if len(normalized) <= _SHORT_MESSAGE_LENGTH and prior_same + 1 >= _SHORT_MESSAGE_REPEAT_LIMIT: suspend( "Automatic suspension: repeated short messages", "auto_suspension_public_spam", occurrences=prior_same + 1, window_seconds=_SPAM_WINDOW_SECONDS, match_type="short", ) return if total_matches >= _SPAM_MESSAGE_LIMIT: suspend( "Automatic suspension: repeated similar public messages", "auto_suspension_public_spam", similar_messages=total_matches, window_seconds=_SPAM_WINDOW_SECONDS, match_type="similar", ) def convert_message(msg: Message) -> dict: # Group reactions by emoji reactions_dict = {} if msg.reactions: for reaction in msg.reactions: emoji = reaction.emoji if emoji not in reactions_dict: reactions_dict[emoji] = { "emoji": emoji, "count": 0, "users": [] } reactions_dict[emoji]["count"] += 1 reactions_dict[emoji]["users"].append({ "id": reaction.user_id, "username": reaction.user.display_name }) # Handle deleted or suspended users if msg.author.deleted or msg.author.suspended: username = f"Deleted User #{msg.author.id}" profile_picture = None verified = False else: username = msg.author.display_name profile_picture = msg.author.profile_picture verified = msg.author.verified return { "id": msg.id, "user_id": msg.author.id, "content": msg.content, "timestamp": msg.timestamp.isoformat(), "is_read": msg.is_read, "is_edited": msg.is_edited, "username": username, "profile_picture": profile_picture, "verified": verified, "reply_to": convert_message(msg.reply_to) if msg.reply_to else None, "reactions": list(reactions_dict.values()), "files": [ { "path": f"/api/uploads/files/normal/{Path(f.path).name}", "id": f.id, "name": f.name, "message_id": f.message_id } for f in (msg.files or []) ] } def convert_dm_envelope(db: Session, envelope: DMEnvelope) -> dict: # Group reactions by emoji reactions_dict = {} if envelope.reactions: for reaction in envelope.reactions: emoji = reaction.emoji if emoji not in reactions_dict: reactions_dict[emoji] = { "emoji": emoji, "count": 0, "users": [] } reactions_dict[emoji]["count"] += 1 reactions_dict[emoji]["users"].append({ "id": reaction.user_id, "username": reaction.user.display_name }) # Get sender info for verified status sender = db.query(User).filter(User.id == envelope.sender_id).first() # Handle deleted or suspended users if sender and (sender.deleted or sender.suspended): sender_verified = False else: sender_verified = sender.verified if sender else False return { "id": envelope.id, "senderId": envelope.sender_id, "recipientId": envelope.recipient_id, "iv": envelope.iv_b64, "ciphertext": envelope.ciphertext_b64, "salt": envelope.salt_b64, "iv2": envelope.iv2_b64, "wrappedMk": envelope.wrapped_mk_b64, "timestamp": envelope.timestamp.isoformat(), "verified": sender_verified, "reactions": list(reactions_dict.values()), "files": [ { "path": f"/api/uploads/files/encrypted/{Path(f.path).name}", "id": f.id, "name": f.name, "dm_envelope_id": f.dm_envelope_id } for f in (envelope.files or []) ] } async def _send_message_internal( message_request: SendMessageRequest, current_user: User, db: Session, files: list[UploadFile] = [], ) -> dict: """Internal function to send a message without requiring a Request object. This can be called from both HTTP endpoints and WebSocket handlers. """ if message_request.reply_to_id: # Check if the message being replied to exists original_message = db.query(Message).filter(Message.id == message_request.reply_to_id).first() if not original_message: raise HTTPException(status_code=404, detail="Original message not found") raw_content = message_request.content.strip() if not raw_content: raise HTTPException( status_code=400, detail="No content provided" ) # Apply profanity filter before storing filtered_content = censor_text(raw_content) escaped_content = html.escape(filtered_content, quote=False) if len(escaped_content) > 4096: raise HTTPException( status_code=400, detail="Message too long" ) new_message = Message( content=escaped_content, user_id=current_user.id, reply_to_id=message_request.reply_to_id, timestamp=datetime.now() ) db.add(new_message) db.commit() db.refresh(new_message) # Handle files if provided (normal, not encrypted) if files: total_size = 0 for up in files: # Accumulate size if available if hasattr(up, "size") and up.size is not None: total_size += int(up.size) else: # If size unknown, read into memory to determine data = await up.read() up.file.seek(0) total_size += len(data) if total_size > MAX_TOTAL_SIZE: raise HTTPException(status_code=400, detail="Total attachments size exceeds 4GB") for up in files: # Sanitize filename original_name = Path(up.filename or "file").name ext = Path(original_name).suffix.lower() uid = uuid.uuid4().hex safe_name = f"{new_message.id}_{uid}{ext or ''}" out_path = FILES_NORMAL_DIR / safe_name content = await up.read() up.file.seek(0) # If image, try lossless optimization try: if up.content_type and up.content_type.startswith("image/"): image = Image.open(io.BytesIO(content)) img_format = image.format or ("PNG" if ext == ".png" else "JPEG") buf = io.BytesIO() save_kwargs = {"optimize": True} if img_format.upper() == "JPEG": # Use quality=95 with optimize to keep high quality (not truly lossless but near) save_kwargs["quality"] = 95 image.save(buf, format=img_format, **save_kwargs) buf.seek(0) content = buf.read() except Exception: # Fallback to original content pass with open(out_path, "wb") as f: f.write(content) mf = MessageFile( message_id=new_message.id, name=original_name, path=str(out_path) ) db.add(mf) db.commit() db.refresh(new_message) # Send push notifications for public messages try: await push_service.send_public_message_notification(db, new_message, exclude_user_id=current_user.id) except Exception as e: logger.error(f"Failed to send push notification for message {new_message.id}: {e}") # Realtime broadcast for HTTP uploads as well try: await messagingManager.broadcast({ "type": "newMessage", "data": convert_message(new_message) }, db) except Exception: pass _monitor_public_message_activity(current_user, filtered_content, db) message_payload = convert_message(new_message) log_public_chat( "message_created", message_id=new_message.id, user_id=current_user.id, username=current_user.username, reply_to=new_message.reply_to_id, attachments=len(new_message.files or []), length=len(new_message.content), suspended=current_user.suspended, content=new_message.content, ) return {"status": "success", "message": message_payload} @router.post("/send_message") @rate_limit_per_ip("30/minute") async def send_message( request: Request, message_request: SendMessageRequest | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db), # Optional multipart form support payload: str | None = Form(default=None), files: list[UploadFile] = File(default=[]), ): # If payload is provided, prefer it for multipart requests if payload and message_request is None: # Expect JSON: {"type":"text","data":{"content": str}, "reply_to_id": number|null} try: obj = json.loads(payload) content = obj.get("content", "") reply_to_id = obj.get("reply_to_id", None) message_request = SendMessageRequest(content=content, reply_to_id=reply_to_id) except Exception: raise HTTPException(status_code=400, detail="Invalid payload JSON") if not message_request: raise HTTPException(status_code=400, detail="Missing request data") return await _send_message_internal(message_request, current_user, db, files) @router.get("/get_messages") @rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse async def get_messages(request: Request, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): messages = db.query(Message).order_by(Message.timestamp.asc()).all() messages_data = [] for msg in messages: messages_data.append(convert_message(msg)) return { "status": "success", "messages": messages_data } @router.post("/dm/send") @rate_limit_per_ip("20/minute") async def dm_send( request: Request, payload: dict | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db), # Multipart support dm_payload: str | None = Form(default=None), files: list[UploadFile] = File(default=[]), fileNames: str | None = Form(default=None), # JSON array of filenames corresponding to files ): if dm_payload and payload is None: try: payload = json.loads(dm_payload) except Exception: raise HTTPException(status_code=400, detail="Invalid dm_payload JSON") if payload is None: raise HTTPException(status_code=400, detail="Missing payload") 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}") try: recipient_id = int(payload["recipientId"]) except (ValueError, TypeError): raise HTTPException(status_code=400, detail="Invalid recipientId") if recipient_id <= 0: raise HTTPException(status_code=400, detail="Invalid recipientId") if recipient_id == current_user.id: raise HTTPException(status_code=400, detail="Cannot send DM to yourself") # Verify recipient exists recipient = db.query(User).filter(User.id == recipient_id).first() if not recipient or recipient.deleted or recipient.suspended: raise HTTPException(status_code=404, detail="Recipient not found") env = DMEnvelope( sender_id=current_user.id, recipient_id=recipient_id, 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) # Save encrypted files if any (no processing) if files: # Validate total size total_size = 0 for file in files: if hasattr(file, "size") and file.size is not None: total_size += int(file.size) else: data = await file.read() file.file.seek(0) total_size += len(data) if total_size > MAX_TOTAL_SIZE: raise HTTPException(status_code=400, detail="Total attachments size exceeds 4GB") names: list[str] = [] if fileNames: try: decoded = json.loads(fileNames) if isinstance(decoded, list): names = [str(x) for x in decoded] except Exception: names = [] for i, file in enumerate(files): provided = names[i] if i < len(names) else None # Sanitize provided name to avoid path traversal if provided and not re.match(r"^[A-Za-z0-9._-]{1,200}$", provided): provided = None original_name = provided or Path(file.filename or "file").name # Save using provided/original name to allow client to reference path directly safe_name = uid = uuid.uuid4().hex out_name = f"{current_user.id}_{env.recipient_id}_{env.id}_{safe_name}" out_path = FILES_ENCRYPTED_DIR / out_name content = await file.read() with open(out_path, "wb") as f: f.write(content) # Save DM file record df = DMFile( message_id=env.id, sender_id=current_user.id, recipient_id=env.recipient_id, path=f"/api/uploads/files/encrypted/{out_name}", name=original_name ) db.add(df) db.commit() db.refresh(env) # 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}") # Realtime notify both users for HTTP requests try: 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, } } await messagingManager.send_to_user(env.recipient_id, payload_ws) await messagingManager.send_to_user(env.sender_id, payload_ws) except Exception: pass log_dm( "message_sent", dm_envelope_id=env.id, sender_id=current_user.id, sender_username=current_user.username, recipient_id=env.recipient_id, attachment_count=len(env.files or []), reply_to=env.reply_to_id, ) return {"status": "ok", "id": env.id} def convert_envelopes(envs: list[DMEnvelope]): return { "status": "ok", "messages": [ { "id": e.id, "senderId": e.sender_id, "recipientId": e.recipient_id, "iv": e.iv_b64, "ciphertext": e.ciphertext_b64, "salt": e.salt_b64, "iv2": e.iv2_b64, "wrappedMk": e.wrapped_mk_b64, "timestamp": e.timestamp.isoformat(), "files": [{"name": file.name, "path": file.path, "id": file.id} for file in e.files] } for e in envs ] } @router.get("/dm/fetch") @rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse async def dm_fetch(request: Request, since: int | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): q = db.query(DMEnvelope).filter(DMEnvelope.recipient_id == current_user.id) if since: q = q.filter(DMEnvelope.id > since) return convert_envelopes(q.order_by(DMEnvelope.id.asc()).all()) @router.get("/dm/history/{other_user_id}") @rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse async def dm_history(request: Request, other_user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): if other_user_id <= 0: raise HTTPException(status_code=400, detail="Invalid user ID") if other_user_id == current_user.id: raise HTTPException(status_code=400, detail="Cannot get history with yourself") # Verify other user exists other_user = db.query(User).filter(User.id == other_user_id).first() if not other_user or other_user.deleted or other_user.suspended: raise HTTPException(status_code=404, detail="User not found") return convert_envelopes( db.query(DMEnvelope) .filter( ((DMEnvelope.sender_id == current_user.id) & (DMEnvelope.recipient_id == other_user_id)) | ((DMEnvelope.sender_id == other_user_id) & (DMEnvelope.recipient_id == current_user.id)) ) .order_by(DMEnvelope.id.asc()) .all() ) @router.get("/dm/conversations") @rate_limit_per_ip("60/minute") # Per-IP limit to prevent abuse async def get_dm_conversations(request: Request, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): # Get all DM conversations where current user is involved conversations_query = db.query(DMEnvelope).filter( (DMEnvelope.sender_id == current_user.id) | (DMEnvelope.recipient_id == current_user.id) ).order_by(DMEnvelope.timestamp.desc()) # Group by the "other user" (not current user) and get latest message conversations = {} for envelope in conversations_query: other_user_id = envelope.recipient_id if envelope.sender_id == current_user.id else envelope.sender_id if other_user_id not in conversations: conversations[other_user_id] = envelope # Get user info for each conversation result = [] for other_user_id, latest_message in conversations.items(): other_user = db.query(User).filter(User.id == other_user_id).first() if other_user: # Calculate unread count for this conversation unread_count = db.query(DMEnvelope).filter( DMEnvelope.sender_id == other_user_id, DMEnvelope.recipient_id == current_user.id, DMEnvelope.id > getattr(latest_message, 'last_read_id', 0) # This would need to be stored somewhere ).count() result.append({ "user": convert_user(other_user), "lastMessage": convert_dm_envelope(db, latest_message), "unreadCount": unread_count }) # Sort by latest message timestamp result.sort(key=lambda x: x["lastMessage"]["timestamp"], reverse=True) return { "status": "success", "conversations": result } @router.put("/edit_message/{message_id}") @rate_limit_per_ip("20/minute") async def edit_message( request: Request, message_id: int, edit_request: EditMessageRequest, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): message = db.query(Message).filter(Message.id == message_id).first() if not message: raise HTTPException(status_code=404, detail="Message not found") if message.user_id != current_user.id: raise HTTPException(status_code=403, detail="You can only edit your own messages") raw_content = edit_request.content.strip() if not raw_content: raise HTTPException(status_code=400, detail="Message content cannot be empty") original_content = message.content sanitized_content = censor_text(raw_content) escaped_content = html.escape(sanitized_content, quote=False) if len(escaped_content) > 4096: raise HTTPException(status_code=400, detail="Message too long") message.content = escaped_content message.is_edited = True db.commit() db.refresh(message) payload = convert_message(message) log_public_chat( "message_edited", message_id=message.id, user_id=current_user.id, username=current_user.username, reply_to=message.reply_to_id, content=message.content, previous_content=original_content, ) return {"status": "success", "message": payload} @router.delete("/delete_message/{message_id}") async def delete_message( message_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): message = db.query(Message).filter(Message.id == message_id).first() if not message: raise HTTPException(status_code=404, detail="Message not found") # Allow owner to delete any message if current_user.username != OWNER_USERNAME and message.user_id != current_user.id: raise HTTPException(status_code=403, detail="You can only delete your own messages") original_content = message.content db.delete(message) db.commit() log_public_chat( "message_deleted", message_id=message_id, actor_id=current_user.id, actor_username=current_user.username, original_author_id=message.user_id, content=original_content, ) return {"status": "success", "message_id": message_id} @router.post("/add_reaction") @rate_limit_per_ip("50/minute") async def add_reaction( request: Request, reaction_request: ReactionRequest, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): # Check if message exists message = db.query(Message).filter(Message.id == reaction_request.message_id).first() if not message: raise HTTPException(status_code=404, detail="Message not found") # Check if reaction already exists existing_reaction = db.query(Reaction).filter( Reaction.message_id == reaction_request.message_id, Reaction.user_id == current_user.id, Reaction.emoji == reaction_request.emoji ).first() if existing_reaction: # Remove existing reaction (toggle off) db.delete(existing_reaction) action = "removed" else: # Add new reaction new_reaction = Reaction( message_id=reaction_request.message_id, user_id=current_user.id, emoji=reaction_request.emoji ) db.add(new_reaction) action = "added" db.commit() # Refresh message to get updated reactions db.refresh(message) message_data = convert_message(message) # Broadcast reaction update try: await messagingManager.broadcast({ "type": "reactionUpdate", "data": { "message_id": reaction_request.message_id, "emoji": reaction_request.emoji, "action": action, "user_id": current_user.id, "username": current_user.username, "reactions": message_data["reactions"] } }, db) except Exception: pass log_public_chat( "reaction_update", message_id=reaction_request.message_id, user_id=current_user.id, username=current_user.username, action=action, emoji=reaction_request.emoji, ) return {"status": "success", "action": action, "reactions": message_data["reactions"]} @router.post("/dm/add_reaction") @rate_limit_per_ip("50/minute") async def add_dm_reaction( request: Request, reaction_request: DMReactionRequest, current_user: User = Depends(get_current_user), db: Session = Depends(get_db) ): # Check if DM envelope exists envelope = db.query(DMEnvelope).filter(DMEnvelope.id == reaction_request.dm_envelope_id).first() if not envelope: raise HTTPException(status_code=404, detail="DM envelope not found") # Check if user is part of this DM conversation if current_user.id not in [envelope.sender_id, envelope.recipient_id]: raise HTTPException(status_code=403, detail="Not authorized to react to this message") # Check if reaction already exists existing_reaction = db.query(DMReaction).filter( DMReaction.dm_envelope_id == reaction_request.dm_envelope_id, DMReaction.user_id == current_user.id, DMReaction.emoji == reaction_request.emoji ).first() if existing_reaction: # Remove existing reaction (toggle off) db.delete(existing_reaction) action = "removed" else: # Add new reaction new_reaction = DMReaction( dm_envelope_id=reaction_request.dm_envelope_id, user_id=current_user.id, emoji=reaction_request.emoji ) db.add(new_reaction) action = "added" db.commit() # Refresh envelope to get updated reactions db.refresh(envelope) envelope_data = convert_dm_envelope(db, envelope) # Broadcast reaction update to both participants try: await messagingManager.broadcast({ "type": "dmReactionUpdate", "data": { "dm_envelope_id": reaction_request.dm_envelope_id, "emoji": reaction_request.emoji, "action": action, "user_id": current_user.id, "username": current_user.username, "reactions": envelope_data["reactions"] } }, db) except Exception: pass log_dm( "reaction_update", dm_envelope_id=reaction_request.dm_envelope_id, user_id=current_user.id, username=current_user.username, action=action, emoji=reaction_request.emoji, ) return {"status": "success", "action": action, "reactions": envelope_data["reactions"]} class MessaggingSocketManager: def __init__(self) -> None: self.connections: list[WebSocket] = [] self.user_by_ws: dict[WebSocket, int] = {} self.online_users: set[int] = set() self.typing_users: dict[int, float] = {} # user_id -> timestamp self.dm_typing_users: dict[int, dict[int, float]] = {} # user_id -> {recipient_id -> timestamp} self.typing_state: dict[int, bool] = {} # user_id -> is_typing (for public chat) self.dm_typing_state: dict[int, dict[int, bool]] = {} # user_id -> {recipient_id -> is_typing} self.ws_subscriptions: dict[WebSocket, set[int]] = {} # websocket -> set of subscribed user_ids self._cleanup_task = None # Update system: sequence numbers and batching self.sequence_numbers: dict[int, int] = {} # user_id -> current sequence number self.pending_updates: dict[WebSocket, list[dict]] = {} # websocket -> list of pending updates self.update_batch_tasks: dict[WebSocket, asyncio.Task] = {} # websocket -> batch task self.last_seq_by_ws: dict[WebSocket, int] = {} # websocket -> last received sequence number self.stored_sequences: dict[tuple[int, int], bool] = {} # (user_id, sequence) -> stored flag self.recent_updates: dict[WebSocket, set[str]] = {} # websocket -> set of recent update signatures self._sequence_lock: dict[int, asyncio.Lock] = {} # user_id -> lock for sequence generation async def send_error(self, websocket: WebSocket, type: str, e: HTTPException): await websocket.send_json({"type": type, "error": {"code": e.status_code, "detail": e.detail}}) async def _get_next_sequence(self, user_id: int) -> int: """Get the next sequence number for a user (shared across all their connections) - thread-safe""" if user_id not in self._sequence_lock: self._sequence_lock[user_id] = asyncio.Lock() async with self._sequence_lock[user_id]: if user_id not in self.sequence_numbers: self.sequence_numbers[user_id] = 0 self.sequence_numbers[user_id] += 1 return self.sequence_numbers[user_id] def _get_update_signature(self, update: dict) -> str: """Generate a unique signature for an update to detect duplicates""" import hashlib import json update_type = update.get("type", "") data = update.get("data", {}) # Create signature based on update type and key identifying fields if update_type == "newMessage": # Deduplicate by message ID sig_data = {"type": update_type, "id": data.get("id")} elif update_type == "messageEdited": # Deduplicate by message ID sig_data = {"type": update_type, "id": data.get("id")} elif update_type == "messageDeleted": # Deduplicate by message ID sig_data = {"type": update_type, "id": data.get("id") or data.get("message_id")} elif update_type == "dmNew": # Deduplicate by envelope ID sig_data = {"type": update_type, "id": data.get("id")} elif update_type == "dmEdited": # Deduplicate by envelope ID sig_data = {"type": update_type, "id": data.get("id")} elif update_type == "dmDeleted": # Deduplicate by envelope ID sig_data = {"type": update_type, "id": data.get("id")} elif update_type == "reactionUpdate": # Deduplicate by message ID + emoji + user ID sig_data = {"type": update_type, "messageId": data.get("message_id"), "emoji": data.get("emoji"), "userId": data.get("userId")} elif update_type == "dmReactionUpdate": # Deduplicate by envelope ID + emoji + user ID sig_data = {"type": update_type, "dmEnvelopeId": data.get("dm_envelope_id"), "emoji": data.get("emoji"), "userId": data.get("userId")} elif update_type == "typing" or update_type == "stopTyping": # Deduplicate by user ID (state tracking already handles this, but extra protection) sig_data = {"type": update_type, "userId": data.get("userId")} elif update_type == "dmTyping" or update_type == "stopDmTyping": # Deduplicate by user ID (recipient ID is implicit - this update is sent TO the recipient) sig_data = {"type": update_type, "userId": data.get("userId")} elif update_type == "statusUpdate": # Deduplicate by user ID sig_data = {"type": update_type, "userId": data.get("userId")} else: # For unknown types, use full data (less efficient but safe) sig_data = {"type": update_type, "data": data} # Create hash of signature data sig_json = json.dumps(sig_data, sort_keys=True) return hashlib.md5(sig_json.encode()).hexdigest() def _add_update(self, websocket: WebSocket, update: dict): """Add an update to the pending batch for a WebSocket (with deduplication)""" if websocket not in self.pending_updates: self.pending_updates[websocket] = [] # Check for duplicates signature = self._get_update_signature(update) if websocket not in self.recent_updates: self.recent_updates[websocket] = set() # Skip if this exact update was recently added if signature in self.recent_updates[websocket]: return # Add to pending updates and track signature self.pending_updates[websocket].append(update) self.recent_updates[websocket].add(signature) # Limit recent updates cache size (keep last 100 signatures per websocket) if len(self.recent_updates[websocket]) > 100: # Remove oldest entries (simple FIFO by converting to list and keeping last 100) # Actually, we'll just clear and rebuild on next flush - simpler approach pass async def _flush_updates(self, websocket: WebSocket, db: Session | None = None): """Flush pending updates for a WebSocket connection""" if websocket not in self.pending_updates or not self.pending_updates[websocket]: return updates = self.pending_updates[websocket] self.pending_updates[websocket] = [] # Clear recent updates cache after flushing (updates are now sent, can be re-added if needed) if websocket in self.recent_updates: # Keep only the last 50 signatures to allow some deduplication across batches recent_list = list(self.recent_updates[websocket]) if len(recent_list) > 50: self.recent_updates[websocket] = set(recent_list[-50:]) else: # Keep all if under limit pass if updates: user_id = self.user_by_ws.get(websocket) if not user_id: # No user associated - this shouldn't happen for authenticated connections # Skip sending to avoid seq: 0 issues logger.warning(f"Attempted to flush updates for unauthenticated websocket, skipping") return seq = await self._get_next_sequence(user_id) # Store updates in database for gap detection (only once per user per sequence) if db: sequence_key = (user_id, seq) # Double-check pattern: check again after getting sequence (in case another connection got the same sequence) if sequence_key not in self.stored_sequences: try: import json # Store the entire batch as a single record update_log = UpdateLog( user_id=user_id, sequence=seq, updates=json.dumps(updates) ) db.add(update_log) db.commit() self.stored_sequences[sequence_key] = True except Exception as e: # Always rollback on error to reset session state try: db.rollback() except Exception: pass # Ignore rollback errors # If we get a UNIQUE constraint error, it means another connection already stored this sequence if "UNIQUE constraint" in str(e) or "IntegrityError" in str(e.__class__.__name__): # Mark as stored to prevent future attempts self.stored_sequences[sequence_key] = True logger.debug(f"Update sequence {seq} for user {user_id} already stored by another connection") else: logger.error(f"Failed to store updates in database: {e}") else: # Already stored, skip logger.debug(f"Update sequence {seq} for user {user_id} already marked as stored") await websocket.send_json({ "type": "updates", "seq": seq, "updates": updates }) async def _schedule_batch_flush(self, websocket: WebSocket, db: Session | None = None): """Schedule a batch flush after a delay (50-100ms)""" if websocket in self.update_batch_tasks: self.update_batch_tasks[websocket].cancel() async def flush_after_delay(): await asyncio.sleep(0.075) # 75ms delay for batching await self._flush_updates(websocket, db) if websocket in self.update_batch_tasks: del self.update_batch_tasks[websocket] self.update_batch_tasks[websocket] = asyncio.create_task(flush_after_delay()) async def _send_update(self, websocket: WebSocket, update_type: str, update_data: dict, db: Session | None = None): """Send an update (will be batched)""" self._add_update(websocket, {"type": update_type, "data": update_data}) await self._schedule_batch_flush(websocket, db) 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, ) while True: try: data = await websocket.receive_json() except Exception as e: logger.error(f"Error receiving WebSocket message: {e}") break type = data["type"] def get_current_user_inner() -> User | None: try: # Ensure session is in a usable state before querying try: db.rollback() except Exception: pass 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 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) else: await websocket.send_json({"type": type, "error": {"code": 400, "detail": "Invalid type"}}) async def disconnect(self, websocket: WebSocket, code: int = 1000, message: str | None = None): try: await websocket.close(code=code, reason=message) finally: self.connections.remove(websocket) async def connect(self, websocket: WebSocket, db: Session): await websocket.accept() client_ip = websocket.client.host if websocket.client else None log_access( "ws_connect", path=str(websocket.url.path), ip=client_ip, ) self.connections.append(websocket) # Initialize update system for this connection self.pending_updates[websocket] = [] self.last_seq_by_ws[websocket] = 0 try: await self.handle_connection(websocket, db) except WebSocketDisconnect as e: logger.info(f"WebSocket disconnected with code {e.code}: {e.reason}") log_access( "ws_disconnect", severity="warning" if e.code != 1000 else "info", path=str(websocket.url.path), ip=client_ip, code=e.code, reason=e.reason, ) finally: # Flush any pending updates before disconnecting if websocket in self.pending_updates: await self._flush_updates(websocket, db) # Cancel any pending batch tasks if websocket in self.update_batch_tasks: self.update_batch_tasks[websocket].cancel() del self.update_batch_tasks[websocket] # Cleanup connection self.connections.remove(websocket) if websocket in self.user_by_ws: user_id = self.user_by_ws[websocket] # Set user offline in DB try: # Ensure session is in a usable state try: db.rollback() except Exception: pass user = db.query(User).filter(User.id == user_id).first() if user: user.online = False user.last_seen = datetime.now() db.commit() # Remove from online users self.online_users.discard(user_id) # Broadcast status change await self.broadcast_status_change(user_id, False, user.last_seen.isoformat(), db) except Exception as e: logger.error(f"Failed to set user offline during cleanup: {e}") try: db.rollback() except Exception: pass finally: del self.user_by_ws[websocket] # Cleanup subscriptions if websocket in self.ws_subscriptions: del self.ws_subscriptions[websocket] # Cleanup update system if websocket in self.pending_updates: del self.pending_updates[websocket] if websocket in self.last_seq_by_ws: del self.last_seq_by_ws[websocket] if websocket in self.recent_updates: del self.recent_updates[websocket] async def broadcast(self, message: dict, db: Session | None = None): """Broadcast a message to all authenticated connections as an update (batched)""" message_type = message.get("type", "") update_data = message.get("data", {}) for websocket in self.connections: # Only send to authenticated websockets (those with user_id set) if websocket in self.user_by_ws: await self._send_update(websocket, message_type, update_data, db) async def send_update_to_user(self, user_id: int, update_type: str, update_data: dict, db: Session | None = None): """Send an update to a specific user (batched)""" for websocket in self.connections: if self.user_by_ws.get(websocket) == user_id: await self._send_update(websocket, update_type, update_data, db) async def send_to_user(self, user_id: int, message: dict): """Send a direct WebSocket message to a specific user (not batched)""" for websocket in self.connections: if self.user_by_ws.get(websocket) == user_id: await websocket.send_json(message) async def send_suspension_to_user(self, user_id: int, reason: str): """Send suspension message to user's WebSocket connections (as batched update)""" await self.send_update_to_user(user_id, "suspended", { "reason": reason }) async def send_deletion_to_user(self, user_id: int): """Send account deletion message to user's WebSocket connections (as batched update)""" await self.send_update_to_user(user_id, "account_deleted", {}) async def broadcast_status_change(self, user_id: int, online: bool, last_seen: str, db: Session | None = None): """Broadcast status change to all connections that are subscribed to this user""" # Send to all connections that have this user in their subscriptions for websocket in self.connections: if websocket in self.ws_subscriptions and user_id in self.ws_subscriptions[websocket]: await self._send_update(websocket, "statusUpdate", { "userId": user_id, "online": online, "lastSeen": last_seen }, db) async def cleanup_stale_typing_indicators(self, db: Session): """Periodically cleanup typing indicators that haven't been updated in 3+ seconds""" while True: try: current_time = time.time() stale_threshold = 3.0 # 3 seconds # Cleanup public chat typing indicators stale_public_typing = [ user_id for user_id, timestamp in self.typing_users.items() if current_time - timestamp > stale_threshold ] for user_id in stale_public_typing: was_typing = self.typing_state.get(user_id, False) del self.typing_users[user_id] # Only send update if state changed (stopped typing) if was_typing: self.typing_state[user_id] = False # Get username from database user = db.query(User).filter(User.id == user_id).first() username = user.username if user else "Unknown" # Broadcast stop typing await self.broadcast({ "type": "stopTyping", "data": { "userId": user_id, "username": username } }, db) # Cleanup DM typing indicators stale_dm_typing = [] for user_id, recipients in self.dm_typing_users.items(): for recipient_id, timestamp in list(recipients.items()): if current_time - timestamp > stale_threshold: stale_dm_typing.append((user_id, recipient_id)) for user_id, recipient_id in stale_dm_typing: was_typing = False if user_id in self.dm_typing_state: was_typing = self.dm_typing_state[user_id].get(recipient_id, False) if user_id in self.dm_typing_users and recipient_id in self.dm_typing_users[user_id]: del self.dm_typing_users[user_id][recipient_id] if not self.dm_typing_users[user_id]: del self.dm_typing_users[user_id] # Only send update if state changed (stopped typing) if was_typing: if user_id in self.dm_typing_state: self.dm_typing_state[user_id][recipient_id] = False # Get username from database user = db.query(User).filter(User.id == user_id).first() username = user.username if user else "Unknown" # Send stop typing to recipient await self.send_update_to_user(recipient_id, "stopDmTyping", { "userId": user_id, "username": username }, db) # Wait 1 second before next cleanup await asyncio.sleep(1.0) except Exception as e: logger.error(f"Error in typing cleanup task: {e}") await asyncio.sleep(1.0) def start_cleanup_task(self): """Start the cleanup task if not already running""" if self._cleanup_task is None or self._cleanup_task.done(): from db import SessionLocal async def cleanup_with_db(): while True: try: with SessionLocal() as db: await self.cleanup_stale_typing_indicators(db) except Exception as e: logger.error(f"Error in cleanup task wrapper: {e}") await asyncio.sleep(1.0) self._cleanup_task = asyncio.create_task(cleanup_with_db()) messagingManager = MessaggingSocketManager() @router.websocket("/chat/ws") async def chat_websocket( websocket: WebSocket, db: Session = Depends(get_db) ): await messagingManager.connect(websocket, db) # File serving endpoints @router.get("/uploads/files/normal/{filename}") async def get_file_normal(filename: str): if not re.match(r"^[A-Za-z0-9._-]+$", filename): raise HTTPException(status_code=400, detail="Invalid file name") path = FILES_NORMAL_DIR / filename if not path.exists(): raise HTTPException(status_code=404, detail="File not found") return FileResponse(str(path)) @router.get("/uploads/files/encrypted/{filename}") async def get_file_encrypted(filename: str, current_user: User = Depends(get_current_user)): if not re.match(r"^[A-Za-z0-9._-]+$", filename): raise HTTPException(status_code=400, detail="Invalid file name") path = FILES_ENCRYPTED_DIR / filename if not path.exists(): raise HTTPException(status_code=404, detail="File not found") match = re.match(r"^(\d+)_(\d+)_(\d+)_.*$", path.resolve().name) if match: sender_id = int(match.group(1)) recipient_id = int(match.group(2)) if not current_user.id in [sender_id, recipient_id]: raise HTTPException(403) else: raise HTTPException(500) return FileResponse(str(path))