From 1f4ca74cf9114beee462e1b3aa7a239b41cfc8ec Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Sun, 9 Mar 2025 16:39:46 -0400
Subject: [PATCH] [DNM] Write to sockets in order, instead of out of order

---
 src/io_util.rs | 22 ++++++++++++++++++++++
 src/main.rs    | 27 ++++++++++++---------------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/src/io_util.rs b/src/io_util.rs
index da1aa5b..c4455b4 100644
--- a/src/io_util.rs
+++ b/src/io_util.rs
@@ -145,6 +145,28 @@ impl<'a> WlMsgWriter<'a> {
         self.write_queue.push(msg);
     }
 
+    pub async fn write(&mut self, msg: WlRawMsg) -> io::Result<()> {
+        let (buf, fds) = msg.into_parts();
+
+        let mut written = 0usize;
+
+        while written < buf.len() {
+            self.egress.writable().await?;
+
+            let res = self
+                .egress
+                .send_with_fd(&buf[written..], unsafe { std::mem::transmute(fds.deref()) });
+
+            match res {
+                Ok(new_written) => written += new_written,
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => return Err(e),
+            }
+        }
+
+        Ok(())
+    }
+
     /// Try to make progress by flushing some of the queued up messages into the stream.
     /// When this resolves, note that we might have only partially written. In that
     /// case the buffer is saved internally in this structure.
diff --git a/src/main.rs b/src/main.rs
index 95da1e7..0cddd42 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -141,7 +141,7 @@ impl<'a> ConnDuplex<'a> {
 
                 match verdict {
                     WlMitmVerdict::Allowed => {
-                        self.downstream_write.queue_write(wl_raw_msg);
+                        self.downstream_write.write(wl_raw_msg).await?;
                     }
                     WlMitmVerdict::Terminate => {
                         return Err(io::Error::new(
@@ -180,18 +180,20 @@ impl<'a> ConnDuplex<'a> {
 
                 match verdict {
                     WlMitmVerdict::Allowed => {
-                        self.upstream_write.queue_write(wl_raw_msg);
+                        self.upstream_write.write(wl_raw_msg).await?;
                     }
                     WlMitmVerdict::Rejected(error_code) => {
-                        self.downstream_write.queue_write(
-                            WlDisplayErrorEvent::new(
-                                WL_DISPLAY_OBJECT_ID,
-                                wl_raw_msg.obj_id,
-                                error_code,
-                                "Rejected by wl-mitm",
+                        self.downstream_write
+                            .write(
+                                WlDisplayErrorEvent::new(
+                                    WL_DISPLAY_OBJECT_ID,
+                                    wl_raw_msg.obj_id,
+                                    error_code,
+                                    "Rejected by wl-mitm",
+                                )
+                                .build(),
                             )
-                            .build(),
-                        );
+                            .await?;
                     }
                     WlMitmVerdict::Terminate => {
                         return Err(io::Error::new(
@@ -213,11 +215,6 @@ impl<'a> ConnDuplex<'a> {
     pub async fn run_to_completion(mut self) -> io::Result<()> {
         loop {
             tokio::select! {
-                biased;
-
-                res = self.downstream_write.dequeue_write() => res?,
-                res = self.upstream_write.dequeue_write() => res?,
-
                 msg = self.upstream_read.read() => {
                     control_flow!(self.handle_s2c_event(msg?).await?);
                 }