Disconnect interfaces on exit (#68)

* Disconnect all interfaces on exit
DefGuard · Nov 24, 2023 · 729c8b5 · 729c8b5
1 parent ebd4094
commit 729c8b5
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 3 deletions.
diff --git a/src-tauri/src/appstate.rs b/src-tauri/src/appstate.rs
@@ -1,7 +1,11 @@
 use crate::{
-    database::{ActiveConnection, DbPool},
+    database::{ActiveConnection, Connection, DbPool},
+    error::Error,
     service::{
-        proto::desktop_daemon_service_client::DesktopDaemonServiceClient, utils::setup_client,
+        proto::{
+            desktop_daemon_service_client::DesktopDaemonServiceClient, RemoveInterfaceRequest,
+        },
+        utils::setup_client,
     },
 };
 use std::sync::{Arc, Mutex};
@@ -57,6 +61,31 @@ impl AppState {
             None // Connection not found
         }
     }
+
+    pub async fn close_all_connections(&self) -> Result<(), crate::error::Error> {
+        for connection in self.get_connections() {
+            debug!("Found active connection");
+            trace!("Connection: {:#?}", connection);
+            debug!("Removing interface");
+            let mut client = self.client.clone();
+            let request = RemoveInterfaceRequest {
+                interface_name: connection.interface_name.clone(),
+            };
+            if let Err(error) = client.remove_interface(request).await {
+                error!("Failed to remove interface: {error}");
+                return Err(Error::InternalError);
+            }
+            debug!("Removed interface");
+            debug!("Saving connection");
+            trace!("Connection: {:#?}", connection);
+            let mut connection: Connection = connection.into();
+            connection.save(&self.get_pool()).await?;
+            debug!("Connection saved");
+            trace!("Saved connection: {:#?}", connection);
+            info!("Location {} disconnected", connection.location_id);
+        }
+        Ok(())
+    }
     pub fn find_connection(&self, location_id: i64) -> Option<ActiveConnection> {
         let connections = self.active_connections.lock().unwrap();
         debug!("Checking for active connection with location id: {location_id} in active connections: {:#?}", connections);

diff --git a/src-tauri/src/bin/defguard-client.rs b/src-tauri/src/bin/defguard-client.rs
@@ -90,7 +90,13 @@ async fn main() {
         .on_system_tray_event(|app, event| match event {
             SystemTrayEvent::MenuItemClick { id, .. } => match id.as_str() {
                 "quit" => {
-                    std::process::exit(0x0);
+                    let app_state: State<AppState> = app.state();
+                    tokio::task::block_in_place(|| {
+                        tokio::runtime::Handle::current().block_on(async {
+                            let _ = app_state.close_all_connections().await;
+                            std::process::exit(0);
+                        });
+                    });
                 }
                 "show" => {
                     if let Some(main_window) = app.get_window("main") {