Route production control through lifecycle ownership

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

@@ -1,6 +1,7 @@
 use std::fmt;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::sync::Mutex;
 
 use lazaret_core::configuration::ResolvedDomainConfiguration;
 use lazaret_core::identity::{BundleDigest, ContainerId, MountIdentity};
@@ -15,7 +16,9 @@ use lazaret_state_anchor::StateAnchorLayout;
 use lazaret_storage_catalog::Catalog;
 
 use crate::control::ControlSessionAdmission;
-use crate::control::{PublishedControlFace, SessionRegistry, StagedControlFace, TrustAnchor};
+use crate::control::{
+    ControlError, PublishedControlFace, SessionRegistry, StagedControlFace, TrustAnchor,
+};
 use crate::domain_store::HostCoordinatorDomainStore;
 use crate::lifecycle::{
     HostCoordinatorLifecycle, LifecycleAction, LifecycleAdapter, SessionAdmissionIdentity,
@@ -63,7 +66,7 @@ pub struct ProductionLifecycleAdapter<'a> {
     registry: Arc<SessionRegistry>,
     trust_anchor: Option<TrustAnchor>,
     staged_control: Option<StagedControlFace>,
-    published_control: Option<PublishedControlFace>,
+    published_control: Option<Arc<PublishedControlFace>>,
     awaiting_runtime: Option<AwaitingAttachedRunsc>,
     runtime: Option<CommittedSandboxRuntime>,
     state_anchor_mount: Option<MountIdentity>,
@@ -169,8 +172,8 @@ impl<'a> ProductionLifecycleAdapter<'a> {
     }
 
     #[must_use]
-    pub fn published_control(&self) -> Option<&PublishedControlFace> {
+    pub fn published_control(&self) -> Option<Arc<PublishedControlFace>> {
-        self.published_control.as_ref()
+        self.published_control.as_ref().map(Arc::clone)
     }
 
     fn require_generation(&self, generation: &Generation) -> Result<(), ProductionLifecycleError> {
@@ -347,13 +350,13 @@ impl<'a> ProductionLifecycleAdapter<'a> {
     }
 
     fn expose_publication(&mut self) -> Result<(), ProductionLifecycleError> {
-        self.published_control = Some(
+        self.published_control = Some(Arc::new(
             self.staged_control
                 .take()
                 .ok_or_else(|| ProductionLifecycleError("staged control face is absent".into()))?
                 .expose()
                 .map_err(failure)?,
-        );
+        ));
         Ok(())
     }
 
@@ -445,6 +448,12 @@ impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
             LifecycleAction::UnmountStateAnchor => self.unmount_state_anchor(),
             LifecycleAction::ReleasePreparedLowerLease => self.release_prepared_lower(),
             LifecycleAction::GracefulSupervisorExit => self.graceful_exit(),
+            LifecycleAction::BeginTeardown => {
+                if let Some(control) = self.published_control.take() {
+                    control.withdraw();
+                }
+                Ok(())
+            }
             LifecycleAction::FinishTeardown => {
                 self.published_control.take();
                 Ok(())
@@ -457,8 +466,7 @@ impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
             | LifecycleAction::AdvanceWaitingSession
             | LifecycleAction::RejectConfigurationMismatch
             | LifecycleAction::RegisterSession
-            | LifecycleAction::DisconnectSession
+            | LifecycleAction::DisconnectSession => Ok(()),
-            | LifecycleAction::BeginTeardown => Ok(()),
         }
     }
 
@@ -476,6 +484,37 @@ impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
     }
 }
 
+/// Serves one production control connection through the exact retained lifecycle owner.
+///
+/// # Errors
+/// Refuses absent publication, poisoned lifecycle ownership, or control/session failure.
+pub fn serve_production_control(
+    lifecycle: &Mutex<ProductionHostCoordinatorLifecycle<'_, '_>>,
+) -> Result<(), ControlError> {
+    let face = lifecycle
+        .lock()
+        .map_err(|_| ControlError::LockPoisoned)?
+        .adapter()
+        .published_control()
+        .ok_or(ControlError::StaleGeneration)?;
+    face.serve_next_coordinated(
+        |admission| {
+            lifecycle
+                .lock()
+                .map_err(|_| ControlError::LockPoisoned)?
+                .admit_session(admission)
+                .map_err(|error| ControlError::Lifecycle(error.to_string()))
+        },
+        |session_id| {
+            lifecycle
+                .lock()
+                .map_err(|_| ControlError::LockPoisoned)?
+                .disconnect_session(session_id)
+                .map_err(|error| ControlError::Lifecycle(error.to_string()))
+        },
+    )
+}
+
 fn path_from_host(bytes: &[u8]) -> &Path {
     use std::os::unix::ffi::OsStrExt;
 

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

@@ -69,6 +69,7 @@ pub enum ControlError {
     InvalidRequest,
     OversizeRequest { length: usize, maximum: usize },
     LockPoisoned,
+    Lifecycle(String),
 }
 
 impl fmt::Display for ControlError {
@@ -100,6 +101,7 @@ impl fmt::Display for ControlError {
                 )
             }
             Self::LockPoisoned => formatter.write_str("control registry lock is poisoned"),
+            Self::Lifecycle(detail) => write!(formatter, "control lifecycle failed: {detail}"),
         }
     }
 }
@@ -1007,12 +1009,42 @@ impl PublishedControlFace {
         &self.paths.identity
     }
 
+    /// Withdraws exact discovery and connection paths while retaining live accepted streams.
+    pub fn withdraw(&self) {
+        remove_if_identity_matches(&self.paths.socket, self.socket_identity);
+        remove_if_identity_matches(&self.paths.identity, self.identity_file_identity);
+    }
+
     /// Accepts and serves one TLS connection. Registered connections remain live until EOF.
     ///
     /// # Errors
     /// Refuses peer UID, publication, TLS, ALPN, certificate, framing, tuple, duplicate, and
     /// persistence failures.
     pub fn serve_next(&self) -> Result<(), ControlError> {
+        self.serve_next_coordinated(
+            |admission| {
+                self.registry
+                    .register(admission.request, &admission.credential)
+                    .map(|_| ())
+            },
+            |session_id| self.registry.disconnect(session_id),
+        )
+    }
+
+    /// Serves one connection while routing authenticated session state through its lifecycle.
+    ///
+    /// # Errors
+    /// Refuses the same conditions as [`Self::serve_next`] plus lifecycle admission or disconnect
+    /// failure.
+    pub fn serve_next_coordinated<Admit, Disconnect>(
+        &self,
+        mut admit: Admit,
+        mut disconnect: Disconnect,
+    ) -> Result<(), ControlError>
+    where
+        Admit: FnMut(ControlSessionAdmission) -> Result<(), ControlError>,
+        Disconnect: FnMut(SessionId) -> Result<(), ControlError>,
+    {
         let (socket, _) = self.listener.accept()?;
         let peer = socket_peercred(&socket).map_err(io::Error::from)?;
         require_peer_uid(peer.uid.as_raw(), self.uid)?;
@@ -1034,7 +1066,9 @@ impl PublishedControlFace {
         let mut stream = rustls::StreamOwned::new(connection, socket);
         match peer_certificate {
             None => self.serve_enrollment(&mut stream),
-            Some(certificates) => self.serve_registration(&mut stream, &certificates),
+            Some(certificates) => {
+                self.serve_registration(&mut stream, &certificates, &mut admit, &mut disconnect)
+            }
         }
     }
 
@@ -1057,11 +1091,17 @@ impl PublishedControlFace {
         )
     }
 
-    fn serve_registration(
+    fn serve_registration<Admit, Disconnect>(
         &self,
         stream: &mut rustls::StreamOwned<ServerConnection, std::os::unix::net::UnixStream>,
         certificates: &[Vec<u8>],
-    ) -> Result<(), ControlError> {
+        admit: &mut Admit,
+        disconnect: &mut Disconnect,
+    ) -> Result<(), ControlError>
+    where
+        Admit: FnMut(ControlSessionAdmission) -> Result<(), ControlError>,
+        Disconnect: FnMut(SessionId) -> Result<(), ControlError>,
+    {
         let [leaf] = certificates else {
             return Err(ControlError::InvalidCertificateProfile);
         };
@@ -1076,7 +1116,11 @@ impl PublishedControlFace {
             self.anchor.sandbox_protocol_credential.credential_id,
             &self.anchor.coordinator_protocol_key,
         )?;
-        let registered = self.registry.register(request, &credential)?;
+        let session_id = request.session_binding.session_id;
+        admit(ControlSessionAdmission {
+            request,
+            credential,
+        })?;
         let result = (|| {
             write_frame(
                 stream,
@@ -1093,8 +1137,8 @@ impl PublishedControlFace {
                 Err(error) => Err(error.into()),
             }
         })();
-        let disconnect = self.registry.disconnect(registered.binding.session_id);
+        let disconnected = disconnect(session_id);
-        result.and(disconnect)
+        result.and(disconnected)
     }
 
     fn validate_publication(&self) -> Result<(), ControlError> {
@@ -1216,8 +1260,7 @@ impl PublishedControlFace {
 
 impl Drop for PublishedControlFace {
     fn drop(&mut self) {
-        remove_if_identity_matches(&self.paths.socket, self.socket_identity);
+        self.withdraw();
-        remove_if_identity_matches(&self.paths.identity, self.identity_file_identity);
     }
 }
 

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,
+    GenerationIdentity, ProtocolVersion, Session, SessionBinding, SessionPhase,
-    SessionRegistrationResponse, TlsCredential,
+    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());
+}