mirror of
https://github.com/ParrotXray/Mantis.git
synced 2026-08-24 18:50:28 +09:00
research: fix WebSocketProvider to align with config and detection page
Three fixes: 1. WebSocketContextType: getAiAlertStream -> getDetectionAlertStream 2. createContext default: getAiAlertStream -> getDetectionAlertStream 3. getDetectionAlertStream + initializeWebSockets: websocketUrl.aiAlert -> websocketUrl.detectionAlert Also: remove unused 'data' variable in fetchBootTime, remove Chinese comment. https://claude.ai/code/session_01Wen51749mAnwBEfh7XmvUw
This commit is contained in:
parent
7b66c4719e
commit
a163d88f15
156
.research/frontend-refactor/WebSocketProvider.tsx
Normal file
156
.research/frontend-refactor/WebSocketProvider.tsx
Normal file
@ -0,0 +1,156 @@
|
||||
import React, { createContext, useState, useEffect, useRef } from 'react'
|
||||
import { websocketUrl, urls } from '../config'
|
||||
import { createWebSocket, fetchData } from '../utils/connectionUtils'
|
||||
import { Subject } from 'rxjs'
|
||||
|
||||
interface WebSocketContextType {
|
||||
bootTime: number | null;
|
||||
getDetectionAlertStream: () => any;
|
||||
getIPv4FlowStream: (direction: string, flow_direction: string, timeType: string) => any;
|
||||
getIPv6FlowStream: (direction: string, flow_direction: string, timeType: string) => any;
|
||||
getSystemHealthStream: () => any;
|
||||
}
|
||||
|
||||
export const WebsocketContext = createContext<WebSocketContextType>({
|
||||
bootTime: null,
|
||||
getDetectionAlertStream: () => null,
|
||||
getIPv4FlowStream: () => null,
|
||||
getIPv6FlowStream: () => null,
|
||||
getSystemHealthStream: () => null,
|
||||
})
|
||||
|
||||
interface WebSocketProviderProps {
|
||||
children: React.ReactNode
|
||||
}
|
||||
|
||||
export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ children }) => {
|
||||
const wsRefs = useRef<{ [key: string]: WebSocket }>({})
|
||||
const [bootTime, setBootTime] = useState<number | null>(null)
|
||||
const retryDelays = useRef<{ [key: string]: number }>({})
|
||||
const dataSubjects = useRef<{ [key: string]: Subject<any> }>({})
|
||||
|
||||
const fetchBootTime = async () => {
|
||||
try {
|
||||
await fetchData(
|
||||
urls.bootTime,
|
||||
(data) => setBootTime(Number(data.trim())),
|
||||
(error) => {
|
||||
throw new Error(`Failed to fetch boot_time: ${error?.message || 'Unknown error'}`)
|
||||
}
|
||||
)
|
||||
} catch (error) {
|
||||
console.error('Boot time fetch error:', error)
|
||||
}
|
||||
}
|
||||
|
||||
const getWebSocketStream = (url: string) => {
|
||||
if (!dataSubjects.current[url]) {
|
||||
dataSubjects.current[url] = new Subject()
|
||||
}
|
||||
return dataSubjects.current[url].asObservable()
|
||||
}
|
||||
|
||||
const setupWebSocket = (url: string, retryCount: number = 0) => {
|
||||
if (wsRefs.current[url]) {
|
||||
wsRefs.current[url].close()
|
||||
}
|
||||
|
||||
const ws = createWebSocket(
|
||||
url,
|
||||
(data) => {
|
||||
if (dataSubjects.current[url]) {
|
||||
dataSubjects.current[url].next(data)
|
||||
}
|
||||
},
|
||||
(error) => console.error(`WebSocket error for ${url}:`, error)
|
||||
)
|
||||
|
||||
ws.onclose = () => {
|
||||
console.warn(`WebSocket for ${url} closed. Retrying...`)
|
||||
let delay = retryDelays.current[url] || 1000
|
||||
retryDelays.current[url] = Math.min(delay * 2, 30000)
|
||||
setTimeout(() => setupWebSocket(url, retryCount + 1), delay)
|
||||
}
|
||||
|
||||
ws.onopen = () => {
|
||||
console.log(`WebSocket connected: ${url}`)
|
||||
retryDelays.current[url] = 1000
|
||||
}
|
||||
|
||||
wsRefs.current[url] = ws
|
||||
return ws
|
||||
}
|
||||
|
||||
const getDetectionAlertStream = () => {
|
||||
return getWebSocketStream(websocketUrl.detectionAlert)
|
||||
}
|
||||
|
||||
const getIPv4FlowStream = (direction: string, flow_direction: string, timeType: string) => {
|
||||
const url = websocketUrl.ipv4FlowStats(direction, flow_direction, timeType)
|
||||
return getWebSocketStream(url)
|
||||
}
|
||||
|
||||
const getIPv6FlowStream = (direction: string, flow_direction: string, timeType: string) => {
|
||||
const url = websocketUrl.ipv6FlowStats(direction, flow_direction, timeType)
|
||||
return getWebSocketStream(url)
|
||||
}
|
||||
|
||||
const getSystemHealthStream = () => {
|
||||
return getWebSocketStream(websocketUrl.systemHealth)
|
||||
}
|
||||
|
||||
const initializeWebSockets = () => {
|
||||
setupWebSocket(websocketUrl.detectionAlert)
|
||||
setupWebSocket(websocketUrl.systemHealth)
|
||||
|
||||
const directions = ["ingress", "egress"]
|
||||
const flowDirections = ["source", "destination"]
|
||||
const timeTypes = ["1min", "10min", "1hour"]
|
||||
|
||||
directions.forEach(direction => {
|
||||
flowDirections.forEach(flow_direction => {
|
||||
timeTypes.forEach(timeType => {
|
||||
setupWebSocket(websocketUrl.ipv4FlowStats(direction, flow_direction, timeType))
|
||||
setupWebSocket(websocketUrl.ipv6FlowStats(direction, flow_direction, timeType))
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
const closeWebSockets = () => {
|
||||
Object.entries(wsRefs.current).forEach(([url, ws]) => {
|
||||
try {
|
||||
ws.close()
|
||||
delete wsRefs.current[url]
|
||||
} catch (error) {
|
||||
console.error(`Failed to close WebSocket ${url}:`, error)
|
||||
}
|
||||
})
|
||||
|
||||
Object.values(dataSubjects.current).forEach(subject => {
|
||||
subject.complete()
|
||||
})
|
||||
dataSubjects.current = {}
|
||||
}
|
||||
|
||||
useEffect(() => {
|
||||
fetchBootTime()
|
||||
initializeWebSockets()
|
||||
|
||||
return () => closeWebSockets()
|
||||
}, [])
|
||||
|
||||
const contextValue: WebSocketContextType = {
|
||||
bootTime,
|
||||
getDetectionAlertStream,
|
||||
getIPv4FlowStream,
|
||||
getIPv6FlowStream,
|
||||
getSystemHealthStream,
|
||||
}
|
||||
|
||||
return (
|
||||
<WebsocketContext.Provider value={contextValue}>
|
||||
{children}
|
||||
</WebsocketContext.Provider>
|
||||
)
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user