fix: address CodeRabbit review feedback on auth security and persistence

Author: Rozerxshashank

src/routes/settings.tsx

@@ -57,6 +57,8 @@ function SettingsPage() {
     useEffect(() => {
         if (typeof window === 'undefined') return;
 
+        let isMounted = true;
+
         // If we already have a token, no need to generate
         const existing = localStorage.getItem('rein_auth_token');
         if (existing) return;
@@ -66,15 +68,19 @@ function SettingsPage() {
         const socket = new WebSocket(wsUrl);
 
         socket.onopen = () => {
-            socket.send(JSON.stringify({ type: 'generate-token' }));
+            if (socket.readyState === WebSocket.OPEN) {
+                socket.send(JSON.stringify({ type: 'generate-token' }));
+            }
         };
 
         socket.onmessage = (event) => {
             try {
                 const data = JSON.parse(event.data);
                 if (data.type === 'token-generated' && data.token) {
-                    setAuthToken(data.token);
-                    localStorage.setItem('rein_auth_token', data.token);
+                    if (isMounted) {
+                        setAuthToken(data.token);
+                        localStorage.setItem('rein_auth_token', data.token);
+                    }
                     socket.close();
                 }
             } catch (e) {
@@ -83,7 +89,10 @@ function SettingsPage() {
         };
 
         return () => {
-            if (socket.readyState === WebSocket.OPEN) socket.close();
+            isMounted = false;
+            if (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING) {
+                socket.close();
+            }
         };
     }, []);
 

src/server/tokenStore.ts

@@ -1,17 +1,22 @@
 import crypto from 'crypto';
 import fs from 'fs';
 import path from 'path';
+import { fileURLToPath } from 'url';
 
 interface TokenEntry {
     token: string;
     createdAt: number;
     lastUsed: number;
 }
 
-const TOKENS_FILE = path.resolve('./src/tokens.json');
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+const TOKENS_FILE = path.resolve(__dirname, '../tokens.json');
 const EXPIRY_MS = 10 * 24 * 60 * 60 * 1000; // 10 days
 
 let tokens: TokenEntry[] = [];
+let lastSaveTime = 0;
+const SAVE_THROTTLE_MS = 60 * 1000; // 1 minute
 
 function load(): void {
     try {
@@ -23,9 +28,13 @@ function load(): void {
     }
 }
 
-function save(): void {
+function save(force = false): void {
+    const now = Date.now();
+    if (!force && (now - lastSaveTime) < SAVE_THROTTLE_MS) return;
+
     try {
         fs.writeFileSync(TOKENS_FILE, JSON.stringify(tokens, null, 2));
+        lastSaveTime = now;
     } catch (e) {
         console.error('Failed to persist tokens:', e);
     }
@@ -35,7 +44,7 @@ function purgeExpired(): void {
     const now = Date.now();
     const before = tokens.length;
     tokens = tokens.filter(t => (now - t.lastUsed) < EXPIRY_MS);
-    if (tokens.length !== before) save();
+    if (tokens.length !== before) save(true);
 }
 
 /** Constant-time string comparison to prevent timing attacks. */
@@ -61,7 +70,7 @@ export function storeToken(token: string): void {
         const now = Date.now();
         tokens.push({ token, createdAt: now, lastUsed: now });
     }
-    save();
+    save(true);
 }
 
 /** Check if a token is already known/stored on the server. */
@@ -75,11 +84,19 @@ export function touchToken(token: string): void {
     const entry = tokens.find(t => timingSafeEqual(t.token, token));
     if (entry) {
         entry.lastUsed = Date.now();
-        // Persist periodically — save only if >60s since last save
-        // to avoid excessive disk I/O on every message
+        save(); // Throttled internally
     }
 }
 
+/** Returns the most recently used active token, if any. */
+export function getActiveToken(): string | null {
+    purgeExpired();
+    if (tokens.length === 0) return null;
+    // Return the one used most recently
+    const sorted = [...tokens].sort((a, b) => b.lastUsed - a.lastUsed);
+    return sorted[0].token;
+}
+
 /** Check if any tokens exist yet (first-run detection). */
 export function hasTokens(): boolean {
     purgeExpired();

src/server/websocket.ts

@@ -1,6 +1,6 @@
 import { WebSocketServer, WebSocket } from 'ws';
 import { InputHandler, InputMessage } from './InputHandler';
-import { storeToken, isKnownToken, touchToken, hasTokens, generateToken } from './tokenStore';
+import { storeToken, isKnownToken, touchToken, generateToken, getActiveToken } from './tokenStore';
 import os from 'os';
 import fs from 'fs';
 import { IncomingMessage } from 'http';
@@ -53,14 +53,6 @@ export function createWsServer(server: any) {
             return;
         }
 
-        // First remote connection ever — accept and store the token
-        if (!hasTokens()) {
-            storeToken(token);
-            wss.handleUpgrade(request, socket, head, (ws) => {
-                wss.emit('connection', ws, request, token, false);
-            });
-            return;
-        }
 
         // Validate against known tokens
         if (!isKnownToken(token)) {
@@ -103,9 +95,15 @@ export function createWsServer(server: any) {
                         ws.send(JSON.stringify({ type: 'auth-error', error: 'Only localhost can generate tokens' }));
                         return;
                     }
-                    const newToken = generateToken();
-                    storeToken(newToken);
-                    ws.send(JSON.stringify({ type: 'token-generated', token: newToken }));
+
+                    // Idempotent: return active token if one exists
+                    let tokenToReturn = getActiveToken();
+                    if (!tokenToReturn) {
+                        tokenToReturn = generateToken();
+                        storeToken(tokenToReturn);
+                    }
+
+                    ws.send(JSON.stringify({ type: 'token-generated', token: tokenToReturn }));
                     return;
                 }