git.delta.rocks / remowt / refs/commits / 5cb6be498b69

difftreelog

feat properly cancel agent task

wplzsrwuYaroslav Bolyukin2024-08-12parent: #9ee216a.patch.diff
in: trunk

3 files changed

modifiedCargo.lockdiffbeforeafterboth
306 "piper",306 "piper",
307]307]
308
309[[package]]
310name = "byteorder"
311version = "1.5.0"
312source = "registry+https://github.com/rust-lang/crates.io-index"
313checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
308314
309[[package]]315[[package]]
310name = "bytes"316name = "bytes"
311version = "1.6.1"317version = "1.7.1"
312source = "registry+https://github.com/rust-lang/crates.io-index"318source = "registry+https://github.com/rust-lang/crates.io-index"
313checksum = "a12916984aab3fa6e39d655a33e09c0071eb36d6ab3aea5c2d78551f1df6d952"319checksum = "8318a53db07bb3f8dca91a600466bdb3f2eaadeedfdbcf02e1accbad9271ba50"
314320
315[[package]]321[[package]]
316name = "cc"322name = "cc"
317version = "1.1.6"323version = "1.1.7"
318source = "registry+https://github.com/rust-lang/crates.io-index"324source = "registry+https://github.com/rust-lang/crates.io-index"
319checksum = "2aba8f4e9906c7ce3c73463f62a7f0c65183ada1a2d47e397cc8810827f9694f"325checksum = "26a5c3fd7bfa1ce3897a3a3501d362b2d87b7f2583ebcb4a949ec25911025cbc"
320326
321[[package]]327[[package]]
322name = "cexpr"328name = "cexpr"
644650
645[[package]]651[[package]]
646name = "indexmap"652name = "indexmap"
647version = "2.2.6"653version = "2.3.0"
648source = "registry+https://github.com/rust-lang/crates.io-index"654source = "registry+https://github.com/rust-lang/crates.io-index"
649checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26"655checksum = "de3fc2e30ba82dd1b3911c8de1ffc143c74a914a14e99514d7637e3099df5ea0"
650dependencies = [656dependencies = [
651 "equivalent",657 "equivalent",
652 "hashbrown",658 "hashbrown",
905911
906[[package]]912[[package]]
907name = "ppv-lite86"913name = "ppv-lite86"
908version = "0.2.17"914version = "0.2.20"
909source = "registry+https://github.com/rust-lang/crates.io-index"915source = "registry+https://github.com/rust-lang/crates.io-index"
910checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"916checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04"
917dependencies = [
918 "zerocopy",
919]
911920
912[[package]]921[[package]]
913name = "proc-macro-crate"922name = "proc-macro-crate"
968977
969[[package]]978[[package]]
970name = "regex"979name = "regex"
971version = "1.10.5"980version = "1.10.6"
972source = "registry+https://github.com/rust-lang/crates.io-index"981source = "registry+https://github.com/rust-lang/crates.io-index"
973checksum = "b91213439dad192326a0d7c6ee3955910425f441d7038e0d6933b0aec5c4517f"982checksum = "4219d74c6b67a3654a9fbebc4b419e22126d13d2f3c4a07ee0cb61ff79a79619"
974dependencies = [983dependencies = [
975 "aho-corasick",984 "aho-corasick",
976 "memchr",985 "memchr",
10891098
1090[[package]]1099[[package]]
1091name = "serde_json"1100name = "serde_json"
1092version = "1.0.120"1101version = "1.0.122"
1093source = "registry+https://github.com/rust-lang/crates.io-index"1102source = "registry+https://github.com/rust-lang/crates.io-index"
1094checksum = "4e0d21c9a8cae1235ad58a00c11cb40d4b1e5c784f1ef2c537876ed6ffd8b7c5"1103checksum = "784b6203951c57ff748476b126ccb5e8e2959a5c19e5c617ab1956be3dbc68da"
1095dependencies = [1104dependencies = [
1096 "itoa",1105 "itoa",
1106 "memchr",
1097 "ryu",1107 "ryu",
1098 "serde",1108 "serde",
1099]1109]
12051215
1206[[package]]1216[[package]]
1207name = "tempfile"1217name = "tempfile"
1208version = "3.10.1"1218version = "3.11.0"
1209source = "registry+https://github.com/rust-lang/crates.io-index"1219source = "registry+https://github.com/rust-lang/crates.io-index"
1210checksum = "85b77fafb263dd9d05cbeac119526425676db3784113aa9295c88498cbf8bff1"1220checksum = "b8fcd239983515c23a32fb82099f97d0b11b8c72f654ed659363a95c3dad7a53"
1211dependencies = [1221dependencies = [
1212 "cfg-if",1222 "cfg-if",
1213 "fastrand",1223 "fastrand",
1224 "once_cell",
1214 "rustix",1225 "rustix",
1215 "windows-sys",1226 "windows-sys",
1216]1227]
12761287
1277[[package]]1288[[package]]
1278name = "toml_datetime"1289name = "toml_datetime"
1279version = "0.6.7"1290version = "0.6.8"
1280source = "registry+https://github.com/rust-lang/crates.io-index"1291source = "registry+https://github.com/rust-lang/crates.io-index"
1281checksum = "f8fb9f64314842840f1d940ac544da178732128f1c78c21772e876579e0da1db"1292checksum = "0dd7358ecb8fc2f8d014bf86f6f638ce72ba252a2c3a2572f2a795f1d23efb41"
12821293
1283[[package]]1294[[package]]
1284name = "toml_edit"1295name = "toml_edit"
1605 "zbus",1616 "zbus",
1606]1617]
1618
1619[[package]]
1620name = "zerocopy"
1621version = "0.7.35"
1622source = "registry+https://github.com/rust-lang/crates.io-index"
1623checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
1624dependencies = [
1625 "byteorder",
1626 "zerocopy-derive",
1627]
1628
1629[[package]]
1630name = "zerocopy-derive"
1631version = "0.7.35"
1632source = "registry+https://github.com/rust-lang/crates.io-index"
1633checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
1634dependencies = [
1635 "proc-macro2",
1636 "quote",
1637 "syn 2.0.72",
1638]
16071639
1608[[package]]1640[[package]]
1609name = "zvariant"1641name = "zvariant"
modifiedcmds/remowt-agent/src/main.rsdiffbeforeafterboth
--- a/cmds/remowt-agent/src/main.rs
+++ b/cmds/remowt-agent/src/main.rs
@@ -2,14 +2,14 @@
 use std::collections::{BTreeMap, HashMap};
 use std::io::{stdout, Write};
 use std::marker::PhantomData;
-use std::sync::{Mutex, RwLock};
+use std::sync::{Arc, Mutex, OnceLock};
 use std::{future, process};
 
 use clap::Parser;
 use polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};
 use tokio::runtime::Handle;
 use tokio::task::{AbortHandle, JoinHandle, LocalSet};
-use tracing::trace;
+use tracing::{info, trace};
 use ui_prompt::dbus::DbusPrompterInterface;
 use ui_prompt::rofi::RofiPrompter;
 use ui_prompt::{PrependSourcePrompter, Prompter, Source};
@@ -58,16 +58,34 @@
     }
 }
 
+struct CancelTaskOnDrop {
+    tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,
+    handle: String,
+}
+impl Drop for CancelTaskOnDrop {
+    fn drop(&mut self) {
+        info!("cancel on drop");
+        if let Some(task) = self
+            .tasks
+            .lock()
+            .expect("not poisoned")
+            .remove(&self.handle)
+        {
+            task.abort();
+        }
+    }
+}
+
 struct Agent {
     helper: PolkitHelperProxy<'static>,
-    tasks: Mutex<HashMap<String, AbortHandle>>,
+    tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,
     connection: Connection,
 }
 impl Agent {
     async fn new(connection: Connection) -> anyhow::Result<Self> {
         Ok(Self {
             helper: PolkitHelperProxy::new(&connection).await?,
-            tasks: Mutex::new(HashMap::new()),
+            tasks: Arc::new(Mutex::new(HashMap::new())),
             connection,
         })
     }
@@ -78,7 +96,7 @@
     /// BeginAuthentication method
     #[allow(clippy::too_many_arguments)]
     async fn begin_authentication(
-        &mut self,
+        &self,
         action_id: String,
         message: String,
         icon_name: String,
@@ -87,12 +105,15 @@
         identities: Vec<Identity>,
     ) -> zbus::fdo::Result<()> {
         use std::fmt::Write;
-        trace!("begin auth");
+        info!("begin auth");
+        let _cancel_guard = Arc::new(OnceLock::new());
         let task = {
             let connection = self.connection.clone();
             let helper = self.helper.clone();
             let cookie = cookie.clone();
+            let _cancel_guard = _cancel_guard.clone();
             tokio::task::spawn(async move {
+                let _cancel_guard = _cancel_guard.clone();
                 trace!("conversation task");
                 let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);
                 if let Some(subject) = details.remove("polkit.caller-pid") {
@@ -121,6 +142,7 @@
                     identities.iter().map(|v| v.to_string()).collect();
                 let identity_displays: Vec<&str> =
                     identity_displays.iter().map(|v| v.as_str()).collect();
+                info!("choose identity");
                 let choosen_identity = match identity_displays.len() {
                     0 => {
                         return Err(fdo::Error::AuthFailed(
@@ -139,6 +161,7 @@
                             .await?
                     }
                 };
+                info!("identity chosen");
 
                 let _ = write!(
                     description,
@@ -148,7 +171,10 @@
                 prompter.description = description;
 
                 prompter.source.push(Source(Cow::Borrowed("polkit daemon")));
+                // let connection = Connection::system().await?;
+                // let helper = PolkitHelperProxy::new(&connection).await?;
                 let prompter = TemporaryPrompterInterface::new(connection, prompter).await;
+                info!("init conv");
                 helper
                     .init_conversation(
                         BackendRequest {
@@ -166,24 +192,26 @@
                 Ok(())
             })
         };
-
         self.tasks
             .lock()
             .unwrap()
             .insert(cookie.clone(), task.abort_handle());
-        let result = task.await.expect("join error");
-        // The only way to no reach this line, is to either panic in previous line, or if authorization cancelled,
-        // while cancellation will remove task by itself.
-        // TODO: But still it would be better to have abort guard, which will remove it from HashMap
-        self.tasks.lock().unwrap().remove(&cookie);
+        info!("abort handle stored");
+        let _ = _cancel_guard.set(CancelTaskOnDrop {
+            tasks: self.tasks.clone(),
+            handle: cookie.clone(),
+        });
+
+        let _ = task.await;
 
-        result
+        Ok(())
     }
 
     /// CancelAuthentication method
     async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {
-        trace!("cancel auth");
+        info!("auth cancelled");
         if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {
+            info!("abort handle found");
             abort.abort();
         }
         // debug!("Authentication cancled ! {cookie}");
modifiednix/nixos-modules.nixdiffbeforeafterboth
--- a/nix/nixos-modules.nix
+++ b/nix/nixos-modules.nix
@@ -10,6 +10,9 @@
       ];
       systemd.services.remowt-polkit-helper = {
         aliases = ["dbus-lach.polkit.helper1.service"];
+        # Restarting can kill in-progress auth requests.
+        # It is good to have it restarted for security, but I didn't decided on the flow yet, graceful shutdown?..
+        unitConfig.X-RestartIfChanged = false;
       };
     };
   };