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

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};
use lazaret_core::network::{CanonicalDerivedNetworkPolicy, ReachabilityExpectation};
use lazaret_core::records::Session;
use lazaret_core::records::{
    Generation, GenerationIdentity, Participant, ParticipantOwnership, ParticipantRole,
    ParticipantState, PreparedArtifact,
};
use lazaret_linux::process::current_process_identity;
use lazaret_runtime_supervisor::network::NetworkProbeRequest;
use lazaret_runtime_supervisor::sandbox::SandboxBundle;
use lazaret_runtime_supervisor::supervision::RuntimeSupervisorLaunch;
use lazaret_state_anchor::StateAnchorLayout;
use lazaret_storage_catalog::Catalog;

use crate::control::ControlSessionAdmission;
use crate::control::{
    ControlError, InteractiveSessionExecutor, NonInteractiveSessionExecutor, PublishedControlFace,
    SessionRegistry, StagedControlFace, TrustAnchor,
};
use crate::domain_store::{DomainStoreMutation, HostCoordinatorDomainStore};
use crate::lifecycle::{
    HostCoordinatorLifecycle, LifecycleAction, LifecycleAdapter, SessionAdmissionIdentity,
    SessionLifecycleAdapter,
};
use crate::network::NetworkCoordinator;
use crate::runtime::{
    AwaitingAttachedRunsc, CommittedSandboxRuntime, SandboxCoordinator,
    StateAnchorAbsenceObservation,
};
use crate::state_anchor::{StateAnchorCoordinator, StateAnchorReleaseAuthority};
use crate::terminal::SessionPromptTransport;

#[derive(Debug)]
pub struct ProductionLifecycleError(String);

impl fmt::Display for ProductionLifecycleError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl std::error::Error for ProductionLifecycleError {}

impl ProductionLifecycleError {
    pub(crate) fn service(error: impl fmt::Display) -> Self {
        Self(error.to_string())
    }
}

fn failure(error: impl fmt::Display) -> ProductionLifecycleError {
    ProductionLifecycleError(error.to_string())
}

pub struct ProductionBootstrapPlan<'a> {
    pub store: Arc<HostCoordinatorDomainStore>,
    pub catalog: &'a Catalog,
    pub configuration: &'a ResolvedDomainConfiguration,
    pub generation_identity: GenerationIdentity,
    pub prepared_artifact: &'a PreparedArtifact,
    pub state_anchor_layout: &'a StateAnchorLayout<'a>,
    pub runtime_launch: RuntimeSupervisorLaunch,
    pub sandbox_bundle: &'a SandboxBundle,
    pub network_policy: &'a CanonicalDerivedNetworkPolicy,
    pub network_probe: NetworkProbeRequest,
    pub reachability: &'a ReachabilityExpectation,
    pub container_id: ContainerId,
    pub runtime_root: PathBuf,
    pub bundle_digest: BundleDigest,
    pub invoking_uid: u32,
    pub non_interactive_executor: Arc<dyn NonInteractiveSessionExecutor>,
    pub interactive_executor: Arc<dyn InteractiveSessionExecutor>,
    pub prompt_transport: Arc<SessionPromptTransport>,
    pub(crate) runtime_staging: &'a crate::production::RuntimeStaging,
}

pub struct ProductionLifecycleAdapter<'a> {
    plan: ProductionBootstrapPlan<'a>,
    registry: Arc<SessionRegistry>,
    trust_anchor: Option<TrustAnchor>,
    staged_control: Option<StagedControlFace>,
    published_control: Option<Arc<PublishedControlFace>>,
    awaiting_runtime: Option<AwaitingAttachedRunsc>,
    runtime: Option<CommittedSandboxRuntime>,
    state_anchor_mount: Option<MountIdentity>,
    state_anchor_release: Option<StateAnchorReleaseAuthority>,
    state_anchor_absence: Option<StateAnchorAbsenceObservation>,
}

pub type ProductionHostCoordinatorLifecycle<'store, 'resources> = HostCoordinatorLifecycle<
    'store,
    ProductionLifecycleAdapter<'resources>,
    ProductionSessionRegistry,
>;

/// Composes the cold-start lifecycle with the exact shared production session registry.
///
/// # Errors
/// Refuses a plan whose Domain Store is not the caller-retained lock-owning store.
pub fn production_lifecycle<'store, 'resources>(
    store: &'store Arc<HostCoordinatorDomainStore>,
    plan: ProductionBootstrapPlan<'resources>,
) -> Result<ProductionHostCoordinatorLifecycle<'store, 'resources>, ProductionLifecycleError> {
    if !Arc::ptr_eq(store, &plan.store) {
        return Err(ProductionLifecycleError(
            "production bootstrap Domain Store authority differs".into(),
        ));
    }
    let (adapter, registry) = ProductionLifecycleAdapter::new(plan);
    Ok(HostCoordinatorLifecycle::new(store, adapter, registry))
}

#[derive(Clone)]
pub struct ProductionSessionRegistry {
    registry: Arc<SessionRegistry>,
}

impl ProductionSessionRegistry {
    #[must_use]
    pub fn registry(&self) -> Arc<SessionRegistry> {
        Arc::clone(&self.registry)
    }
}

impl SessionLifecycleAdapter for ProductionSessionRegistry {
    type Error = ProductionLifecycleError;
    type Admission = ControlSessionAdmission;

    fn admission_identity(&self, admission: &Self::Admission) -> SessionAdmissionIdentity {
        <Arc<SessionRegistry> as SessionLifecycleAdapter>::admission_identity(
            &self.registry,
            admission,
        )
    }

    fn register_session(
        &self,
        store: &HostCoordinatorDomainStore,
        entry_action: LifecycleAction,
        admission: Self::Admission,
    ) -> Result<Session, Self::Error> {
        <Arc<SessionRegistry> as SessionLifecycleAdapter>::register_session(
            &self.registry,
            store,
            entry_action,
            admission,
        )
        .map_err(failure)
    }

    fn disconnect_session(
        &self,
        store: &HostCoordinatorDomainStore,
        session: &Session,
    ) -> Result<(), Self::Error> {
        <Arc<SessionRegistry> as SessionLifecycleAdapter>::disconnect_session(
            &self.registry,
            store,
            session,
        )
        .map_err(failure)
    }
}

impl<'a> ProductionLifecycleAdapter<'a> {
    #[must_use]
    pub fn new(plan: ProductionBootstrapPlan<'a>) -> (Self, ProductionSessionRegistry) {
        let registry = Arc::new(SessionRegistry::new(
            Arc::clone(&plan.store),
            plan.generation_identity.clone(),
        ));
        (
            Self {
                plan,
                registry: Arc::clone(&registry),
                trust_anchor: None,
                staged_control: None,
                published_control: None,
                awaiting_runtime: None,
                runtime: None,
                state_anchor_mount: None,
                state_anchor_release: None,
                state_anchor_absence: None,
            },
            ProductionSessionRegistry { registry },
        )
    }

    #[must_use]
    pub fn published_control(&self) -> Option<Arc<PublishedControlFace>> {
        self.published_control.as_ref().map(Arc::clone)
    }

    fn require_generation(&self, generation: &Generation) -> Result<(), ProductionLifecycleError> {
        if generation.domain_id == self.plan.generation_identity.domain_id
            && generation.generation_id == self.plan.generation_identity.generation_id
            && generation.configuration_fingerprint
                == self.plan.generation_identity.configuration_fingerprint
            && generation.coordinator_id == self.plan.generation_identity.coordinator_id
            && generation.sandbox_id == self.plan.generation_identity.sandbox_id
        {
            Ok(())
        } else {
            Err(ProductionLifecycleError(
                "production bootstrap generation identity changed".into(),
            ))
        }
    }

    fn store(&self) -> &HostCoordinatorDomainStore {
        &self.plan.store
    }

    fn sandbox(&self) -> SandboxCoordinator<'_> {
        SandboxCoordinator::new(self.store())
    }

    fn network(&self) -> NetworkCoordinator<'_> {
        NetworkCoordinator::new(self.store())
    }

    fn state_anchor(&self) -> StateAnchorCoordinator<'_> {
        StateAnchorCoordinator::new(self.store(), self.plan.catalog)
    }

    fn create_trust_anchor(&mut self) -> Result<(), ProductionLifecycleError> {
        if self.trust_anchor.is_some() {
            return Err(ProductionLifecycleError(
                "Trust Anchor already exists".into(),
            ));
        }
        self.trust_anchor =
            Some(TrustAnchor::create(self.plan.generation_identity.clone()).map_err(failure)?);
        Ok(())
    }

    fn start_runtime_supervisor(&mut self) -> Result<(), ProductionLifecycleError> {
        let sandbox = self.sandbox();
        sandbox
            .begin_intent(
                self.plan.generation_identity.generation_id,
                self.plan.container_id,
                &self.plan.runtime_root,
                self.plan
                    .runtime_staging
                    .runsc_root_identity()
                    .map_err(failure)?,
                self.plan.bundle_digest,
                current_process_identity().map_err(failure)?,
            )
            .map_err(failure)?;
        let launch = &self.plan.runtime_launch;
        self.awaiting_runtime = Some(
            AwaitingAttachedRunsc::launch_networked(&sandbox, &self.network(), launch)
                .map_err(failure)?,
        );
        Ok(())
    }

    fn lease_prepared_lower(&self) -> Result<(), ProductionLifecycleError> {
        self.state_anchor()
            .begin_acquisition(
                self.plan.generation_identity.generation_id,
                self.plan.prepared_artifact,
                self.plan.state_anchor_layout,
            )
            .map(|_| ())
            .map_err(failure)
    }

    fn mount_state_anchor(&mut self) -> Result<(), ProductionLifecycleError> {
        let store = Arc::clone(&self.plan.store);
        let coordinator = StateAnchorCoordinator::new(&store, self.plan.catalog);
        let mount = self
            .awaiting_runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("Runtime Supervisor is absent".into()))?
            .mount_state_anchor(
                &coordinator,
                self.plan.prepared_artifact.prepared_artifact_id,
            )
            .map_err(|error| failure(format!("mounting State Anchor: {error}")))?;
        self.state_anchor_mount = Some(mount);
        Ok(())
    }

    fn start_connectivity_helper(&mut self) -> Result<(), ProductionLifecycleError> {
        let store = Arc::clone(&self.plan.store);
        let network = NetworkCoordinator::new(&store);
        self.awaiting_runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("Runtime Supervisor is absent".into()))?
            .receive_helper_staged(&network)
            .map(|_| ())
            .map_err(|error| failure(format!("starting connectivity helper: {error}")))
    }

    fn start_sidecar(&mut self) -> Result<(), ProductionLifecycleError> {
        self.awaiting_runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("Runtime Supervisor is absent".into()))?
            .start_sidecar()
            .map_err(|error| failure(format!("starting sidecar: {error}")))
    }

    fn mark_network_ready(&mut self) -> Result<(), ProductionLifecycleError> {
        let store = Arc::clone(&self.plan.store);
        let network = NetworkCoordinator::new(&store);
        let mut probe = self.plan.network_probe.clone();
        self.awaiting_runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("Runtime Supervisor is absent".into()))?
            .receive_network_readiness(
                &network,
                self.plan.network_policy,
                |namespace| {
                    probe.enforcement_namespace = namespace;
                    probe
                },
                self.plan.reachability,
            )
            .map(|_| ())
            .map_err(|error| failure(format!("marking network ready: {error}")))
    }

    fn start_sandbox(&mut self) -> Result<(), ProductionLifecycleError> {
        let store = Arc::clone(&self.plan.store);
        let sandbox = SandboxCoordinator::new(&store);
        let mut pending = self
            .awaiting_runtime
            .take()
            .ok_or_else(|| ProductionLifecycleError("Runtime Supervisor is absent".into()))?;
        let staged = pending
            .receive_staged(&sandbox)
            .map_err(|error| failure(format!("receiving attached runsc staging: {error}")))?;
        let handoff = self
            .plan
            .sandbox_bundle
            .capture_handoff_authority(self.plan.configuration)
            .map_err(|error| failure(format!("capturing sandbox handoff: {error}")))?;
        let (runtime, _) = pending
            .acknowledge(staged, &handoff)
            .map_err(|error| failure(format!("releasing sandbox handoff: {error}")))?
            .commit_networked(&sandbox)
            .map_err(|error| failure(format!("committing sandbox readiness: {error}")))?;
        self.runtime = Some(runtime);
        Ok(())
    }

    fn stage_publication(&mut self) -> Result<(), ProductionLifecycleError> {
        let anchor = self
            .trust_anchor
            .take()
            .ok_or_else(|| ProductionLifecycleError("Trust Anchor is absent".into()))?;
        let generation_id = self.plan.generation_identity.generation_id;
        self.plan
            .store
            .apply(&[DomainStoreMutation::PutParticipant(
                Participant::bare_intent(generation_id, ParticipantRole::Publication),
            )])
            .map_err(failure)?;
        let staged = StagedControlFace::stage_with_executors(
            path_from_host(
                self.plan
                    .configuration
                    .runtime_domain_root()
                    .canonical_label()
                    .as_bytes(),
            ),
            self.plan.invoking_uid,
            anchor,
            Arc::clone(&self.registry),
            Arc::clone(&self.plan.non_interactive_executor),
            Arc::clone(&self.plan.interactive_executor),
        )
        .map_err(failure)?;
        let ownership = staged.recovery_ownership().map_err(failure)?;
        self.plan
            .store
            .apply(&[DomainStoreMutation::PutParticipant(
                Participant::snapshot(
                    generation_id,
                    ParticipantRole::Publication,
                    ParticipantState::Intent,
                    Some(ParticipantOwnership::Publication(ownership)),
                )
                .map_err(failure)?,
            )])
            .map_err(failure)?;
        self.staged_control = Some(staged);
        Ok(())
    }

    fn expose_publication(&mut self) -> Result<(), ProductionLifecycleError> {
        let published = self
            .staged_control
            .take()
            .ok_or_else(|| ProductionLifecycleError("staged control face is absent".into()))?
            .expose()
            .map_err(failure)?;
        let ownership = published.recovery_ownership().map_err(failure)?;
        self.plan
            .store
            .apply(&[DomainStoreMutation::PutParticipant(
                Participant::snapshot(
                    self.plan.generation_identity.generation_id,
                    ParticipantRole::Publication,
                    ParticipantState::Committed,
                    Some(ParticipantOwnership::Publication(ownership)),
                )
                .map_err(failure)?,
            )])
            .map_err(failure)?;
        self.published_control = Some(Arc::new(published));
        Ok(())
    }

    fn clean_publication_after_withdrawal(&self) -> Result<(), ProductionLifecycleError> {
        let generation_id = self.plan.generation_identity.generation_id;
        let current = self
            .plan
            .store
            .snapshot()
            .map_err(failure)?
            .participants
            .into_iter()
            .find(|participant| {
                participant.generation_id() == generation_id
                    && participant.role() == ParticipantRole::Publication
            })
            .ok_or_else(|| ProductionLifecycleError("Publication participant is absent".into()))?;
        let release = Participant::snapshot(
            generation_id,
            ParticipantRole::Publication,
            ParticipantState::ReleaseIntent,
            current.ownership().cloned(),
        )
        .map_err(failure)?;
        let clean = Participant::snapshot(
            generation_id,
            ParticipantRole::Publication,
            ParticipantState::Clean,
            Some(ParticipantOwnership::empty(ParticipantRole::Publication)),
        )
        .map_err(failure)?;
        self.plan
            .store
            .apply(&[
                DomainStoreMutation::PutParticipant(release),
                DomainStoreMutation::CleanParticipantAfterAbsence(clean),
            ])
            .map_err(failure)
    }

    fn stop_sandbox(&mut self) -> Result<(), ProductionLifecycleError> {
        let store = Arc::clone(&self.plan.store);
        let sandbox = SandboxCoordinator::new(&store);
        self.runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("committed runtime is absent".into()))?
            .stop_sandbox(&sandbox)
            .map_err(failure)?;
        self.plan.runtime_staging.cleanup_overlay().map_err(failure)
    }

    fn stop_sidecar(&mut self) -> Result<(), ProductionLifecycleError> {
        let store = Arc::clone(&self.plan.store);
        let network = NetworkCoordinator::new(&store);
        network
            .begin_release(self.plan.generation_identity.generation_id)
            .map_err(failure)?;
        self.runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("committed runtime is absent".into()))?
            .stop_sidecar()
            .map_err(failure)?;
        network
            .complete_clean_after_absence(self.plan.generation_identity.generation_id)
            .map_err(failure)
    }

    fn unmount_state_anchor(&mut self) -> Result<(), ProductionLifecycleError> {
        let coordinator = self.state_anchor();
        self.state_anchor_release = Some(
            coordinator
                .begin_release(
                    self.plan.generation_identity.generation_id,
                    self.plan.state_anchor_layout,
                )
                .map_err(failure)?,
        );
        let runtime = self
            .runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("committed runtime is absent".into()))?;
        let release = self
            .state_anchor_release
            .as_ref()
            .ok_or_else(|| ProductionLifecycleError("State Anchor release is absent".into()))?;
        runtime
            .unmount_state_anchor(release.mount())
            .map_err(failure)?;
        self.state_anchor_absence = Some(
            runtime
                .observe_state_anchor_absence(
                    release.user_namespace(),
                    release.mount_namespace(),
                    release.mount(),
                )
                .map_err(failure)?,
        );
        self.state_anchor_mount = None;
        Ok(())
    }

    fn release_prepared_lower(&mut self) -> Result<(), ProductionLifecycleError> {
        let release = self
            .state_anchor_release
            .as_ref()
            .ok_or_else(|| ProductionLifecycleError("State Anchor release is absent".into()))?;
        let first_observation = self.state_anchor_absence.ok_or_else(|| {
            ProductionLifecycleError("State Anchor absence evidence is absent".into())
        })?;
        let runtime = self
            .runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("committed runtime is absent".into()))?;
        let coordinator = StateAnchorCoordinator::new(&self.plan.store, self.plan.catalog);
        coordinator
            .complete_release_after_unmount(
                release,
                self.plan.state_anchor_layout,
                first_observation,
                || {
                    runtime
                        .observe_state_anchor_absence(
                            release.user_namespace(),
                            release.mount_namespace(),
                            release.mount(),
                        )
                        .map_err(|error| {
                            crate::state_anchor::StateAnchorCoordinatorError::InvalidState(
                                match error {
                                    crate::runtime::SandboxCoordinatorError::InvalidState(
                                        detail,
                                    ) => detail,
                                    _ => "Runtime Supervisor absence observation failed",
                                },
                            )
                        })
                },
            )
            .map_err(failure)?;
        self.state_anchor_release = None;
        self.state_anchor_absence = None;
        Ok(())
    }

    fn graceful_exit(&mut self) -> Result<(), ProductionLifecycleError> {
        self.runtime
            .take()
            .ok_or_else(|| ProductionLifecycleError("committed runtime is absent".into()))?
            .graceful_exit()
            .map_err(failure)
    }
}

impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
    type Error = ProductionLifecycleError;

    fn execute(
        &mut self,
        action: LifecycleAction,
        generation: &Generation,
    ) -> Result<(), Self::Error> {
        self.require_generation(generation)?;
        match action {
            LifecycleAction::CreateTrustAnchor => self.create_trust_anchor(),
            LifecycleAction::StartRuntimeSupervisor => self.start_runtime_supervisor(),
            LifecycleAction::LeasePreparedLower => self.lease_prepared_lower(),
            LifecycleAction::MountStateAnchor => self.mount_state_anchor(),
            LifecycleAction::StartSidecar => self.start_sidecar(),
            LifecycleAction::StartConnectivityHelper => self.start_connectivity_helper(),
            LifecycleAction::MarkNetworkReady => self.mark_network_ready(),
            LifecycleAction::StartSandbox => self.start_sandbox(),
            LifecycleAction::PublishDomain => self.stage_publication(),
            LifecycleAction::StopSandbox => self.stop_sandbox(),
            LifecycleAction::StopSidecar => self.stop_sidecar(),
            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.as_ref() {
                    control.withdraw().map_err(failure)?;
                    self.clean_publication_after_withdrawal()?;
                }
                self.published_control.take();
                Ok(())
            }
            LifecycleAction::FinishTeardown => {
                self.published_control.take();
                Ok(())
            }
            LifecycleAction::StartNewDomain
            | LifecycleAction::StartReadySession
            | LifecycleAction::WaitForDomain
            | LifecycleAction::RestartWaitingDomain
            | LifecycleAction::ArmRuntimeCausalEdge
            | LifecycleAction::AdvanceWaitingSession
            | LifecycleAction::RejectConfigurationMismatch
            | LifecycleAction::RegisterSession
            | LifecycleAction::DisconnectSession => Ok(()),
        }
    }

    fn after_commit(
        &mut self,
        action: LifecycleAction,
        generation: &Generation,
    ) -> Result<(), Self::Error> {
        self.require_generation(generation)?;
        if action == LifecycleAction::PublishDomain {
            self.expose_publication()
        } else {
            Ok(())
        }
    }
}

/// 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| {
            let mut lifecycle = lifecycle.lock().map_err(|_| ControlError::LockPoisoned)?;
            lifecycle
                .disconnect_session(session_id)
                .map_err(|error| ControlError::Lifecycle(error.to_string()))?;
            lifecycle
                .teardown_if_idle()
                .map(|_| ())
                .map_err(|error| ControlError::Lifecycle(error.to_string()))
        },
    )
}

fn path_from_host(bytes: &[u8]) -> &Path {
    use std::os::unix::ffi::OsStrExt;

    Path::new(std::ffi::OsStr::from_bytes(bytes))
}