diff --git a/migrations/0007_add_chat_message.sql b/migrations/0007_add_chat_message.sql new file mode 100644 index 0000000..ecddf6a --- /dev/null +++ b/migrations/0007_add_chat_message.sql @@ -0,0 +1,15 @@ +-- Migration: 0007_add_chat_message +-- Adds the chat_message table used by ChatDO for encrypted message persistence. +-- Content is encrypted with AES-256-GCM before storage (via ChatDO._persist_message). + +CREATE TABLE IF NOT EXISTS chat_message ( + id TEXT PRIMARY KEY, + classroom_id TEXT NOT NULL, + user_id TEXT NOT NULL, + display_name TEXT, + content TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (datetime('now')) +); + +CREATE INDEX IF NOT EXISTS idx_chat_classroom ON chat_message(classroom_id); +CREATE INDEX IF NOT EXISTS idx_chat_created ON chat_message(classroom_id, created_at); diff --git a/public/classroom_poc.html b/public/classroom_poc.html index 66e2773..ce19fa9 100644 --- a/public/classroom_poc.html +++ b/public/classroom_poc.html @@ -180,7 +180,7 @@

Room Chat

-
+
No messages yet!
@@ -284,6 +284,9 @@

// State let ws = null; let presenceWs = null; + let chatWs = null; + let chatReconnectAttempts = 0; + let chatReconnectTimer = null; let roomId = ''; let myParticipantId = ''; let myDisplayName = ''; @@ -536,6 +539,7 @@

reconnectAttempts = 0; setConnStatus(true); connectPresenceWS(); + connectChatWS(); showToast('Connected to classroom', 'success'); }; @@ -559,8 +563,10 @@

intentionalDisconnect = true; if (reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; } if (presenceReconnectTimer) { clearTimeout(presenceReconnectTimer); presenceReconnectTimer = null; } + if (chatReconnectTimer) { clearTimeout(chatReconnectTimer); chatReconnectTimer = null; } if (ws) { ws.close(); ws = null; } if (presenceWs) { presenceWs.close(); presenceWs = null; } + if (chatWs) { chatWs.close(); chatWs = null; } setConnStatus(false); } @@ -626,6 +632,79 @@

}; } + // Chat WebSocket (ChatDO — /ws/chat/:id) + function connectChatWS() { + if (chatWs && (chatWs.readyState === WebSocket.OPEN || chatWs.readyState === WebSocket.CONNECTING)) return; + const proto = location.protocol === 'https:' ? 'wss:' : 'ws:'; + const params = new URLSearchParams(); + const token = localStorage.getItem('auth_token'); + if (token) { + // Prefer authenticated token so the server derives identity server-side. + params.set('token', token); + } else { + // Fallback to anonymous identity only when explicitly enabled server-side. + params.set('user_id', myParticipantId); + params.set('display_name', myDisplayName); + } + const qs = params.toString(); + const url = qs + ? `${proto}//${location.host}/ws/chat/${encodeURIComponent(roomId)}?${qs}` + : `${proto}//${location.host}/ws/chat/${encodeURIComponent(roomId)}`; + chatWs = new WebSocket(url); + chatWs.onopen = () => { + chatReconnectAttempts = 0; + if (chatReconnectTimer) { clearTimeout(chatReconnectTimer); chatReconnectTimer = null; } + }; + chatWs.onmessage = (ev) => { + try { + const data = JSON.parse(ev.data); + if (data.type === 'chat_history') { + // Always clear the placeholder, even if history is empty + chatMessages.innerHTML = ''; + chatFirstMessage = false; + (data.messages || []).forEach((msg) => appendChatMessage( + msg.display_name || msg.user_id, + msg.text, + msg.user_id === myParticipantId + )); + } else if (data.type === 'chat_error') { + showToast(data.message || 'Failed to send chat message', 'error'); + } else if (data.type === 'chat_message') { + appendChatMessage( + data.display_name || data.user_id, + data.text, + data.user_id === myParticipantId + ); + } + } catch (err) { + console.warn('[chat] dropped malformed ws frame', err, ev.data); + } + }; + chatWs.onerror = () => { + console.warn('[chat] WebSocket error'); + }; + chatWs.onclose = () => { + if (intentionalDisconnect) return; + if (chatReconnectAttempts >= MAX_RECONNECT) { + showToast('Lost connection to chat — please rejoin.', 'warning'); + return; + } + const delay = Math.min(BASE_DELAY * Math.pow(2, chatReconnectAttempts), 30000); + chatReconnectAttempts++; + chatReconnectTimer = setTimeout(() => { + if (!chatWs || chatWs.readyState === WebSocket.CLOSED) connectChatWS(); + }, delay); + }; + } + + function chatSend(obj) { + if (chatWs && chatWs.readyState === WebSocket.OPEN) { + chatWs.send(JSON.stringify(obj)); + return true; + } + return false; + } + function presenceSend(obj) { if (presenceWs && presenceWs.readyState === WebSocket.OPEN) { presenceWs.send(JSON.stringify(obj)); @@ -769,10 +848,6 @@

} break; - case 'chat_message': - appendChatMessage(data.display_name, data.text, data.participant_id === myParticipantId); - break; - case 'seat_updated': if (!participants[data.participant_id] && data.participant_id !== myParticipantId) { participants[data.participant_id] = { @@ -997,13 +1072,17 @@

presenceSend({ type: 'presence', hand_raised: nextHandRaised }); }); - // Chat + // Chat chatForm.addEventListener('submit', (e) => { e.preventDefault(); const text = chatInput.value.trim(); if (!text) return; - wsSend({ type: 'chat_message', text, timestamp: new Date().toISOString() }); - chatInput.value = ''; + const ok = chatSend({ type: 'chat_message', text }); + if (ok) { + chatInput.value = ''; + } else { + showToast('Chat is not connected. Retry in a moment…', 'warning'); + } }); function appendChatMessage(name, text, isSelf) { diff --git a/src/chat_do.py b/src/chat_do.py new file mode 100644 index 0000000..3b4c13c --- /dev/null +++ b/src/chat_do.py @@ -0,0 +1,364 @@ +import base64 +import datetime +import json +import os +import re +import uuid +from urllib.parse import urlparse, parse_qs + +from workers import Response, DurableObject + +import js +from pyodide.ffi import to_js +from js import WebSocketPair, WebSocketRequestResponsePair + + +# Minimal local helpers to avoid circular imports with worker.py + +def _derive_key(secret: str) -> bytes: + import hashlib + return hashlib.sha256(secret.encode("utf-8")).digest() + + +def _encrypt_xor(plaintext: str, secret: str) -> str: + if not plaintext: + return "" + key = _derive_key(secret) + data = plaintext.encode("utf-8") + ks = (key * (len(data) // len(key) + 1))[: len(data)] + return base64.b64encode(bytes(a ^ b for a, b in zip(data, ks))).decode("ascii") + + +def _decrypt_xor(ciphertext: str, secret: str) -> str: + if not ciphertext: + return "" + try: + key = _derive_key(secret) + raw = base64.b64decode(ciphertext) + ks = (key * (len(raw) // len(key) + 1))[: len(raw)] + return bytes(a ^ b for a, b in zip(raw, ks)).decode("utf-8") + except Exception: + return "[decryption error]" + + +def _pbkdf2_key(secret: str) -> bytes: + import hashlib + salt = hashlib.sha256(b"aol-edu-aes-salt-v1" + secret.encode()).digest() + return hashlib.pbkdf2_hmac("sha256", secret.encode("utf-8"), salt, 100_000) + + +async def _import_aes_key(key_bytes: bytes) -> object: + key_buf = to_js(key_bytes, create_pyproxies=False) + algo = to_js({"name": "AES-GCM"}, dict_converter=js.Object.fromEntries) + usages = to_js(["encrypt", "decrypt"]) + return await js.crypto.subtle.importKey("raw", key_buf, algo, False, usages) + + +async def encrypt_aes(plaintext: str, secret: str) -> str: + if not plaintext: + return "" + key_bytes = _pbkdf2_key(secret) + crypto_key = await _import_aes_key(key_bytes) + iv_array = js.Uint8Array.new(12) + js.crypto.getRandomValues(iv_array) + algo = to_js({"name": "AES-GCM", "iv": iv_array}, dict_converter=js.Object.fromEntries) + data = to_js(plaintext.encode("utf-8"), create_pyproxies=False) + ct_buf = await js.crypto.subtle.encrypt(algo, crypto_key, data) + ct = bytes(js.Uint8Array.new(ct_buf)) + return "v1:" + base64.b64encode(bytes(iv_array) + ct).decode("ascii") + + +async def decrypt_aes(ciphertext: str, secret: str) -> str: + if not ciphertext: + return "" + if not ciphertext.startswith("v1:"): + return _decrypt_xor(ciphertext, secret) + try: + raw = base64.b64decode(ciphertext[3:]) + iv, ct = raw[:12], raw[12:] + except Exception: + return "[decryption error]" + try: + key_bytes = _pbkdf2_key(secret) + crypto_key = await _import_aes_key(key_bytes) + iv_array = to_js(iv, create_pyproxies=False) + algo = to_js({"name": "AES-GCM", "iv": iv_array}, dict_converter=js.Object.fromEntries) + data = to_js(ct, create_pyproxies=False) + pt_buf = await js.crypto.subtle.decrypt(algo, crypto_key, data) + return bytes(js.Uint8Array.new(pt_buf)).decode("utf-8") + except Exception: + return "[decryption error]" + + +def new_id() -> str: + b = bytearray(os.urandom(16)) + b[6] = (b[6] & 0x0F) | 0x40 + b[8] = (b[8] & 0x3F) | 0x80 + h = b.hex() + return f"{h[:8]}-{h[8:12]}-{h[12:16]}-{h[16:20]}-{h[20:]}" + + +def verify_token(raw: str, secret: str): + import hashlib, hmac as _hmac + if not raw: + return None + try: + token = raw.removeprefix("Bearer ").strip() + dot = token.rfind(".") + if dot == -1: + return None + p, sig = token[:dot], token[dot + 1 :] + exp = _hmac.new(secret.encode("utf-8"), p.encode("utf-8"), hashlib.sha256).hexdigest() + if not _hmac.compare_digest(sig, exp): + return None + padding = (4 - len(p) % 4) % 4 + return json.loads(base64.b64decode(p + "=" * padding).decode("utf-8")) + except Exception: + return None + + +class ChatDO(DurableObject): + """Room-scoped real-time chat Durable Object.""" + + _MAX_BUFFER = 200 + + def __init__(self, ctx, env): + super().__init__(ctx, env) + self.sessions = {} + self.messages = [] + self._history_loaded = False + self._room_id = getattr(self, "_room_id", "") + + for ws in self.ctx.getWebSockets(): + try: + attachment = ws.deserializeAttachment() + if not attachment: + continue + data = json.loads(attachment) if isinstance(attachment, str) else attachment + session_id = data.get("session_id", str(uuid.uuid4())) + user_id = str(data.get("user_id", ""))[:64] + display_name = str(data.get("display_name", user_id or "Unknown"))[:64] + if not self._room_id: + rid = str(data.get("classroom_id", "")) + if rid: + self._room_id = rid + if not user_id: + continue + self.sessions[session_id] = { + "ws": ws, + "user_id": user_id, + "display_name": display_name, + } + except Exception as exc: + print(f"[ChatDO.__init__.restore] error={exc!r}") + + self.ctx.setWebSocketAutoResponse( + WebSocketRequestResponsePair.new("ping", "pong") + ) + + async def _load_history(self, classroom_id: str): + if self._history_loaded: + return + try: + enc_key = getattr(self.env, "CHAT_ENCRYPTION_KEY", None) or getattr(self.env, "ENCRYPTION_KEY", "") + rows = await self.env.DB.prepare( + "SELECT id, user_id, display_name, content, created_at FROM chat_message WHERE classroom_id = ? ORDER BY created_at DESC LIMIT ?" + ).bind(classroom_id, self._MAX_BUFFER).all() + loaded = [] + for row in (rows.results or []): + try: + text = await decrypt_aes(row.content or "", enc_key) if enc_key else (row.content or "") + except Exception: + text = "[decryption error]" + name = getattr(row, "display_name", None) or row.user_id + loaded.append({ + "id": row.id, + "user_id": row.user_id, + "display_name": name, + "text": text, + "timestamp": row.created_at or "", + }) + # reverse to chronological order for playback + self.messages = list(reversed(loaded)) + self._history_loaded = True + except Exception as exc: + print(f"[ChatDO._load_history] error={exc!r}") + + async def on_fetch(self, request): + upgrade = request.headers.get("Upgrade") or "" + if upgrade.lower() != "websocket": + return Response(json.dumps({"error": "Expected WebSocket upgrade"}), status=426, headers={"Content-Type": "application/json"}) + + parsed = urlparse(request.url) + qs = parse_qs(parsed.query) + token_param = (qs.get("token") or [None])[0] + user_param = (qs.get("user_id") or [None])[0] + display_param = (qs.get("display_name") or [None])[0] + # Derive classroom_id from path segment /ws/chat/:id + m = re.search(r"/ws/chat/([A-Za-z0-9_-]+)", parsed.path or "") + classroom_id = m.group(1) if m else "" + + authenticated_user = verify_token(token_param or "", self.env.JWT_SECRET) if token_param else None + if authenticated_user: + user_id = str(authenticated_user.get("id", "")) + display_name = str(authenticated_user.get("username") or user_id) + else: + if not user_param: + return Response(json.dumps({"error": "Invalid user"}), status=400, headers={"Content-Type": "application/json"}) + user_id = str(user_param) + display_name = str(display_param or user_id) + + user_id = user_id[:64] + display_name = display_name[:64] + if not user_id: + return Response(json.dumps({"error": "Invalid user_id"}), status=400, headers={"Content-Type": "application/json"}) + + client, server = WebSocketPair.new().object_values() + self.ctx.acceptWebSocket(server) + + session_id = str(uuid.uuid4()) + # Include classroom_id in server-side attachment to support hibernation resume + self._room_id = classroom_id + attachment = json.dumps({ + "session_id": session_id, + "user_id": user_id, + "display_name": display_name, + "classroom_id": classroom_id, + }) + server.serializeAttachment(attachment) + + self.sessions[session_id] = { + "ws": server, + "user_id": user_id, + "display_name": display_name, + } + + await self._load_history(classroom_id) + try: + server.send(json.dumps({"type": "chat_history", "messages": list(self.messages)})) + except Exception as exc: + print(f"[ChatDO.on_fetch.send_history] error={exc!r}") + + return Response(None, status=101, web_socket=client) + + async def on_webSocketMessage(self, ws, message): + try: + raw = message if isinstance(message, str) else message.decode("utf-8") + if len(raw) > 4096: + return + data = json.loads(raw) + except Exception: + return + if not isinstance(data, dict): + return + + # Lookup session + sid = None + info = None + for s_id, s_info in self.sessions.items(): + try: + if s_info["ws"] == ws: + sid, info = s_id, s_info + break + except Exception: + pass + if not info: + try: + raw_att = ws.deserializeAttachment() + if raw_att: + att = json.loads(raw_att) if isinstance(raw_att, str) else raw_att + sid = att.get("session_id") + info = self.sessions.get(sid) + except Exception: + pass + if not info: + return + + if data.get("type") != "chat_message": + return + + raw_text = data.get("text", "") + if not isinstance(raw_text, str): + return + text = raw_text.strip()[:500] + if not text: + return + + classroom_id = getattr(self, "_room_id", "") + if not classroom_id: + try: + raw_att = ws.deserializeAttachment() + if raw_att: + att = json.loads(raw_att) if isinstance(raw_att, str) else raw_att + classroom_id = str(att.get("classroom_id", "")) + if classroom_id: + self._room_id = classroom_id + except Exception: + pass + + timestamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + msg_id = new_id() + entry = { + "id": msg_id, + "user_id": info["user_id"], + "display_name": info["display_name"], + "text": text, + "timestamp": timestamp, + } + # Persist first; only broadcast if durable + ok = await self._persist_message(msg_id, classroom_id, info["user_id"], info["display_name"], text) + if not ok: + try: + ws.send(json.dumps({"type": "chat_error", "message": "Failed to persist message"})) + except Exception: + pass + return + + self.messages.append(entry) + if len(self.messages) > self._MAX_BUFFER: + self.messages = self.messages[-self._MAX_BUFFER :] + + self._broadcast(json.dumps({ + "type": "chat_message", + "id": msg_id, + "user_id": info["user_id"], + "display_name": info["display_name"], + "text": text, + "timestamp": timestamp, + })) + + async def on_webSocketClose(self, ws, _code, _reason, _was_clean): + for sid, s_info in list(self.sessions.items()): + try: + if s_info["ws"] == ws: + self.sessions.pop(sid, None) + break + except Exception: + pass + + async def on_webSocketError(self, _ws, error): + print(f"[ChatDO.on_webSocketError] error={error!r}") + + def _broadcast(self, payload, exclude_session_id=None): + for sid, info in self.sessions.items(): + if sid == exclude_session_id: + continue + try: + info["ws"].send(payload) + except Exception as exc: + print(f"[ChatDO._broadcast] sid={sid} user_id={info.get('user_id')} error={exc!r}") + + async def _persist_message(self, msg_id: str, classroom_id: str, user_id: str, display_name: str, text: str) -> bool: + try: + enc_key = getattr(self.env, "CHAT_ENCRYPTION_KEY", None) or getattr(self.env, "ENCRYPTION_KEY", "") + if not enc_key: + return False + encrypted_content = await encrypt_aes(text, enc_key) + await self.env.DB.prepare( + "INSERT INTO chat_message (id, classroom_id, user_id, display_name, content, created_at) VALUES (?, ?, ?, ?, ?, datetime('now'))" + ).bind(msg_id, classroom_id, user_id, display_name, encrypted_content).run() + return True + except Exception as exc: + print(f"[ChatDO._persist_message] error={exc!r}") + return False + diff --git a/src/worker.py b/src/worker.py index f274a07..49aef40 100644 --- a/src/worker.py +++ b/src/worker.py @@ -50,6 +50,18 @@ from pyodide.ffi import to_js from js import WebSocketPair, WebSocketRequestResponsePair import uuid +try: + from chat_do import ChatDO as _ChatDO +except ModuleNotFoundError: + try: + from src.chat_do import ChatDO as _ChatDO + except ModuleNotFoundError: + # Fallback stub for local tooling environments where chat_do is absent. + class _ChatDO(DurableObject): + pass + +# Re-export as ChatDO so wrangler binding resolves in production +ChatDO = _ChatDO _SENTRY_INITIALIZED = False _SENTRY_DSN: str = "" @@ -700,7 +712,7 @@ async def send_password_reset_email(to_email: str, _username: str, token: str, e updated_at TEXT NOT NULL DEFAULT (datetime('now')), FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE )""", - # Email verification tokens (token_hash = SHA-256 of plaintext token) + # Email verification tokens """CREATE TABLE IF NOT EXISTS email_verification_tokens ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, @@ -710,7 +722,7 @@ async def send_password_reset_email(to_email: str, _username: str, token: str, e FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE )""", "CREATE INDEX IF NOT EXISTS idx_evtoken_user ON email_verification_tokens(user_id)", - # Password reset tokens (token_hash = SHA-256 of plaintext token) + # Password reset tokens """CREATE TABLE IF NOT EXISTS password_reset_tokens ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, @@ -720,6 +732,16 @@ async def send_password_reset_email(to_email: str, _username: str, token: str, e FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE )""", "CREATE INDEX IF NOT EXISTS idx_prtoken_user ON password_reset_tokens(user_id)", + # Chat messages + """CREATE TABLE IF NOT EXISTS chat_message ( + id TEXT PRIMARY KEY, + classroom_id TEXT NOT NULL, + user_id TEXT NOT NULL, + display_name TEXT, + content TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (datetime('now')) + )""", + "CREATE INDEX IF NOT EXISTS idx_chat_created ON chat_message(classroom_id, created_at)", ] @@ -2103,23 +2125,9 @@ def _valid_norm_position(value): }), exclude_session_id=sid) elif msg_type == "chat_message": - raw_text = data.get("text", "") - if not isinstance(raw_text, str): - return - text = raw_text.strip()[:500] - if not text: - return - raw_timestamp = data.get("timestamp", "") - timestamp = raw_timestamp[:64] if isinstance(raw_timestamp, str) else "" - self._broadcast(json.dumps({ - "type": "chat_message", - "participant_id": info["participant_id"], - "display_name": info["display_name"], - "text": text, - "timestamp": timestamp, - })) + # chat_message is handled by ChatDO (/ws/chat/:id). + pass - elif msg_type == "update_seat": # Classroom layout, keep in sync with DESK_ROWS/DESK_COLS in classroom_poc.html DESK_ROWS = 3 @@ -2623,6 +2631,17 @@ async def _dispatch(request, env): await capture_exception(e, request, env, "presence_do_dispatch") return err("Failed to connect to presence channel", 500) + m_chat = re.fullmatch(r"/ws/chat/([A-Za-z0-9_-]+)", path) + if m_chat: + room_id = m_chat.group(1) + try: + do_id = env.CHAT_DO.idFromName(room_id) + stub = env.CHAT_DO.get(do_id) + return await stub.fetch(request) + except Exception as e: + await capture_exception(e, request, env, "chat_do_dispatch") + return err("Failed to connect to chat channel", 500) + if path.startswith("/api/"): if path == "/api/init" and method == "POST": try: diff --git a/wrangler.toml b/wrangler.toml index 761e130..c1a98a0 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -6,7 +6,7 @@ compatibility_flags = ["python_workers"] [assets] directory = "./public" binding = "ASSETS" -run_worker_first = ["/api/*"] +run_worker_first = ["/api/*", "/ws/*"] html_handling = "auto-trailing-slash" not_found_handling = "404-page" @@ -28,6 +28,10 @@ class_name = "ClassroomDO" name = "PRESENCE_DO" class_name = "PresenceDO" +[[durable_objects.bindings]] +name = "CHAT_DO" +class_name = "ChatDO" + [[migrations]] tag = "v1" new_sqlite_classes = ["ClassroomDO"] @@ -36,6 +40,10 @@ new_sqlite_classes = ["ClassroomDO"] tag = "v2" new_sqlite_classes = ["PresenceDO"] +[[migrations]] +tag = "v3" +new_sqlite_classes = ["ChatDO"] + [observability] enabled = false