xembed-sni-proxy/src/sni_proxy.rs

use tokio::sync::mpsc;
use tracing::{debug, info};
use x11rb::connection::Connection;
use x11rb::protocol::composite::ConnectionExt as _;
use x11rb::protocol::damage::{self, ConnectionExt as _};
use x11rb::protocol::xproto::*;
use x11rb::protocol::xtest::ConnectionExt as _;
use x11rb::rust_connection::RustConnection;
use x11rb::CURRENT_TIME;

use crate::sni_dbus::{self, SniHandle};
use crate::tray_manager::Atoms;
use crate::DbusCommand;

const DEFAULT_ICON_SIZE: u16 = 64;

// XEMBED message constants.
const XEMBED_EMBEDDED_NOTIFY: u32 = 0;
const XEMBED_VERSION: u32 = 0;

pub struct SniProxy {
    conn: *const RustConnection,
    container: Window,
    client_window: Window,
    icon_size: u16,
    damage: Option<damage::Damage>,
    sni_handle: SniHandle,
}

// Safety: RustConnection is Send+Sync, and we only use it on the single-threaded runtime.
unsafe impl Send for SniProxy {}

/// Determine the icon size for a client window.
/// Reads the client's actual geometry and WM_NORMAL_HINTS to find its preferred size.
/// Falls back to DEFAULT_ICON_SIZE, overridable via TRAY_ICON_SIZE env var.
fn determine_icon_size(conn: &RustConnection, client_window: Window) -> u16 {
    // Env var override takes priority.
    if let Ok(val) = std::env::var("TRAY_ICON_SIZE") {
        if let Ok(size) = val.parse::<u16>() {
            if size > 0 {
                debug!("using TRAY_ICON_SIZE={size} from environment");
                return size;
            }
        }
    }

    // Try reading the client's current geometry.
    if let Some(geom) = conn
        .get_geometry(client_window)
        .ok()
        .and_then(|c| c.reply().ok())
    {
        let size = geom.width.max(geom.height);
        // Only trust it if it's a reasonable icon size (not 1x1 placeholder).
        if size >= 8 {
            debug!(
                "client 0x{client_window:x} native size {w}x{h}, using {size}",
                w = geom.width,
                h = geom.height,
            );
            return size;
        }
    }

    DEFAULT_ICON_SIZE
}

impl SniProxy {
    pub async fn new(
        conn: &RustConnection,
        screen_num: usize,
        client_window: Window,
        atoms: &Atoms,
        cmd_tx: mpsc::UnboundedSender<DbusCommand>,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let screen = &conn.setup().roots[screen_num];

        let icon_size = determine_icon_size(conn, client_window);
        info!("icon size for 0x{client_window:x}: {icon_size}x{icon_size}");

        // Create container window for the client. KDE's approach: use MANUAL
        // composite redirect so the X server never renders the window on screen —
        // pixels only exist in the offscreen pixmap. Plus stack below + opacity 0
        // as defense in depth.
        let container = conn.generate_id()?;
        conn.create_window(
            0,
            container,
            screen.root,
            -2000,
            -2000,
            icon_size,
            icon_size,
            0,
            WindowClass::INPUT_OUTPUT,
            screen.root_visual,
            &CreateWindowAux::default()
                .override_redirect(1)
                .background_pixel(0),
        )?;

        // Stack below all other windows.
        conn.configure_window(
            container,
            &ConfigureWindowAux::default().stack_mode(StackMode::BELOW),
        )?;

        // Map the container (must be mapped for the client to be viewable).
        conn.map_window(container)?;

        // Reparent client into our container.
        conn.reparent_window(client_window, container, 0, 0)?;

        // Resize client to fill container.
        conn.configure_window(
            client_window,
            &ConfigureWindowAux::default()
                .width(u32::from(icon_size))
                .height(u32::from(icon_size)),
        )?;

        // Map the client window.
        conn.map_window(client_window)?;

        // MANUAL redirect: the window is NOT rendered on screen by the X server.
        // Pixels only go to the offscreen composite pixmap, which we read via
        // get_image. This is how KDE's xembedsniproxy hides the container.
        conn.composite_redirect_window(
            client_window,
            x11rb::protocol::composite::Redirect::MANUAL,
        )?;

        // Create damage tracking on the client.
        let damage_id = conn.generate_id()?;
        conn.damage_create(
            damage_id,
            client_window,
            damage::ReportLevel::NON_EMPTY,
        )?;

        // Send XEMBED_EMBEDDED_NOTIFY to the client.
        let event = ClientMessageEvent::new(
            32,
            client_window,
            atoms.xembed,
            [
                CURRENT_TIME,
                XEMBED_EMBEDDED_NOTIFY,
                0, // detail
                container,
                XEMBED_VERSION,
            ],
        );
        conn.send_event(false, client_window, EventMask::NO_EVENT, event)?;
        conn.flush()?;

        info!(
            "created proxy: client=0x{client_window:x} container=0x{container:x} damage={damage_id}"
        );

        // Do initial icon capture.
        let icon_data = capture_icon(conn, client_window, icon_size);

        // Spawn a D-Bus connection for this icon.
        let sni_handle =
            sni_dbus::register_sni(client_window, icon_data, cmd_tx).await?;

        let proxy = Self {
            conn: conn as *const RustConnection,
            container,
            client_window,
            icon_size,
            damage: Some(damage_id),
            sni_handle,
        };

        Ok(proxy)
    }

    fn conn(&self) -> &RustConnection {
        // Safety: lifetime tied to TrayManager which holds the connection.
        unsafe { &*self.conn }
    }

    pub fn container(&self) -> Window {
        self.container
    }

    pub fn damage_id(&self) -> Option<u32> {
        self.damage
    }

    /// Re-capture the icon pixels and emit a NewIcon D-Bus signal.
    pub async fn update_icon(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        let conn = self.conn();

        // Subtract damage so we get notified again next time.
        if let Some(damage_id) = self.damage {
            conn.damage_subtract(damage_id, 0u32, 0u32)?;
        }

        let icon_data = capture_icon(conn, self.client_window, self.icon_size);
        self.sni_handle.update_icon(icon_data).await;
        Ok(())
    }

    /// Inject a click using XTest fake_input. On XWayland, send_event synthetic
    /// events aren't delivered to Wine apps. XTest generates events
    /// indistinguishable from real hardware input.
    ///
    /// Because the container lives at (-2000, -2000), XWayland has no valid
    /// Wayland surface for it and pointer warping is a no-op. We temporarily
    /// move the container to a real screen position, inject the click, then
    /// move it back.
    pub fn inject_click(
        &self,
        _x: i32,
        _y: i32,
        button: u8,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let conn = self.conn();
        let root = conn.setup().roots[0].root;

        // Save current pointer position.
        let pointer = conn.query_pointer(root)?.reply()?;

        // Move container to the cursor position so XWayland has a valid
        // surface for pointer input, and so Wine opens menus at the cursor.
        let half = i32::from(self.icon_size) / 2;
        conn.configure_window(
            self.container,
            &ConfigureWindowAux::default()
                .x(i32::from(pointer.root_x) - half)
                .y(i32::from(pointer.root_y) - half)
                .stack_mode(StackMode::ABOVE),
        )?;

        // Round-trip: ensure XWayland has processed the configure before
        // we try to warp the pointer to the new surface position.
        conn.get_input_focus()?.reply()?;

        // Warp pointer to the center of the container (the Wayland surface).
        let half = (self.icon_size / 2) as i16;
        conn.warp_pointer(0u32, self.container, 0, 0, 0, 0, half, half)?;

        // XTest fake button press + release.
        conn.xtest_fake_input(
            BUTTON_PRESS_EVENT,
            button,
            CURRENT_TIME,
            root,
            0,
            0,
            0,
        )?;
        conn.xtest_fake_input(
            BUTTON_RELEASE_EVENT,
            button,
            CURRENT_TIME,
            root,
            0,
            0,
            0,
        )?;

        // Warp pointer back to original position.
        conn.warp_pointer(0u32, root, 0, 0, 0, 0, pointer.root_x, pointer.root_y)?;

        // Move container back off-screen and re-lower it.
        conn.configure_window(
            self.container,
            &ConfigureWindowAux::default()
                .x(-2000)
                .y(-2000)
                .stack_mode(StackMode::BELOW),
        )?;

        conn.flush()?;
        debug!(
            "xtest button {button} on 0x{:x} via container (pointer restored to {}, {})",
            self.client_window, pointer.root_x, pointer.root_y
        );
        Ok(())
    }

    pub async fn shutdown(self) {
        let conn = self.conn();

        if let Some(damage_id) = self.damage {
            let _ = conn.damage_destroy(damage_id);
        }
        let _ = conn.reparent_window(
            self.client_window,
            conn.setup().roots[0].root,
            0,
            0,
        );
        let _ = conn.destroy_window(self.container);
        let _ = conn.flush();

        self.sni_handle.shutdown().await;
        info!("proxy for 0x{:x} shut down", self.client_window);
    }
}

/// Capture the client window's pixels and convert BGRA → ARGB (network byte order).
/// Returns (width, height, argb_data) or transparent fallback on failure.
fn capture_icon(conn: &RustConnection, window: Window, icon_size: u16) -> (i32, i32, Vec<u8>) {
    let size = i32::from(icon_size);

    let cookie = match conn.get_image(
        ImageFormat::Z_PIXMAP,
        window,
        0,
        0,
        icon_size,
        icon_size,
        !0, // all planes
    ) {
        Ok(c) => c,
        Err(e) => {
            debug!("get_image request failed for 0x{window:x}: {e}");
            return (size, size, vec![0u8; (size * size * 4) as usize]);
        }
    };
    let reply = match cookie.reply() {
        Ok(reply) => reply,
        Err(e) => {
            debug!("get_image failed for 0x{window:x}: {e}");
            return (size, size, vec![0u8; (size * size * 4) as usize]);
        }
    };

    let data = &reply.data;
    let pixel_count = (size * size) as usize;
    let mut argb = Vec::with_capacity(pixel_count * 4);

    // X11 ZPixmap on little-endian: bytes are B, G, R, A (or X).
    // SNI IconPixmap expects ARGB in network byte order (big-endian): A, R, G, B.
    for i in 0..pixel_count {
        let off = i * 4;
        if off + 3 < data.len() {
            let b = data[off];
            let g = data[off + 1];
            let r = data[off + 2];
            let a = data[off + 3];
            argb.push(a);
            argb.push(r);
            argb.push(g);
            argb.push(b);
        } else {
            argb.extend_from_slice(&[0, 0, 0, 0]);
        }
    }

    (size, size, argb)
}