Implement WebSocket Authentication for Enhanced Security

src/hooks/useRemoteConnection.ts

@@ -3,6 +3,9 @@ import { useState, useEffect, useCallback } from 'react';
 export const useRemoteConnection = () => {
     const [ws, setWs] = useState<WebSocket | null>(null);
     const [status, setStatus] = useState<'connecting' | 'connected' | 'disconnected'>('disconnected');
+    const [isAuthenticated, setIsAuthenticated] = useState(false);
+    const [authRequired, setAuthRequired] = useState(false);
+    const [authError, setAuthError] = useState(false);
 
     useEffect(() => {
         const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
@@ -16,9 +19,38 @@ export const useRemoteConnection = () => {
             setStatus('connecting');
             const socket = new WebSocket(wsUrl);
 
-            socket.onopen = () => setStatus('connected');
+            socket.onopen = () => {
+                setStatus('connected');
+                // Auth state is determined by the 'connected' message
+            };
+
+            socket.onmessage = (event) => {
+                try {
+                    const msg = JSON.parse(event.data);
+                    if (msg.type === 'connected') {
+                        if (msg.authRequired) {
+                            setAuthRequired(true);
+                            setIsAuthenticated(false);
+                        } else {
+                            setIsAuthenticated(true);
+                            setAuthRequired(false);
+                        }
+                    } else if (msg.type === 'auth-success') {
+                        setIsAuthenticated(true);
+                        setAuthError(false);
+                    } else if (msg.type === 'auth-failed') {
+                        setAuthError(true);
+                        setIsAuthenticated(false);
+                    }
+                } catch (e) {
+                    console.error('WS Message Error', e);
+                }
+            };
+
             socket.onclose = () => {
                 setStatus('disconnected');
+                setIsAuthenticated(false);
+                setAuthRequired(false);
                 reconnectTimer = setTimeout(connect, 3000);
             };
             socket.onerror = (e) => {
@@ -42,5 +74,11 @@ export const useRemoteConnection = () => {
         }
     }, [ws]);
 
-    return { status, send };
+    const authenticate = useCallback((pin: string) => {
+        if (ws?.readyState === WebSocket.OPEN) {
+            ws.send(JSON.stringify({ type: 'authenticate', pin }));
+        }
+    }, [ws]);
+
+    return { status, send, isAuthenticated, authRequired, authenticate, authError };
 };

src/routes/settings.tsx

@@ -13,6 +13,7 @@ function SettingsPage() {
     const [invertScroll, setInvertScroll] = useState(CONFIG.MOUSE_INVERT);
     const [sensitivity, setSensitivity] = useState(CONFIG.MOUSE_SENSITIVITY);
     const [qrData, setQrData] = useState('');
+    const [pin, setPin] = useState<string | null>(null);
 
     // Load initial state
     useEffect(() => {
@@ -53,12 +54,24 @@ function SettingsPage() {
 
         socket.onopen = () => {
             console.log('Connected to local server for IP detection');
-            socket.send(JSON.stringify({ type: 'get-ip' }));
+            // socket.send(JSON.stringify({ type: 'get-ip' })); // Not strictly needed as connected sends logic
         };
 
         socket.onmessage = (event) => {
             try {
                 const data = JSON.parse(event.data);
+
+                if (data.type === 'connected') {
+                    if (data.serverIp) {
+                        console.log('Auto-detected IP:', data.serverIp);
+                        setIp(data.serverIp);
+                    }
+                    if (data.pin) {
+                        console.log('Received session PIN');
+                        setPin(data.pin);
+                    }
+                }
+
                 if (data.type === 'server-ip' && data.ip) {
                     console.log('Auto-detected IP:', data.ip);
                     setIp(data.ip);
@@ -84,6 +97,19 @@ function SettingsPage() {
             <div className="p-6 pb-safe max-w-md mx-auto space-y-8 min-h-full">
                 <h1 className="text-3xl font-bold pt-4">Settings</h1>
 
+                {pin && (
+                    <div className="alert alert-info shadow-lg">
+                        <div>
+                            <svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" className="stroke-current flex-shrink-0 w-6 h-6"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth="2" d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z"></path></svg>
+                            <div>
+                                <h3 className="font-bold">Session PIN</h3>
+                                <div className="text-xs">Enter this PIN on your mobile device</div>
+                            </div>
+                        </div>
+                        <div className="font-mono text-4xl font-black">{pin}</div>
+                    </div>
+                )}
+
                 <div className="form-control w-full">
                     <label className="label">
                         <span className="label-text">Server IP (for Remote)</span>