fix(backend+frontend): WebSocket auth, Yjs persistence fix, keep docs in memory (Issues #2, #4, #5)

backend/src/server.ts

@@ -102,7 +102,7 @@ export async function createServer(opts: ServerOptions) {
   registerShareRoutes(fastify, opts.db, authService);
 
   // Yjs collaboration WebSocket
-  registerYjsWebSocket(fastify);
+  registerYjsWebSocket(fastify, authService);
 
   return fastify;
 }

backend/src/websocket/yjsServer.ts

@@ -5,6 +5,7 @@ import * as Y from 'yjs';
 import { LeveldbPersistence } from 'y-leveldb';
 import type { FastifyInstance } from 'fastify';
 import type { WebSocket } from '@fastify/websocket';
+import type { AuthService } from '../auth/AuthService.js';
 
 const PERSISTENCE_DIR = process.env.YJS_PERSISTENCE_DIR || '/tmp/yjs-documents';
 
@@ -18,6 +19,12 @@ let persistenceReady: Promise<void> | null = null;
 export async function initYjsPersistence(): Promise<void> {
   try {
     persistence = new LeveldbPersistence(PERSISTENCE_DIR);
+    // Set persistenceReady so getOrCreateDoc can await it
+    persistenceReady = (async () => {
+      // Wait for LevelDB to be ready
+      await (persistence as any).getAllDocNames();
+    })();
+    await persistenceReady;
     console.log(`Yjs persistence initialized at ${PERSISTENCE_DIR}`);
   } catch (err) {
     console.warn('Yjs persistence failed to init (non-fatal):', err);
@@ -80,7 +87,7 @@ function broadcastUpdate(docName: string, update: Uint8Array, exclude?: WebSocke
   }
 }
 
-export function registerYjsWebSocket(fastify: FastifyInstance): void {
+export function registerYjsWebSocket(fastify: FastifyInstance, authService: AuthService): void {
   fastify.get('/ws/collab/:docName', { websocket: true }, async (socket: WebSocket, request: any) => {
     const docName = request.params.docName as string;
     if (!docName) {
@@ -88,6 +95,18 @@ export function registerYjsWebSocket(fastify: FastifyInstance): void {
       return;
     }
 
+    // Issue #2: Authenticate WebSocket connection via query param token
+    const token = request.query.token as string;
+    if (!token) {
+      socket.close(4001, 'Authentication required');
+      return;
+    }
+    const user = authService.getUserFromSession(token);
+    if (!user) {
+      socket.close(4001, 'Invalid or expired session');
+      return;
+    }
+
     const doc = await getOrCreateDoc(docName);
 
     // Track connection
@@ -113,16 +132,14 @@ export function registerYjsWebSocket(fastify: FastifyInstance): void {
       }
     });
 
-    // Clean up on disconnect
+    // Clean up on disconnect (Issue #5: keep doc in memory for reconnection)
     socket.on('close', () => {
       const conns = connections.get(docName);
       if (conns) {
         conns.delete(socket);
         if (conns.size === 0) {
           connections.delete(docName);
-          // Optionally keep doc in memory for a while
-          // For now, remove it to free memory
-          docs.delete(docName);
+          // Keep doc in memory for reconnection (persistence handles durability)
         }
       }
     });