Refactor window creation functions for improved environment variable handling

desktop/tauri/src-tauri/src/window.rs

@@ -3,6 +3,7 @@ use tauri::{
     image::Image, AppHandle, Listener, Manager, Result, Theme, UserAttentionType, WebviewUrl,
     WebviewWindow, WebviewWindowBuilder,
 };
+use std::sync::{atomic::{AtomicBool, Ordering}};
 
 use crate::{portmaster::PortmasterExt, traymenu};
 
@@ -11,6 +12,8 @@ const DARK_PM_ICON: &[u8] = include_bytes!("../../../../assets/data/icons/pm_dar
 
 const CUSTOM_ENVVAR_FOR_WEBVIEW_PROCESS: &str = "PORTMASTER_UI_WEBVIEW_PROCESS";
 
+static UI_PROCESS_ENV_VAR_DEFINED_FLAG: AtomicBool = AtomicBool::new(false);
+
 /// Either returns the existing "main" window or creates a new one.
 ///
 /// The window is not automatically shown (i.e it starts hidden).
@@ -26,15 +29,18 @@ pub fn create_main_window(app: &AppHandle) -> Result<WebviewWindow> {
     } else {
         debug!("[tauri] creating main window");
 
-        do_before_window_create(); // required operations before window creation
+        do_before_any_window_create(); // required operations before window creation
         let res = WebviewWindowBuilder::new(app, "main", WebviewUrl::App("index.html".into()))
             .title("Portmaster")
             .visible(false)
             .inner_size(1200.0, 700.0)
             .min_inner_size(800.0, 600.0)
             .theme(Some(Theme::Dark))
+            .on_page_load(|_window, _event| {
+                debug!("[tauri] main window page loaded: {}", _event.url());
+                do_after_main_window_created(); // required operations after Main window creation
+            })
             .build();
-        do_after_window_create(); // required operations after window creation
 
         match res {
             Ok(win) => {
@@ -74,7 +80,7 @@ pub fn create_splash_window(app: &AppHandle) -> Result<WebviewWindow> {
         Ok(window)
     } else {
 
-        do_before_window_create(); // required operations before window creation
+        do_before_any_window_create(); // required operations before window creation
         let window = WebviewWindowBuilder::new(app, "splash", WebviewUrl::App("index.html".into()))
             .center()
             .closable(false)
@@ -84,7 +90,6 @@ pub fn create_splash_window(app: &AppHandle) -> Result<WebviewWindow> {
             .title("Portmaster")
             .inner_size(600.0, 250.0)
             .build()?;
-        do_after_window_create(); // required operations after window creation
         set_window_icon(&window);
 
         let _ = window.request_user_attention(Some(UserAttentionType::Informational));
@@ -125,26 +130,36 @@ pub fn set_window_icon(window: &WebviewWindow) {
     };
 }
 
-/// This function must be called before the window is created.
+/// This function must be called before any window is created.
 /// 
 /// Temporarily sets the environment variable `PORTMASTER_WEBVIEW_UI_PROCESS` to "true".
 /// This ensures that any child process (i.e., the WebView process) spawned during window creation
 /// will inherit this environment variable. This allows portmaster-core to detect that the process
 /// is a child WebView of the main process.
 /// 
-/// IMPORTANT: After the window is created, you must call `do_after_window_create()` to remove
+/// IMPORTANT: After the 'Main' window is created, you must call `do_after_main_window_created()` to remove
 /// the environment variable from the main process environment.
-pub fn do_before_window_create() {
+/// This ensures that any subsequent child processes (such as those created by "open external" functionality)
+/// will not inherit this environment variable, correctly indicating that they are not part of the
+/// Portmaster UI WebView process.
+pub fn do_before_any_window_create() {
+    UI_PROCESS_ENV_VAR_DEFINED_FLAG.store(true, Ordering::SeqCst);
     std::env::set_var(CUSTOM_ENVVAR_FOR_WEBVIEW_PROCESS, "true");
 }
 
-/// This function must be called after the window is created.
+/// This function must be called after the Main window is created.
 /// 
 /// Removes the `PORTMASTER_WEBVIEW_UI_PROCESS` environment variable from the main process.
 /// This ensures that only the child WebView process has the variable set, and the main process
 /// does not retain it.
-pub fn do_after_window_create() {
-    std::env::remove_var(CUSTOM_ENVVAR_FOR_WEBVIEW_PROCESS);
+pub fn do_after_main_window_created() {
+     let flag_was_set = UI_PROCESS_ENV_VAR_DEFINED_FLAG.compare_exchange(
+        true, false, Ordering::SeqCst, Ordering::SeqCst
+    ).is_ok();
+
+    if flag_was_set {
+        std::env::remove_var(CUSTOM_ENVVAR_FOR_WEBVIEW_PROCESS);
+    }
 } 
 
 /// Opens a window for the tauri application.