Route production control through lifecycle ownership

crates/lazaret-host-coordinator/tests/control.rs

@@ -2,25 +2,28 @@ use std::fs::{File, OpenOptions};
 use std::io::Write;
 use std::os::unix::fs::{MetadataExt, PermissionsExt};
 use std::path::PathBuf;
-use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
 use std::thread::{self, JoinHandle};
 
 use ed25519_dalek::SigningKey;
 use lazaret_core::identity::{
     ClientId, ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId, HostPath,
-    NetworkSidecarId, SandboxId, SessionId,
+    NetworkSidecarId, SandboxId, SessionAnchorId, SessionId,
 };
 use lazaret_core::records::{
     ClientEnrollmentResponse, ControlRequest, ControlResponse, Domain, DurableDecision, Generation,
-    GenerationIdentity, ProtocolVersion, SessionBinding, SessionPhase, SessionRegistrationRequest,
-    SessionRegistrationResponse, TlsCredential,
+    GenerationIdentity, ProtocolVersion, Session, SessionBinding, SessionPhase,
+    SessionRegistrationRequest, SessionRegistrationResponse, TlsCredential,
 };
 use lazaret_host_coordinator::control::{
     CONTROL_ALPN, ControlError, PendingClientEnrollment, PublishedControlFace, SessionRegistry,
     StagedControlFace, TrustAnchor,
 };
 use lazaret_host_coordinator::domain_store::{DomainStoreMutation, HostCoordinatorDomainStore};
+use lazaret_host_coordinator::lifecycle::{
+    HostCoordinatorLifecycle, LifecycleAction, LifecycleAdapter,
+};
 use lazaret_linux::filesystem::try_exclusive_flock;
 use rustix::process::getuid;
 use x509_parser::extensions::GeneralName;
@@ -28,6 +31,28 @@ use x509_parser::parse_x509_certificate;
 
 static NEXT_DIRECTORY: AtomicU64 = AtomicU64::new(0);
 
+#[derive(Clone, Default)]
+struct ControlLifecycleAdapter {
+    publication: Arc<Mutex<Option<Arc<PublishedControlFace>>>>,
+}
+
+impl LifecycleAdapter for ControlLifecycleAdapter {
+    type Error = ControlError;
+
+    fn execute(
+        &mut self,
+        action: LifecycleAction,
+        _generation: &Generation,
+    ) -> Result<(), Self::Error> {
+        if action == LifecycleAction::BeginTeardown
+            && let Some(face) = self.publication.lock().unwrap().take()
+        {
+            face.withdraw();
+        }
+        Ok(())
+    }
+}
+
 struct TestDirectory(PathBuf);
 
 impl TestDirectory {
@@ -152,6 +177,48 @@ fn open_preparing_store(
     store
 }
 
+fn open_empty_store(directory: &TestDirectory) -> Arc<HostCoordinatorDomainStore> {
+    let lock = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create_new(true)
+        .open(directory.path("domain.lock"))
+        .unwrap();
+    Arc::new(
+        HostCoordinatorDomainStore::open(
+            try_exclusive_flock(lock).unwrap(),
+            &directory.path("domain.redb"),
+            &domain(),
+        )
+        .unwrap(),
+    )
+}
+
+fn start_sandbox(
+    lifecycle: &mut HostCoordinatorLifecycle<'_, ControlLifecycleAdapter, Arc<SessionRegistry>>,
+    identity: &GenerationIdentity,
+) {
+    let owner = Session {
+        generation_id: identity.generation_id,
+        session_id: SessionId::new(bytes(80)),
+        cli_id: ClientId::new(bytes(81)),
+        session_anchor_id: SessionAnchorId::new(bytes(82)),
+        session_phase: SessionPhase::New,
+    };
+    lifecycle
+        .start_new_domain(generation(identity, DurableDecision::None), owner)
+        .unwrap();
+    lifecycle.create_trust_anchor().unwrap();
+    lifecycle.start_runtime_supervisor().unwrap();
+    lifecycle.lease_prepared_lower().unwrap();
+    lifecycle.mount_state_anchor().unwrap();
+    lifecycle.start_sidecar().unwrap();
+    lifecycle.start_connectivity_helper().unwrap();
+    lifecycle.mark_network_ready().unwrap();
+    lifecycle.arm_runtime_causal_edge().unwrap();
+    lifecycle.start_sandbox().unwrap();
+}
+
 fn publish(
     directory: &TestDirectory,
     identity: &GenerationIdentity,
@@ -708,3 +775,94 @@ fn prepared_control_face_admits_only_after_generation_publication() {
     let client = enroll(&face, ClientId::new(bytes(90)));
     assert_eq!(client.credential.generation_identity, identity);
 }
+
+#[test]
+fn coordinated_control_session_tracks_lifecycle_and_withdraws_before_teardown() {
+    let directory = TestDirectory::new("coordinated-lifecycle");
+    let identity = generation_identity(7);
+    let store = open_empty_store(&directory);
+    let registry = Arc::new(SessionRegistry::new(Arc::clone(&store), identity.clone()));
+    let adapter = ControlLifecycleAdapter::default();
+    let publication_slot = Arc::clone(&adapter.publication);
+    let mut lifecycle = HostCoordinatorLifecycle::new(&store, adapter, Arc::clone(&registry));
+    start_sandbox(&mut lifecycle, &identity);
+    let staged = StagedControlFace::stage(
+        &directory.0,
+        getuid().as_raw(),
+        TrustAnchor::create(identity.clone()).unwrap(),
+        Arc::clone(&registry),
+    )
+    .unwrap();
+    lifecycle.publish_domain().unwrap();
+    lifecycle
+        .disconnect_session(SessionId::new(bytes(80)))
+        .unwrap();
+    let face = Arc::new(staged.expose().unwrap());
+    *publication_slot.lock().unwrap() = Some(Arc::clone(&face));
+    let client = enroll(&face, ClientId::new(bytes(90)));
+    let binding = registration_binding(&identity, ClientId::new(bytes(90)), 91, 92);
+    let session_id = binding.session_id;
+    let lifecycle = Arc::new(Mutex::new(lifecycle));
+
+    thread::scope(|scope| {
+        let serving_face = Arc::clone(&face);
+        let serving_lifecycle = Arc::clone(&lifecycle);
+        let serving = scope.spawn(move || {
+            serving_face.serve_next_coordinated(
+                |admission| {
+                    serving_lifecycle
+                        .lock()
+                        .unwrap()
+                        .admit_session(admission)
+                        .map_err(|error| ControlError::Lifecycle(error.to_string()))
+                },
+                |session_id| {
+                    serving_lifecycle
+                        .lock()
+                        .unwrap()
+                        .disconnect_session(session_id)
+                        .map_err(|error| ControlError::Lifecycle(error.to_string()))
+                },
+            )
+        });
+        let config = PublishedControlFace::client_config(
+            face.publication_identity(),
+            Some((&client.credential, client.pending.private_key())),
+        )
+        .unwrap();
+        let mut stream =
+            PublishedControlFace::connect(face.socket_path(), config, face.publication_identity())
+                .unwrap();
+        PublishedControlFace::send(
+            &mut stream,
+            &ControlRequest::Register(SessionRegistrationRequest {
+                protocol_version: ProtocolVersion::V1,
+                session_binding: binding,
+            }),
+        )
+        .unwrap();
+        let _: ControlResponse = PublishedControlFace::receive(&mut stream).unwrap();
+        {
+            let lifecycle = lifecycle.lock().unwrap();
+            assert_eq!(
+                lifecycle.session_phase(session_id),
+                Some(SessionPhase::Attached)
+            );
+            assert_eq!(lifecycle.registered_sessions(), 1);
+            assert!(!lifecycle.is_teardown_ready());
+        }
+        drop(stream);
+        serving.join().unwrap().unwrap();
+    });
+
+    let mut lifecycle = lifecycle.lock().unwrap();
+    assert_eq!(
+        lifecycle.session_phase(session_id),
+        Some(SessionPhase::Closed)
+    );
+    assert!(lifecycle.is_teardown_ready());
+    lifecycle.begin_teardown().unwrap();
+    assert!(!face.socket_path().exists());
+    assert!(!face.identity_path().exists());
+    assert!(std::os::unix::net::UnixStream::connect(face.socket_path()).is_err());
+}