RedBear-OS/local/recipes/system/redbear-polkit/source/src/main.rs

use std::{collections::HashMap, env, error::Error, process, time::Duration};

use tokio::runtime::Builder as RuntimeBuilder;
use zbus::{
    Address,
    connection::Builder as ConnectionBuilder,
    interface,
    zvariant::{ObjectPath, OwnedObjectPath},
};

const BUS_NAME: &str = "org.freedesktop.PolicyKit1";
const AUTHORITY_PATH: &str = "/org/freedesktop/PolicyKit1/Authority";

type AuthorizationDetails = HashMap<String, String>;
type EnumeratedAction = (
    String,
    String,
    String,
    String,
    String,
    u32,
    AuthorizationDetails,
);

#[derive(Debug, Default)]
struct PolicyKitAuthority;

enum Command {
    Run,
    Help,
}

fn usage() -> &'static str {
    "Usage: redbear-polkit [--help]"
}

async fn wait_for_dbus_socket() {
    let socket_path = env::var("DBUS_STARTER_ADDRESS")
        .ok()
        .and_then(|addr| addr.strip_prefix("unix:path=").map(String::from))
        .unwrap_or_else(|| "/run/dbus/system_bus_socket".to_string());

    for _ in 0..30 {
        if tokio::net::UnixStream::connect(&socket_path).await.is_ok() {
            return;
        }
        tokio::time::sleep(Duration::from_secs(1)).await;
    }
    eprintln!("redbear-polkit: timed out waiting for D-Bus socket at {socket_path}");
}

fn parse_args() -> Result<Command, String> {
    let mut args = env::args().skip(1);

    match args.next() {
        None => Ok(Command::Run),
        Some(arg) if arg == "--help" || arg == "-h" => {
            if args.next().is_some() {
                return Err(String::from("unexpected extra arguments after --help"));
            }

            Ok(Command::Help)
        }
        Some(arg) => Err(format!("unrecognized argument '{arg}'")),
    }
}

fn parse_object_path(path: &str) -> Result<OwnedObjectPath, Box<dyn Error>> {
    Ok(OwnedObjectPath::try_from(path.to_owned())?)
}

fn system_connection_builder() -> Result<ConnectionBuilder<'static>, Box<dyn Error>> {
    if let Ok(address) = env::var("DBUS_STARTER_ADDRESS") {
        Ok(ConnectionBuilder::address(Address::try_from(address.as_str())?)?)
    } else {
        Ok(ConnectionBuilder::address(Address::try_from("unix:path=/run/dbus/system_bus_socket")?)?)
    }
}

fn spawn_signal_handler(shutdown_tx: tokio::sync::watch::Sender<bool>) {
    tokio::spawn(async move {
        #[cfg(unix)]
        {
            use tokio::signal::unix::{SignalKind, signal};
            if let Ok(mut sigterm) = signal(SignalKind::terminate()) {
                tokio::select! {
                    _ = sigterm.recv() => {},
                    _ = tokio::signal::ctrl_c() => {},
                }
            } else {
                let _ = tokio::signal::ctrl_c().await;
            }
        }
        #[cfg(not(unix))]
        {
            let _ = tokio::signal::ctrl_c().await;
        }
        let _ = shutdown_tx.send(true);
    });
}

#[interface(name = "org.freedesktop.PolicyKit1.Authority")]
impl PolicyKitAuthority {
    #[zbus(name = "CheckAuthorization")]
    fn check_authorization(
        &self,
        _action_id: &str,
        _details: AuthorizationDetails,
        _flags: u32,
        _cancellation_id: &str,
    ) -> (bool, bool, AuthorizationDetails) {
        (true, false, AuthorizationDetails::new())
    }

    #[zbus(name = "RegisterAuthenticationAgent")]
    fn register_authentication_agent(
        &self,
        _session: (&str, ObjectPath<'_>),
        _locale: &str,
        _object_path: &str,
    ) {
    }

    #[zbus(name = "UnregisterAuthenticationAgent")]
    fn unregister_authentication_agent(
        &self,
        _session: (&str, ObjectPath<'_>),
        _object_path: &str,
    ) {
    }

    #[zbus(name = "EnumerateActions")]
    fn enumerate_actions(&self, _locale: &str) -> Vec<EnumeratedAction> {
        Vec::new()
    }

    #[zbus(property(emits_changed_signal = "const"), name = "BackendName")]
    fn backend_name(&self) -> String {
        String::from("redbear-permit-all")
    }

    #[zbus(property(emits_changed_signal = "const"), name = "BackendVersion")]
    fn backend_version(&self) -> String {
        String::from("0.1.0")
    }
}

async fn run_daemon() -> Result<(), Box<dyn Error>> {
    wait_for_dbus_socket().await;

    let (shutdown_tx, mut shutdown_rx) = tokio::sync::watch::channel(false);
    spawn_signal_handler(shutdown_tx);

    let mut last_err = None;
    for attempt in 1..=5 {
        let _authority_path = parse_object_path(AUTHORITY_PATH)?;

        match system_connection_builder()?
            .name(BUS_NAME)?
            .serve_at(AUTHORITY_PATH, PolicyKitAuthority)?
            .build()
            .await
        {
            Ok(connection) => {
                eprintln!("redbear-polkit: registered {BUS_NAME} on the system bus");
                let _ = shutdown_rx.changed().await;
                eprintln!("redbear-polkit: shutdown signal received, exiting cleanly");
                drop(connection);
                return Ok(());
            }
            Err(err) => {
                if attempt < 5 {
                    eprintln!(
                        "redbear-polkit: attempt {attempt}/5 failed ({err}), retrying in 2s..."
                    );
                    tokio::time::sleep(Duration::from_secs(2)).await;
                }
                last_err = Some(err.into());
            }
        }
    }
    Err(last_err.unwrap())
}

fn main() {
    match parse_args() {
        Ok(Command::Help) => {
            println!("{}", usage());
        }
        Ok(Command::Run) => {
            let runtime = match RuntimeBuilder::new_multi_thread().enable_all().build() {
                Ok(runtime) => runtime,
                Err(err) => {
                    eprintln!("redbear-polkit: failed to create tokio runtime: {err}");
                    process::exit(1);
                }
            };

            if let Err(err) = runtime.block_on(run_daemon()) {
                eprintln!("redbear-polkit: fatal error: {err}");
                process::exit(1);
            }
        }
        Err(err) => {
            eprintln!("redbear-polkit: {err}");
            eprintln!("{}", usage());
            process::exit(1);
        }
    }
}