fix: guard mihomo websocket reconnects across restart

Closes #1841

src/main/core/mihomoApi.ts

@@ -14,16 +14,134 @@ const mihomoApiLogger = createLogger('MihomoApi')
 
 let axiosIns: AxiosInstance | null = null
 let currentIpcPath: string = ''
-let mihomoTrafficWs: WebSocket | null = null
-let trafficRetry = 10
-let mihomoMemoryWs: WebSocket | null = null
-let memoryRetry = 10
-let mihomoLogsWs: WebSocket | null = null
-let logsRetry = 10
-let mihomoConnectionsWs: WebSocket | null = null
-let connectionsRetry = 10
 
 const MAX_RETRY = 10
+const RECONNECT_INTERVAL_MS = 1000
+
+interface MihomoStreamState {
+  ws: WebSocket | null
+  retry: number
+  active: boolean
+  generation: number
+  reconnectTimer: NodeJS.Timeout | null
+}
+
+const trafficStream: MihomoStreamState = {
+  ws: null,
+  retry: MAX_RETRY,
+  active: false,
+  generation: 0,
+  reconnectTimer: null
+}
+const memoryStream: MihomoStreamState = {
+  ws: null,
+  retry: MAX_RETRY,
+  active: false,
+  generation: 0,
+  reconnectTimer: null
+}
+const logsStream: MihomoStreamState = {
+  ws: null,
+  retry: MAX_RETRY,
+  active: false,
+  generation: 0,
+  reconnectTimer: null
+}
+const connectionsStream: MihomoStreamState = {
+  ws: null,
+  retry: MAX_RETRY,
+  active: false,
+  generation: 0,
+  reconnectTimer: null
+}
+
+function clearStreamReconnect(stream: MihomoStreamState): void {
+  if (!stream.reconnectTimer) return
+  clearTimeout(stream.reconnectTimer)
+  stream.reconnectTimer = null
+}
+
+function disposeStreamSocket(ws: WebSocket): void {
+  ws.onmessage = null
+  ws.onclose = null
+  ws.onerror = null
+  ws.removeAllListeners()
+
+  if (ws.readyState === WebSocket.OPEN) {
+    ws.close()
+  } else if (ws.readyState === WebSocket.CONNECTING) {
+    ws.terminate()
+  }
+}
+
+function activateStream(stream: MihomoStreamState): void {
+  stream.active = true
+  stream.retry = MAX_RETRY
+  clearStreamReconnect(stream)
+}
+
+function stopStream(stream: MihomoStreamState): void {
+  stream.active = false
+  stream.retry = 0
+  stream.generation++
+  clearStreamReconnect(stream)
+
+  const ws = stream.ws
+  stream.ws = null
+  if (ws) {
+    disposeStreamSocket(ws)
+  }
+}
+
+function beginStreamConnection(stream: MihomoStreamState): number | null {
+  if (!stream.active) return null
+
+  stream.generation++
+  clearStreamReconnect(stream)
+
+  const ws = stream.ws
+  stream.ws = null
+  if (ws) {
+    disposeStreamSocket(ws)
+  }
+
+  return stream.generation
+}
+
+function isCurrentStream(stream: MihomoStreamState, generation: number): boolean {
+  return stream.active && stream.generation === generation
+}
+
+function scheduleStreamReconnect(
+  stream: MihomoStreamState,
+  generation: number,
+  connect: () => Promise<void>
+): void {
+  if (!isCurrentStream(stream, generation) || stream.retry <= 0) return
+
+  stream.retry--
+  clearStreamReconnect(stream)
+  stream.reconnectTimer = setTimeout(() => {
+    stream.reconnectTimer = null
+    if (isCurrentStream(stream, generation)) {
+      void connect()
+    }
+  }, RECONNECT_INTERVAL_MS)
+}
+
+function closeErroredStreamSocket(
+  stream: MihomoStreamState,
+  generation: number,
+  ws: WebSocket
+): void {
+  if (!isCurrentStream(stream, generation)) return
+
+  if (ws.readyState === WebSocket.OPEN) {
+    ws.close()
+  } else if (ws.readyState === WebSocket.CONNECTING) {
+    ws.terminate()
+  }
+}
 
 export const getAxios = async (force: boolean = false): Promise<AxiosInstance> => {
   const dynamicIpcPath = getMihomoIpcPath()
@@ -227,33 +345,31 @@ export const mihomoSmartFlushCache = async (configName?: string): Promise<void>
 }
 
 export const startMihomoTraffic = async (): Promise<void> => {
-  trafficRetry = MAX_RETRY
+  activateStream(trafficStream)
   await mihomoTraffic()
 }
 
 export const stopMihomoTraffic = (): void => {
-  trafficRetry = 0
-
-  if (mihomoTrafficWs) {
-    mihomoTrafficWs.removeAllListeners()
-    if (mihomoTrafficWs.readyState === WebSocket.OPEN) {
-      mihomoTrafficWs.close()
-    }
-    mihomoTrafficWs = null
-  }
+  stopStream(trafficStream)
 }
 
 const mihomoTraffic = async (): Promise<void> => {
+  const generation = beginStreamConnection(trafficStream)
+  if (generation === null) return
+
   const dynamicIpcPath = getMihomoIpcPath()
   const wsUrl = `ws+unix:${dynamicIpcPath}:/traffic`
 
   mihomoApiLogger.info(`Creating traffic WebSocket with URL: ${wsUrl}`)
-  mihomoTrafficWs = new WebSocket(wsUrl)
+  const ws = new WebSocket(wsUrl)
+  trafficStream.ws = ws
+
+  ws.onmessage = async (e): Promise<void> => {
+    if (!isCurrentStream(trafficStream, generation)) return
 
-  mihomoTrafficWs.onmessage = async (e): Promise<void> => {
     const data = e.data as string
     const json = JSON.parse(data) as IMihomoTrafficInfo
-    trafficRetry = MAX_RETRY
+    trafficStream.retry = MAX_RETRY
     try {
       mainWindow?.webContents.send('mihomoTraffic', json)
       if (process.platform !== 'linux') {
@@ -270,47 +386,41 @@ const mihomoTraffic = async (): Promise<void> => {
     }
   }
 
-  mihomoTrafficWs.onclose = (): void => {
-    if (trafficRetry) {
-      trafficRetry--
-      setTimeout(mihomoTraffic, 1000)
-    }
+  ws.onclose = (): void => {
+    if (!isCurrentStream(trafficStream, generation)) return
+    trafficStream.ws = null
+    scheduleStreamReconnect(trafficStream, generation, mihomoTraffic)
   }
 
-  mihomoTrafficWs.onerror = (error): void => {
+  ws.onerror = (error): void => {
     mihomoApiLogger.error('Traffic WebSocket error', error)
-    if (mihomoTrafficWs) {
-      mihomoTrafficWs.close()
-      mihomoTrafficWs = null
-    }
+    closeErroredStreamSocket(trafficStream, generation, ws)
   }
 }
 
 export const startMihomoMemory = async (): Promise<void> => {
-  memoryRetry = MAX_RETRY
+  activateStream(memoryStream)
   await mihomoMemory()
 }
 
 export const stopMihomoMemory = (): void => {
-  memoryRetry = 0
-
-  if (mihomoMemoryWs) {
-    mihomoMemoryWs.removeAllListeners()
-    if (mihomoMemoryWs.readyState === WebSocket.OPEN) {
-      mihomoMemoryWs.close()
-    }
-    mihomoMemoryWs = null
-  }
+  stopStream(memoryStream)
 }
 
 const mihomoMemory = async (): Promise<void> => {
+  const generation = beginStreamConnection(memoryStream)
+  if (generation === null) return
+
   const dynamicIpcPath = getMihomoIpcPath()
   const wsUrl = `ws+unix:${dynamicIpcPath}:/memory`
-  mihomoMemoryWs = new WebSocket(wsUrl)
+  const ws = new WebSocket(wsUrl)
+  memoryStream.ws = ws
+
+  ws.onmessage = (e): void => {
+    if (!isCurrentStream(memoryStream, generation)) return
 
-  mihomoMemoryWs.onmessage = (e): void => {
     const data = e.data as string
-    memoryRetry = MAX_RETRY
+    memoryStream.retry = MAX_RETRY
     try {
       mainWindow?.webContents.send('mihomoMemory', JSON.parse(data) as IMihomoMemoryInfo)
     } catch {
@@ -318,48 +428,42 @@ const mihomoMemory = async (): Promise<void> => {
     }
   }
 
-  mihomoMemoryWs.onclose = (): void => {
-    if (memoryRetry) {
-      memoryRetry--
-      setTimeout(mihomoMemory, 1000)
-    }
+  ws.onclose = (): void => {
+    if (!isCurrentStream(memoryStream, generation)) return
+    memoryStream.ws = null
+    scheduleStreamReconnect(memoryStream, generation, mihomoMemory)
   }
 
-  mihomoMemoryWs.onerror = (): void => {
-    if (mihomoMemoryWs) {
-      mihomoMemoryWs.close()
-      mihomoMemoryWs = null
-    }
+  ws.onerror = (): void => {
+    closeErroredStreamSocket(memoryStream, generation, ws)
   }
 }
 
 export const startMihomoLogs = async (): Promise<void> => {
-  logsRetry = MAX_RETRY
+  activateStream(logsStream)
   await mihomoLogs()
 }
 
 export const stopMihomoLogs = (): void => {
-  logsRetry = 0
-
-  if (mihomoLogsWs) {
-    mihomoLogsWs.removeAllListeners()
-    if (mihomoLogsWs.readyState === WebSocket.OPEN) {
-      mihomoLogsWs.close()
-    }
-    mihomoLogsWs = null
-  }
+  stopStream(logsStream)
 }
 
 const mihomoLogs = async (): Promise<void> => {
+  const generation = beginStreamConnection(logsStream)
+  if (generation === null) return
+
   const { 'log-level': logLevel = 'info' } = await getControledMihomoConfig()
   const dynamicIpcPath = getMihomoIpcPath()
   const wsUrl = `ws+unix:${dynamicIpcPath}:/logs?level=${logLevel}`
 
-  mihomoLogsWs = new WebSocket(wsUrl)
+  const ws = new WebSocket(wsUrl)
+  logsStream.ws = ws
+
+  ws.onmessage = (e): void => {
+    if (!isCurrentStream(logsStream, generation)) return
 
-  mihomoLogsWs.onmessage = (e): void => {
     const data = e.data as string
-    logsRetry = MAX_RETRY
+    logsStream.retry = MAX_RETRY
     try {
       mainWindow?.webContents.send('mihomoLogs', JSON.parse(data) as IMihomoLogInfo)
     } catch {
@@ -367,46 +471,40 @@ const mihomoLogs = async (): Promise<void> => {
     }
   }
 
-  mihomoLogsWs.onclose = (): void => {
-    if (logsRetry) {
-      logsRetry--
-      setTimeout(mihomoLogs, 1000)
-    }
+  ws.onclose = (): void => {
+    if (!isCurrentStream(logsStream, generation)) return
+    logsStream.ws = null
+    scheduleStreamReconnect(logsStream, generation, mihomoLogs)
   }
 
-  mihomoLogsWs.onerror = (): void => {
-    if (mihomoLogsWs) {
-      mihomoLogsWs.close()
-      mihomoLogsWs = null
-    }
+  ws.onerror = (): void => {
+    closeErroredStreamSocket(logsStream, generation, ws)
   }
 }
 
 export const startMihomoConnections = async (): Promise<void> => {
-  connectionsRetry = MAX_RETRY
+  activateStream(connectionsStream)
   await mihomoConnections()
 }
 
 export const stopMihomoConnections = (): void => {
-  connectionsRetry = 0
-
-  if (mihomoConnectionsWs) {
-    mihomoConnectionsWs.removeAllListeners()
-    if (mihomoConnectionsWs.readyState === WebSocket.OPEN) {
-      mihomoConnectionsWs.close()
-    }
-    mihomoConnectionsWs = null
-  }
+  stopStream(connectionsStream)
 }
 
 const mihomoConnections = async (): Promise<void> => {
+  const generation = beginStreamConnection(connectionsStream)
+  if (generation === null) return
+
   const dynamicIpcPath = getMihomoIpcPath()
   const wsUrl = `ws+unix:${dynamicIpcPath}:/connections`
-  mihomoConnectionsWs = new WebSocket(wsUrl)
+  const ws = new WebSocket(wsUrl)
+  connectionsStream.ws = ws
+
+  ws.onmessage = (e): void => {
+    if (!isCurrentStream(connectionsStream, generation)) return
 
-  mihomoConnectionsWs.onmessage = (e): void => {
     const data = e.data as string
-    connectionsRetry = MAX_RETRY
+    connectionsStream.retry = MAX_RETRY
     try {
       mainWindow?.webContents.send('mihomoConnections', JSON.parse(data) as IMihomoConnectionsInfo)
     } catch {
@@ -414,18 +512,14 @@ const mihomoConnections = async (): Promise<void> => {
     }
   }
 
-  mihomoConnectionsWs.onclose = (): void => {
-    if (connectionsRetry) {
-      connectionsRetry--
-      setTimeout(mihomoConnections, 1000)
-    }
+  ws.onclose = (): void => {
+    if (!isCurrentStream(connectionsStream, generation)) return
+    connectionsStream.ws = null
+    scheduleStreamReconnect(connectionsStream, generation, mihomoConnections)
   }
 
-  mihomoConnectionsWs.onerror = (): void => {
-    if (mihomoConnectionsWs) {
-      mihomoConnectionsWs.close()
-      mihomoConnectionsWs = null
-    }
+  ws.onerror = (): void => {
+    closeErroredStreamSocket(connectionsStream, generation, ws)
   }
 }