From eab4adc9917ab5483a8f5b146c368090992a0f5f Mon Sep 17 00:00:00 2001
From: En-En <39373446+En-En-Code@users.noreply.github.com>
Date: Thu, 13 Mar 2025 13:44:41 +0000
Subject: [PATCH] Fix most Broken Pipe IO aborts

Makes sure the xwayland_satellite main thread ends so that the file descriptors don't continue to be used after Fixture drops them
---
 src/lib.rs           | 12 +++++++++++-
 tests/integration.rs | 27 +++++++++++++++++++--------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index a3975eb..8c3982c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -41,6 +41,9 @@ pub trait RunData {
     }
     fn created_server(&self) {}
     fn connected_server(&self) {}
+    fn quit_rx(&self) -> Option<&UnixStream> {
+        None
+    }
     fn xwayland_ready(&self, _display: String, _pid: u32) {}
 }
 
@@ -115,7 +118,6 @@ pub fn main(data: impl RunData) -> Option<()> {
             panic!("first poll failed: {e:?}")
         }
     };
-    drop(finish_rx);
 
     server_state.connect(connection);
     server_state.run();
@@ -127,11 +129,16 @@ pub fn main(data: impl RunData) -> Option<()> {
     let server_fd = unsafe { BorrowedFd::borrow_raw(server_state.clientside_fd().as_raw_fd()) };
     let display_fd = unsafe { BorrowedFd::borrow_raw(display.backend().poll_fd().as_raw_fd()) };
 
+    // `finish_rx` only writes the status code of `Xwayland` exiting, so it is reasonable to use as
+    // the UnixStream of choice when not running the integration tests.
+    let quit_rx = data.quit_rx().unwrap_or(&finish_rx);
+
     let mut fds = [
         PollFd::from_borrowed_fd(server_fd, PollFlags::IN),
         PollFd::new(&xsock_wl, PollFlags::IN),
         PollFd::from_borrowed_fd(display_fd, PollFlags::IN),
         PollFd::new(&ready_rx, PollFlags::IN),
+        PollFd::new(quit_rx, PollFlags::IN),
     ];
 
     let mut ready = false;
@@ -141,6 +148,9 @@ pub fn main(data: impl RunData) -> Option<()> {
                 if !fds[3].revents().is_empty() {
                     ready = true;
                 }
+                if !fds[4].revents().is_empty() {
+                    return None;
+                }
             }
             Err(other) => panic!("Poll failed: {other:?}"),
         }
diff --git a/tests/integration.rs b/tests/integration.rs
index ab0a279..9cad70a 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -1,5 +1,6 @@
 use rustix::event::{poll, PollFd, PollFlags};
 use rustix::process::{Pid, Signal, WaitOptions};
+use std::io::Write;
 use std::mem::ManuallyDrop;
 use std::os::fd::{AsRawFd, BorrowedFd};
 use std::os::unix::net::UnixStream;
@@ -22,15 +23,17 @@ struct TestDataInner {
     display: Mutex<Option<String>>,
     server: Mutex<Option<UnixStream>>,
     pid: Mutex<Option<u32>>,
+    quit_rx: Option<UnixStream>,
 }
 
 #[derive(Default, Clone)]
 struct TestData(Arc<TestDataInner>);
 
 impl TestData {
-    fn new(server: UnixStream) -> Self {
+    fn new(server: UnixStream, quit_rx: UnixStream) -> Self {
         Self(Arc::new(TestDataInner {
             server: Mutex::new(server.into()),
+            quit_rx: Some(quit_rx),
             ..Default::default()
         }))
     }
@@ -52,6 +55,10 @@ impl xwls::RunData for TestData {
         self.server_connected.store(true, Ordering::Relaxed);
     }
 
+    fn quit_rx(&self) -> Option<&UnixStream> {
+        self.quit_rx.as_ref()
+    }
+
     fn xwayland_ready(&self, display: String, pid: u32) {
         *self.display.lock().unwrap() = Some(display);
         *self.pid.lock().unwrap() = Some(pid);
@@ -74,17 +81,18 @@ struct Fixture {
     pollfd: PollFd<'static>,
     display: String,
     pid: Pid,
+    quit_tx: UnixStream,
 }
 
 impl Drop for Fixture {
     fn drop(&mut self) {
         let thread = unsafe { ManuallyDrop::take(&mut self.thread) };
-        if thread.is_finished() {
-            thread.join().expect("Main thread panicked");
-        } else {
-            rustix::process::kill_process(self.pid, Signal::Term).unwrap();
-            rustix::process::waitpid(Some(self.pid), WaitOptions::NOHANG).unwrap();
-        }
+        // Sending anything to the quit receiver to stop the main loop. Then we guarantee a main
+        // thread does not use file descriptors which outlive the Fixture's BorrowedFd
+        self.quit_tx.write_all(&1_usize.to_ne_bytes()).unwrap();
+        thread.join().expect("Main thread panicked");
+        rustix::process::kill_process(self.pid, Signal::Term).unwrap();
+        rustix::process::waitpid(Some(self.pid), WaitOptions::NOHANG).unwrap();
     }
 }
 
@@ -99,11 +107,13 @@ impl Fixture {
                 .init();
         });
 
+        let (quit_tx, quit_rx) = UnixStream::pair().unwrap();
+
         let (a, b) = UnixStream::pair().unwrap();
         let mut testwl = testwl::Server::new(false);
         pre_connect(&mut testwl);
         testwl.connect(a);
-        let our_data = TestData::new(b);
+        let our_data = TestData::new(b, quit_rx);
         let data = our_data.clone();
         let thread = std::thread::spawn(move || xwls::main(data));
 
@@ -155,6 +165,7 @@ impl Fixture {
             pollfd,
             display,
             pid: Pid::from_raw(pid as _).expect("Xwayland PID was invalid?"),
+            quit_tx,
         }
     }
     fn new() -> Self {