From 289dcdf03076402ac7284dcad4b5c60ae9052eab Mon Sep 17 00:00:00 2001
From: song_xiao_lin <song_xiao_lin@126.com>
Date: Wed, 8 May 2024 15:53:46 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E9=87=8D=E8=BF=9E?=
 =?UTF-8?q?=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 public/background.js      |  2 +-
 public/connect.js         |  6 ++---
 src/components/Contro.tsx | 46 +++++++++++++++++++++++++--------------
 3 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/public/background.js b/public/background.js
index 9e600be..4b67a04 100644
--- a/public/background.js
+++ b/public/background.js
@@ -13,7 +13,7 @@ chrome.runtime.onMessage.addListener(function (msg, sender, sendResponse) {
             console.info("Start to connect websocket")
             const host = msg['host'] || "127.0.0.1"
             const port = msg['port'] || 11212
-            websocketManager.connectWebsocket(`ws://${host}:${port}/?token=${"a"}`)
+            websocketManager.connectWebsocket(`ws://${host}:${port}/?token=${"a"}`, port)
             break;
         case ActionType.DISCONNECT:
             websocketManager.disconnectWebsocket();
diff --git a/public/connect.js b/public/connect.js
index 8dc3619..c901b62 100644
--- a/public/connect.js
+++ b/public/connect.js
@@ -13,12 +13,12 @@ export class WebSocketManager {
         this.intervalId = null;
     }
 
-    connectWebsocket(url) {
+    connectWebsocket(url, port) {
         this.disconnectWebsocket();
         this.socket = new WebSocket(url);
 
         this.socket.onopen = () => {
-            chrome.runtime.sendMessage({action: ActionType.STATUS, connected: true});
+            chrome.runtime.sendMessage({action: ActionType.STATUS, connected: true, port: port});
             this.startHeartbeat();
         };
 
@@ -27,7 +27,7 @@ export class WebSocketManager {
         };
 
         this.socket.onclose = () => {
-            chrome.runtime.sendMessage({action: ActionType.STATUS, connected: false});
+            chrome.runtime.sendMessage({action: ActionType.STATUS, connected: false, port: port});
         };
 
         this.socket.onerror = (error) => {
diff --git a/src/components/Contro.tsx b/src/components/Contro.tsx
index 3da355a..29ae888 100644
--- a/src/components/Contro.tsx
+++ b/src/components/Contro.tsx
@@ -10,7 +10,7 @@ import {
 } from "@assets/icon/icon";
 import { wsc } from "@network/chrome";
 import "./Contro.css";
-import {ActionType} from "../../public/connect";
+import { ActionType } from "../../public/connect";
 
 interface ControProps {}
 export const Contro: React.FC<ControProps> = () => {
@@ -21,8 +21,35 @@ export const Contro: React.FC<ControProps> = () => {
   const [enginePortTemp, setEnginePortTemp] = useState<string>(enginePort);
 
   useEffect(() => {
-    findPort(11212, 11222);
+    const yakitConnectInfo = localStorage.getItem("yakit-connect");
+    if (!yakitConnectInfo) {
+      findPort(11212, 11222);
+    } else {
+      const { port, connected } = JSON.parse(yakitConnectInfo);
+      setConnected(connected);
+      setEnginePort(port + "");
+      setAutoFindFailedReason("");
+    }
+
+    wsc.onWSCMessage((message) => {
+      if (message.action === ActionType.STATUS) {
+        setConnected(message.connected);
+        if (message.connected === false) {
+          localStorage.setItem("yakit-connect", "");
+          setAutoFindFailedReason(
+            "Yakit WebSocket Controller Port is not right"
+          );
+        } else {
+          localStorage.setItem(
+            "yakit-connect",
+            JSON.stringify({ connected: true, port: message.port })
+          );
+          setAutoFindFailedReason("");
+        }
+      }
+    });
   }, []);
+
   const findPort = (port: number, max: number) => {
     const ws = new WebSocket(`ws://127.0.0.1:${port}`);
     ws.onclose = (e: CloseEvent) => {
@@ -34,6 +61,7 @@ export const Contro: React.FC<ControProps> = () => {
         setConnected(false);
         setEnginePort("");
         setAutoFindFailedReason("Cannot found Yakit");
+        localStorage.setItem("yakit-connect", "");
       }
     };
     ws.onopen = () => {
@@ -47,20 +75,6 @@ export const Contro: React.FC<ControProps> = () => {
 
   const connectPort = (port: number) => {
     wsc.connect(port);
-    wsc.onWSCMessage((message) => {
-      if (message.action === ActionType.STATUS) {
-        console.log("onWSCMessage", message)
-        if (message.connected === undefined) {
-          return;
-        }
-        setConnected(message.connected);
-        if (message.connected === false) {
-          setAutoFindFailedReason("Yakit WebSocket Controller Port is not right");
-        } else {
-          setAutoFindFailedReason("");
-        }
-      }
-    });
   };
 
   const safeConnected = useMemo(() => {