From 88b9a5ba9dfedb0ccc4904004472aa1c4d24abb0 Mon Sep 17 00:00:00 2001
From: NriotHrreion <nriot233@outlook.com>
Date: Fri, 9 Aug 2024 17:24:58 +0800
Subject: [PATCH] fix(demo): Websocket error in demo mode

---
 app/(pages)/dashboard/page.tsx |  8 ++++++--
 app/api/os/route.ts            | 10 ++--------
 hooks/useOS.ts                 | 12 +++++++++++-
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/app/(pages)/dashboard/page.tsx b/app/(pages)/dashboard/page.tsx
index 7cf89b1..176e81d 100644
--- a/app/(pages)/dashboard/page.tsx
+++ b/app/(pages)/dashboard/page.tsx
@@ -12,6 +12,7 @@ import GPUWidget from "@/components/dashboard/gpu-widget";
 import BatteryWidget from "@/components/dashboard/battery-widget";
 import OSWidget from "@/components/dashboard/os-widget";
 import { WebSocketContext } from "@/hooks/useOS";
+import { isDemo } from "@/lib/global";
 
 export default function Page() {
     const [mounted, setMounted] = useState<boolean>(false);
@@ -22,8 +23,11 @@ export default function Page() {
 
         document.title = "Ferrum - 仪表盘";
 
-        const _ws = new WebSocket(`ws://${window.location.host}/api/os`);
-        setWebSocket(_ws);
+        var _ws: WebSocket;
+        if(!isDemo) {
+            _ws = new WebSocket(`ws://${window.location.host}/api/os`);
+            setWebSocket(_ws);
+        }
 
         return () => _ws?.close();
     }, []);
diff --git a/app/api/os/route.ts b/app/api/os/route.ts
index 46e82c8..b3d0108 100644
--- a/app/api/os/route.ts
+++ b/app/api/os/route.ts
@@ -12,8 +12,6 @@ import cookie from "cookie";
 import { error } from "@/lib/packet";
 import { isDemo, tokenStorageKey } from "@/lib/global";
 import { validateToken } from "@/lib/token";
-// Demo
-import demoOS from "@/lib/demo/os.json";
 
 export function GET() {
     return error(400);
@@ -24,6 +22,8 @@ export function SOCKET(
     req: IncomingMessage,
     _server: WebSocketServer,
 ) {
+    if(isDemo) return;
+
     const token = cookie.parse(req.headers.cookie ?? "")[tokenStorageKey];
 
     if(!token) {
@@ -40,12 +40,6 @@ export function SOCKET(
     console.log("[Server: /api/os] Socket client connected.");
 
     const handleRequest = async () => {
-        if(isDemo) {
-            client.send(JSON.stringify(demoOS));
-
-            return;
-        }
-
         const cpu = await si.cpu();
         const cpuTemp = await si.cpuTemperature();
         const mem = await si.mem();
diff --git a/hooks/useOS.ts b/hooks/useOS.ts
index e87fd8b..233854b 100644
--- a/hooks/useOS.ts
+++ b/hooks/useOS.ts
@@ -1,7 +1,11 @@
 import React, { useContext, useEffect } from "react";
 
+import { isDemo } from "@/lib/global";
+// Demo
+import demoOS from "@/lib/demo/os.json";
+
 interface WebSocketContextType {
-    ws: WebSocket | null;
+    ws: WebSocket | null
 }
 
 export const WebSocketContext = React.createContext<WebSocketContextType>({ ws: null });
@@ -65,6 +69,12 @@ export function useOS(listener: (msg: OSWebSocketMessage) => void) {
     const { ws } = useContext(WebSocketContext);
 
     useEffect(() => {
+        if(isDemo) {
+            listener(demoOS as OSWebSocketMessage);
+
+            return;
+        }
+
         const controller = new AbortController();
 
         ws?.addEventListener("message", (e) => {