diff --git a/desktop/tauri/src-tauri/src/main.rs b/desktop/tauri/src-tauri/src/main.rs index 69a14ee8..9b4b81eb 100644 --- a/desktop/tauri/src-tauri/src/main.rs +++ b/desktop/tauri/src-tauri/src/main.rs @@ -87,10 +87,21 @@ impl portmaster::Handler for WsHandler { error!("failed to close splash window: {}", err.to_string()); } - let handle = self.handle.clone(); - tauri::async_runtime::spawn(async move { - traymenu::tray_handler(cli, handle).await; - }); + // Cancel the previous tray handler task if it exists + let portmaster = self.handle.portmaster(); + if let Ok(mut task_guard) = portmaster.tray_handler_task.lock() { + if let Some(old_task) = task_guard.take() { + debug!("Aborting previous tray handler task"); + old_task.abort(); + } + + // Start new tray handler and store the task handle + let handle = self.handle.clone(); + let task = tauri::async_runtime::spawn(async move { + traymenu::tray_handler(cli, handle).await; + }); + *task_guard = Some(task); + } } fn on_disconnect(&mut self) { @@ -252,40 +263,47 @@ fn main() { .expect("error while running tauri application"); app.run(|handle, e| { - if let RunEvent::WindowEvent { label, event, .. } = e { - if label != "main" { - // We only have one window at most so any other label is unexpected - return; - } + match e { + RunEvent::WindowEvent { label, event, .. } => { + if label != "main" { + // We only have one window at most so any other label is unexpected + return; + } - // Do not let the user close the window, instead send an event to the main - // window so we can show the "will not stop portmaster" dialog and let the window - // close itself using - // - // window.__TAURI__.window.getCurrent().close() - // - // Note: the above javascript does NOT trigger the CloseRequested event so - // there's no need to handle that case here. - if let WindowEvent::CloseRequested { api, .. } = event { - debug!( - "window (label={}) close request received, forwarding to user-interface.", - label - ); + // Do not let the user close the window, instead send an event to the main + // window so we can show the "will not stop portmaster" dialog and let the window + // close itself using + // + // window.__TAURI__.window.getCurrent().close() + // + // Note: the above javascript does NOT trigger the CloseRequested event so + // there's no need to handle that case here. + if let WindowEvent::CloseRequested { api, .. } = event { + debug!( + "window (label={}) close request received, forwarding to user-interface.", + label + ); - // Manually save the window state on close attempt. - // This ensures the state is saved since we prevent the close event. - let _ = handle.save_window_state(WINDOW_STATE_FLAGS_TO_SAVE); + // Manually save the window state on close attempt. + // This ensures the state is saved since we prevent the close event. + let _ = handle.save_window_state(WINDOW_STATE_FLAGS_TO_SAVE); - api.prevent_close(); - if let Some(window) = handle.get_webview_window(label.as_str()) { - let result = window.emit("exit-requested", ""); - if let Err(err) = result { - error!("failed to emit event: {}", err.to_string()); + api.prevent_close(); + if let Some(window) = handle.get_webview_window(label.as_str()) { + let result = window.emit("exit-requested", ""); + if let Err(err) = result { + error!("failed to emit event: {}", err.to_string()); + } + } else { + error!("window was None"); } - } else { - error!("window was None"); } } + RunEvent::ExitRequested { .. } => { + debug!("Application exit requested, shutting down websocket"); + portmaster::websocket::shutdown_websocket(); + } + _ => {} } }); } diff --git a/desktop/tauri/src-tauri/src/portmaster/mod.rs b/desktop/tauri/src-tauri/src/portmaster/mod.rs index 19c33d7a..9ad81f69 100644 --- a/desktop/tauri/src-tauri/src/portmaster/mod.rs +++ b/desktop/tauri/src-tauri/src/portmaster/mod.rs @@ -18,7 +18,7 @@ pub mod commands; // The websocket module spawns an async function on tauri's runtime that manages // a persistent connection to the Portmaster websocket API and updates the tauri Portmaster // Plugin instance. -mod websocket; +pub mod websocket; // The notification module manages system notifications from portmaster. mod notifications; @@ -72,6 +72,9 @@ pub struct PortmasterInterface { // whether or not the angular application should call window.show after it // finished bootstrapping. should_show_after_bootstrap: AtomicBool, + + // handle to the tray handler task so we can abort it when reconnecting + pub tray_handler_task: Mutex>>, } impl PortmasterInterface { @@ -300,6 +303,14 @@ impl PortmasterInterface { fn on_disconnect(&self) { self.is_reachable.store(false, Ordering::Relaxed); + // Abort the tray handler task if it's running + if let Ok(mut task_guard) = self.tray_handler_task.lock() { + if let Some(task) = task_guard.take() { + debug!("Aborting tray handler task"); + task.abort(); + } + } + // clear the current api client reference. { let mut guard = self.api.lock().unwrap(); @@ -337,6 +348,7 @@ pub fn setup(app: AppHandle) { handle_notifications: AtomicBool::new(false), handle_prompts: AtomicBool::new(false), should_show_after_bootstrap: AtomicBool::new(true), + tray_handler_task: Mutex::new(None), }; app.manage(interface); diff --git a/desktop/tauri/src-tauri/src/portmaster/websocket.rs b/desktop/tauri/src-tauri/src/portmaster/websocket.rs index 6dca8519..be2efc90 100644 --- a/desktop/tauri/src-tauri/src/portmaster/websocket.rs +++ b/desktop/tauri/src-tauri/src/portmaster/websocket.rs @@ -1,9 +1,17 @@ use super::PortmasterExt; use crate::portapi::client::connect; use log::{debug, error, info, warn}; +use std::sync::atomic::{AtomicBool, Ordering}; use tauri::{AppHandle, Runtime}; use tokio::time::{sleep, Duration}; +static WEBSOCKET_SHUTDOWN: AtomicBool = AtomicBool::new(false); + +/// Signals the websocket thread to stop reconnecting and shut down gracefully. +pub fn shutdown_websocket() { + WEBSOCKET_SHUTDOWN.store(true, Ordering::Release); +} + /// Starts a backround thread (via tauri::async_runtime) that connects to the Portmaster /// Websocket database API. pub fn start_websocket_thread(app: AppHandle) { @@ -11,6 +19,12 @@ pub fn start_websocket_thread(app: AppHandle) { tauri::async_runtime::spawn(async move { loop { + // Check if we should shutdown before attempting to connect + if WEBSOCKET_SHUTDOWN.load(Ordering::Acquire) { + debug!("WebSocket thread shutting down gracefully"); + break; + } + debug!("Trying to connect to websocket endpoint"); let api = connect("ws://127.0.0.1:817/api/database/v1").await; @@ -23,12 +37,29 @@ pub fn start_websocket_thread(app: AppHandle) { portmaster.on_connect(cli.clone()); - while !cli.is_closed() { - let _ = sleep(Duration::from_secs(1)).await; + // Monitor connection status + loop { + if WEBSOCKET_SHUTDOWN.load(Ordering::Acquire) { + debug!("Shutdown signal received, closing connection"); + break; + } + + if cli.is_closed() { + warn!("Connection to portmaster lost"); + break; + } + + sleep(Duration::from_secs(1)).await; } portmaster.on_disconnect(); + // If shutdown was requested, exit the loop + if WEBSOCKET_SHUTDOWN.load(Ordering::Acquire) { + debug!("Exiting websocket thread after disconnect"); + break; + } + warn!("lost connection to portmaster, retrying ....") } Err(err) => { @@ -36,10 +67,18 @@ pub fn start_websocket_thread(app: AppHandle) { app.portmaster().on_disconnect(); - // sleep and retry + // Check shutdown flag before sleeping + if WEBSOCKET_SHUTDOWN.load(Ordering::Acquire) { + debug!("Shutdown requested, not retrying connection"); + break; + } + + // Sleep and retry with constant 2 second delay sleep(Duration::from_secs(2)).await; } } } + + info!("WebSocket thread terminated"); }); }