// shared/hooks/useWebSocket.ts import { useEffect, useRef, useState, useCallback } from "react"; export interface LogMessage { timestamp: string; service: string; level: "log" | "info" | "success" | "warn" | "error"; message: string; host: string; attributes?: Record; } interface WebSocketOptions { url: string; token?: string; autoConnect?: boolean; reconnectInterval?: number; maxReconnectAttempts?: number; } export const useWebSocket = (options: WebSocketOptions) => { const { url, token, autoConnect = true, reconnectInterval = 3000, maxReconnectAttempts = 5, } = options; const [isConnected, setIsConnected] = useState(false); const [isAuthenticated, setIsAuthenticated] = useState(false); const [lastMessage, setLastMessage] = useState(null); const [error, setError] = useState(null); const wsRef = useRef(null); const reconnectAttemptsRef = useRef(0); const reconnectTimeoutRef = useRef(null); const urlRef = useRef(url); const tokenRef = useRef(token); const authTimeoutRef = useRef(null); const disconnect = useCallback(() => { if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); reconnectTimeoutRef.current = null; } if (authTimeoutRef.current) { clearTimeout(authTimeoutRef.current); authTimeoutRef.current = null; } if (wsRef.current) { const ws = wsRef.current; wsRef.current = null; if ( ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING ) { ws.close(1000, "Disconnecting"); } } setIsConnected(false); setIsAuthenticated(false); }, []); const sendAuth = useCallback(() => { if ( wsRef.current && wsRef.current.readyState === WebSocket.OPEN && tokenRef.current ) { const authMessage = { type: "auth", payload: { token: tokenRef.current, }, }; console.log("Sending auth message..."); wsRef.current.send(JSON.stringify(authMessage)); // Set timeout for auth response authTimeoutRef.current = setTimeout(() => { if (!isAuthenticated) { console.error("Auth timeout"); setError("Authentication timeout"); disconnect(); } }, 5000); } }, [isAuthenticated, disconnect]); const connect = useCallback(() => { // Если URL изменился, пересоздаем соединение if (urlRef.current !== url) { console.log("URL changed, forcing new connection"); disconnect(); urlRef.current = url; tokenRef.current = token; } // Если уже есть открытое соединение, не создаем новое if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) { console.log("WebSocket already connected"); return; } try { console.log("Connecting to WebSocket:", url); wsRef.current = new WebSocket(url); wsRef.current.onopen = () => { console.log("WebSocket connected, sending auth..."); setIsConnected(true); setError(null); reconnectAttemptsRef.current = 0; // Send authentication immediately after connection sendAuth(); }; wsRef.current.onmessage = (event) => { try { const data = JSON.parse(event.data); console.log("WebSocket message received:", data); // Check if it's an auth response if (data.type === "auth") { if (data.success) { console.log("Authentication successful"); setIsAuthenticated(true); setError(null); if (authTimeoutRef.current) { clearTimeout(authTimeoutRef.current); } } else { console.error("Authentication failed:", data.error); setError(data.error || "Authentication failed"); setIsAuthenticated(false); disconnect(); } } else { // Regular log message setLastMessage(data); } } catch (err) { console.error("Failed to parse WebSocket message:", err); } }; wsRef.current.onerror = (event) => { console.error("WebSocket error:", event); setError("Connection error"); }; wsRef.current.onclose = (event) => { console.log( "WebSocket disconnected, code:", event.code, "reason:", event.reason, ); setIsConnected(false); setIsAuthenticated(false); // Если URL изменился, не переподключаемся автоматически if (urlRef.current !== url) { console.log("URL changed, will reconnect manually"); return; } // Не переподключаемся при нормальном закрытии if (event.code === 1000) { console.log("Normal closure, not reconnecting"); return; } // Attempt reconnection if ( reconnectAttemptsRef.current < maxReconnectAttempts && tokenRef.current ) { console.log(`Reconnecting in ${reconnectInterval}ms...`); reconnectTimeoutRef.current = setTimeout(() => { reconnectAttemptsRef.current++; connect(); }, reconnectInterval); } else if (reconnectAttemptsRef.current >= maxReconnectAttempts) { setError("Max reconnection attempts reached"); } }; } catch (err) { setError(err instanceof Error ? err.message : "Failed to connect"); } }, [ url, token, reconnectInterval, maxReconnectAttempts, disconnect, sendAuth, ]); const reconnect = useCallback(() => { console.log("Manual reconnect triggered"); disconnect(); setTimeout(() => { connect(); }, 100); }, [disconnect, connect]); const sendMessage = useCallback( (data: any) => { if ( wsRef.current && wsRef.current.readyState === WebSocket.OPEN && isAuthenticated ) { wsRef.current.send(JSON.stringify(data)); return true; } else { console.warn("WebSocket is not authenticated or not connected"); return false; } }, [isAuthenticated], ); const updateFilter = useCallback( (hosts: string[], services: string[]) => { const message = { type: "filter", payload: { hosts, services }, }; console.log("Sending filter update:", message); sendMessage(message); }, [sendMessage], ); // Подключаемся при монтировании и при изменении URL или токена useEffect(() => { if (autoConnect && url && token) { connect(); } return () => { disconnect(); }; }, [autoConnect, connect, disconnect, url, token]); return { isConnected: isConnected && isAuthenticated, isAuthenticated, lastMessage, error, connect, disconnect, reconnect, sendMessage, updateFilter, }; };