Implement real-time WebSocket

This commit is contained in:
2025-08-25 17:40:47 +03:00
Unverified
parent eda9de699e
commit 8acb2dd507
2 changed files with 130 additions and 3 deletions
+80
View File
@@ -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}") @router.put("/edit_message/{message_id}")
async def edit_message( async def edit_message(
message_id: int, message_id: int,
@@ -193,6 +223,7 @@ async def reply_message(
class MessaggingSocketManager: class MessaggingSocketManager:
def __init__(self) -> None: def __init__(self) -> None:
self.connections: list[WebSocket] = [] self.connections: list[WebSocket] = []
self.user_by_ws: dict[WebSocket, int] = {}
async def send_error(self, websocket: WebSocket, type: str, e: HTTPException): 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}}) 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() current_user = get_current_user_inner()
if not current_user: if not current_user:
raise HTTPException(401) raise HTTPException(401)
self.user_by_ws[websocket] = current_user.id
await websocket.send_json({"type": type, "data": await get_messages(current_user, db)}) await websocket.send_json({"type": type, "data": await get_messages(current_user, db)})
except HTTPException as e: except HTTPException as e:
@@ -230,6 +262,7 @@ class MessaggingSocketManager:
current_user = get_current_user_inner() current_user = get_current_user_inner()
if not current_user: if not current_user:
raise HTTPException(401) raise HTTPException(401)
self.user_by_ws[websocket] = current_user.id
request: SendMessageRequest = SendMessageRequest.model_validate(data["data"]) request: SendMessageRequest = SendMessageRequest.model_validate(data["data"])
@@ -242,6 +275,46 @@ class MessaggingSocketManager:
await websocket.send_json({"type": type, "data": response}) await websocket.send_json({"type": type, "data": response})
except HTTPException as e: except HTTPException as e:
await self.send_error(websocket, type, 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": elif type == "editMessage":
try: try:
current_user = get_current_user_inner() current_user = get_current_user_inner()
@@ -310,11 +383,18 @@ class MessaggingSocketManager:
logger.info(f"WebSocket disconnected with code {e.code}: {e.reason}") logger.info(f"WebSocket disconnected with code {e.code}: {e.reason}")
finally: finally:
self.connections.remove(websocket) self.connections.remove(websocket)
if websocket in self.user_by_ws:
del self.user_by_ws[websocket]
async def broadcast(self, message: dict): async def broadcast(self, message: dict):
for websocket in self.connections: for websocket in self.connections:
await websocket.send_json(message) 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() messagingManager = MessaggingSocketManager()
@router.websocket("/chat/ws") @router.websocket("/chat/ws")
+49 -2
View File
@@ -5,6 +5,8 @@ import { ecdhSharedSecret, deriveWrappingKey } from "../crypto/asymmetric";
import { importAesGcmKey, aesGcmEncrypt, aesGcmDecrypt } from "../crypto/symmetric"; import { importAesGcmKey, aesGcmEncrypt, aesGcmDecrypt } from "../crypto/symmetric";
import { randomBytes } from "../crypto/kdf"; import { randomBytes } from "../crypto/kdf";
import { getCurrentKeys } from "../auth/crypto"; import { getCurrentKeys } from "../auth/crypto";
import { websocket } from "../websocket";
import type { WebSocketMessage } from "../core/types";
import type { Tabs } from "mdui/components/tabs"; import type { Tabs } from "mdui/components/tabs";
function b64(a: Uint8Array): string { return btoa(String.fromCharCode(...a)); } function b64(a: Uint8Array): string { return btoa(String.fromCharCode(...a)); }
@@ -90,11 +92,33 @@ async function loadUsers() {
dmPanel = new DmPanel( dmPanel = new DmPanel(
async (text: string) => { async (text: string) => {
if (activeDm?.publicKey) { 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) => { tabs.addEventListener("change", (e: any) => {
if (e.detail?.value === "dms") { if (e.detail?.value === "dms") {
ensureUsersLoaded(); ensureUsersLoaded();
dmPanel?.activate();
} }
}); });
} }
init(); 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 {}
});