Implement xwayland-satellite integration

4 files changed, 470 insertions, 0 deletions

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 2ac37721..ac7cf6a1 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -37,6 +37,7 @@ pub mod scale;
 pub mod spawning;
 pub mod transaction;
 pub mod watcher;
+pub mod xwayland;
 
 pub static IS_SYSTEMD_SERVICE: AtomicBool = AtomicBool::new(false);
 
diff --git a/src/utils/spawning.rs b/src/utils/spawning.rs
index 286df68f..0e118785 100644
--- a/src/utils/spawning.rs
+++ b/src/utils/spawning.rs
@@ -16,6 +16,7 @@ use crate::utils::expand_home;
 pub static REMOVE_ENV_RUST_BACKTRACE: AtomicBool = AtomicBool::new(false);
 pub static REMOVE_ENV_RUST_LIB_BACKTRACE: AtomicBool = AtomicBool::new(false);
 pub static CHILD_ENV: RwLock<Environment> = RwLock::new(Environment(Vec::new()));
+pub static CHILD_DISPLAY: RwLock<Option<String>> = RwLock::new(None);
 
 static ORIGINAL_NOFILE_RLIMIT_CUR: Atomic<rlim_t> = Atomic::new(0);
 static ORIGINAL_NOFILE_RLIMIT_MAX: Atomic<rlim_t> = Atomic::new(0);
@@ -116,6 +117,14 @@ fn spawn_sync(
         process.env_remove("RUST_LIB_BACKTRACE");
     }
 
+    // Set DISPLAY if needed.
+    let display = CHILD_DISPLAY.read().unwrap();
+    if let Some(display) = &*display {
+        process.env("DISPLAY", display);
+    } else {
+        process.env_remove("DISPLAY");
+    }
+
     // Set configured environment.
     let env = CHILD_ENV.read().unwrap();
     for var in &env.0 {
diff --git a/src/utils/xwayland/mod.rs b/src/utils/xwayland/mod.rs
new file mode 100644
index 00000000..a04b5a53
--- /dev/null
+++ b/src/utils/xwayland/mod.rs
@@ -0,0 +1,151 @@
+use std::os::fd::OwnedFd;
+use std::os::linux::net::SocketAddrExt;
+use std::os::unix::net::{SocketAddr, UnixListener};
+
+use anyhow::{anyhow, ensure, Context as _};
+use rustix::fs::{lstat, mkdir};
+use rustix::io::Errno;
+use rustix::process::getuid;
+use smithay::reexports::rustix::fs::{unlink, OFlags};
+use smithay::reexports::rustix::process::getpid;
+use smithay::reexports::rustix::{self};
+
+pub mod satellite;
+
+const TMP_UNIX_DIR: &str = "/tmp";
+const X11_TMP_UNIX_DIR: &str = "/tmp/.X11-unix";
+
+struct X11Connection {
+    display_name: String,
+    abstract_fd: OwnedFd,
+    unix_fd: OwnedFd,
+    _unix_guard: Unlink,
+    _lock_guard: Unlink,
+}
+
+struct Unlink(String);
+impl Drop for Unlink {
+    fn drop(&mut self) {
+        let _ = unlink(&self.0);
+    }
+}
+
+// Adapted from Mutter code:
+// https://gitlab.gnome.org/GNOME/mutter/-/blob/48.3.1/src/wayland/meta-xwayland.c?ref_type=tags#L513
+fn ensure_x11_unix_dir() -> anyhow::Result<()> {
+    match mkdir(X11_TMP_UNIX_DIR, 0o1777.into()) {
+        Ok(()) => Ok(()),
+        Err(Errno::EXIST) => {
+            ensure_x11_unix_perms().context("wrong X11 directory permissions")?;
+            Ok(())
+        }
+        Err(err) => Err(err).context("error creating X11 directory"),
+    }
+}
+
+fn ensure_x11_unix_perms() -> anyhow::Result<()> {
+    let x11_tmp = lstat(X11_TMP_UNIX_DIR).context("error checking X11 directory permissions")?;
+    let tmp = lstat(TMP_UNIX_DIR).context("error checking /tmp directory permissions")?;
+
+    ensure!(
+        x11_tmp.st_uid == tmp.st_uid || x11_tmp.st_uid == getuid().as_raw(),
+        "wrong ownership for X11 directory"
+    );
+    ensure!(
+        (x11_tmp.st_mode & 0o022) == 0o022,
+        "X11 directory is not writable"
+    );
+    ensure!(
+        (x11_tmp.st_mode & 0o1000) == 0o1000,
+        "X11 directory is missing the sticky bit"
+    );
+
+    Ok(())
+}
+
+fn pick_x11_display(start: u32) -> anyhow::Result<(u32, OwnedFd, Unlink)> {
+    for n in start..start + 50 {
+        let lock_path = format!("/tmp/.X{n}-lock");
+        let flags = OFlags::WRONLY | OFlags::CLOEXEC | OFlags::CREATE | OFlags::EXCL;
+        let Ok(lock_fd) = rustix::fs::open(&lock_path, flags, 0o444.into()) else {
+            // FIXME: check if the target process is dead and reuse the lock.
+            continue;
+        };
+        return Ok((n, lock_fd, Unlink(lock_path)));
+    }
+
+    Err(anyhow!("no free X11 display found after 50 attempts"))
+}
+
+fn bind_to_socket(addr: &SocketAddr) -> anyhow::Result<UnixListener> {
+    let listener = UnixListener::bind_addr(addr).context("error binding socket")?;
+    Ok(listener)
+}
+
+fn bind_to_abstract_socket(display: u32) -> anyhow::Result<UnixListener> {
+    let name = format!("/tmp/.X11-unix/X{display}");
+    let addr = SocketAddr::from_abstract_name(name).unwrap();
+    bind_to_socket(&addr)
+}
+
+fn bind_to_unix_socket(display: u32) -> anyhow::Result<(UnixListener, Unlink)> {
+    let name = format!("/tmp/.X11-unix/X{display}");
+    let addr = SocketAddr::from_pathname(&name).unwrap();
+    // Unlink old leftover socket if any.
+    let _ = unlink(&name);
+    let guard = Unlink(name);
+    bind_to_socket(&addr).map(|listener| (listener, guard))
+}
+
+fn open_display_sockets(display: u32) -> anyhow::Result<(UnixListener, UnixListener, Unlink)> {
+    let a = bind_to_abstract_socket(display).context("error binding to abstract socket")?;
+    let (u, g) = bind_to_unix_socket(display).context("error binding to unix socket")?;
+    Ok((a, u, g))
+}
+
+fn setup_connection() -> anyhow::Result<X11Connection> {
+    let _span = tracy_client::span!("open_x11_sockets");
+
+    ensure_x11_unix_dir()?;
+
+    let mut n = 0;
+    let mut attempt = 0;
+    let (display, lock_guard, a, u, unix_guard) = loop {
+        let (display, lock_fd, lock_guard) = pick_x11_display(n)?;
+
+        // Write our PID into the lock file.
+        let pid_string = format!("{:>10}\n", getpid().as_raw_nonzero());
+        if let Err(err) = rustix::io::write(&lock_fd, pid_string.as_bytes()) {
+            return Err(err).context("error writing PID to X11 lock file");
+        }
+        drop(lock_fd);
+
+        match open_display_sockets(display) {
+            Ok((a, u, g)) => {
+                break (display, lock_guard, a, u, g);
+            }
+            Err(err) => {
+                if attempt == 50 {
+                    return Err(err)
+                        .context("error opening X11 sockets after creating a lock file");
+                }
+
+                n = display + 1;
+                attempt += 1;
+                continue;
+            }
+        }
+    };
+
+    let display_name = format!(":{display}");
+    let abstract_fd = OwnedFd::from(a);
+    let unix_fd = OwnedFd::from(u);
+
+    Ok(X11Connection {
+        display_name,
+        abstract_fd,
+        unix_fd,
+        _unix_guard: unix_guard,
+        _lock_guard: lock_guard,
+    })
+}
diff --git a/src/utils/xwayland/satellite.rs b/src/utils/xwayland/satellite.rs
new file mode 100644
index 00000000..0c5a2d67
--- /dev/null
+++ b/src/utils/xwayland/satellite.rs
@@ -0,0 +1,309 @@
+use std::os::fd::{AsRawFd as _, BorrowedFd, OwnedFd};
+use std::os::unix::net::UnixListener;
+use std::os::unix::process::CommandExt as _;
+use std::path::{Path, PathBuf};
+use std::process::{Command, Stdio};
+use std::thread;
+
+use calloop::channel::Sender;
+use calloop::generic::Generic;
+use calloop::{Interest, Mode, PostAction, RegistrationToken};
+use smithay::reexports::rustix::io::{fcntl_setfd, FdFlags};
+
+use crate::niri::State;
+use crate::utils::expand_home;
+use crate::utils::xwayland::X11Connection;
+
+pub struct Satellite {
+    x11: X11Connection,
+    abstract_token: Option<RegistrationToken>,
+    unix_token: Option<RegistrationToken>,
+    to_main: Sender<ToMain>,
+}
+
+enum ToMain {
+    SetupWatch,
+}
+
+impl Satellite {
+    pub fn display_name(&self) -> &str {
+        &self.x11.display_name
+    }
+}
+
+pub fn setup(state: &mut State) {
+    if state.niri.satellite.is_some() {
+        return;
+    }
+
+    let config = state.niri.config.borrow();
+    let xwls_config = &config.xwayland_satellite;
+    if xwls_config.off {
+        return;
+    }
+
+    if !test_ondemand(&xwls_config.path) {
+        return;
+    }
+    drop(config);
+
+    let x11 = match super::setup_connection() {
+        Ok(x11) => x11,
+        Err(err) => {
+            warn!("error opening X11 sockets, disabling xwayland-satellite integration: {err:?}");
+            return;
+        }
+    };
+
+    let event_loop = &state.niri.event_loop;
+    let (to_main, rx) = calloop::channel::channel();
+    event_loop
+        .insert_source(rx, move |event, _, state| match event {
+            calloop::channel::Event::Msg(msg) => match msg {
+                ToMain::SetupWatch => setup_watch(state),
+            },
+            calloop::channel::Event::Closed => (),
+        })
+        .unwrap();
+
+    state.niri.satellite = Some(Satellite {
+        x11,
+        abstract_token: None,
+        unix_token: None,
+        to_main,
+    });
+
+    setup_watch(state);
+}
+
+fn test_ondemand(path: &str) -> bool {
+    let _span = tracy_client::span!("satellite::test_ondemand");
+
+    // Expand `~` at the start.
+    let mut path = Path::new(path);
+    let expanded = expand_home(path);
+    match &expanded {
+        Ok(Some(expanded)) => path = expanded.as_ref(),
+        Ok(None) => (),
+        Err(err) => {
+            warn!("error expanding ~: {err:?}");
+        }
+    }
+
+    let mut process = Command::new(path);
+    process
+        .args([":0", "--test-listenfd-support"])
+        .stdin(Stdio::null())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .env_remove("DISPLAY")
+        .env_remove("RUST_BACKTRACE")
+        .env_remove("RUST_LIB_BACKTRACE");
+
+    let mut child = match process.spawn() {
+        Ok(child) => child,
+        Err(err) => {
+            info!("error spawning xwayland-satellite at {path:?}, disabling integration: {err}");
+            return false;
+        }
+    };
+
+    let status = match child.wait() {
+        Ok(status) => status,
+        Err(err) => {
+            info!("error waiting for xwayland-satellite, disabling integration: {err}");
+            return false;
+        }
+    };
+
+    if !status.success() {
+        info!("xwayland-satellite doesn't support on-demand activation, disabling integration");
+        return false;
+    }
+
+    true
+}
+
+// When xwayland-satellite fails to start and accept a connection on the socket, the socket will
+// keep triggering our event source, even after the X11 client quits, resulting in a busyloop of
+// trying to start xwayland-satellite. This function will clear out (accept and drop) all pending
+// connections on the socket before registering a new event source, working around this problem.
+// When the problem happens, it's very likely that xwayland-satellite won't be able to accept the
+// pending client (since it had just failed to do so), so it's fine to drop the connections.
+fn clear_out_pending_connections(fd: OwnedFd) -> OwnedFd {
+    let listener = UnixListener::from(fd);
+
+    if let Err(err) = listener.set_nonblocking(true) {
+        warn!("error setting X11 socket to nonblocking: {err:?}");
+        return OwnedFd::from(listener);
+    }
+
+    while listener.accept().is_ok() {}
+
+    if let Err(err) = listener.set_nonblocking(false) {
+        warn!("error setting X11 socket to blocking: {err:?}");
+    }
+
+    OwnedFd::from(listener)
+}
+
+fn setup_watch(state: &mut State) {
+    let Some(satellite) = state.niri.satellite.as_mut() else {
+        return;
+    };
+
+    let event_loop = &state.niri.event_loop;
+
+    if let Some(token) = satellite.abstract_token.take() {
+        error!("abstract_token must be None in setup_watch()");
+        event_loop.remove(token);
+    }
+    if let Some(token) = satellite.unix_token.take() {
+        error!("unix_token must be None in setup_watch()");
+        event_loop.remove(token);
+    }
+
+    let fd = satellite.x11.abstract_fd.try_clone().unwrap();
+    let fd = clear_out_pending_connections(fd);
+    let source = Generic::new(fd, Interest::READ, Mode::Level);
+    let token = event_loop
+        .insert_source(source, move |_, _, state| {
+            if let Some(satellite) = &mut state.niri.satellite {
+                // Remove the other source.
+                if let Some(token) = satellite.unix_token.take() {
+                    state.niri.event_loop.remove(token);
+                }
+                // Clear this source.
+                satellite.abstract_token = None;
+
+                debug!("connection to X11 abstract socket; spawning xwayland-satellite");
+                let path = state.niri.config.borrow().xwayland_satellite.path.clone();
+                spawn(path, satellite);
+            }
+            Ok(PostAction::Remove)
+        })
+        .unwrap();
+    satellite.abstract_token = Some(token);
+
+    let fd = satellite.x11.unix_fd.try_clone().unwrap();
+    let fd = clear_out_pending_connections(fd);
+    let source = Generic::new(fd, Interest::READ, Mode::Level);
+    let token = event_loop
+        .insert_source(source, move |_, _, state| {
+            if let Some(satellite) = &mut state.niri.satellite {
+                // Remove the other source.
+                if let Some(token) = satellite.abstract_token.take() {
+                    state.niri.event_loop.remove(token);
+                }
+                // Clear this source.
+                satellite.unix_token = None;
+
+                debug!("connection to X11 unix socket; spawning xwayland-satellite");
+                let path = state.niri.config.borrow().xwayland_satellite.path.clone();
+                spawn(path, satellite);
+            }
+            Ok(PostAction::Remove)
+        })
+        .unwrap();
+    satellite.unix_token = Some(token);
+}
+
+fn spawn(path: String, xwl: &Satellite) {
+    let _span = tracy_client::span!("satellite::spawn");
+
+    let abstract_fd = xwl.x11.abstract_fd.try_clone().unwrap();
+    let unix_fd = xwl.x11.unix_fd.try_clone().unwrap();
+    let to_main = xwl.to_main.clone();
+
+    // Expand `~` at the start.
+    let mut path = PathBuf::from(path);
+    let expanded = expand_home(&path);
+    match expanded {
+        Ok(Some(expanded)) => path = expanded,
+        Ok(None) => (),
+        Err(err) => {
+            warn!("error expanding ~: {err:?}");
+        }
+    }
+
+    let mut process = Command::new(&path);
+    process.arg(&xwl.x11.display_name).env_remove("DISPLAY");
+
+    // We don't want it spamming the niri output.
+    process
+        .env_remove("RUST_BACKTRACE")
+        .env_remove("RUST_LIB_BACKTRACE");
+    process
+        .stdin(Stdio::null())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null());
+
+    // Spawning and waiting takes some milliseconds, so do it in a thread.
+    let res = thread::Builder::new()
+        .name("Xwl-s Spawner".to_owned())
+        .spawn(move || {
+            spawn_and_wait(&path, process, abstract_fd, unix_fd);
+
+            // Once xwayland-satellite crashes or fails to spawn, re-establish our X11 socket watch
+            // to try again next time.
+            let _ = to_main.send(ToMain::SetupWatch);
+        });
+
+    if let Err(err) = res {
+        warn!("error spawning a thread to spawn xwayland-satellite: {err:?}");
+        let _ = xwl.to_main.send(ToMain::SetupWatch);
+    }
+}
+
+fn spawn_and_wait(path: &Path, mut process: Command, abstract_fd: OwnedFd, unix_fd: OwnedFd) {
+    let abstract_raw = abstract_fd.as_raw_fd();
+    let unix_raw = unix_fd.as_raw_fd();
+
+    process
+        .arg("-listenfd")
+        .arg(abstract_raw.to_string())
+        .arg("-listenfd")
+        .arg(unix_raw.to_string());
+
+    unsafe {
+        process.pre_exec(move || {
+            // We're about to exec xwl-s; perfect time to clear CLOEXEC on the file descriptors
+            // that we want to pass it.
+
+            // We're not dropping these until after spawn().
+            let abstract_fd = BorrowedFd::borrow_raw(abstract_raw);
+            let unix_fd = BorrowedFd::borrow_raw(unix_raw);
+
+            fcntl_setfd(abstract_fd, FdFlags::empty())?;
+            fcntl_setfd(unix_fd, FdFlags::empty())?;
+
+            Ok(())
+        })
+    };
+
+    let mut child = {
+        let _span = tracy_client::span!();
+        match process.spawn() {
+            Ok(child) => child,
+            Err(err) => {
+                warn!("error spawning {path:?}: {err:?}");
+                return;
+            }
+        }
+    };
+
+    // The process spawned, we can drop our fds.
+    drop(abstract_fd);
+    drop(unix_fd);
+
+    let status = match child.wait() {
+        Ok(status) => status,
+        Err(err) => {
+            warn!("error waiting for xwayland-satellite: {err:?}");
+            return;
+        }
+    };
+
+    // This is most likely a crash, hence warn!().
+    warn!("xwayland-satellite exited with: {status}");
+}