From c7003326dbcb37194c00afcce2225a028153f619 Mon Sep 17 00:00:00 2001
From: Huakun Shen <huakun.shen@huakunshen.com>
Date: Tue, 12 Nov 2024 12:25:24 -0500
Subject: [PATCH] fix: implemented file server for custom UI commands on
 Windows (#24)

---
 apps/desktop/src-tauri/src/lib.rs             |   8 +-
 apps/desktop/src/lib/cmds/ext.ts              |  31 ++++-
 apps/desktop/src/routes/+page.svelte          |   7 +-
 .../routes/extension/ui-iframe/+page.svelte   |   8 +-
 apps/desktop/src/routes/settings/+page.svelte |   4 +-
 packages/api/src/commands/extension.ts        |   6 +
 packages/tauri-plugins/jarvis/build.rs        |   1 +
 .../tauri-plugins/jarvis/permissions/all.toml |   1 +
 .../commands/spawn_extension_file_server.toml |  13 ++
 .../permissions/autogenerated/reference.md    |  26 ++++
 .../jarvis/permissions/schemas/schema.json    |  10 ++
 .../jarvis/src/commands/extension.rs          | 128 ++++++++++++++----
 packages/tauri-plugins/jarvis/src/lib.rs      |   1 +
 .../jarvis/src/model/extension.rs             |  15 +-
 14 files changed, 215 insertions(+), 44 deletions(-)
 create mode 100644 packages/tauri-plugins/jarvis/permissions/autogenerated/commands/spawn_extension_file_server.toml

diff --git a/apps/desktop/src-tauri/src/lib.rs b/apps/desktop/src-tauri/src/lib.rs
index 05ed166..eb57386 100644
--- a/apps/desktop/src-tauri/src/lib.rs
+++ b/apps/desktop/src-tauri/src/lib.rs
@@ -113,8 +113,12 @@ pub fn run() {
                 Some(ext) => {
                     // let app_state = app_handle.state::<tauri_plugin_jarvis::model::app_state::AppState>();
                     // let extension_path = app_state.extension_path.lock().unwrap().clone();
-                    // tauri_file_server(app_handle, request, extension_path)
-                    tauri_file_server(app_handle, request, ext.path.clone(), ext.dist.clone())
+                    tauri_file_server(
+                        app_handle,
+                        request,
+                        ext.info.path.clone(),
+                        ext.info.dist.clone(),
+                    )
                 }
                 None => tauri::http::Response::builder()
                     .status(tauri::http::StatusCode::NOT_FOUND)
diff --git a/apps/desktop/src/lib/cmds/ext.ts b/apps/desktop/src/lib/cmds/ext.ts
index 1aede5c..b29abe9 100644
--- a/apps/desktop/src/lib/cmds/ext.ts
+++ b/apps/desktop/src/lib/cmds/ext.ts
@@ -2,10 +2,12 @@ import { appState } from "@/stores"
 import { winExtMap } from "@/stores/winExtMap"
 import { trimSlash } from "@/utils/url"
 import { constructExtensionSupportDir } from "@kksh/api"
+import { spawnExtensionFileServer } from "@kksh/api/commands"
 import { CustomUiCmd, ExtPackageJsonExtra, TemplateUiCmd } from "@kksh/api/models"
 import { launchNewExtWindow } from "@kksh/extension"
 import { convertFileSrc } from "@tauri-apps/api/core"
 import * as fs from "@tauri-apps/plugin-fs"
+import { platform } from "@tauri-apps/plugin-os"
 import { goto } from "$app/navigation"
 
 export async function createExtSupportDir(extPath: string) {
@@ -44,28 +46,43 @@ export async function onCustomUiCmdSelect(
 	// console.log("onCustomUiCmdSelect", ext, cmd, isDev, hmr)
 	await createExtSupportDir(ext.extPath)
 	let url = cmd.main
-
-	if (hmr && isDev && cmd.devMain) {
+	const useDevMain = hmr && isDev && cmd.devMain
+	if (useDevMain) {
 		url = cmd.devMain
 	} else {
 		url = decodeURIComponent(convertFileSrc(`${trimSlash(cmd.main)}`, "ext"))
 	}
-	const url2 = `/extension/ui-iframe?url=${encodeURIComponent(url)}&extPath=${encodeURIComponent(ext.extPath)}`
+	let url2 = `/extension/ui-iframe?url=${encodeURIComponent(url)}&extPath=${encodeURIComponent(ext.extPath)}`
 	if (cmd.window) {
 		const winLabel = await winExtMap.registerExtensionWithWindow({
 			extPath: ext.extPath,
 			dist: cmd.dist
 		})
-		console.log("Launch new window, ", winLabel)
+		if (platform() === "windows" && !useDevMain) {
+			const addr = await spawnExtensionFileServer(winLabel)
+			const newUrl = `http://${addr}`
+			url2 = `/extension/ui-iframe?url=${encodeURIComponent(newUrl)}&extPath=${encodeURIComponent(ext.extPath)}`
+		}
+		console.log("URL 2", url2)
 		const window = launchNewExtWindow(winLabel, url2, cmd.window)
 		window.onCloseRequested(async (event) => {
 			await winExtMap.unregisterExtensionFromWindow(winLabel)
 		})
 	} else {
 		console.log("Launch main window")
-		return winExtMap
-			.registerExtensionWithWindow({ windowLabel: "main", extPath: ext.extPath, dist: cmd.dist })
-			.then(() => goto(url2))
+		const winLabel = await winExtMap.registerExtensionWithWindow({
+			windowLabel: "main",
+			extPath: ext.extPath,
+			dist: cmd.dist
+		})
+		if (platform() === "windows" && !useDevMain) {
+			const addr = await spawnExtensionFileServer(winLabel) // addr has format "127.0.0.1:<port>"
+			console.log("Extension file server address: ", addr)
+			const newUrl = `http://${addr}`
+			url2 = `/extension/ui-iframe?url=${encodeURIComponent(newUrl)}&extPath=${encodeURIComponent(ext.extPath)}`
+		}
+		console.log("URL 2", url2)
+		goto(url2)
 	}
 	appState.clearSearchTerm()
 }
diff --git a/apps/desktop/src/routes/+page.svelte b/apps/desktop/src/routes/+page.svelte
index ec0a512..d64616a 100644
--- a/apps/desktop/src/routes/+page.svelte
+++ b/apps/desktop/src/routes/+page.svelte
@@ -24,7 +24,7 @@
 	import { cn, commandScore } from "@kksh/ui/utils"
 	import { getCurrentWebviewWindow } from "@tauri-apps/api/webviewWindow"
 	import { exit } from "@tauri-apps/plugin-process"
-	import { CircleXIcon, EllipsisVerticalIcon, RefreshCcwIcon } from "lucide-svelte"
+	import { ArrowBigUpIcon, CircleXIcon, EllipsisVerticalIcon, RefreshCcwIcon } from "lucide-svelte"
 
 	let inputEle: HTMLInputElement | null = null
 	function onKeyDown(event: KeyboardEvent) {
@@ -100,7 +100,6 @@
 				</DropdownMenu.Trigger>
 				<DropdownMenu.Content class="w-80">
 					<DropdownMenu.Group>
-						<DropdownMenu.Separator />
 						<DropdownMenu.Item onclick={() => exit()}>
 							<CircleXIcon class="h-4 w-4 text-red-500" />
 							Quit
@@ -116,12 +115,12 @@
 						<DropdownMenu.Item onclick={toggleDevTools}>
 							<Icon icon="mingcute:code-fill" class="mr-2 h-5 w-5 text-green-500" />
 							Toggle Devtools
-							<DropdownMenu.Shortcut>⌃+Shift+I</DropdownMenu.Shortcut>
+							<DropdownMenu.Shortcut><span class="flex items-center">⌃+<ArrowBigUpIcon class="w-4 h-4" />+I</span></DropdownMenu.Shortcut>
 						</DropdownMenu.Item>
 						<DropdownMenu.Item onclick={() => location.reload()}>
 							<RefreshCcwIcon class="mr-2 h-4 w-4 text-green-500" />
 							Reload Window
-							<DropdownMenu.Shortcut>⌃+Shift+R</DropdownMenu.Shortcut>
+							<DropdownMenu.Shortcut><span class="flex items-center">⌃+<ArrowBigUpIcon class="w-4 h-4" />+R</span></DropdownMenu.Shortcut>
 						</DropdownMenu.Item>
 						<DropdownMenu.Item
 							onclick={() => {
diff --git a/apps/desktop/src/routes/extension/ui-iframe/+page.svelte b/apps/desktop/src/routes/extension/ui-iframe/+page.svelte
index 071f3bd..d0665bf 100644
--- a/apps/desktop/src/routes/extension/ui-iframe/+page.svelte
+++ b/apps/desktop/src/routes/extension/ui-iframe/+page.svelte
@@ -68,6 +68,7 @@
 		// 	navigateTo(localePath("/"))
 		// },
 		goBack: async () => {
+			console.log("goBack iframe ui API called")
 			if (isInMainWindow()) {
 				goto("/")
 			} else {
@@ -132,6 +133,7 @@
 	} satisfies IApp
 
 	function onBackBtnClicked() {
+		console.log("onBackBtnClicked")
 		if (isInMainWindow()) {
 			goHome()
 		} else {
@@ -173,13 +175,12 @@
 		class={cn("absolute", positionToTailwindClasses(uiControl.backBtnPosition))}
 		size="icon"
 		variant="outline"
-		data-tauri-drag-region
 		onclick={onBackBtnClicked}
 	>
 		{#if appWin.label === "main"}
-			<ArrowLeftIcon class="w-4" data-tauri-drag-region />
+			<ArrowLeftIcon class="w-4" />
 		{:else}
-			<XIcon class="w-4" data-tauri-drag-region />
+			<XIcon class="w-4" />
 		{/if}
 	</Button>
 {/if}
@@ -198,7 +199,6 @@
 		class={cn("absolute", positionToTailwindClasses(uiControl.refreshBtnPosition))}
 		size="icon"
 		variant="outline"
-		data-tauri-drag-region
 		onclick={iframeUiAPI.reloadPage}
 	>
 		<RefreshCcwIcon class="w-4" />
diff --git a/apps/desktop/src/routes/settings/+page.svelte b/apps/desktop/src/routes/settings/+page.svelte
index cf5544f..9032f1e 100644
--- a/apps/desktop/src/routes/settings/+page.svelte
+++ b/apps/desktop/src/routes/settings/+page.svelte
@@ -47,9 +47,9 @@
 				<Switch bind:checked={$appConfig} />
 			</li> -->
 	</ul>
-	{#if dev}
+	<!-- {#if dev}
 		<Shiki class="w-full overflow-x-auto" lang="json" code={JSON.stringify($appConfig, null, 2)} />
-	{/if}
+	{/if} -->
 </main>
 
 <style scoped>
diff --git a/packages/api/src/commands/extension.ts b/packages/api/src/commands/extension.ts
index 66481bb..7aaa734 100644
--- a/packages/api/src/commands/extension.ts
+++ b/packages/api/src/commands/extension.ts
@@ -31,6 +31,12 @@ export function unregisterExtensionWindow(label: string): Promise<void> {
 	})
 }
 
+export function spawnExtensionFileServer(windowLabel: string): Promise<string> {
+	return invoke<string>(generateJarvisPluginCommand("spawn_extension_file_server"), {
+		windowLabel
+	})
+}
+
 export function registerExtensionSpawnedProcess(windowLabel: string, pid: number): Promise<void> {
 	return invoke(generateJarvisPluginCommand("register_extension_spawned_process"), {
 		windowLabel,
diff --git a/packages/tauri-plugins/jarvis/build.rs b/packages/tauri-plugins/jarvis/build.rs
index c3f87ca..812db0a 100644
--- a/packages/tauri-plugins/jarvis/build.rs
+++ b/packages/tauri-plugins/jarvis/build.rs
@@ -62,6 +62,7 @@ const COMMANDS: &[&str] = &[
     "register_extension_spawned_process",
     "unregister_extension_spawned_process",
     "unregister_extension_window",
+    "spawn_extension_file_server",
     "get_ext_label_map",
     // "ext_store_wrapper_set",
     // "ext_store_wrapper_get",
diff --git a/packages/tauri-plugins/jarvis/permissions/all.toml b/packages/tauri-plugins/jarvis/permissions/all.toml
index 059f06a..e2dd0d9 100644
--- a/packages/tauri-plugins/jarvis/permissions/all.toml
+++ b/packages/tauri-plugins/jarvis/permissions/all.toml
@@ -58,6 +58,7 @@ commands.allow = [
     "is_window_label_registered",
     "register_extension_window",
     "unregister_extension_window",
+    "spawn_extension_file_server",
     "register_extension_spawned_process",
     "unregister_extension_spawned_process",
     "get_ext_label_map",
diff --git a/packages/tauri-plugins/jarvis/permissions/autogenerated/commands/spawn_extension_file_server.toml b/packages/tauri-plugins/jarvis/permissions/autogenerated/commands/spawn_extension_file_server.toml
new file mode 100644
index 0000000..f798860
--- /dev/null
+++ b/packages/tauri-plugins/jarvis/permissions/autogenerated/commands/spawn_extension_file_server.toml
@@ -0,0 +1,13 @@
+# Automatically generated - DO NOT EDIT!
+
+"$schema" = "../../schemas/schema.json"
+
+[[permission]]
+identifier = "allow-spawn-extension-file-server"
+description = "Enables the spawn_extension_file_server command without any pre-configured scope."
+commands.allow = ["spawn_extension_file_server"]
+
+[[permission]]
+identifier = "deny-spawn-extension-file-server"
+description = "Denies the spawn_extension_file_server command without any pre-configured scope."
+commands.deny = ["spawn_extension_file_server"]
diff --git a/packages/tauri-plugins/jarvis/permissions/autogenerated/reference.md b/packages/tauri-plugins/jarvis/permissions/autogenerated/reference.md
index 27858dc..1a00216 100644
--- a/packages/tauri-plugins/jarvis/permissions/autogenerated/reference.md
+++ b/packages/tauri-plugins/jarvis/permissions/autogenerated/reference.md
@@ -1610,6 +1610,32 @@ Denies the sleep_displays command without any pre-configured scope.
 <tr>
 <td>
 
+`jarvis:allow-spawn-extension-file-server`
+
+</td>
+<td>
+
+Enables the spawn_extension_file_server command without any pre-configured scope.
+
+</td>
+</tr>
+
+<tr>
+<td>
+
+`jarvis:deny-spawn-extension-file-server`
+
+</td>
+<td>
+
+Denies the spawn_extension_file_server command without any pre-configured scope.
+
+</td>
+</tr>
+
+<tr>
+<td>
+
 `jarvis:allow-start-server`
 
 </td>
diff --git a/packages/tauri-plugins/jarvis/permissions/schemas/schema.json b/packages/tauri-plugins/jarvis/permissions/schemas/schema.json
index 28823bb..89bae7c 100644
--- a/packages/tauri-plugins/jarvis/permissions/schemas/schema.json
+++ b/packages/tauri-plugins/jarvis/permissions/schemas/schema.json
@@ -909,6 +909,16 @@
           "type": "string",
           "const": "deny-sleep-displays"
         },
+        {
+          "description": "Enables the spawn_extension_file_server command without any pre-configured scope.",
+          "type": "string",
+          "const": "allow-spawn-extension-file-server"
+        },
+        {
+          "description": "Denies the spawn_extension_file_server command without any pre-configured scope.",
+          "type": "string",
+          "const": "deny-spawn-extension-file-server"
+        },
         {
           "description": "Enables the start_server command without any pre-configured scope.",
           "type": "string",
diff --git a/packages/tauri-plugins/jarvis/src/commands/extension.rs b/packages/tauri-plugins/jarvis/src/commands/extension.rs
index 10af74b..db5f9cd 100644
--- a/packages/tauri-plugins/jarvis/src/commands/extension.rs
+++ b/packages/tauri-plugins/jarvis/src/commands/extension.rs
@@ -1,14 +1,20 @@
 use crate::JarvisState;
 use crate::{
     model::{
-        extension::Extension,
+        extension::{Extension, ExtensionInfo},
         manifest::{ExtPackageJsonExtra, MANIFEST_FILE_NAME},
     },
     utils::manifest::load_jarvis_ext_manifest,
 };
 use std::collections::HashMap;
-use std::{fmt::format, path::PathBuf};
-use tauri::{command, AppHandle, Runtime, State, Window};
+use std::net::{SocketAddr, TcpListener};
+use std::{
+    path::PathBuf,
+    sync::{Arc, Mutex},
+};
+use tauri::{AppHandle, Runtime, State, Window};
+use tower_http::cors::CorsLayer;
+use tower_http::services::ServeDir;
 
 /// manifest_path can be folder of package.json
 /// If it's a folder, join it with package.json
@@ -61,6 +67,9 @@ pub async fn is_window_label_registered<R: Runtime>(
         .contains_key(label.as_str()))
 }
 
+/// When ui iframe ext loaded, register the window containing the iframe with the extension info
+/// Extension info contains path to extension command and dist path, so Kunkun's custom file server knows where to load
+/// static files from when a request is received from a window.
 #[tauri::command]
 pub async fn register_extension_window<R: Runtime>(
     _app: AppHandle<R>,
@@ -74,17 +83,14 @@ pub async fn register_extension_window<R: Runtime>(
         None => format!("main:ext:{}", uuid::Uuid::new_v4()),
     };
     let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
-    // if label_ext_map.contains_key(window_label_2.as_str()) {
-    //     return Err(format!(
-    //         "Window with label {} is already registered",
-    //         &window_label_2
-    //     ));
-    // }
     let ext = Extension {
-        path: extension_path,
-        processes: vec![],
-        dist: dist,
-        // identifier: manifest.kunkun.identifier,
+        info: ExtensionInfo {
+            path: extension_path,
+            processes: vec![],
+            dist: dist,
+        },
+        shutdown_handle: Arc::new(Mutex::new(None)),
+        server_handle: Arc::new(Mutex::new(None)),
     };
     label_ext_map.insert(window_label_2.clone(), ext);
     Ok(window_label_2)
@@ -106,7 +112,7 @@ pub async fn register_extension_spawned_process<R: Runtime>(
         ));
     }
     let ext = label_ext_map.get_mut(window_label.as_str()).unwrap();
-    ext.processes.push(pid);
+    ext.info.processes.push(pid);
     Ok(())
 }
 
@@ -121,6 +127,7 @@ pub async fn unregister_extension_spawned_process<R: Runtime>(
     label_ext_map
         .get_mut(window_label.as_str())
         .unwrap()
+        .info
         .processes
         .retain(|p| *p != pid);
     Ok(())
@@ -131,20 +138,95 @@ pub async fn get_ext_label_map<R: Runtime>(
     _app: AppHandle<R>,
     _window: Window<R>,
     state: State<'_, JarvisState>,
-) -> Result<HashMap<String, Extension>, String> {
-    Ok(state.window_label_ext_map.lock().unwrap().clone())
+) -> Result<HashMap<String, ExtensionInfo>, String> {
+    let label_ext_map = state.window_label_ext_map.lock().unwrap();
+    // turn label_ext_map from HashMap<String, Extension> to HashMap<String, ExtensionInfo>
+    let label_ext_map_info: HashMap<String, ExtensionInfo> = label_ext_map
+        .iter()
+        .map(|(label, ext)| (label.clone(), ext.info.clone()))
+        .collect();
+    Ok(label_ext_map_info)
 }
 
+/// unregister extension window
 #[tauri::command]
 pub async fn unregister_extension_window<R: Runtime>(
     _app: AppHandle<R>,
     state: State<'_, JarvisState>,
     label: String,
-) -> Result<bool, String> {
-    Ok(state
-        .window_label_ext_map
-        .lock()
-        .unwrap()
-        .remove(label.as_str())
-        .is_some())
+) -> Result<(), String> {
+    // find extension, if there is shutdown handle, shutdown it
+    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
+    // find extension info with window_label
+    let ext = label_ext_map.get(label.as_str()).cloned();
+    // drop(label_ext_map);
+    match ext {
+        Some(ext) => {
+            let shutdown_handle = ext.shutdown_handle.lock().unwrap();
+            if let Some(shutdown_handle) = shutdown_handle.as_ref() {
+                shutdown_handle.shutdown();
+                log::info!("Shutdown extension file server with label {}", label);
+            }
+            let server_handle = ext.server_handle.lock().unwrap();
+            if let Some(server_handle) = server_handle.as_ref() {
+                server_handle.abort();
+                log::info!("Abort extension file server with label {}", label);
+            }
+            label_ext_map.remove(label.as_str());
+            log::info!("Unregistered extension window with label {}", label);
+        }
+        None => return Err(format!("Extension with label {} not found", label)),
+    }
+    // state
+    //     .window_label_ext_map
+    //     .lock()
+    //     .unwrap()
+    //     .remove(label.as_str())
+    //     .is_some()
+    Ok(())
+}
+
+/// spawn extension file server, only for Windows
+#[tauri::command]
+pub async fn spawn_extension_file_server<R: Runtime>(
+    _app: AppHandle<R>,
+    state: State<'_, JarvisState>,
+    window_label: String,
+) -> Result<SocketAddr, String> {
+    let mut label_ext_map = state.window_label_ext_map.lock().unwrap();
+    // find extension info with window_label
+    let ext = label_ext_map.get(window_label.as_str());
+    if ext.is_none() {
+        return Err(format!("Extension with label {} not found", window_label));
+    }
+    let ext = ext.unwrap();
+    let mut ext_path = ext.info.path.clone();
+    if let Some(dist) = ext.info.dist.clone() {
+        ext_path = ext_path.join(dist);
+    }
+    // TODO: spawn file server
+    let shutdown_handle = axum_server::Handle::new();
+    let shutdown_handle_clone = shutdown_handle.clone();
+    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+    let addr = listener.local_addr().unwrap();
+    let rest_router = axum::Router::new()
+        .layer(CorsLayer::permissive())
+        .nest_service("/", ServeDir::new(ext_path));
+
+    let server_handle = tauri::async_runtime::spawn(async move {
+        // axum_server::bind(addr)
+        axum_server::from_tcp(listener)
+            .handle(shutdown_handle_clone)
+            .serve(rest_router.into_make_service())
+            // .serve(combined_router.into_make_service())
+            .await
+            .unwrap();
+    });
+    // add server handle and shutdown handle to extension
+    let mut ext = label_ext_map.get_mut(window_label.as_str()).unwrap();
+    ext.server_handle.lock().unwrap().replace(server_handle);
+    ext.shutdown_handle.lock().unwrap().replace(shutdown_handle);
+
+    // TODO: add server handle and shutdown handle
+    Ok(addr)
 }
diff --git a/packages/tauri-plugins/jarvis/src/lib.rs b/packages/tauri-plugins/jarvis/src/lib.rs
index ca9dbdb..e52b2cc 100644
--- a/packages/tauri-plugins/jarvis/src/lib.rs
+++ b/packages/tauri-plugins/jarvis/src/lib.rs
@@ -125,6 +125,7 @@ pub fn init<R: Runtime>() -> TauriPlugin<R> {
             commands::extension::register_extension_spawned_process,
             commands::extension::unregister_extension_spawned_process,
             commands::extension::get_ext_label_map,
+            commands::extension::spawn_extension_file_server,
             /* ---------------------- extension storage API wrapper --------------------- */
             // commands::storage::ext_store_wrapper_set,
             // commands::storage::ext_store_wrapper_get,
diff --git a/packages/tauri-plugins/jarvis/src/model/extension.rs b/packages/tauri-plugins/jarvis/src/model/extension.rs
index 65da38b..47e145d 100644
--- a/packages/tauri-plugins/jarvis/src/model/extension.rs
+++ b/packages/tauri-plugins/jarvis/src/model/extension.rs
@@ -1,11 +1,22 @@
 use super::manifest::Permissions;
 use serde::{Deserialize, Serialize};
-use std::path::PathBuf;
+use std::{
+    net::SocketAddr,
+    path::PathBuf,
+    sync::{Arc, Mutex},
+};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Extension {
+pub struct ExtensionInfo {
     pub path: PathBuf,
     pub processes: Vec<u32>,
     pub dist: Option<String>,
     // pub identifier: String,
 }
+
+#[derive(Debug, Clone)]
+pub struct Extension {
+    pub info: ExtensionInfo,
+    pub shutdown_handle: Arc<Mutex<Option<axum_server::Handle>>>,
+    pub server_handle: Arc<std::sync::Mutex<Option<tauri::async_runtime::JoinHandle<()>>>>,
+}