git.delta.rocks / remowt / refs/commits / 7c2fb57d5b71

1use std::collections::HashMap;2use std::future;3use std::marker::PhantomData;4use std::sync::{Mutex, RwLock};56use polkit_shared::{BackendRequest, Identity};7use tokio::runtime::Handle;8use tokio::task::{AbortHandle, JoinHandle, LocalSet};9use tracing::trace;10use ui_prompt::dbus::DbusPrompterInterface;11use ui_prompt::rofi::RofiPrompter;12use ui_prompt::Prompter;13use zbus::zvariant::{OwnedValue, Str};14use zbus::ObjectServer;15use zbus::{interface, proxy, Connection};16use zbus_polkit::policykit1::Subject;1718struct TemporaryPrompterInterface<P: Prompter + Send + Sync + 'static> {19    connection: Connection,20    path: String,21    _marker: PhantomData<P>,22}23impl<P: Prompter + Send + Sync + 'static> TemporaryPrompterInterface<P> {24    async fn new(connection: Connection, prompter: P) -> Self {25        let path = format!(26            "/remowt/prompters/{}",27            uuid::Uuid::new_v4().to_string().replace("-", "_")28        );29        let _ = connection30            .object_server()31            .at(path.clone(), DbusPrompterInterface(prompter))32            .await;33        Self {34            connection,35            path,36            _marker: PhantomData,37        }38    }39}40impl<P: Prompter + Send + Sync + 'static> Drop for TemporaryPrompterInterface<P> {41    fn drop(&mut self) {42        // FIXME: block_in_place prevents to moving to current_thread runtime43        // There should be a blocking way to remove ObjectServer listener.44        // As far as I can see, it is only async because of async RwLock, shouldn't it be45        // just a sync lock?46        tokio::task::block_in_place(move || {47            Handle::current().block_on(async {48                let _ = self49                    .connection50                    .object_server()51                    .remove::<DbusPrompterInterface<P>, String>(self.path.clone())52                    .await;53            });54        });55    }56}5758struct Agent {59    helper: PolkitHelperProxy<'static>,60    tasks: Mutex<HashMap<String, AbortHandle>>,61    connection: Connection,62}63impl Agent {64    async fn new(connection: Connection) -> anyhow::Result<Self> {65        Ok(Self {66            helper: PolkitHelperProxy::new(&connection).await?,67            tasks: Mutex::new(HashMap::new()),68            connection,69        })70    }71}7273#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]74impl Agent {75    /// BeginAuthentication method76    #[allow(clippy::too_many_arguments)]77    async fn begin_authentication(78        &mut self,79        action_id: String,80        message: String,81        icon_name: String,82        details: HashMap<String, String>,83        cookie: String,84        identities: Vec<Identity>,85    ) -> zbus::fdo::Result<()> {86        trace!("begin auth");87        let task = {88            let connection = self.connection.clone();89            let helper = self.helper.clone();90            let cookie = cookie.clone();91            tokio::task::spawn(async move {92                trace!("conversation task");93                let prompter = TemporaryPrompterInterface::new(connection, RofiPrompter).await;94                helper95                    .init_conversation(96                        BackendRequest {97                            cookie: cookie.to_owned(),98                            environment: HashMap::new(),99                            prompter_path: prompter.path.clone(),100                            // TODO: Let user choose101                            identity: identities.get(0).expect("first always exists").clone(),102                        }, // cookie.to_owned(), HashMap::new(), prompter.path.clone()103                    )104                    .await?;105                println!("ASKED");106                dbg!(action_id, message, icon_name, details, cookie, identities);107108                Ok(())109            })110        };111112        self.tasks113            .lock()114            .unwrap()115            .insert(cookie.clone(), task.abort_handle());116        let result = task.await.expect("join error");117        // The only way to no reach this line, is to either panic in previous line, or if authorization cancelled,118        // while cancellation will remove task by itself.119        // TODO: But still it would be better to have abort guard, which will remove it from HashMap120        self.tasks.lock().unwrap().remove(&cookie);121122        result123    }124125    /// CancelAuthentication method126    async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {127        trace!("cancel auth");128        if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {129            abort.abort();130        }131        // debug!("Authentication cancled ! {cookie}");132        Ok(())133    }134}135136const OBJ_PATH: &str = "/0lach/polkitAgent";137#[tokio::main]138async fn main() -> anyhow::Result<()> {139    tracing_subscriber::fmt::init();140141    trace!("started");142    let conn = Connection::system().await?;143144    let proxy = zbus_polkit::policykit1::AuthorityProxy::new(&conn).await?;145    conn.object_server()146        .at(OBJ_PATH, Agent::new(conn.clone()).await?)147        .await?;148149    let session_id = std::env::var("XDG_SESSION_ID")?;150    let mut details = HashMap::new();151    let val: OwnedValue = {152        let wrapped: Str<'_> = session_id.into();153        wrapped.into()154    };155    details.insert("session-id".to_string(), val);156    proxy157        .register_authentication_agent(158            &Subject {159                subject_kind: "unix-session".to_string(),160                subject_details: details,161            },162            "C",163            OBJ_PATH,164        )165        .await?;166167    future::pending().await168}169170#[proxy(171    interface = "lach.PolkitHelper",172    default_service = "lach.polkit.helper1",173    default_path = "/lach/PolkitHelper"174)]175trait PolkitHelper {176    fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;177}