Stop terminal grace polling after authority loss

crates/lazaret-host-coordinator/src/terminal.rs

@@ -63,6 +63,12 @@ enum GuestReleaseAuthority {
     HostOnly,
 }
 
+enum GuestWaitOutcome {
+    Absent,
+    TimedOut,
+    Revoked,
+}
+
 #[derive(Debug)]
 pub enum TerminalError {
     Io(io::Error),
@@ -788,10 +794,19 @@ fn release_physical(
                 failure.get_or_insert(error);
                 break;
             }
-            if let Err(error) = wait_until(SIGNAL_GRACE, || process.group_absent()) {
+            match wait_for_group_absence(process, authorize_guest) {
-                failure.get_or_insert(error);
+                Ok(GuestWaitOutcome::Absent) => break,
+                Ok(GuestWaitOutcome::TimedOut) => {}
+                Ok(GuestWaitOutcome::Revoked) => {
+                    guest_authorized = false;
                     break;
                 }
+                Err(error) => {
+                    failure.get_or_insert(error);
+                    guest_authorized = false;
+                    break;
+                }
+            }
         }
     }
     match process.exec_exited() {
@@ -812,29 +827,7 @@ fn release_physical(
             None
         }
     };
-    let mut output = Vec::new();
+    let output = drain_output(process, session_id, connected);
-    if connected {
-        let deadline = Instant::now() + DRAIN_DEADLINE;
-        let mut bytes = [0_u8; MAX_OUTPUT_BYTES];
-        let mut drained = 0;
-        while Instant::now() < deadline && drained < MAX_DRAIN_BYTES {
-            match process.read_output(&mut bytes) {
-                Ok(0) => break,
-                Ok(length) => {
-                    let length = length.min(MAX_DRAIN_BYTES - drained);
-                    drained += length;
-                    output.push(SessionTerminalMessage::Output(TerminalOutput {
-                        session_id,
-                        bytes: bytes[..length].to_vec(),
-                    }));
-                }
-                Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
-                    std::thread::sleep(POLL_INTERVAL);
-                }
-                Err(_) => break,
-            }
-        }
-    }
     process.close_pty();
     if guest_authorized && authorize_guest() {
         match process.group_absent() {
@@ -860,6 +853,38 @@ fn release_physical(
     ))
 }
 
+fn drain_output(
+    process: &mut dyn SessionTerminalProcess,
+    session_id: SessionId,
+    connected: bool,
+) -> Vec<SessionTerminalMessage> {
+    let mut output = Vec::new();
+    if !connected {
+        return output;
+    }
+    let deadline = Instant::now() + DRAIN_DEADLINE;
+    let mut bytes = [0_u8; MAX_OUTPUT_BYTES];
+    let mut drained = 0;
+    while Instant::now() < deadline && drained < MAX_DRAIN_BYTES {
+        match process.read_output(&mut bytes) {
+            Ok(0) => break,
+            Ok(length) => {
+                let length = length.min(MAX_DRAIN_BYTES - drained);
+                drained += length;
+                output.push(SessionTerminalMessage::Output(TerminalOutput {
+                    session_id,
+                    bytes: bytes[..length].to_vec(),
+                }));
+            }
+            Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
+                std::thread::sleep(POLL_INTERVAL);
+            }
+            Err(_) => break,
+        }
+    }
+    output
+}
+
 fn with_cleanup<T>(original: TerminalError, cleanup: Result<T, TerminalError>) -> TerminalError {
     match cleanup {
         Ok(_) => original,
@@ -870,6 +895,23 @@ fn with_cleanup<T>(original: TerminalError, cleanup: Result<T, TerminalError>) -
     }
 }
 
+fn wait_for_group_absence(
+    process: &dyn SessionTerminalProcess,
+    authorize_guest: &mut dyn FnMut() -> bool,
+) -> Result<GuestWaitOutcome, TerminalError> {
+    let deadline = Instant::now() + SIGNAL_GRACE;
+    while Instant::now() < deadline {
+        if !authorize_guest() {
+            return Ok(GuestWaitOutcome::Revoked);
+        }
+        if process.group_absent()? {
+            return Ok(GuestWaitOutcome::Absent);
+        }
+        std::thread::sleep(POLL_INTERVAL);
+    }
+    Ok(GuestWaitOutcome::TimedOut)
+}
+
 fn flush_input(terminal: &mut ActiveTerminal) -> Result<(), TerminalError> {
     while !terminal.pending_input.is_empty() {
         let (first, second) = terminal.pending_input.as_slices();
@@ -888,20 +930,6 @@ fn flush_input(terminal: &mut ActiveTerminal) -> Result<(), TerminalError> {
     Ok(())
 }
 
-fn wait_until(
-    duration: Duration,
-    mut predicate: impl FnMut() -> Result<bool, TerminalError>,
-) -> Result<(), TerminalError> {
-    let deadline = Instant::now() + duration;
-    while Instant::now() < deadline {
-        if predicate()? {
-            return Ok(());
-        }
-        std::thread::sleep(POLL_INTERVAL);
-    }
-    Ok(())
-}
-
 fn exit_message(
     session_id: SessionId,
     loss: Option<TerminalLoss>,

crates/lazaret-host-coordinator/src/terminal/tests.rs

@@ -167,25 +167,7 @@ impl Fixture {
     }
 
     fn release_sandbox(&self) {
-        let sandbox = self
+        release_sandbox(&self.store);
-            .store
-            .snapshot()
-            .unwrap()
-            .participants
-            .into_iter()
-            .find(|participant| participant.role() == ParticipantRole::Sandbox)
-            .unwrap();
-        self.store
-            .apply(&[DomainStoreMutation::PutParticipant(
-                Participant::snapshot(
-                    sandbox.generation_id(),
-                    sandbox.role(),
-                    ParticipantState::ReleaseIntent,
-                    sandbox.ownership().cloned(),
-                )
-                .unwrap(),
-            )])
-            .unwrap();
     }
 
     fn add_newer_generation(&self) {
@@ -203,6 +185,27 @@ impl Fixture {
     }
 }
 
+fn release_sandbox(store: &HostCoordinatorDomainStore) {
+    let sandbox = store
+        .snapshot()
+        .unwrap()
+        .participants
+        .into_iter()
+        .find(|participant| participant.role() == ParticipantRole::Sandbox)
+        .unwrap();
+    store
+        .apply(&[DomainStoreMutation::PutParticipant(
+            Participant::snapshot(
+                sandbox.generation_id(),
+                sandbox.role(),
+                ParticipantState::ReleaseIntent,
+                sandbox.ownership().cloned(),
+            )
+            .unwrap(),
+        )])
+        .unwrap();
+}
+
 fn insert_attached_sessions(store: &HostCoordinatorDomainStore, session_ids: &[SessionId]) {
     for (index, session_id) in session_ids.iter().enumerate() {
         let mut session = Session {
@@ -222,12 +225,13 @@ fn insert_attached_sessions(store: &HostCoordinatorDomainStore, session_ids: &[S
     }
 }
 
-#[derive(Debug)]
 struct MockState {
     output: VecDeque<Vec<u8>>,
     log: Arc<Mutex<Vec<String>>>,
     guest: MockGuestState,
     resources: MockResourceState,
+    group_queries: usize,
+    group_query_hook: Option<(usize, Arc<dyn Fn() + Send + Sync>)>,
 }
 
 #[derive(Debug)]
@@ -367,11 +371,29 @@ impl SessionTerminalProcess for MockProcess {
     }
 
     fn group_absent(&self) -> Result<bool, TerminalError> {
-        let state = self.state.lock().unwrap();
+        let (identity_mismatch, group_absent, hook) = {
-        if state.guest.identity_mismatch {
+            let mut state = self.state.lock().unwrap();
+            state.group_queries += 1;
+            let group_queries = state.group_queries;
+            state.log.lock().unwrap().push("query:group".into());
+            let hook = state
+                .group_query_hook
+                .as_ref()
+                .filter(|(query, _)| *query == group_queries)
+                .map(|(_, hook)| Arc::clone(hook));
+            (
+                state.guest.identity_mismatch,
+                !state.guest.group_present,
+                hook,
+            )
+        };
+        if let Some(hook) = hook {
+            hook();
+        }
+        if identity_mismatch {
             return Err(TerminalError::GuestIdentityMismatch);
         }
-        Ok(!state.guest.group_present)
+        Ok(group_absent)
     }
 
     fn terminate_exec(&self) -> io::Result<()> {
@@ -468,6 +490,8 @@ impl SessionTerminalFactory for MockFactory {
                 exec_reaped: false,
                 pty_open: true,
             },
+            group_queries: 0,
+            group_query_hook: None,
         }));
         self.states.lock().unwrap().push(Arc::clone(&state));
         Ok((running, Box::new(MockProcess { state })))
@@ -537,6 +561,16 @@ fn assert_no_group_signal(factory: &MockFactory) {
     );
 }
 
+fn assert_no_guest_action_after(factory: &MockFactory, marker: &str) {
+    let log = factory.log.lock().unwrap();
+    let marker = log.iter().position(|action| action == marker).unwrap();
+    assert!(
+        log[marker + 1..]
+            .iter()
+            .all(|action| !action.starts_with("query:") && !action.starts_with("signal:"))
+    );
+}
+
 #[test]
 fn acquisition_refuses_authority_outside_the_committed_sandbox() {
     let session = id(6);
@@ -694,6 +728,68 @@ fn explicit_close_and_loss_revalidate_before_guest_cleanup() {
     );
 }
 
+#[test]
+fn grace_wait_revocation_and_read_error_stop_all_later_guest_actions() {
+    let session = id(6);
+    let (fixture, factory, mut engine) = setup(&[session]);
+    engine.acquire(&acquisition(session)).unwrap();
+    let store = Arc::clone(&fixture.store);
+    let log = Arc::clone(&factory.log);
+    factory.state(0).lock().unwrap().group_query_hook = Some((
+        2,
+        Arc::new(move || {
+            log.lock().unwrap().push("revoke:sandbox".into());
+            release_sandbox(&store);
+        }),
+    ));
+    engine
+        .handle(
+            GenerationId::new(7),
+            SessionTerminalEvent::Close(lazaret_core::records::SessionClose {
+                session_id: session,
+            }),
+        )
+        .unwrap();
+    assert_no_guest_action_after(&factory, "revoke:sandbox");
+    let state = factory.state(0);
+    assert!(state.lock().unwrap().resources.exec_reaped);
+    assert!(!state.lock().unwrap().resources.pty_open);
+    assert_eq!(
+        fixture.store.snapshot().unwrap().session_terminals[0].state(),
+        SessionTerminalState::ReleaseIntent
+    );
+
+    let fixture = Fixture::new(&[session]);
+    let store = Arc::new(FailingRunningStore::new(Arc::clone(&fixture.store)));
+    let (factory, mut engine) = engine_for_store(store.clone());
+    engine.acquire(&acquisition(session)).unwrap();
+    let failing_store = Arc::clone(&store);
+    let log = Arc::clone(&factory.log);
+    factory.state(0).lock().unwrap().group_query_hook = Some((
+        2,
+        Arc::new(move || {
+            log.lock().unwrap().push("revoke:read-error".into());
+            failing_store.arm_snapshot_failure();
+        }),
+    ));
+    engine
+        .handle(
+            GenerationId::new(7),
+            SessionTerminalEvent::Close(lazaret_core::records::SessionClose {
+                session_id: session,
+            }),
+        )
+        .unwrap();
+    assert_no_guest_action_after(&factory, "revoke:read-error");
+    let state = factory.state(0);
+    assert!(state.lock().unwrap().resources.exec_reaped);
+    assert!(!state.lock().unwrap().resources.pty_open);
+    assert_eq!(
+        fixture.store.snapshot().unwrap().session_terminals[0].state(),
+        SessionTerminalState::ReleaseIntent
+    );
+}
+
 #[test]
 fn stale_caller_and_superseded_generation_refuse_without_cross_generation_action() {
     let session = id(6);
@@ -923,6 +1019,8 @@ fn running_store_failures_release_every_live_partial_handle() {
                     exec_reaped: false,
                     pty_open: true,
                 },
+                group_queries: 0,
+                group_query_hook: None,
             }));
             let mut process_handle = MockProcess {
                 state: Arc::clone(&state),