Skip to content

Fix stdin relay for pipes in IsoSession #1450

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Dom Giandinoto (daamenik) wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Fix stdin relay for pipes in IsoSession #1450

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 97 additions & 44 deletions src/mxc-sdk/src/backends/isolation_session/common/manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,11 @@ use crate::isolation_session_bindings::bindings::{
IsoSessionUserResult,
};
use windows::Win32::Foundation::{HANDLE, WAIT_OBJECT_0};
use windows::Win32::Storage::FileSystem::{GetFileType, FILE_TYPE_CHAR};
use windows::Win32::System::Com::{CoDecrementMTAUsage, CoIncrementMTAUsage, CO_MTA_USAGE_COOKIE};
use windows::Win32::System::Console::{
GetStdHandle, STD_ERROR_HANDLE, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE,
GetConsoleMode, GetStdHandle, CONSOLE_MODE, STD_ERROR_HANDLE, STD_INPUT_HANDLE,
STD_OUTPUT_HANDLE,
};
use windows::Win32::System::Threading::{CreateEventW, SetEvent, WaitForSingleObject};
use windows::UI::WindowId;
Expand All @@ -42,6 +44,23 @@ use super::process_options::{build_iso_process_options, ProcessOptions};
// The OS account name is lifecycle addressing data, not diagnostic identity.
const AGENT_IDENTITY_MARKER: &str = "isolation-session-agent";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StdinHandleKind {
Console,
PipeOrFile,
}

fn classify_stdin_handle(handle: HANDLE) -> StdinHandleKind {
let file_type = unsafe { GetFileType(handle) };
if file_type == FILE_TYPE_CHAR {
let mut mode = CONSOLE_MODE(0);
if unsafe { GetConsoleMode(handle, &mut mode) }.is_ok() {
return StdinHandleKind::Console;
}
}
StdinHandleKind::PipeOrFile
}

/// Keeps the process's MTA alive for as long as the lifecycle's WinRT objects,
/// which outlive the call that created them.
/// `CoIncrementMTAUsage` holds a reference releasable from any thread;
Expand Down Expand Up @@ -418,10 +437,7 @@ impl IsolationSessionManager {
let wxc_stdin = unsafe { GetStdHandle(STD_INPUT_HANDLE) }
.map_err(|e| lifecycle_err(format!("GetStdHandle(stdin) failed: {}", e)))?;

// Manual-reset stop event for the stdin relay. Effective for a waitable
// `h_read` (console = TTY mode); for pipe handles it has no effect on a
// blocked `ReadFile`, so that relay misses its join and ends only when
// the calling process exits. Only `wxc-exec` reaches the non-TTY case.
// Manual-reset stop event for console-backed stdin relays.
let stdin_stop_event = unsafe {
CreateEventW(None, true, false, PCWSTR::null())
.map_err(|e| lifecycle_err(format!("CreateEventW(stdin stop): {}", e)))?
Expand Down Expand Up @@ -516,58 +532,81 @@ impl IsolationSessionManager {
} else {
None
};
// Stdin: in interactive mode use the console-aware relay so
// `WINDOW_BUFFER_SIZE_EVENT` records propagate as
// `ResizeConsole(cols, rows)` calls on the agent's inner ConPTY.
// In non-interactive mode the agent's stdin is plain byte-oriented
// and the simpler stop-aware pipe relay is appropriate.
// Stdin relay selection is based on the actual local stdin handle,
// not on `options.interactive` (which is derived from stdout). A
// pipe/file read must be cancelled with `PipeReadCanceller`; a
// console handle can be interrupted by the stop event and, when the
// agent side is interactive, needs the console-aware relay for
// resize records.
enum StdinRelayKind {
None,
CancellablePipe,
Pipe(PipeRelayWithStopParams),
Console(ConsoleRelayParams),
}

let stdin_relay_state = if stdin_handle_val == 0 {
StdinRelayKind::None
} else {
// Owned, like the output relays': `process.Close()` releases the
// original while this relay may still be writing.
let stdin_h_write =
duplicate_handle(HANDLE(stdin_handle_val as *mut core::ffi::c_void))
.map_err(|e| lifecycle_err(format!("duplicate stdin handle: {}", e)))?;
// Shares the event object, so `signal_stop` still reaches the relay.
let stdin_h_stop = duplicate_handle(stdin_stop_owned.get())
.map_err(|e| lifecycle_err(format!("duplicate stdin stop event: {}", e)))?;
if options.interactive {
// Clone the WinRT process handle so the relay thread holds
// its own ref-counted reference (WinRT clone = AddRef),
// released when the thread frees the params it owns.
let process_for_resize = process.clone();
let impersonation = self.impersonation.clone();
StdinRelayKind::Console(ConsoleRelayParams {
h_read: wxc_stdin,
h_write: stdin_h_write,
h_stop_event: stdin_h_stop,
resize_callback: Box::new(move |cols, rows| {
let _ = owned_thread::call(&impersonation, || {
let _ = process_for_resize.ResizeConsole(cols, rows);
Ok(())
});
}),
})
} else {
StdinRelayKind::Pipe(PipeRelayWithStopParams {
h_read: wxc_stdin,
h_write: stdin_h_write,
h_stop_event: stdin_h_stop,
})
match classify_stdin_handle(wxc_stdin) {
StdinHandleKind::PipeOrFile => StdinRelayKind::CancellablePipe,
StdinHandleKind::Console => {
// Owned, like the output relays': `process.Close()`
// releases the original while this relay may still be
// writing.
let stdin_h_write =
duplicate_handle(HANDLE(stdin_handle_val as *mut core::ffi::c_void))
.map_err(|e| {
lifecycle_err(format!("duplicate stdin handle: {}", e))
})?;
// Shares the event object, so `signal_stop` still reaches the relay.
let stdin_h_stop =
duplicate_handle(stdin_stop_owned.get()).map_err(|e| {
lifecycle_err(format!("duplicate stdin stop event: {}", e))
})?;
if options.interactive {
// Clone the WinRT process handle so the relay thread holds
// its own ref-counted reference (WinRT clone = AddRef),
// released when the thread frees the params it owns.
let process_for_resize = process.clone();
let impersonation = self.impersonation.clone();
StdinRelayKind::Console(ConsoleRelayParams {
h_read: wxc_stdin,
h_write: stdin_h_write,
h_stop_event: stdin_h_stop,
resize_callback: Box::new(move |cols, rows| {
let _ = owned_thread::call(&impersonation, || {
let _ = process_for_resize.ResizeConsole(cols, rows);
Ok(())
});
}),
})
} else {
StdinRelayKind::Pipe(PipeRelayWithStopParams {
h_read: wxc_stdin,
h_write: stdin_h_write,
h_stop_event: stdin_h_stop,
})
}
}
}
};

// Matched by value: the params move into the spawn call, which puts
// them on the heap under the relay thread's ownership.
scope.stdin = match stdin_relay_state {
StdinRelayKind::None => None,
StdinRelayKind::CancellablePipe => {
let (thread, canceller) = unsafe {
create_relay_thread(
wxc_stdin,
HANDLE(stdin_handle_val as *mut core::ffi::c_void),
)
}
Comment on lines +599 to +605
.map_err(|e| lifecycle_err(format!("create stdin relay: {}", e)))?;
scope.stdin_canceller = Some(canceller);
Some(thread)
}
StdinRelayKind::Pipe(params) => Some(
unsafe { create_relay_thread_with_stop(params) }
.map_err(|e| lifecycle_err(format!("create stdin relay: {}", e)))?,
Expand Down Expand Up @@ -1035,6 +1074,10 @@ struct RelayScope<'a> {
/// ends keeps them open past the workload's exit, so EOF alone is not a
/// bound.
output_cancellers: Vec<PipeReadCanceller>,
/// Ends a pipe/file-backed stdin relay. Console stdin relays use the stop
/// event instead because `CancelSynchronousIo` has console-handle edge
/// cases.
stdin_canceller: Option<PipeReadCanceller>,
finished: bool,
}

Expand All @@ -1051,6 +1094,7 @@ impl<'a> RelayScope<'a> {
stderr: None,
stdin: None,
output_cancellers: Vec::new(),
stdin_canceller: None,
finished: false,
}
}
Expand All @@ -1068,6 +1112,12 @@ impl<'a> RelayScope<'a> {
}
}

fn cancel_stdin_read(&self) {
if let Some(canceller) = &self.stdin_canceller {
canceller.close();
}
}

/// Whether the relay exited. An absent relay was never spawned.
fn join(relay: &Option<OwnedHandle>, timeout_ms: u32) -> bool {
match relay {
Expand All @@ -1085,12 +1135,14 @@ impl<'a> RelayScope<'a> {

fn stop_stdin(&self) {
self.signal_stop();
self.cancel_stdin_read();
Self::join(&self.stdin, Self::JOIN_MS);
}

/// Teardown for a workload that has already exited.
fn finish(&mut self) {
self.signal_stop();
self.cancel_stdin_read();
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) {
Expand All @@ -1110,6 +1162,7 @@ impl Drop for RelayScope<'_> {
return;
}
self.signal_stop();
self.cancel_stdin_read();
self.cancel_output_reads();
let _ = self.process.Terminate();
Self::join(&self.stdout, Self::JOIN_MS);
Expand Down Expand Up @@ -1296,9 +1349,9 @@ fn plan_wait(
}

/// Three-tier graceful shutdown for an `IsoSessionProcess` that's still
/// running after `WaitForExit(timeout_ms)` returns. Tier 1: close stdin —
/// many REPLs exit on EOF alone. Tier 2: `SendCtrlClose` — ConPTY-only;
/// `E_NOTIMPL` outside ConPTY, benign. Tier 3: force-terminate, wait
/// running after `WaitForExit(timeout_ms)` returns. Tier 1 asks the platform
/// to close stdin, tier 2 tries `SendCtrlClose` (ConPTY-only; `E_NOTIMPL`
/// outside ConPTY is benign), and tier 3 force-terminates, then waits
/// infinitely (`WaitForExit(0)` = INFINITE) for the kill to land.
///
/// The first `ExitCode()` query is `?`-propagated: a failure there means
Expand Down
53 changes: 43 additions & 10 deletions src/mxc-sdk/src/backends/isolation_session/common/pipe_relay.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,12 @@
//! desktop-session boundary.
//!
//! Two relay variants:
//! - [`create_relay_thread`] — cancellable; used for stdout / stderr. Ends at
//! EOF or when its [`PipeReadCanceller`] fires.
//! - [`create_relay_thread_with_stop`] — stop-event-aware; used for stdin
//! in TTY (ConPTY) mode where the agent can exit while the local stdin
//! handle remains open.
//! - [`create_relay_thread`] — cancellable; used for stdout / stderr and for
//! pipe/file-backed stdin. Ends at EOF or when its [`PipeReadCanceller`]
//! fires.
//! - [`create_relay_thread_with_stop`] — stop-event-aware; used for
//! console-backed stdin when the agent can exit while the local stdin handle
//! remains open.

use std::io::Read;

Expand Down Expand Up @@ -132,10 +133,9 @@ pub(super) unsafe fn create_relay_thread(

// ── Stop-event-aware pipe relay ────────────────────────────────────────────
//
// Used for the IsolationSession stdin relay in TTY (ConPTY) mode, where the
// agent process can exit naturally while wxc-exec's stdin remains open (held
// by the parent: node-pty, shell, etc.). Without an external signal the relay
// sits in `ReadFile` forever.
// Used for the IsolationSession stdin relay when wxc-exec's stdin is a console
// handle and the agent process can exit naturally while that handle remains
// open. Without an external signal the relay sits in `ReadFile` forever.
//
// `CancelSynchronousIo` is the obvious alternative but has documented edge
// cases on console handles, and wxc-exec's stdin is a console handle in the
Expand All @@ -149,7 +149,7 @@ pub(super) unsafe fn create_relay_thread(
// signalled when open", so the wait returns immediately, the relay enters
// `ReadFile`, and from there the stop event cannot interrupt it. For
// pipe-backed stdin (the non-TTY case), use the simpler EOF-driven
// `create_relay_thread` and rely on natural EOF or process exit for cleanup.
// `create_relay_thread` and rely on natural EOF or its canceller for cleanup.

/// Parameters for a stop-event-aware relay thread. The thread loops
/// `WaitForMultipleObjects({h_stop_event, h_read})`; copies a chunk when
Expand Down Expand Up @@ -407,6 +407,39 @@ mod tests {
);
}

/// Pipe-backed stdin uses `create_relay_thread` so shutdown can cancel the
/// blocked source read. Once cancelled, the relay must also release its
/// duplicate of the agent stdin write handle so the reader observes EOF.
#[test]
fn a_cancelled_stdin_pipe_relay_releases_its_destination_handle() {
let (source_read, _source_write) = create_test_pipes();
let (mut dest_read, dest_write) = create_test_pipes();

let (relay_thread, canceller) =
unsafe { create_relay_thread(source_read.get(), dest_write.get()).unwrap() };

// Mirror `process.Close()`: the caller drops its original agent-stdin
// write handle while the relay's duplicate is still live.
drop(dest_write);

assert_eq!(
unsafe { WaitForSingleObject(relay_thread.get(), 250) },
WAIT_TIMEOUT,
"the relay should still be blocked before the cancel"
);

canceller.close();

assert_eq!(
unsafe { WaitForSingleObject(relay_thread.get(), 5000) },
WAIT_OBJECT_0,
"the cancelled stdin relay should exit promptly"
);

let dest_send = SendOwnedHandle::take(&mut dest_read);
assert_eq!(read_from_pipe(dest_send.get()), "");
}

/// Closing the source handle the relay was built from does not disturb it:
/// it reads a duplicate. This is what lets `IsoSessionProcess::Close` run
/// unconditionally.
Expand Down
Loading