Guard more effectively against stdout truncation

docs/development/architecture/backends/isolation-session/state-aware-rust.md

@@ -421,6 +421,22 @@ waiter, so the caller decides when to block. Either way, two concurrent exec
 calls against the same `sandboxId` are not coordinated by MXC; the OS-side
 service serialises (or rejects, depending on session state) at its own layer.
 
+### Output draining after relayed execution
+
+After waiting for the workload and stopping stdin forwarding, `create_process`
+copies remaining stdout and stderr until both finish, neither successfully
+writes output for five seconds, or 60 seconds pass, whichever comes first.
+These limits do not apply to SDK `Piped` execution.
+
+On timeout, MXC cancels reads and briefly waits for the copying threads to exit.
+Cancellation can discard data already read, but cannot interrupt blocked writes
+or flushes. Those threads may keep pipe handles open and finish writing after
+cleanup returns. Descendant processes and stdin handling are unchanged.
+
+Reaching the drain limit does not change the workload's exit code or emit a
+truncation warning. The drain fix reduces premature cancellation of progressing
+output; it does not guarantee complete output when a limit is reached.
+
 ### Deprovision and concurrent sandboxes
 
 `deprovision` removes only its own agent user (`deprovision_agent_user`).

src/mxc-sdk/src/backends/isolation_session/common/manager.rs

@@ -7,11 +7,11 @@
 //! against the local console.
 
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::mxc_common::audit::{AuditEvent, AuditEventName, KillMethod, TeardownStatus};
 use crate::mxc_common::logger::Logger;
-use crate::mxc_common::process_util::{OwnedHandle, PipeReadCanceller};
-use crate::mxc_common::sandbox_process::StreamCloser;
+use crate::mxc_common::process_util::OwnedHandle;
 use crate::mxc_common::state_aware_backend::{ExecHandle, ExecOutcome};
 
 use crate::isolation_session_bindings::bindings::{
@@ -35,7 +35,8 @@ use super::error::{
 };
 use super::owned_thread::{self, Impersonation};
 use super::pipe_relay::{
-    create_relay_thread, create_relay_thread_with_stop, duplicate_handle, PipeRelayWithStopParams,
+    create_relay_thread, create_relay_thread_with_stop, drain_output_relays_within,
+    duplicate_handle, OutputRelayDrain, PipeRelay, PipeRelayWithStopParams,
 };
 use super::process_options::{build_iso_process_options, ProcessOptions};
 
@@ -491,28 +492,28 @@ impl IsolationSessionManager {
             }
 
             scope.stdout = if stdout_handle_val != 0 {
-                let (thread, canceller) = unsafe {
-                    create_relay_thread(
-                        HANDLE(stdout_handle_val as *mut core::ffi::c_void),
-                        wxc_stdout,
-                    )
-                }
-                .map_err(|e| lifecycle_err(format!("create stdout relay: {}", e)))?;
-                scope.output_cancellers.push(canceller);
-                Some(thread)
+                Some(
+                    unsafe {
+                        create_relay_thread(
+                            HANDLE(stdout_handle_val as *mut core::ffi::c_void),
+                            wxc_stdout,
+                        )
+                    }
+                    .map_err(|e| lifecycle_err(format!("create stdout relay: {}", e)))?,
+                )
             } else {
                 None
             };
             scope.stderr = if stderr_handle_val != 0 {
-                let (thread, canceller) = unsafe {
-                    create_relay_thread(
-                        HANDLE(stderr_handle_val as *mut core::ffi::c_void),
-                        wxc_stderr,
-                    )
-                }
-                .map_err(|e| lifecycle_err(format!("create stderr relay: {}", e)))?;
-                scope.output_cancellers.push(canceller);
-                Some(thread)
+                Some(
+                    unsafe {
+                        create_relay_thread(
+                            HANDLE(stderr_handle_val as *mut core::ffi::c_void),
+                            wxc_stderr,
+                        )
+                    }
+                    .map_err(|e| lifecycle_err(format!("create stderr relay: {}", e)))?,
+                )
             } else {
                 None
             };
@@ -581,7 +582,7 @@ impl IsolationSessionManager {
 
             let outcome = started.wait(options.timeout_ms)?;
             scope.stop_stdin();
-            scope.finish();
+            scope.finish()?;
 
             Ok(match outcome {
                 ExecOutcome::Exited(exit_code) => exit_code,
@@ -1022,26 +1023,23 @@ impl StartedProcess {
 /// Owns the relay threads and the workload for the span of `create_process`.
 ///
 /// [`finish`](Self::finish) tears down a workload that has already exited.
-/// Dropping without it abandons the run: the workload is terminated so the
-/// output relays reach EOF. Every join is bounded, so a relay whose read never
-/// returns cannot wedge the caller's thread.
+/// Dropping without it cancels output reads and terminates the workload.
+/// Every join is bounded; cancellation does not interrupt destination writes.
 struct RelayScope<'a> {
     process: &'a IsoSessionProcess,
     stop_event: HANDLE,
-    stdout: Option<OwnedHandle>,
-    stderr: Option<OwnedHandle>,
+    stdout: Option<PipeRelay>,
+    stderr: Option<PipeRelay>,
     stdin: Option<OwnedHandle>,
-    /// Ends the output relays' reads. A descendant that inherited the write
-    /// ends keeps them open past the workload's exit, so EOF alone is not a
-    /// bound.
-    output_cancellers: Vec<PipeReadCanceller>,
     finished: bool,
 }
 
 impl<'a> RelayScope<'a> {
     const JOIN_MS: u32 = 1000;
-    /// Expiry cancels the read, dropping an in-flight chunk.
-    const DRAIN_MS: u32 = 5000;
+    /// Stop waiting only after all output relays are idle for this long.
+    const IDLE_MS: u64 = 5000;
+    /// Overall cap so a pathological one-byte-forever trickle still returns.
+    const ABSOLUTE_CAP_MS: u64 = 60_000;
 
     fn new(process: &'a IsoSessionProcess, stop_event: HANDLE) -> Self {
         Self {
@@ -1050,7 +1048,6 @@ impl<'a> RelayScope<'a> {
             stdout: None,
             stderr: None,
             stdin: None,
-            output_cancellers: Vec::new(),
             finished: false,
         }
     }
@@ -1062,45 +1059,55 @@ impl<'a> RelayScope<'a> {
         }
     }
 
-    fn cancel_output_reads(&self) {
-        for canceller in &self.output_cancellers {
-            canceller.close();
-        }
-    }
-
     /// Whether the relay exited. An absent relay was never spawned.
-    fn join(relay: &Option<OwnedHandle>, timeout_ms: u32) -> bool {
+    fn join(relay: Option<&OwnedHandle>, timeout_ms: u32) -> bool {
         match relay {
             // SAFETY: an owned thread handle held by this scope.
             Some(t) => (unsafe { WaitForSingleObject(t.get(), timeout_ms) }) == WAIT_OBJECT_0,
             None => true,
         }
     }
 
-    /// Each relay reads its own duplicate, so one still running does not keep
-    /// the pipe handles alive.
+    fn join_output(relay: &Option<PipeRelay>, timeout_ms: u32) -> bool {
+        Self::join(relay.as_ref().map(PipeRelay::thread), timeout_ms)
+    }
+
+    fn output_relays(&self) -> Vec<&PipeRelay> {
+        [&self.stdout, &self.stderr]
+            .into_iter()
+            .filter_map(Option::as_ref)
+            .collect()
+    }
+
+    /// Release the process's handles. A still-running relay owns its duplicates
+    /// until it exits.
     fn close_handles(&self) {
         let _ = self.process.Close();
     }
 
     fn stop_stdin(&self) {
         self.signal_stop();
-        Self::join(&self.stdin, Self::JOIN_MS);
+        Self::join(self.stdin.as_ref(), Self::JOIN_MS);
     }
 
     /// Teardown for a workload that has already exited.
-    fn finish(&mut self) {
+    fn finish(&mut self) -> Result<(), IsolationSessionError> {
         self.signal_stop();
-        let stdout_drained = Self::join(&self.stdout, Self::DRAIN_MS);
-        let stderr_drained = Self::join(&self.stderr, Self::DRAIN_MS);
-        if !(stdout_drained && stderr_drained) {
-            self.cancel_output_reads();
-            Self::join(&self.stdout, Self::JOIN_MS);
-            Self::join(&self.stderr, Self::JOIN_MS);
+        let output_relays = self.output_relays();
+        let output_drain = drain_output_relays_within(
+            &output_relays,
+            Duration::from_millis(Self::IDLE_MS),
+            Duration::from_millis(Self::ABSOLUTE_CAP_MS),
+        )
+        .map_err(|error| lifecycle_err(format!("drain output relays: {error}")))?;
+        if output_drain == OutputRelayDrain::GaveUp {
+            Self::join_output(&self.stdout, Self::JOIN_MS);
+            Self::join_output(&self.stderr, Self::JOIN_MS);
         }
-        Self::join(&self.stdin, Self::JOIN_MS);
+        Self::join(self.stdin.as_ref(), Self::JOIN_MS);
         self.close_handles();
         self.finished = true;
+        Ok(())
     }
 }
 
@@ -1110,11 +1117,13 @@ impl Drop for RelayScope<'_> {
             return;
         }
         self.signal_stop();
-        self.cancel_output_reads();
+        for relay in self.output_relays() {
+            relay.cancel_read();
+        }
         let _ = self.process.Terminate();
-        Self::join(&self.stdout, Self::JOIN_MS);
-        Self::join(&self.stderr, Self::JOIN_MS);
-        Self::join(&self.stdin, Self::JOIN_MS);
+        Self::join_output(&self.stdout, Self::JOIN_MS);
+        Self::join_output(&self.stderr, Self::JOIN_MS);
+        Self::join(self.stdin.as_ref(), Self::JOIN_MS);
         self.close_handles();
     }
 }