Fix race between parent and monitor in use_pty (#1130)

Author: squell

src/exec/event.rs

@@ -91,6 +91,11 @@ impl EventHandle {
             }
         }
     }
+
+    /// Is this event handle ready to be processed?
+    pub(super) fn is_active(&self) -> bool {
+        self.should_poll
+    }
 }
 
 /// The kind of event that will be monitored for a file descriptor.

src/exec/use_pty/backchannel.rs

@@ -199,31 +199,31 @@ impl AsFd for ParentBackchannel {
 /// Different messages exchanged between the monitor and the parent process using a [`ParentBackchannel`].
 #[derive(PartialEq, Eq)]
 pub(super) enum MonitorMessage {
-    ExecCommand,
+    Edge,
     Signal(c_int),
 }
 
 impl MonitorMessage {
     const LEN: usize = PREFIX_LEN + MONITOR_DATA_LEN;
-    const EXEC_CMD: Prefix = 0;
+    const EDGE_CMD: Prefix = 0;
     const SIGNAL: Prefix = 1;
 
     fn from_parts(prefix: Prefix, data: MonitorData) -> Self {
         match prefix {
-            Self::EXEC_CMD => Self::ExecCommand,
+            Self::EDGE_CMD => Self::Edge,
             Self::SIGNAL => Self::Signal(data),
             _ => unreachable!(),
         }
     }
 
     fn to_parts(&self) -> (Prefix, MonitorData) {
         let prefix = match self {
-            MonitorMessage::ExecCommand => Self::EXEC_CMD,
+            MonitorMessage::Edge => Self::EDGE_CMD,
             MonitorMessage::Signal(_) => Self::SIGNAL,
         };
 
         let data = match self {
-            MonitorMessage::ExecCommand => 0,
+            MonitorMessage::Edge => 0,
             MonitorMessage::Signal(data) => *data,
         };
 
@@ -234,7 +234,7 @@ impl MonitorMessage {
 impl std::fmt::Debug for MonitorMessage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::ExecCommand => "ExecCommand".fmt(f),
+            Self::Edge => "Edge".fmt(f),
             &Self::Signal(signal) => write!(f, "Signal({})", signal_fmt(signal)),
         }
     }

src/exec/use_pty/monitor.rs

@@ -75,8 +75,8 @@ pub(super) fn exec_monitor(
         err
     })?;
     // Given that `UnixStream` delivers messages in order it shouldn't be possible to
-    // receive an event different to `ExecCommand` at the beginning.
-    debug_assert_eq!(event, MonitorMessage::ExecCommand);
+    // receive an event different to `Edge` at the beginning.
+    debug_assert_eq!(event, MonitorMessage::Edge);
 
     // FIXME (ogsudo): Some extra config happens here if selinux is available.
 
@@ -186,6 +186,14 @@ pub(super) fn exec_monitor(
         }
     }
 
+    // Wait for the parent to give us red light before shutting down. This avoids missing
+    // output when the monitor exits too quickly.
+    let event = retry_while_interrupted(|| backchannel.recv()).map_err(|err| {
+        dev_warn!("cannot receive red light from parent: {err}");
+        err
+    })?;
+    debug_assert_eq!(event, MonitorMessage::Edge);
+
     // FIXME (ogsudo): The tty is restored here if selinux is available.
 
     // We call `_exit` instead of `exit` to avoid flushing the parent's IO streams by accident.
@@ -271,8 +279,8 @@ impl<'a> MonitorClosure<'a> {
             }
             Ok(event) => {
                 match event {
-                    // We shouldn't receive this event more than once.
-                    MonitorMessage::ExecCommand => unreachable!(),
+                    // We shouldn't receive this event at this point in the protocol
+                    MonitorMessage::Edge => unreachable!(),
                     // Forward signal to the command.
                     MonitorMessage::Signal(signal) => {
                         if let Some(command_pid) = self.command_pid {

src/exec/use_pty/parent.rs

@@ -217,12 +217,10 @@ pub(in crate::exec) fn exec_pty(
     drop(backchannels.monitor);
 
     // Send green light to the monitor after closing the follower.
-    retry_while_interrupted(|| backchannels.parent.send(&MonitorMessage::ExecCommand)).map_err(
-        |err| {
-            dev_error!("cannot send green light to monitor: {err}");
-            err
-        },
-    )?;
+    retry_while_interrupted(|| backchannels.parent.send(&MonitorMessage::Edge)).map_err(|err| {
+        dev_error!("cannot send green light to monitor: {err}");
+        err
+    })?;
 
     let mut closure = ParentClosure::new(
         monitor_pid,
@@ -247,6 +245,7 @@ pub(in crate::exec) fn exec_pty(
     // FIXME (ogsudo): Retry if `/dev/tty` is revoked.
 
     // Flush the terminal
+    closure.tty_pipe.right().set_nonblocking()?;
     closure.tty_pipe.flush_left().ok();
 
     // Restore the terminal settings
@@ -359,10 +358,19 @@ impl ParentClosure {
     }
 
     fn run(&mut self, registry: EventRegistry<Self>) -> io::Result<ExitReason> {
-        match registry.event_loop(self) {
+        let result = match registry.event_loop(self) {
             StopReason::Break(err) | StopReason::Exit(ParentExit::Backchannel(err)) => Err(err),
             StopReason::Exit(ParentExit::Command(exit_reason)) => Ok(exit_reason),
-        }
+        };
+        // Send red light to the monitor after processing all events
+        retry_while_interrupted(|| self.backchannel.send(&MonitorMessage::Edge)).map_err(
+            |err| {
+                dev_error!("cannot send red light to monitor: {err}");
+                err
+            },
+        )?;
+
+        result
     }
 
     /// Read an event from the backchannel and return the event if it should break the event loop.

src/exec/use_pty/pipe/mod.rs

@@ -119,9 +119,27 @@ impl<L: Read + Write + AsFd, R: Read + Write + AsFd> Pipe<L, R> {
         }
     }
 
-    /// Ensure that all the contents of the pipe's internal buffer are written to the left side.
+    /// Flush the pipe, ensuring that all the contents are written to the left side.
     pub(super) fn flush_left(&mut self) -> io::Result<()> {
-        self.buffer_rl.flush(&mut self.left)
+        let buffer = &mut self.buffer_rl;
+        let source = &mut self.right;
+        let sink = &mut self.left;
+
+        // Flush the ring buffer, then process any eventual bytes still in-flight.
+        buffer.internal.remove(sink)?;
+
+        if buffer.write_handle.is_active() {
+            let mut buf = [0u8; RingBuffer::LEN];
+            loop {
+                match source.read(&mut buf) {
+                    Ok(read_bytes) => sink.write_all(&buf[..read_bytes])?,
+                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => break,
+                    Err(e) => return Err(e),
+                }
+            }
+        }
+
+        sink.flush()
     }
 }
 
@@ -198,12 +216,4 @@ impl<R: Read, W: Write> Buffer<R, W> {
         // Return whether we actually freed up some buffer space
         Ok(removed_len > 0)
     }
-
-    /// Flush this buffer, ensuring that all the contents of its internal buffer are written.
-    fn flush(&mut self, write: &mut W) -> io::Result<()> {
-        // Remove bytes from the buffer and write them.
-        self.internal.remove(write)?;
-
-        write.flush()
-    }
 }

src/exec/use_pty/pipe/ring_buffer.rs

@@ -10,7 +10,7 @@ pub(super) struct RingBuffer {
 
 impl RingBuffer {
     /// The size of the internal storage of the ring buffer.
-    const LEN: usize = 8 * 1024;
+    pub(super) const LEN: usize = 8 * 1024;
 
     /// Create a new, empty buffer.
     pub(super) fn new() -> Self {

src/system/term/mod.rs

@@ -97,6 +97,22 @@ impl PtyLeader {
 
         Ok(())
     }
+
+    pub(crate) fn set_nonblocking(&self) -> io::Result<()> {
+        let fd = self.file.as_fd();
+        // SAFETY: these two calls to fcntl are memory safe (and the file descriptor is valid as well)
+        unsafe {
+            let flags = cerr(libc::fcntl(fd.as_raw_fd(), libc::F_GETFL))?;
+
+            // Set the O_NONBLOCK flag
+            cerr(libc::fcntl(
+                fd.as_raw_fd(),
+                libc::F_SETFL,
+                flags | libc::O_NONBLOCK,
+            ))?;
+        }
+        Ok(())
+    }
 }
 
 impl io::Read for PtyLeader {