From 8acb2dd50777d724de90239a0e77ee6b68bd90e1 Mon Sep 17 00:00:00 2001 From: denis0001-dev Date: Mon, 25 Aug 2025 17:40:47 +0300 Subject: [PATCH] Implement real-time WebSocket --- backend/routes/messaging.py | 80 +++++++++++++++++++++++++++++++++++++ frontend/src/chat/dm.ts | 53 ++++++++++++++++++++++-- 2 files changed, 130 insertions(+), 3 deletions(-) diff --git a/backend/routes/messaging.py b/backend/routes/messaging.py index 57fc406..f009f03 100644 --- a/backend/routes/messaging.py +++ b/backend/routes/messaging.py @@ -114,6 +114,36 @@ async def dm_fetch(since: int | None = None, current_user: User = Depends(get_cu } +@router.get("/dm/history/{other_user_id}") +async def dm_history(other_user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): + envs = ( + 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() + ) + 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(), + } + for e in envs + ] + } + + @router.put("/edit_message/{message_id}") async def edit_message( message_id: int, @@ -193,6 +223,7 @@ async def reply_message( class MessaggingSocketManager: def __init__(self) -> None: self.connections: list[WebSocket] = [] + self.user_by_ws: dict[WebSocket, int] = {} 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}}) @@ -221,6 +252,7 @@ class MessaggingSocketManager: 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)}) except HTTPException as e: @@ -230,6 +262,7 @@ class MessaggingSocketManager: current_user = get_current_user_inner() if not current_user: raise HTTPException(401) + self.user_by_ws[websocket] = current_user.id request: SendMessageRequest = SendMessageRequest.model_validate(data["data"]) @@ -242,6 +275,46 @@ class MessaggingSocketManager: await websocket.send_json({"type": type, "data": response}) except HTTPException as e: await self.send_error(websocket, type, e) + elif type == "dmSend": + 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"], + ) + db.add(env) + db.commit() + db.refresh(env) + await self.send_to_user(env.recipient_id, { + "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(), + } + }) + await websocket.send_json({"type": type, "data": {"status": "ok", "id": env.id}}) + except HTTPException as e: + await self.send_error(websocket, type, e) elif type == "editMessage": try: current_user = get_current_user_inner() @@ -310,11 +383,18 @@ class MessaggingSocketManager: logger.info(f"WebSocket disconnected with code {e.code}: {e.reason}") finally: self.connections.remove(websocket) + if websocket in self.user_by_ws: + del self.user_by_ws[websocket] async def broadcast(self, message: dict): for websocket in self.connections: await websocket.send_json(message) + async def send_to_user(self, user_id: int, message: dict): + for websocket in self.connections: + if self.user_by_ws.get(websocket) == user_id: + await websocket.send_json(message) + messagingManager = MessaggingSocketManager() @router.websocket("/chat/ws") diff --git a/frontend/src/chat/dm.ts b/frontend/src/chat/dm.ts index 4b8dd96..6d3e720 100644 --- a/frontend/src/chat/dm.ts +++ b/frontend/src/chat/dm.ts @@ -5,6 +5,8 @@ import { ecdhSharedSecret, deriveWrappingKey } from "../crypto/asymmetric"; import { importAesGcmKey, aesGcmEncrypt, aesGcmDecrypt } from "../crypto/symmetric"; import { randomBytes } from "../crypto/kdf"; import { getCurrentKeys } from "../auth/crypto"; +import { websocket } from "../websocket"; +import type { WebSocketMessage } from "../core/types"; import type { Tabs } from "mdui/components/tabs"; function b64(a: Uint8Array): string { return btoa(String.fromCharCode(...a)); } @@ -90,11 +92,33 @@ async function loadUsers() { dmPanel = new DmPanel( async (text: string) => { if (activeDm?.publicKey) { - try { await sendDm(activeDm.userId, activeDm.publicKey, text); } catch {} + // WebSocket realtime send + const keys = getCurrentKeys(); + if (!keys) return; + const mk = randomBytes(32); + const wkSalt = randomBytes(16); + const shared = ecdhSharedSecret(keys.privateKey, ub64(activeDm.publicKey)); + const wkRaw = await deriveWrappingKey(shared, wkSalt, new Uint8Array([1])); + const wk = await importAesGcmKey(wkRaw); + const encMsg = await aesGcmEncrypt(await importAesGcmKey(mk), new TextEncoder().encode(text)); + const wrap = await aesGcmEncrypt(wk, mk); + const payload: WebSocketMessage = { + type: "dmSend", + credentials: { scheme: "Bearer", credentials: (await import("../auth/api")).authToken! }, + data: { + recipientId: activeDm.userId, + iv: b64(encMsg.iv), + ciphertext: b64(encMsg.ciphertext), + salt: b64(wkSalt), + iv2: b64(wrap.iv), + wrappedMk: b64(wrap.ciphertext) + } + }; + websocket.send(JSON.stringify(payload)); } }, () => { - // For v1, no DM history yet; just clear + // For v1, load history for this DM } ); } @@ -124,8 +148,31 @@ function init() { tabs.addEventListener("change", (e: any) => { if (e.detail?.value === "dms") { ensureUsersLoaded(); + dmPanel?.activate(); } }); } -init(); \ No newline at end of file +init(); + +// realtime incoming DMs +websocket.addEventListener("message", async (e) => { + try { + const msg = JSON.parse((e as MessageEvent).data); + if (msg?.type === "dmNew" && activeDm && msg.data.senderId === activeDm.userId) { + const plaintext = await decryptDm(msg.data, activeDm.publicKey!); + const container = document.getElementById("chat-messages")!; + const div = document.createElement("div"); + div.className = "message received"; + const inner = document.createElement("div"); + inner.className = "message-inner"; + const content = document.createElement("div"); + content.className = "message-content"; + content.textContent = plaintext; + inner.appendChild(content); + div.appendChild(inner); + container.appendChild(div); + container.scrollTop = container.scrollHeight; + } + } catch {} +}); \ No newline at end of file