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

use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::Arc;

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, 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::{PublishedControlFace, SessionRegistry, StagedControlFace, TrustAnchor};
use crate::domain_store::HostCoordinatorDomainStore;
use crate::lifecycle::{
    HostCoordinatorLifecycle, LifecycleAction, LifecycleAdapter, SessionAdmissionIdentity,
    SessionLifecycleAdapter,
};
use crate::network::NetworkCoordinator;
use crate::runtime::{AwaitingAttachedRunsc, CommittedSandboxRuntime, SandboxCoordinator};
use crate::state_anchor::{StateAnchorCoordinator, StateAnchorReleaseAuthority};

#[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 {}

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 struct ProductionLifecycleAdapter<'a> {
    plan: ProductionBootstrapPlan<'a>,
    registry: Arc<SessionRegistry>,
    trust_anchor: Option<TrustAnchor>,
    staged_control: Option<StagedControlFace>,
    published_control: Option<PublishedControlFace>,
    awaiting_runtime: Option<AwaitingAttachedRunsc>,
    runtime: Option<CommittedSandboxRuntime>,
    state_anchor_mount: Option<MountIdentity>,
    state_anchor_release: Option<StateAnchorReleaseAuthority>,
}

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,
            },
            ProductionSessionRegistry { registry },
        )
    }

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

    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.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(failure)?;
        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(failure)
    }

    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(failure)
    }

    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(failure)
    }

    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(failure)?;
        let handoff = self
            .plan
            .sandbox_bundle
            .capture_handoff_authority(self.plan.configuration)
            .map_err(failure)?;
        let (runtime, _) = pending
            .acknowledge(staged, &handoff)
            .map_err(failure)?
            .commit_networked(&sandbox)
            .map_err(failure)?;
        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()))?;
        self.staged_control = Some(
            StagedControlFace::stage(
                path_from_host(
                    self.plan
                        .configuration
                        .runtime_domain_root()
                        .canonical_label()
                        .as_bytes(),
                ),
                self.plan.invoking_uid,
                anchor,
                Arc::clone(&self.registry),
            )
            .map_err(failure)?,
        );
        Ok(())
    }

    fn expose_publication(&mut self) -> Result<(), ProductionLifecycleError> {
        self.published_control = Some(
            self.staged_control
                .take()
                .ok_or_else(|| ProductionLifecycleError("staged control face is absent".into()))?
                .expose()
                .map_err(failure)?,
        );
        Ok(())
    }

    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)
    }

    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)?,
        );
        self.runtime
            .as_mut()
            .ok_or_else(|| ProductionLifecycleError("committed runtime is absent".into()))?
            .unmount_state_anchor()
            .map_err(failure)?;
        self.state_anchor_mount = None;
        Ok(())
    }

    fn release_prepared_lower(&mut self) -> Result<(), ProductionLifecycleError> {
        let release = self
            .state_anchor_release
            .take()
            .ok_or_else(|| ProductionLifecycleError("State Anchor release is absent".into()))?;
        self.state_anchor()
            .complete_release_after_unmount(&release, self.plan.state_anchor_layout)
            .map_err(failure)
    }

    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::FinishTeardown => {
                self.published_control.take();
                Ok(())
            }
            LifecycleAction::StartNewDomain
            | LifecycleAction::StartReadySession
            | LifecycleAction::WaitForDomain
            | LifecycleAction::RestartWaitingDomain
            | LifecycleAction::ArmRuntimeCausalEdge
            | LifecycleAction::AdvanceWaitingSession
            | LifecycleAction::RejectConfigurationMismatch
            | LifecycleAction::RegisterSession
            | LifecycleAction::DisconnectSession
            | LifecycleAction::BeginTeardown => 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(())
        }
    }
}

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

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