wry/src/cli/clients.rs

use {
    crate::{
        cli::{
            GlobalArgs,
            json::{JsonClient, jsonl},
        },
        ipc::{self, Client as IpcClient},
        tools::tool_client::{Handle, ToolClient, with_tool_client},
        utils::errorfmt::ErrorFmt,
        wire::{JayClientQueryId, jay_client_query, jay_compositor},
    },
    ahash::AHashMap,
    clap::{Args, Subcommand},
    std::{cell::RefCell, mem, rc::Rc},
    uapi::c,
};

#[derive(Args, Debug)]
pub struct ClientsArgs {
    #[clap(subcommand)]
    cmd: Option<ClientsCmd>,
}

#[derive(Subcommand, Debug)]
enum ClientsCmd {
    /// Show information about clients.
    Show(ShowArgs),
    /// Disconnect a client.
    Kill(KillArgs),
}

#[derive(Args, Debug)]
struct ShowArgs {
    #[clap(subcommand)]
    cmd: ShowCmd,
}

#[derive(Subcommand, Debug)]
enum ShowCmd {
    /// Show all clients.
    All,
    /// Show a client with a given ID.
    Id(ShowIdArgs),
    /// Interactively select a window and show information about its client.
    SelectWindow,
}

#[derive(Args, Debug)]
struct ShowIdArgs {
    /// The ID of the client.
    id: u64,
}

#[derive(Args, Debug)]
struct KillArgs {
    #[clap(subcommand)]
    cmd: KillCmd,
}

#[derive(Subcommand, Debug)]
enum KillCmd {
    /// Kill the client with a given ID.
    Id(KillIdArgs),
    /// Interactively select a window and kill its client.
    SelectWindow,
}

#[derive(Args, Debug)]
struct KillIdArgs {
    /// The ID of the client.
    id: u64,
}

pub fn main(global: GlobalArgs, clients_args: ClientsArgs) {
    with_tool_client(global.log_level, |tc| async move {
        let clients = Rc::new(Clients { tc: tc.clone() });
        clients.run(&global, clients_args).await;
    });
}

struct Clients {
    tc: Rc<ToolClient>,
}

impl Clients {
    async fn run(&self, global: &GlobalArgs, args: ClientsArgs) {
        let tc = &self.tc;
        let cmd = args
            .cmd
            .unwrap_or(ClientsCmd::Show(ShowArgs { cmd: ShowCmd::All }));
        if self.run_ipc(global, &cmd) {
            return;
        }
        let comp = tc.jay_compositor().await;
        match cmd {
            ClientsCmd::Show(a) => {
                let id = tc.id();
                tc.send(jay_compositor::CreateClientQuery { self_id: comp, id });
                match a.cmd {
                    ShowCmd::All => {
                        tc.send(jay_client_query::AddAll { self_id: id });
                    }
                    ShowCmd::Id(a) => {
                        tc.send(jay_client_query::AddId {
                            self_id: id,
                            id: a.id,
                        });
                    }
                    ShowCmd::SelectWindow => {
                        let client_id = tc.select_toplevel_client().await;
                        if client_id == 0 {
                            fatal!("Did not select a window");
                        }
                        tc.send(jay_client_query::AddId {
                            self_id: id,
                            id: client_id,
                        });
                    }
                }
                tc.send(jay_client_query::Execute { self_id: id });
                let clients = handle_client_query(tc, id).await;
                let mut clients = clients.values().collect::<Vec<_>>();
                clients.sort_by_key(|c| c.id);
                if global.json {
                    for client in clients {
                        let client = make_json_client(client);
                        jsonl(&client);
                    }
                } else {
                    let mut prefix = "    ".to_string();
                    let mut printer = ClientPrinter {
                        prefix: &mut prefix,
                    };
                    for client in clients {
                        println!("- client:");
                        printer.print_client(client);
                    }
                }
            }
            ClientsCmd::Kill(a) => match a.cmd {
                KillCmd::Id(id) => {
                    tc.send(jay_compositor::KillClient {
                        self_id: comp,
                        id: id.id,
                    });
                }
                KillCmd::SelectWindow => {
                    let client_id = tc.select_toplevel_client().await;
                    if client_id == 0 {
                        fatal!("Did not select a window");
                    }
                    tc.send(jay_compositor::KillClient {
                        self_id: comp,
                        id: client_id,
                    });
                }
            },
        }
        tc.round_trip().await;
    }

    fn run_ipc(&self, global: &GlobalArgs, cmd: &ClientsCmd) -> bool {
        match cmd {
            ClientsCmd::Show(a) => {
                let id = match &a.cmd {
                    ShowCmd::All => None,
                    ShowCmd::Id(a) => Some(a.id),
                    ShowCmd::SelectWindow => return false,
                };
                let clients: Vec<IpcClient> = match ipc::request(&ipc::Request::ClientsGet { id }) {
                    Ok(clients) => clients,
                    Err(e) if e.can_fallback() => return false,
                    Err(e) => fatal!("Could not query clients over IPC: {}", ErrorFmt(e)),
                };
                let mut clients = clients.into_iter().map(Client::from).collect::<Vec<_>>();
                clients.sort_by_key(|c| c.id);
                if global.json {
                    for client in &clients {
                        jsonl(&make_json_client(client));
                    }
                } else {
                    let mut prefix = "    ".to_string();
                    let mut printer = ClientPrinter {
                        prefix: &mut prefix,
                    };
                    for client in &clients {
                        println!("- client:");
                        printer.print_client(client);
                    }
                }
                true
            }
            ClientsCmd::Kill(a) => match &a.cmd {
                KillCmd::Id(id) => {
                    match ipc::request_unit(&ipc::Request::ClientsKill { id: id.id }) {
                        Ok(()) => true,
                        Err(e) if e.can_fallback() => false,
                        Err(e) => fatal!("Could not kill client over IPC: {}", ErrorFmt(e)),
                    }
                }
                KillCmd::SelectWindow => false,
            },
        }
    }
}

#[derive(Default)]
pub struct Client {
    pub id: u64,
    pub sandboxed: bool,
    pub sandbox_engine: Option<String>,
    pub sandbox_app_id: Option<String>,
    pub sandbox_instance_id: Option<String>,
    pub uid: Option<c::uid_t>,
    pub pid: Option<c::pid_t>,
    pub is_xwayland: bool,
    pub comm: Option<String>,
    pub exe: Option<String>,
}

impl From<IpcClient> for Client {
    fn from(client: IpcClient) -> Self {
        Self {
            id: client.client_id,
            sandboxed: client.sandboxed,
            sandbox_engine: client.sandbox_engine,
            sandbox_app_id: client.sandbox_app_id,
            sandbox_instance_id: client.sandbox_instance_id,
            uid: client.uid,
            pid: client.pid,
            is_xwayland: client.is_xwayland,
            comm: client.comm,
            exe: client.exe,
        }
    }
}

pub async fn handle_client_query(
    tl: &Rc<ToolClient>,
    id: JayClientQueryId,
) -> AHashMap<u64, Client> {
    use jay_client_query::*;
    let c = Rc::new(RefCell::new(Vec::<Client>::new()));
    macro_rules! last {
        ($c:ident) => {
            $c.borrow_mut().last_mut().unwrap()
        };
    }
    Start::handle(tl, id, c.clone(), |c, event| {
        c.borrow_mut().push(Client::default());
        last!(c).id = event.id;
    });
    Sandboxed::handle(tl, id, c.clone(), |c, _event| {
        last!(c).sandboxed = true;
    });
    SandboxEngine::handle(tl, id, c.clone(), |c, event| {
        last!(c).sandbox_engine = Some(event.engine.to_string());
    });
    SandboxAppId::handle(tl, id, c.clone(), |c, event| {
        last!(c).sandbox_app_id = Some(event.app_id.to_string());
    });
    SandboxInstanceId::handle(tl, id, c.clone(), |c, event| {
        last!(c).sandbox_instance_id = Some(event.instance_id.to_string());
    });
    Uid::handle(tl, id, c.clone(), |c, event| {
        last!(c).uid = Some(event.uid);
    });
    Pid::handle(tl, id, c.clone(), |c, event| {
        last!(c).pid = Some(event.pid);
    });
    IsXwayland::handle(tl, id, c.clone(), |c, _event| {
        last!(c).is_xwayland = true;
    });
    Comm::handle(tl, id, c.clone(), |c, event| {
        last!(c).comm = Some(event.comm.to_string());
    });
    Exe::handle(tl, id, c.clone(), |c, event| {
        last!(c).exe = Some(event.exe.to_string());
    });
    tl.round_trip().await;
    mem::take(&mut *c.borrow_mut())
        .into_iter()
        .map(|c| (c.id, c))
        .collect()
}

pub struct ClientPrinter<'a> {
    pub prefix: &'a mut String,
}

impl ClientPrinter<'_> {
    pub fn print_client(&mut self, c: &Client) {
        let p = &self.prefix;
        macro_rules! opt {
            ($field:ident, $pretty:expr) => {
                if let Some(v) = &c.$field {
                    println!("{p}{}: {}", $pretty, v);
                }
            };
        }
        macro_rules! bol {
            ($field:ident, $pretty:expr) => {
                if c.$field {
                    println!("{p}{}", $pretty);
                }
            };
        }
        println!("{p}id: {}", c.id);
        bol!(sandboxed, "sandboxed");
        opt!(sandbox_engine, "sandbox engine");
        opt!(sandbox_app_id, "sandbox app id");
        opt!(sandbox_instance_id, "sandbox instance id");
        opt!(uid, "uid");
        opt!(pid, "pid");
        bol!(is_xwayland, "xwayland");
        opt!(comm, "comm");
        opt!(exe, "exe");
    }
}

pub fn make_json_client(client: &Client) -> JsonClient<'_> {
    JsonClient {
        client_id: client.id,
        sandboxed: client.sandboxed,
        sandbox_engine: client.sandbox_engine.as_deref(),
        sandbox_app_id: client.sandbox_app_id.as_deref(),
        sandbox_instance_id: client.sandbox_instance_id.as_deref(),
        uid: client.uid,
        pid: client.pid,
        is_xwayland: client.is_xwayland,
        comm: client.comm.as_deref(),
        exe: client.exe.as_deref(),
    }
}