pixelflux/wayland/

vkclient.rs

//! `zwp_virtual_keyboard_v1` client for typing Unicode text into a nested
//! compositor's socket.
//!
//! pixelflux is normally the compositor, but in a nested deployment (a labwc/kwin
//! session running as a client of pixelflux) the apps live on that inner
//! compositor's socket, and keys injected into pixelflux's own seat resolve
//! against pixelflux's keymap — an overlay the inner compositor never sees. So
//! text is typed here as a client of whichever compositor the apps live under —
//! by Computer-Use actions and by selkies over the `type_text_wayland` ABI —
//! reusing the seat's [`KeymapPolicy`] over a US base: base-reachable characters
//! press their ordinary keycodes, everything else is overlay-bound. The client
//! is PERSISTENT per socket: the connection, virtual-keyboard device and its
//! uploaded keymap live across calls, so a flush re-uploads (and settles) only
//! when the accumulated keymap actually changed, and key events ride the
//! protocol's ordering in one batch with a single closing round-trip. Any
//! failure drops the cached connection and one reconnect is attempted before
//! the error propagates, so a restarted app compositor heals on the next call.
//! Blocking, off the compositor thread, with every round-trip deadline-bounded so
//! a wedged compositor cannot hang the caller forever.

use std::collections::{HashMap, HashSet};
use std::os::fd::AsFd;
use std::os::unix::net::UnixStream;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};
use std::time::Duration;

use wayland_client::protocol::{wl_registry, wl_seat};
use wayland_client::{delegate_noop, Connection, Dispatch, EventQueue, QueueHandle};
use wayland_protocols_misc::zwp_virtual_keyboard_v1::client::{
    zwp_virtual_keyboard_manager_v1::ZwpVirtualKeyboardManagerV1,
    zwp_virtual_keyboard_v1::ZwpVirtualKeyboardV1,
};

use crate::wayland::keymap::{compile_rmlvo, KeymapPolicy};
use crate::wayland::wlclient::{bounded_roundtrip, impl_sync_callback, memfd_with, SyncState};

/// Overlay keycodes stay under the X11 255 ceiling so XWayland apps under the app
/// compositor can still receive them (the seat's own overlay sits above 255).
const OVERLAY_FIRST_XKB: u32 = 150;
const OVERLAY_LAST_XKB: u32 = 255;
const OVERLAY_SLOTS: usize = (OVERLAY_LAST_XKB - OVERLAY_FIRST_XKB + 1) as usize;
/// wl_keyboard / zwp_virtual_keyboard key events carry evdev codes (xkb - 8).
const EVDEV_OFFSET: u32 = 8;
const KEYMAP_FORMAT_XKB_V1: u32 = 1;

#[derive(Default)]
struct Globals {
    seat: Option<wl_seat::WlSeat>,
    manager: Option<ZwpVirtualKeyboardManagerV1>,
    sync_done: bool,
}

impl SyncState for Globals {
    fn sync_done_mut(&mut self) -> &mut bool {
        &mut self.sync_done
    }
}
impl_sync_callback!(Globals);

impl Dispatch<wl_registry::WlRegistry, ()> for Globals {
    fn event(
        state: &mut Self,
        registry: &wl_registry::WlRegistry,
        event: wl_registry::Event,
        _: &(),
        _: &Connection,
        qh: &QueueHandle<Self>,
    ) {
        if let wl_registry::Event::Global { name, interface, .. } = event {
            // Version 1 of each suffices: the seat is only the manager argument.
            match interface.as_str() {
                "wl_seat" if state.seat.is_none() => {
                    state.seat = Some(registry.bind(name, 1, qh, ()));
                }
                "zwp_virtual_keyboard_manager_v1" if state.manager.is_none() => {
                    state.manager = Some(registry.bind(name, 1, qh, ()));
                }
                _ => {}
            }
        }
    }
}

delegate_noop!(Globals: ignore wl_seat::WlSeat);
delegate_noop!(Globals: ZwpVirtualKeyboardManagerV1);
delegate_noop!(Globals: ZwpVirtualKeyboardV1);

/// The US base keymap text, compiled once per process: selkies types per text
/// commit, and xkbcommon compilation is the expensive part of a call.
pub(crate) fn us_base_text() -> Option<&'static str> {
    static CACHE: OnceLock<Option<String>> = OnceLock::new();
    CACHE.get_or_init(|| compile_rmlvo("", "", "us", "", "")).as_deref()
}

/// The keymap policy over that base, shared across calls: rebuilding it costs a
/// keymap compile per typed batch, and its accumulated overlay assignments are
/// device-independent, so repeat batches reuse their slots. Locked — callers
/// run off-thread and Computer-Use may type concurrently with selkies.
fn shared_policy() -> Option<&'static Mutex<KeymapPolicy>> {
    static POLICY: OnceLock<Option<Mutex<KeymapPolicy>>> = OnceLock::new();
    POLICY
        .get_or_init(|| {
            let mut policy =
                KeymapPolicy::with_overlay_range(OVERLAY_FIRST_XKB, OVERLAY_LAST_XKB);
            policy.rebuild_base(us_base_text()?.to_string()).then(|| Mutex::new(policy))
        })
        .as_ref()
}

fn upload_keymap(
    vk: &ZwpVirtualKeyboardV1,
    queue: &mut EventQueue<Globals>,
    text: &str,
) -> Result<(), String> {
    let mut data = text.as_bytes().to_vec();
    // Compositors parse the mapping as a NUL-terminated string.
    data.push(0);
    let fd = memfd_with(&data)?;
    vk.keymap(KEYMAP_FORMAT_XKB_V1, fd.as_fd(), data.len() as u32);
    queue.flush().map_err(|e| format!("flush keymap: {e}"))
}

/// Universal keysym for a character: Latin-1 keysyms are their codepoint and
/// everything above rides the Unicode plane — the two encodings every compositor
/// and toolkit translate algorithmically. Editing controls map to their keys;
/// other controls have no keysym. Fixed grammar only: which keysym spells a
/// char in any richer sense is selkies' policy, delivered via `type_keysyms_to`.
fn universal_keysym(c: char) -> Option<u32> {
    match c {
        '\n' | '\r' => Some(0xFF0D),
        '\t' => Some(0xFF09),
        '\x1b' => Some(0xFF1B),
        _ => match c as u32 {
            cp @ 0x20..=0xFF => Some(cp),
            cp @ 0x100.. => Some(0x0100_0000 | cp),
            _ => None,
        },
    }
}

/// Bumped whenever the shared policy's keymap changes; each cached typer
/// re-uploads (and settles) only when it is behind this.
static KEYMAP_GENERATION: AtomicU64 = AtomicU64::new(1);

struct Typer {
    conn: Connection,
    queue: EventQueue<Globals>,
    state: Globals,
    vk: ZwpVirtualKeyboardV1,
    uploaded_generation: u64,
}

fn typers() -> &'static Mutex<HashMap<String, Typer>> {
    static TYPERS: OnceLock<Mutex<HashMap<String, Typer>>> = OnceLock::new();
    TYPERS.get_or_init(|| Mutex::new(HashMap::new()))
}

fn connect_typer(socket_path: &str) -> Result<Typer, String> {
    let stream =
        UnixStream::connect(socket_path).map_err(|e| format!("connect {socket_path}: {e}"))?;
    let conn = Connection::from_socket(stream).map_err(|e| format!("wayland setup: {e}"))?;
    let mut queue = conn.new_event_queue();
    let qh = queue.handle();
    let _registry = conn.display().get_registry(&qh, ());
    let mut state = Globals::default();
    bounded_roundtrip(&conn, &mut queue, &mut state)?;
    let seat = state.seat.take().ok_or("app compositor advertises no wl_seat")?;
    let manager = state
        .manager
        .take()
        .ok_or("app compositor does not advertise zwp_virtual_keyboard_manager_v1")?;
    let vk = manager.create_virtual_keyboard(&seat, &qh, ());
    // Surfaces an "unauthorized" bind error before the first keymap upload.
    bounded_roundtrip(&conn, &mut queue, &mut state)?;
    Ok(Typer { conn, queue, state, vk, uploaded_generation: 0 })
}

fn flush_keysyms(
    typer: &mut Typer,
    policy: &mut KeymapPolicy,
    keysyms: &[u32],
) -> Result<(), String> {
    let none = HashSet::new();
    // Chunk bound: at most OVERLAY_SLOTS keysyms per bind call, so a batch can
    // never recycle a slot it assigned earlier in the same batch.
    for chunk in keysyms.chunks(OVERLAY_SLOTS) {
        let (keycodes, changed) = policy.bind_many_plain(chunk, &none);
        if changed {
            KEYMAP_GENERATION.fetch_add(1, Ordering::Relaxed);
        }
        // The protocol requires a keymap before the first key event even when
        // the whole text resolves in the base; after that, only a changed
        // keymap costs an upload, its compositor-side compile and the settle.
        let generation = KEYMAP_GENERATION.load(Ordering::Relaxed);
        if typer.uploaded_generation != generation {
            upload_keymap(&typer.vk, &mut typer.queue, &policy.keymap_text())?;
            bounded_roundtrip(&typer.conn, &mut typer.queue, &mut typer.state)?;
            std::thread::sleep(Duration::from_millis(10));
            typer.uploaded_generation = generation;
        }
        // One protocol-ordered batch, one closing round-trip: the compositor
        // serializes per client, so per-key pacing bought nothing but latency.
        for &kc in &keycodes {
            // Below the evdev offset the keysym has no bindable keycode.
            if kc < EVDEV_OFFSET {
                continue;
            }
            typer.vk.key(0, kc - EVDEV_OFFSET, 1);
            typer.vk.key(0, kc - EVDEV_OFFSET, 0);
        }
        typer.queue.flush().map_err(|e| format!("flush keys: {e}"))?;
        bounded_roundtrip(&typer.conn, &mut typer.queue, &mut typer.state)?;
    }
    Ok(())
}

/// Type `text` in order through the persistent client for `socket_path`.
/// Codepoints with no keysym are skipped. Blocking; call off the compositor's
/// calloop thread.
pub fn type_text_to(socket_path: &str, text: &str) -> Result<(), String> {
    let syms: Vec<u32> = text.chars().filter_map(universal_keysym).collect();
    type_keysyms_to(socket_path, &syms)
}

/// Like [`type_text_to`], but taps the given keysyms verbatim: the caller owns
/// which keysym spells which character; this owns delivery (base-reachable
/// keysyms press their ordinary keycodes, the rest overlay-bind).
pub fn type_keysyms_to(socket_path: &str, keysyms: &[u32]) -> Result<(), String> {
    let mut policy = shared_policy()
        .ok_or("us base keymap failed to compile")?
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    let mut typers = typers().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
    let mut last_err = None;
    for _ in 0..2 {
        if !typers.contains_key(socket_path) {
            typers.insert(socket_path.to_string(), connect_typer(socket_path)?);
        }
        let typer = typers.get_mut(socket_path).expect("typer just ensured");
        match flush_keysyms(typer, &mut policy, keysyms) {
            Ok(()) => return Ok(()),
            Err(e) => {
                // Dropping the typer closes the socket; the compositor tears the
                // device down and releases anything it held. Reconnect once —
                // a restarted app compositor reappears under the same name.
                typers.remove(socket_path);
                last_err = Some(e);
            }
        }
    }
    Err(last_err.unwrap_or_else(|| "virtual-keyboard flush failed".to_string()))
}