feat: add token-based WebSocket authentication (#25)

Rozerxshashank · Rozerxshashank · commit f25e711a86ee · 2026-02-17T20:55:06.000+05:30
diff --git a/.gitignore b/.gitignore
@@ -11,4 +11,5 @@ count.txt
 .output
 .vinxi
 todos.json
+tokens.json
 SPEC.md
diff --git a/src/hooks/useRemoteConnection.ts b/src/hooks/useRemoteConnection.ts
@@ -8,7 +8,22 @@ export const useRemoteConnection = () => {
         let isMounted = true;
         const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
         const host = window.location.host;
-        const wsUrl = `${protocol}//${host}/ws`;
+
+        // Get token from URL params (passed via QR code) or localStorage
+        const urlParams = new URLSearchParams(window.location.search);
+        const urlToken = urlParams.get('token');
+        const storedToken = localStorage.getItem('rein_auth_token');
+        const token = urlToken || storedToken;
+
+        // Persist URL token to localStorage for future reconnections
+        if (urlToken && urlToken !== storedToken) {
+            localStorage.setItem('rein_auth_token', urlToken);
+        }
+
+        let wsUrl = `${protocol}//${host}/ws`;
+        if (token) {
+            wsUrl += `?token=${encodeURIComponent(token)}`;
+        }
 
         let reconnectTimer: NodeJS.Timeout;
 
@@ -24,7 +39,6 @@ export const useRemoteConnection = () => {
                 wsRef.current = null;
             }
 
-            console.log(`Connecting to ${wsUrl}`);
             setStatus('connecting');
             const socket = new WebSocket(wsUrl);
 
@@ -37,8 +51,7 @@ export const useRemoteConnection = () => {
                     reconnectTimer = setTimeout(connect, 3000);
                 }
             };
-            socket.onerror = (e) => {
-                console.error("WS Error", e);
+            socket.onerror = () => {
                 socket.close();
             };
 
diff --git a/src/routes/settings.tsx b/src/routes/settings.tsx
@@ -42,12 +42,51 @@ function SettingsPage() {
     const [qrData, setQrData] = useState('');
 
     // Load initial state (IP is not stored in localStorage; only sensitivity, invert, theme are client settings)
+    const [authToken, setAuthToken] = useState(() => {
+        if (typeof window === 'undefined') return '';
+        return localStorage.getItem('rein_auth_token') || '';
+    });
+
     useEffect(() => {
         const defaultIp = typeof window !== 'undefined' ? window.location.hostname : 'localhost';
         setIp(defaultIp);
         setFrontendPort(String(CONFIG.FRONTEND_PORT));
     }, []);
 
+    // Auto-generate token on settings page load (localhost only)
+    useEffect(() => {
+        if (typeof window === 'undefined') return;
+
+        // If we already have a token, no need to generate
+        const existing = localStorage.getItem('rein_auth_token');
+        if (existing) return;
+
+        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+        const wsUrl = `${protocol}//${window.location.host}/ws`;
+        const socket = new WebSocket(wsUrl);
+
+        socket.onopen = () => {
+            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);
+                    socket.close();
+                }
+            } catch (e) {
+                console.error(e);
+            }
+        };
+
+        return () => {
+            if (socket.readyState === WebSocket.OPEN) socket.close();
+        };
+    }, []);
+
     // Effect: Update LocalStorage when settings change
     useEffect(() => {
         localStorage.setItem('rein_sensitivity', String(sensitivity));
@@ -57,43 +96,47 @@ function SettingsPage() {
         localStorage.setItem('rein_invert', JSON.stringify(invertScroll));
     }, [invertScroll]);
 
+    // Effect: Theme
     useEffect(() => {
         if (typeof window === 'undefined') return;
         localStorage.setItem(APP_CONFIG.THEME_STORAGE_KEY, theme);
         document.documentElement.setAttribute('data-theme', theme);
     }, [theme]);
 
-    // Generate QR when IP changes (IP is not persisted to localStorage)
+    // Generate QR when IP changes or Token changes
     useEffect(() => {
         if (!ip || typeof window === 'undefined') return;
         const appPort = String(CONFIG.FRONTEND_PORT);
         const protocol = window.location.protocol;
-        const shareUrl = `${protocol}//${ip}:${appPort}/trackpad`;
+        let shareUrl = `${protocol}//${ip}:${appPort}/trackpad`;
+
+        // Token is embedded in the URL — QR automatically includes it
+        if (authToken) {
+            shareUrl += `?token=${encodeURIComponent(authToken)}`;
+        }
+
         QRCode.toDataURL(shareUrl)
             .then(setQrData)
             .catch((e) => console.error('QR Error:', e));
-    }, [ip]);
+    }, [ip, authToken]);
 
     // Effect: Auto-detect LAN IP from Server (only if on localhost)
     useEffect(() => {
         if (typeof window === 'undefined') return;
         if (window.location.hostname !== 'localhost') return;
 
-        console.log('Attempting to auto-detect IP...');
         const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
         const wsUrl = `${protocol}//${window.location.host}/ws`;
         const socket = new WebSocket(wsUrl);
 
         socket.onopen = () => {
-            console.log('Connected to local server for IP detection');
             socket.send(JSON.stringify({ type: 'get-ip' }));
         };
 
         socket.onmessage = (event) => {
             try {
                 const data = JSON.parse(event.data);
                 if (data.type === 'server-ip' && data.ip) {
-                    console.log('Auto-detected IP:', data.ip);
                     setIp(data.ip);
                     socket.close();
                 }
@@ -161,6 +204,7 @@ function SettingsPage() {
                                 <span>Slow</span>
                                 <span>Default</span>
                                 <span>Fast</span>
+
                             </div>
                         </div>
 
diff --git a/src/server/tokenStore.ts b/src/server/tokenStore.ts
@@ -0,0 +1,95 @@
+import crypto from 'crypto';
+import fs from 'fs';
+import path from 'path';
+
+interface TokenEntry {
+    token: string;
+    createdAt: number;
+    lastUsed: number;
+}
+
+const TOKENS_FILE = path.resolve('./src/tokens.json');
+const EXPIRY_MS = 10 * 24 * 60 * 60 * 1000; // 10 days
+
+let tokens: TokenEntry[] = [];
+
+function load(): void {
+    try {
+        if (fs.existsSync(TOKENS_FILE)) {
+            tokens = JSON.parse(fs.readFileSync(TOKENS_FILE, 'utf-8'));
+        }
+    } catch {
+        tokens = [];
+    }
+}
+
+function save(): void {
+    try {
+        fs.writeFileSync(TOKENS_FILE, JSON.stringify(tokens, null, 2));
+    } catch (e) {
+        console.error('Failed to persist tokens:', e);
+    }
+}
+
+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();
+}
+
+/** Constant-time string comparison to prevent timing attacks. */
+function timingSafeEqual(a: string, b: string): boolean {
+    if (a.length !== b.length) return false;
+    try {
+        return crypto.timingSafeEqual(Buffer.from(a), Buffer.from(b));
+    } catch {
+        return false;
+    }
+}
+
+/**
+ * Store a token upon successful connection.
+ * If it already exists, refresh its lastUsed timestamp.
+ */
+export function storeToken(token: string): void {
+    purgeExpired();
+    const existing = tokens.find(t => timingSafeEqual(t.token, token));
+    if (existing) {
+        existing.lastUsed = Date.now();
+    } else {
+        const now = Date.now();
+        tokens.push({ token, createdAt: now, lastUsed: now });
+    }
+    save();
+}
+
+/** Check if a token is already known/stored on the server. */
+export function isKnownToken(token: string): boolean {
+    purgeExpired();
+    return tokens.some(t => timingSafeEqual(t.token, token));
+}
+
+/** Refresh the lastUsed timestamp for a token. */
+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
+    }
+}
+
+/** Check if any tokens exist yet (first-run detection). */
+export function hasTokens(): boolean {
+    purgeExpired();
+    return tokens.length > 0;
+}
+
+/** Generate a cryptographically random token. */
+export function generateToken(): string {
+    return crypto.randomUUID();
+}
+
+// Load persisted tokens on startup
+load();
diff --git a/src/server/websocket.ts b/src/server/websocket.ts
@@ -1,12 +1,12 @@
 import { WebSocketServer, WebSocket } from 'ws';
 import { InputHandler, InputMessage } from './InputHandler';
+import { storeToken, isKnownToken, touchToken, hasTokens, generateToken } from './tokenStore';
 import os from 'os';
 import fs from 'fs';
-import { Server, IncomingMessage } from 'http';
+import { IncomingMessage } from 'http';
 import { Socket } from 'net';
 
-// Helper to find LAN IP
-function getLocalIp() {
+function getLocalIp(): string {
     const nets = os.networkInterfaces();
     for (const name of Object.keys(nets)) {
         for (const net of nets[name]!) {
@@ -18,27 +18,64 @@ function getLocalIp() {
     return 'localhost';
 }
 
-export function createWsServer(server: Server) {
+function isLocalhost(request: IncomingMessage): boolean {
+    const addr = request.socket.remoteAddress;
+    if (!addr) return false;
+    return addr === '127.0.0.1' || addr === '::1' || addr === '::ffff:127.0.0.1';
+}
+
+export function createWsServer(server: any) {
     const wss = new WebSocketServer({ noServer: true });
     const inputHandler = new InputHandler();
     const LAN_IP = getLocalIp();
     const MAX_PAYLOAD_SIZE = 10 * 1024; // 10KB limit
 
-    console.log(`WebSocket Server initialized (Upgrade mode)`);
-    console.log(`WS LAN IP: ${LAN_IP}`);
-
     server.on('upgrade', (request: IncomingMessage, socket: Socket, head: Buffer) => {
-        const pathname = request.url;
+        const url = new URL(request.url || '', `http://${request.headers.host}`);
+
+        if (url.pathname !== '/ws') return;
+
+        const token = url.searchParams.get('token');
+        const local = isLocalhost(request);
+
+        // Localhost is always allowed (settings page, IP detection, etc.)
+        if (local) {
+            wss.handleUpgrade(request, socket, head, (ws) => {
+                wss.emit('connection', ws, request, token, true);
+            });
+            return;
+        }
+
+        // Remote connections require a token
+        if (!token) {
+            socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+            socket.destroy();
+            return;
+        }
 
-        if (pathname === '/ws') {
+        // First remote connection ever — accept and store the token
+        if (!hasTokens()) {
+            storeToken(token);
             wss.handleUpgrade(request, socket, head, (ws) => {
-                wss.emit('connection', ws, request);
+                wss.emit('connection', ws, request, token, false);
             });
+            return;
         }
+
+        // Validate against known tokens
+        if (!isKnownToken(token)) {
+            socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+            socket.destroy();
+            return;
+        }
+
+        wss.handleUpgrade(request, socket, head, (ws) => {
+            wss.emit('connection', ws, request, token, false);
+        });
     });
 
-    wss.on('connection', (ws: WebSocket) => {
-        console.log('Client connected to /ws');
+    wss.on('connection', (ws: WebSocket, _request: IncomingMessage, token: string | null, isLocal: boolean) => {
+        if (token) storeToken(token);
 
         ws.send(JSON.stringify({ type: 'connected', serverIp: LAN_IP }));
 
@@ -54,22 +91,33 @@ export function createWsServer(server: Server) {
 
                 const msg = JSON.parse(raw);
 
+                if (token) touchToken(token);
+
                 if (msg.type === 'get-ip') {
                     ws.send(JSON.stringify({ type: 'server-ip', ip: LAN_IP }));
                     return;
                 }
 
+                if (msg.type === 'generate-token') {
+                    if (!isLocal) {
+                        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 }));
+                    return;
+                }
+
                 if (msg.type === 'update-config') {
-                    console.log('Updating config:', msg.config);
                     try {
                         const configPath = './src/server-config.json';
-                        // eslint-disable-next-line @typescript-eslint/no-require-imports
-                        const current = fs.existsSync(configPath) ? JSON.parse(fs.readFileSync(configPath, 'utf-8')) : {};
+                        const current = fs.existsSync(configPath)
+                            ? JSON.parse(fs.readFileSync(configPath, 'utf-8'))
+                            : {};
                         const newConfig = { ...current, ...msg.config };
-
                         fs.writeFileSync(configPath, JSON.stringify(newConfig, null, 2));
                         ws.send(JSON.stringify({ type: 'config-updated', success: true }));
-                        console.log('Config updated. Vite should auto-restart.');
                     } catch (e) {
                         console.error('Failed to update config:', e);
                         ws.send(JSON.stringify({ type: 'config-updated', success: false, error: String(e) }));
@@ -83,9 +131,7 @@ export function createWsServer(server: Server) {
             }
         });
 
-        ws.on('close', () => {
-            console.log('Client disconnected');
-        });
+        ws.on('close', () => { /* client disconnected */ });
 
         ws.onerror = (error) => {
             console.error('WebSocket error:', error);
