add verify before ws connection

renderer/src/App.vue

@@ -45,14 +45,12 @@ onMounted(async () => {
 	setDefaultCss();
 
 	pinkLog('OpenMCP Client 启动');
-	const platform = getPlatform();
 
 	// 跳转到首页
-	if (platform !== 'web') {
+	if (route.name !== 'debug') {
-		if (route.name !== 'debug') {
+		const targetRoute = import.meta.env.BASE_URL + 'debug';
-			router.replace('/debug');
+		console.log('go to ' + targetRoute);
-			router.push('/debug');
+		router.push(targetRoute);
-		}
 	}
 
 	// 进行桥接

renderer/src/api/message-bridge.ts

@@ -1,12 +1,10 @@
 import { pinkLog, redLog } from '@/views/setting/util';
 import { acquireVsCodeApi, electronApi, getPlatform } from './platform';
-import { privilegeStatus } from '@/components/password-dialog/status';
 
 export interface VSCodeMessage {
 	command: string;
 	data?: unknown;
 	callbackId?: string;
-	password?: string;
 }
 
 export interface RestFulResponse {
@@ -67,51 +65,57 @@ export class MessageBridge {
 	}
 
 	// WebSocket 环境连接
-	private setupWebSocket() {
+	public setupWebSocket(setupSignature?: string) {
-
+		const wsUrl = setupSignature || this.setupSignature;
-		const wsUrl = this.setupSignature;
 
 		if (typeof wsUrl !== 'string') {
 			throw new Error('setupSignature must be a string');
 		}
-
+		
 		this.ws = new WebSocket(wsUrl);
-
-		this.ws.onmessage = (event) => {
-			try {				
-				const message = JSON.parse(event.data) as VSCodeMessage;
-				this.dispatchMessage(message);
-			} catch (err) {
-				console.error('Message parse error:', err);
-				console.log(event);
-			}
-		};
-
-		this.ws.onclose = () => {
-			redLog('WebSocket connection closed');
-		};
-
-		this.postMessage = (message) => {
-			if (this.ws?.readyState === WebSocket.OPEN) {
-				console.log('send', message);
-				message.password = privilegeStatus.password;
-				this.ws.send(JSON.stringify(message));
-			}
-		};
-
 		const ws = this.ws;
 
 		this.isConnected = new Promise<boolean>((resolve, reject) => {
 			ws.onopen = () => {
 				resolve(true);
 			};
+
+			ws.onmessage = (event) => {
+				try {				
+					const message = JSON.parse(event.data) as VSCodeMessage;
+					this.dispatchMessage(message);
+				} catch (err) {
+					console.error('Message parse error:', err);
+					console.log(event);
+				}
+			};
+	
+			ws.onerror = (err) => {
+				redLog('WebSocket error:');				
+				resolve(false);
+			};
+	
+			ws.onclose = () => {
+				redLog('WebSocket connection closed');
+				resolve(false);
+			};
+	
+			this.postMessage = (message) => {
+				if (this.ws?.readyState === WebSocket.OPEN) {
+					console.log('send', message);
+					this.ws.send(JSON.stringify(message));
+				}
+			};
 		});
+
+		return this.isConnected;
 	}
 
 	public async awaitForWebsocket() {
 		if (this.isConnected) {
-			await this.isConnected;
+			return await this.isConnected;
 		}
+		return false;
 	}
 
 	private setupElectronListener() {

renderer/src/components/guide/tour.vue

@@ -46,7 +46,7 @@
         <el-tour-step
             target="#sidebar-connect"
             :prev-button-props="{ children: '上一步' }"
-            :next-button-props="{ children: '下一步', onClick: () => router.push('/connect') }"
+            :next-button-props="{ children: '下一步', onClick: () => router.push(baseUrl + 'connect') }"
             :show-close="false"
         >
             <template #header>
@@ -95,7 +95,7 @@
         <el-tour-step
             target="#sidebar-debug"
             :prev-button-props="{ children: '上一步' }"
-            :next-button-props="{ children: '下一步', onClick: () => router.push('/debug') }"
+            :next-button-props="{ children: '下一步', onClick: () => router.push(baseUrl + 'debug') }"
             :show-close="false"
         >
             <template #header>
@@ -129,8 +129,8 @@
 
         <el-tour-step
             target="#sidebar-setting"
-            :prev-button-props="{ children: '上一步', onClick: () => router.push('/debug') }"
+            :prev-button-props="{ children: '上一步', onClick: () => router.push(baseUrl + 'debug') }"
-            :next-button-props="{ children: '下一步', onClick: () => router.push('/setting') }"
+            :next-button-props="{ children: '下一步', onClick: () => router.push(baseUrl + 'setting') }"
             :show-close="false"
         >
             <template #header>
@@ -205,7 +205,7 @@
 
 
         <el-tour-step
-            :prev-button-props="{ children: '上一步', onClick: () => router.push('/setting') }"
+            :prev-button-props="{ children: '上一步', onClick: () => router.push(baseUrl + 'setting') }"
             :next-button-props="{ children: '下一步' }"
             :show-close="false"
         >
@@ -224,7 +224,7 @@
 
 
         <el-tour-step
-            :prev-button-props="{ children: '上一步', onClick: () => router.push('/setting') }"
+            :prev-button-props="{ children: '上一步', onClick: () => router.push(baseUrl + 'setting') }"
             :next-button-props="{ children: '结束', onClick: () => finishTour() }"
             :show-close="false"
         >
@@ -260,6 +260,8 @@ const { t } = useI18n();
 
 const router = useRouter();
 
+const baseUrl = import.meta.env.BASE_URL;
+
 function finishTour() {
     openTour.value = false;
     userHasReadGuide.value = true;

renderer/src/components/main-panel/chat/message/toolcall.vue

@@ -127,7 +127,7 @@ import { useI18n } from 'vue-i18n';
 import MessageMeta from './message-meta.vue';
 import { markdownToHtml } from '@/components/main-panel/chat/markdown/markdown';
 import { createTest } from '@/views/setting/llm';
-import { IToolRenderMessage, MessageState } from '../chat-box/chat';
+import { type IToolRenderMessage, MessageState } from '../chat-box/chat';
 import type { ToolCallContent } from '@/hook/type';
 
 import ToolcallResultItem from './toolcall-result-item.vue';

renderer/src/components/password-dialog/index.vue

@@ -24,9 +24,10 @@ const dialogVisible = ref(true);
 
 const handleSubmit = async () => {
     const bridge = useMessageBridge();
-    const res = await bridge.commandRequest('ciallo', { password: privilegeStatus.password });
+    
-
+    const ok = await bridge.setupWebSocket(import.meta.env.VITE_WEBSOCKET_URL + '?t=' + privilegeStatus.password);
-    if (res.code === 200) {
+    
+    if (ok) {
         ElMessage.success('密码验证成功，欢迎回来锦恢');
         dialogVisible.value = false;
 

renderer/src/components/sidebar/sidebar-item-container.vue

@@ -29,8 +29,10 @@ function isActive(name: string) {
 	return route.name === name;
 }
 
+const baseUrl = import.meta.env.BASE_URL;
+
 function gotoOption(ident: string) {
-	router.push('/' + ident);
+	router.push(baseUrl + ident);
 }
 
 </script>

renderer/src/router/index.ts

@@ -1,31 +1,33 @@
 import { createRouter, createWebHistory, type RouteRecordRaw } from "vue-router";
 
+const baseURL = import.meta.env.BASE_URL;
+
 const routes: Array<RouteRecordRaw> = [
 	{
 		name : "default",
 		path : "/",
-		redirect : "/debug"
+		redirect : baseURL + "/debug"
 	},
 	{
-		path: "/debug",
+		path: baseURL + "/debug",
 		name: "debug",
 		component: () => import( /* webpackMode: "eager" */ "@/views/debug/index.vue"),
 		meta: { title: "Debug" }
 	},
 	{
-		path: "/connect",
+		path: baseURL + "/connect",
 		name: "connect",
 		component: () => import( /* webpackMode: "eager" */ "@/views/connect/index.vue"),
 		meta: { title: "Connect" }
 	},
 	{
-		path: "/setting",
+		path: baseURL + "/setting",
 		name: "setting",
 		component: () => import( /* webpackMode: "eager" */ "@/views/setting/index.vue"),
 		meta: { title: "Setting" }
 	},
 	{
-		path: "/about",
+		path: baseURL + "/about",
 		name: "about",
 		component: () => import( /* webpackMode: "eager" */ "@/views/about/index.vue"),
 		meta: { title: "Tools" }

renderer/tsconfig.app.json

@@ -10,6 +10,10 @@
     ],
     "compilerOptions": {
         "tsBuildInfoFile": "./node_modules/.tmp/tsconfig.app.tsbuildinfo",
+        "lib": [
+            "es2022",
+			"dom"
+        ],
         "paths": {
             "@/*": [
                 "./src/*"

renderer/tsconfig.node.json

@@ -11,7 +11,11 @@
 	"compilerOptions": {
 		"noEmit": true,
 		"tsBuildInfoFile": "./node_modules/.tmp/tsconfig.node.tsbuildinfo",
-		"module": "ESNext",
+		"module": "es2022",
+		"lib": [
+            "es2022",
+			"dom"
+        ],
 		"moduleResolution": "Bundler",
 		"types": [
 			"node"

service/src/main.ts

@@ -1,4 +1,4 @@
-import WebSocket from 'ws';
+import { WebSocketServer } from 'ws';
 import pino from 'pino';
 
 import { routeMessage } from './common/router';
@@ -26,7 +26,6 @@ const logger = pino({
 });
 
 export type MessageHandler = (message: VSCodeMessage) => void;
-const wss = new (WebSocket as any).Server({ port: 8282 });
 
 interface IStdioLaunchSignature {
     type: 'stdio';
@@ -98,6 +97,10 @@ function updateConnectionOption(data: any) {
 const devHome = path.join(__dirname, '..', '..');
 setRunningCWD(devHome);
 
+const wss = new WebSocketServer({ port: 8282 });
+
+console.log('listen on ws://localhost:8282');
+
 wss.on('connection', (ws: any) => {
 
     // 仿造 webview 进行统一接口访问

service/src/server.ts

@@ -1,4 +1,4 @@
-import WebSocket from 'ws';
+import { WebSocketServer } from 'ws';
 import pino from 'pino';
 
 import { routeMessage } from './common/router';
@@ -27,7 +27,6 @@ const logger = pino({
 });
 
 export type MessageHandler = (message: VSCodeMessage) => void;
-const wss = new (WebSocket as any).Server({ port: 8282 });
 
 interface IStdioLaunchSignature {
     type: 'stdio';
@@ -36,7 +35,7 @@ interface IStdioLaunchSignature {
 }
 
 interface ISSELaunchSignature {
-    type:'sse';
+    type: 'sse';
     url: string;
     oauth: string;
 }
@@ -45,13 +44,13 @@ export type ILaunchSigature = IStdioLaunchSignature | ISSELaunchSignature;
 
 function refreshConnectionOption(envPath: string) {
     const defaultOption = {
-        type:'stdio',
+        type: 'stdio',
         command: 'mcp',
         args: ['run', 'main.py'],
         cwd: '../server'
     };
 
-    fs.writeFileSync(envPath, JSON.stringify(defaultOption, null, 4));   
+    fs.writeFileSync(envPath, JSON.stringify(defaultOption, null, 4));
 
     return defaultOption;
 }
@@ -82,7 +81,7 @@ const authPassword = JSON.parse(fs.readFileSync(path.join(__dirname, '..', '.env
 
 function updateConnectionOption(data: any) {
     const envPath = path.join(__dirname, '..', '.env');
-    
+
     if (data.connectionType === 'STDIO') {
         const connectionItem = {
             type: 'stdio',
@@ -106,6 +105,37 @@ function updateConnectionOption(data: any) {
 const devHome = path.join(__dirname, '..', '..');
 setRunningCWD(devHome);
 
+function verifyToken(url: string) {
+    try {
+        const token = url.split('=')[1];
+        return token === authPassword.toString();
+    } catch (error) {
+        return false;
+    }
+}
+
+const wss = new WebSocketServer(
+    {
+        port: 8282,
+        verifyClient: (info, callback) => {
+            console.log(info.req.url);
+
+            const ok = verifyToken(info.req.url || '');
+
+            console.log(ok);
+            
+
+            if (!ok) {
+                callback(false, 401, 'Unauthorized: Invalid token');
+            } else {
+                callback(true); // 允许连接
+            }
+        }
+    },
+);
+
+console.log('listen on ws://localhost:8282');
+
 wss.on('connection', (ws: any) => {
 
     // 仿造 webview 进行统一接口访问
@@ -125,42 +155,7 @@ wss.on('connection', (ws: any) => {
     // 注册消息接受的管线
     webview.onDidReceiveMessage(message => {
         logger.info(`command: [${message.command || 'No Command'}]`);
-        const { command, data, password } = message;
+        const { command, data } = message;
-
-        console.log(command, data);
-        
-
-        if (command === 'ciallo') {
-            if (data.password === authPassword) {
-                webview.postMessage({
-                    command,
-                    data: {
-                        code: 200,
-                        msg: 'ciallo'
-                    }
-                });
-            } else {
-                webview.postMessage({
-                    command,
-                    data: {
-                        code: 403,
-                        msg: '没有权限'
-                    }
-                });
-            }
-            return;
-        }
-
-        if (password !== authPassword) {
-            webview.postMessage({
-                command,
-                data: {
-                    code: 403,
-                    msg: '没有权限'
-                }
-            });
-            return;
-        }
 
         switch (command) {
             case 'web/launch-signature':
@@ -175,7 +170,7 @@ wss.on('connection', (ws: any) => {
                         url: option.url,
                         oauth: option.oauth || ''
                     };
-            
+
                 const launchResult = {
                     code: 200,
                     msg: launchResultMessage
@@ -187,11 +182,11 @@ wss.on('connection', (ws: any) => {
                 });
 
                 break;
-            
+
             case 'web/update-connection-sigature':
                 updateConnectionOption(data);
                 break;
-        
+
             default:
                 routeMessage(command, data, webview);
                 break;

tsconfig.json

@@ -1,19 +1,28 @@
 {
     "compilerOptions": {
-    "module": "commonjs",
+        "module": "commonjs",
-    "target": "es6",
+        "target": "es6",
-    "outDir": "./dist",
+        "outDir": "./dist",
-    "rootDir": "./src",
+        "rootDir": "./src",
-    "strict": true,
+        "strict": true,
-    "esModuleInterop": true,
+        "esModuleInterop": true,
-    "skipLibCheck": true,
+        "skipLibCheck": true,
-    "forceConsistentCasingInFileNames": true,
+        "forceConsistentCasingInFileNames": true,
-    "experimentalDecorators": true,
+        "experimentalDecorators": true,
-    // 允许访问 openmcp-sdk 目录
+        // 允许访问 openmcp-sdk 目录
-    "paths": {
+        "paths": {
-            "@openmcp-sdk/*": ["./openmcp-sdk/*"]
+            "@openmcp-sdk/*": [
+                "./openmcp-sdk/*"
+            ]
         }
     },
-    "include": ["src/**/*"],
+    "include": [
-    "exclude": ["node_modules", "dist", "service", "renderer"]
+        "src/**/*"
+    ],
+    "exclude": [
+        "node_modules",
+        "dist",
+        "service",
+        "renderer"
+    ]
 }