[WSLC] Harden state-aware daemon (PR 2a/3): typed errors, validate-then-admit exec, idle-watchdog

src/backends/wslc/common/src/wsl_container_runner.rs

@@ -29,6 +29,7 @@ use wxc_common::sandbox_process::StdioMode;
 use wxc_common::script_runner::ScriptRunner;
 use wxc_common::string_util::{to_wide, CoTaskMemPWSTR};
 
+use crate::container_steps::sdk_error;
 use crate::policy_mapping;
 use crate::stream_buffer::{stream_pair, StreamReader, StreamWriter};
 use crate::wslc_bindings::*;
@@ -593,16 +594,6 @@ enum TarFormat {
     Unknown,
 }
 
-/// Create a ScriptResponse error from an HRESULT failure with optional SDK error message.
-fn sdk_error(context: &str, hr: HRESULT, sdk_msg: &str) -> ScriptResponse {
-    let msg = if sdk_msg.is_empty() {
-        format!("{}: HRESULT 0x{:08X}", context, hr as u32)
-    } else {
-        format!("{}: {} (HRESULT 0x{:08X})", context, sdk_msg, hr as u32)
-    };
-    ScriptResponse::error(&msg)
-}
-
 /// Builds a user-facing prerequisite error for the components `WslcGetMissingComponents`
 /// reports as missing. `missing` may combine multiple bits, and the guidance is branched
 /// per-component so a user missing only `VirtualMachinePlatform` isn't told to update WSL

src/backends/wslc/daemon/src/control_server.rs

@@ -24,9 +24,9 @@ use std::time::Duration;
 use anyhow::{bail, Context, Result};
 use serde::de::DeserializeOwned;
 use serde::Serialize;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use tokio::net::windows::named_pipe::{NamedPipeServer, ServerOptions};
-use tokio::sync::{Notify, Semaphore};
+use tokio::sync::{oneshot, Semaphore};
 use tokio::task::JoinSet;
 use tokio::time::timeout;
 use windows::core::{PCWSTR, PWSTR};
@@ -41,10 +41,10 @@ use windows::Win32::Security::{
 use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
 
 use wslc_common::daemon_protocol::{
-    encode_frame, DaemonRequest, DaemonResponse, ErrKind, StreamFrame, MAX_FRAME_SIZE,
+    encode_frame, DaemonRequest, DaemonResponse, StreamFrame, MAX_FRAME_SIZE,
 };
 
-use crate::session_manager::SessionHandle;
+use crate::session_manager::{SessionHandle, WorkerError};
 
 /// Upper bound on concurrently-serviced client connections. At capacity the
 /// accept loop applies backpressure (a new connection waits for a slot) instead
@@ -85,15 +85,22 @@ pub fn bind(pipe_name: &str) -> Result<(OwnerOnlySecurity, NamedPipeServer)> {
 /// `active_clients` tracks in-flight requests so the idle watchdog does not tear
 /// the daemon down mid-request. Concurrency is bounded by a semaphore, and on
 /// shutdown all in-flight handlers are drained before returning so the caller
-/// can release the WSLc session without racing a live handler.
+/// can release the WSLc session without racing a live handler. `activity` is a
+/// monotonic connection counter the watchdog compares across polls to catch
+/// bursts that start and finish between two of its samples.
 pub async fn run(
     session: SessionHandle,
     pipe_name: String,
     security: OwnerOnlySecurity,
     first_instance: NamedPipeServer,
-    active_clients: Arc<AtomicUsize>,
-    shutdown: Arc<Notify>,
+    signals: crate::DaemonSignals,
 ) -> Result<()> {
+    let crate::DaemonSignals {
+        active_clients,
+        activity,
+        shutdown,
+        draining,
+    } = signals;
     let mut server = first_instance;
     let limiter = Arc::new(Semaphore::new(MAX_CONCURRENT_CLIENTS));
     let mut clients: JoinSet<()> = JoinSet::new();
@@ -106,6 +113,11 @@ pub async fn run(
 
     loop {
         tokio::select! {
+            // Bias toward shutdown: once the watchdog signals, prefer tearing
+            // down over accepting a connection that raced the final idle sample.
+            biased;
+
+            _ = shutdown.notified() => break,
             connect = server.connect() => {
                 if let Err(e) = connect {
                     eprintln!("[wslc-daemon] pipe connect error: {e}");
@@ -118,7 +130,19 @@ pub async fn run(
                     }
                     continue;
                 }
+                // The watchdog may have entered the draining state after its
+                // final sample but before this connection arrived. Refuse it
+                // rather than provision into a session that is about to be
+                // released, handing the client an ID that teardown invalidates.
+                // The client re-spawns a fresh daemon once our record is gone.
+                if draining.load(Ordering::SeqCst) {
+                    break;
+                }
                 let connected = server;
+                // Record the connection so an idle streak that spans this
+                // request is invalidated even if it completes between polls.
+                activity.fetch_add(1, Ordering::SeqCst);
+
                 // Recreate the next listening instance before servicing this one
                 // so the next client is not refused. Transient failures are
                 // retried with bounded backoff rather than tearing the whole
@@ -147,7 +171,6 @@ pub async fn run(
             }
             // Reap finished handlers so the JoinSet does not accumulate.
             Some(_) = clients.join_next() => {}
-            _ = shutdown.notified() => break,
         }
     }
 
@@ -361,7 +384,7 @@ async fn handle_client(mut pipe: NamedPipeServer, session: SessionHandle) -> Res
         DaemonRequest::Provision(config) => {
             let resp = match session.provision(config).await {
                 Ok(sandbox_id) => DaemonResponse::Provisioned { sandbox_id },
-                Err(e) => err_response(ErrKind::Backend, e),
+                Err(e) => worker_err_response(e),
             };
             write_frame(&mut pipe, &resp).await?;
         }
@@ -384,7 +407,15 @@ async fn handle_client(mut pipe: NamedPipeServer, session: SessionHandle) -> Res
     Ok(())
 }
 
-/// Exec: admit with `Ok`, then stream the run's outcome as [`StreamFrame`]s.
+/// Exec: validate-then-admit, then stream the run's outcome as [`StreamFrame`]s.
+///
+/// The sandbox is validated (exists + started) *before* the `Ok` admission is
+/// written, and — critically — admission is **atomic** with the start of the
+/// run on the worker thread (see [`SessionHandle::exec`]): the worker validates
+/// and begins running within one command handler, so no `Stop`/`Deprovision`
+/// can invalidate the checked state between the admission and the run. An
+/// unknown/not-started sandbox therefore comes back as a pre-admission typed
+/// [`DaemonResponse::Err`] rather than a post-admission stream `Error` frame.
 ///
 /// TODO(fill-in): bidirectional live stdio (client `Stdin` frames -> process,
 /// process stdout/stderr -> `Stdout`/`Stderr` frames). The skeleton runs the
@@ -394,34 +425,60 @@ async fn handle_exec(
     session: SessionHandle,
     config: wslc_common::daemon_protocol::ExecConfig,
 ) -> Result<()> {
-    write_frame(&mut pipe, &DaemonResponse::Ok).await?;
-    let terminal = match session.exec(config).await {
-        Ok(code) => StreamFrame::Exit { code },
-        Err(e) => StreamFrame::Error {
-            message: format!("{e:#}"),
+    // Await the worker's admission decision before writing anything: a rejected
+    // exec is a pre-admission typed error, never a post-admission stream frame.
+    write_exec_result(&mut pipe, session.exec(config).await).await
+}
+
+/// Turn an exec **admission** outcome into the client's frame sequence, generic
+/// over the transport so the protocol can be exercised over an in-memory duplex
+/// in tests. On rejection it writes a single typed [`DaemonResponse::Err`]; on
+/// admission it writes `Ok` then awaits completion and writes exactly one
+/// terminal [`StreamFrame`] — `Exit` on success, `Error` on a run failure or a
+/// dropped completion channel.
+async fn write_exec_result<S: AsyncWrite + Unpin>(
+    pipe: &mut S,
+    admission: Result<oneshot::Receiver<Result<i32, WorkerError>>, WorkerError>,
+) -> Result<()> {
+    let done = match admission {
+        Ok(done) => done,
+        Err(e) => {
+            write_frame(pipe, &worker_err_response(e)).await?;
+            return Ok(());
+        }
+    };
+    write_frame(pipe, &DaemonResponse::Ok).await?;
+    let terminal = match done.await {
+        Ok(Ok(code)) => StreamFrame::Exit { code },
+        Ok(Err(e)) => StreamFrame::Error {
+            message: e.to_string(),
+        },
+        Err(_) => StreamFrame::Error {
+            message: "WSLc worker dropped the exec reply channel".to_string(),
         },
     };
-    write_frame(&mut pipe, &terminal).await?;
+    write_frame(pipe, &terminal).await?;
     Ok(())
 }
 
-/// Map a `Result<()>` to `Ok` / `Err` response.
-fn ok_or_err(result: Result<()>) -> DaemonResponse {
+/// Map a worker `Result<()>` to an `Ok` / typed `Err` response.
+fn ok_or_err(result: Result<(), WorkerError>) -> DaemonResponse {
     match result {
         Ok(()) => DaemonResponse::Ok,
-        Err(e) => err_response(ErrKind::Backend, e),
+        Err(e) => worker_err_response(e),
     }
 }
 
-fn err_response(kind: ErrKind, e: anyhow::Error) -> DaemonResponse {
+/// Build a [`DaemonResponse::Err`] carrying the worker error's protocol `kind`.
+fn worker_err_response(e: WorkerError) -> DaemonResponse {
     DaemonResponse::Err {
-        kind,
-        message: format!("{e:#}"),
+        kind: e.kind(),
+        message: e.to_string(),
     }
 }
 
 /// Read one length-prefixed frame and deserialise it.
-async fn read_frame<T: DeserializeOwned>(pipe: &mut NamedPipeServer) -> Result<T> {
+async fn read_frame<S: AsyncRead + Unpin, T: DeserializeOwned>(pipe: &mut S) -> Result<T> {
     let mut len_buf = [0u8; 4];
     pipe.read_exact(&mut len_buf).await?;
     let len = u32::from_le_bytes(len_buf) as usize;
@@ -434,9 +491,104 @@ async fn read_frame<T: DeserializeOwned>(pipe: &mut NamedPipeServer) -> Result<T
 }
 
 /// Serialise `msg` and write it as a length-prefixed frame.
-async fn write_frame<T: Serialize>(pipe: &mut NamedPipeServer, msg: &T) -> Result<()> {
+async fn write_frame<S: AsyncWrite + Unpin, T: Serialize>(pipe: &mut S, msg: &T) -> Result<()> {
     let frame = encode_frame(msg)?;
     pipe.write_all(&frame).await?;
     pipe.flush().await?;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::io::duplex;
+    use wslc_common::daemon_protocol::ErrKind;
+
+    /// A rejected admission (unknown sandbox) round-trips as a single typed
+    /// `DaemonResponse::Err { NotProvisioned }` through frame encode → transport
+    /// → decode, with no terminal frame following it.
+    #[tokio::test]
+    async fn exec_rejection_round_trips_as_typed_error() {
+        let (mut server, mut client) = duplex(64 * 1024);
+        let admission = Err(WorkerError::NotProvisioned("wslc:nope".to_string()));
+
+        write_exec_result(&mut server, admission).await.unwrap();
+        drop(server);
+
+        let resp: DaemonResponse = read_frame(&mut client).await.unwrap();
+        match resp {
+            DaemonResponse::Err { kind, message } => {
+                assert_eq!(kind, ErrKind::NotProvisioned);
+                assert!(message.contains("wslc:nope"), "message was {message:?}");
+            }
+            other => panic!("expected a typed Err response, got {other:?}"),
+        }
+        // A rejected exec is a single frame: nothing else follows.
+        assert!(read_frame::<_, StreamFrame>(&mut client).await.is_err());
+    }
+
+    /// The `NotStarted` admission contract has an SDK-free regression here: a
+    /// provisioned-but-not-started sandbox surfaces as a typed pre-admission
+    /// error, exercised without constructing a real container handle.
+    #[tokio::test]
+    async fn exec_not_started_round_trips_as_typed_error() {
+        let (mut server, mut client) = duplex(64 * 1024);
+        let admission = Err(WorkerError::NotStarted("wslc:cold".to_string()));
+
+        write_exec_result(&mut server, admission).await.unwrap();
+        drop(server);
+
+        let resp: DaemonResponse = read_frame(&mut client).await.unwrap();
+        assert_eq!(
+            resp,
+            DaemonResponse::Err {
+                kind: ErrKind::NotStarted,
+                message: "sandbox wslc:cold is not started".to_string(),
+            }
+        );
+    }
+
+    /// Dropping the completion sender after admission must produce exactly one
+    /// terminal `StreamFrame::Error` — never a hang or a malformed stream.
+    #[tokio::test]
+    async fn dropped_completion_channel_yields_single_error_terminal() {
+        let (mut server, mut client) = duplex(64 * 1024);
+        let (done_tx, done_rx) = oneshot::channel::<Result<i32, WorkerError>>();
+        drop(done_tx);
+
+        write_exec_result(&mut server, Ok(done_rx)).await.unwrap();
+        drop(server);
+
+        let admit: DaemonResponse = read_frame(&mut client).await.unwrap();
+        assert_eq!(admit, DaemonResponse::Ok);
+        let terminal: StreamFrame = read_frame(&mut client).await.unwrap();
+        match terminal {
+            StreamFrame::Error { message } => {
+                assert!(
+                    message.contains("dropped the exec reply channel"),
+                    "message was {message:?}"
+                );
+            }
+            other => panic!("expected a terminal Error frame, got {other:?}"),
+        }
+        // Exactly one terminal frame is emitted.
+        assert!(read_frame::<_, StreamFrame>(&mut client).await.is_err());
+    }
+
+    /// A successful run writes admission `Ok` then a single `Exit` terminal
+    /// carrying the process exit code.
+    #[tokio::test]
+    async fn successful_exec_writes_ok_then_exit() {
+        let (mut server, mut client) = duplex(64 * 1024);
+        let (done_tx, done_rx) = oneshot::channel::<Result<i32, WorkerError>>();
+        done_tx.send(Ok(7)).unwrap();
+
+        write_exec_result(&mut server, Ok(done_rx)).await.unwrap();
+        drop(server);
+
+        let admit: DaemonResponse = read_frame(&mut client).await.unwrap();
+        assert_eq!(admit, DaemonResponse::Ok);
+        let terminal: StreamFrame = read_frame(&mut client).await.unwrap();
+        assert_eq!(terminal, StreamFrame::Exit { code: 7 });
+    }
+}