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

use std::fmt;
use std::io;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::process::Stdio;
use std::time::{Duration, Instant};

use lazaret_core::identity::{
    BundleDigest, ContainerId, GenerationId, HostPath, PidfsProcessIdentityV1,
};
use lazaret_core::records::{
    HostProcessOwnershipSlot, Participant, ParticipantOwnership, ParticipantRole, ParticipantState,
    PersistedHostProcessOwnership, ProcessOwnershipCommitAck, SandboxOwnership,
    StagedProcessOwnership,
};
use lazaret_linux::ipc::private_stream_pair;
use lazaret_linux::process::{
    ChildStandardIo, RetainedProcess, child_broker, current_process_identity,
};
use lazaret_runtime_supervisor::supervision::{
    AttachedRunscSignal, CoordinatorChannel, RuntimeSupervisorLaunch, SupervisorEvent,
};
use rustix::pipe::{PipeFlags, pipe_with};
use rustix::process::Signal;

use crate::domain_store::{DomainStoreError, DomainStoreMutation, HostCoordinatorDomainStore};

const SUPERVISOR_EXIT_TIMEOUT: Duration = Duration::from_secs(5);

#[derive(Debug)]
pub enum SandboxCoordinatorError {
    Store(DomainStoreError),
    Io(io::Error),
    InvalidState(&'static str),
}

impl fmt::Display for SandboxCoordinatorError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Store(error) => {
                write!(formatter, "Sandbox Domain Store operation failed: {error}")
            }
            Self::Io(error) => write!(formatter, "Sandbox process operation failed: {error}"),
            Self::InvalidState(detail) => write!(formatter, "Sandbox state refused: {detail}"),
        }
    }
}

impl std::error::Error for SandboxCoordinatorError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Store(error) => Some(error),
            Self::Io(error) => Some(error),
            Self::InvalidState(_) => None,
        }
    }
}

impl From<DomainStoreError> for SandboxCoordinatorError {
    fn from(error: DomainStoreError) -> Self {
        Self::Store(error)
    }
}

impl From<io::Error> for SandboxCoordinatorError {
    fn from(error: io::Error) -> Self {
        Self::Io(error)
    }
}

pub struct SandboxCoordinator<'store> {
    domain_store: &'store HostCoordinatorDomainStore,
}

impl<'store> SandboxCoordinator<'store> {
    #[must_use]
    pub const fn new(domain_store: &'store HostCoordinatorDomainStore) -> Self {
        Self { domain_store }
    }

    /// Persists bare Sandbox intent and the immutable physical launch snapshot.
    ///
    /// # Errors
    /// Refuses a missing generation, changed physical identity, or existing nonmatching state.
    pub fn begin_intent(
        &self,
        generation_id: GenerationId,
        container_id: ContainerId,
        runtime_root: &Path,
        bundle_digest: BundleDigest,
        host_coordinator: PidfsProcessIdentityV1,
    ) -> Result<(), SandboxCoordinatorError> {
        let snapshot = self.domain_store.snapshot()?;
        if !snapshot
            .generations
            .iter()
            .any(|generation| generation.generation_id == generation_id)
        {
            return Err(SandboxCoordinatorError::InvalidState(
                "Sandbox generation is absent",
            ));
        }
        let ownership = SandboxOwnership {
            container_id: Some(container_id),
            runtime_root: Some(HostPath::new(runtime_root.as_os_str().as_bytes().to_vec())),
            bundle_digest: Some(bundle_digest),
            host_coordinator: PersistedHostProcessOwnership::committed(host_coordinator),
            runtime_supervisor: PersistedHostProcessOwnership::unassigned(),
            attached_runsc: PersistedHostProcessOwnership::unassigned(),
        };
        let intent = sandbox_participant(generation_id, ParticipantState::Intent, ownership)?;
        match sandbox_role(&snapshot.participants, generation_id) {
            None => self.domain_store.apply(&[
                DomainStoreMutation::PutParticipant(Participant::bare_intent(
                    generation_id,
                    ParticipantRole::Sandbox,
                )),
                DomainStoreMutation::PutParticipant(intent),
            ])?,
            Some(current) if current == &intent => {}
            Some(_) => {
                return Err(SandboxCoordinatorError::InvalidState(
                    "Sandbox intent already differs",
                ));
            }
        }
        Ok(())
    }

    /// # Errors
    /// Refuses a mismatched RS tuple or any skipped ownership phase.
    pub fn stage_runtime_supervisor(
        &self,
        staged: StagedProcessOwnership,
    ) -> Result<(), SandboxCoordinatorError> {
        self.require_staged(staged, HostProcessOwnershipSlot::SandboxRuntimeSupervisor)?;
        let (current, mut ownership) = self.intent_ownership(staged.generation_id())?;
        let expected = PersistedHostProcessOwnership::staged(staged.identity());
        if ownership.runtime_supervisor == expected {
            return Ok(());
        }
        if !ownership.runtime_supervisor.is_unassigned() {
            return Err(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor ownership changed",
            ));
        }
        ownership.runtime_supervisor = expected;
        self.put_intent(current.generation_id(), ownership)
    }

    /// # Errors
    /// Refuses an RS staged identity that is not durably exact.
    pub fn acknowledge_runtime_supervisor(
        &self,
        staged: StagedProcessOwnership,
    ) -> Result<ProcessOwnershipCommitAck, SandboxCoordinatorError> {
        self.require_staged(staged, HostProcessOwnershipSlot::SandboxRuntimeSupervisor)?;
        let (_, ownership) = self.intent_ownership(staged.generation_id())?;
        if ownership.runtime_supervisor != PersistedHostProcessOwnership::staged(staged.identity())
        {
            return Err(SandboxCoordinatorError::InvalidState(
                "staged Runtime Supervisor ownership is not durable",
            ));
        }
        Ok(ProcessOwnershipCommitAck::v1(staged))
    }

    /// Commits the executing RS and durably stages its exact attached-runsc child.
    ///
    /// # Errors
    /// Refuses any tuple, phase, or immutable-evidence mismatch.
    pub fn stage_attached_runsc(
        &self,
        staged: StagedProcessOwnership,
    ) -> Result<ProcessOwnershipCommitAck, SandboxCoordinatorError> {
        self.require_staged(staged, HostProcessOwnershipSlot::SandboxAttachedRunsc)?;
        let (current, mut ownership) = self.intent_ownership(staged.generation_id())?;
        let rs = *ownership
            .runtime_supervisor
            .authoritative_pidfs()
            .map_err(|_| {
                SandboxCoordinatorError::InvalidState("Runtime Supervisor is not staged")
            })?;
        if !ownership.runtime_supervisor.is_staged() {
            return Err(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor readiness cannot skip staging",
            ));
        }
        if !ownership.attached_runsc.is_unassigned() {
            return Err(SandboxCoordinatorError::InvalidState(
                "attached runsc ownership changed",
            ));
        }
        ownership.runtime_supervisor = PersistedHostProcessOwnership::committed(rs);
        ownership.attached_runsc = PersistedHostProcessOwnership::staged(staged.identity());
        self.put_intent(current.generation_id(), ownership)?;
        Ok(ProcessOwnershipCommitAck::v1(staged))
    }

    /// Promotes the exact attached runsc only after the RS reports target readiness.
    ///
    /// # Errors
    /// Refuses any incomplete physical snapshot or changed staged identity.
    pub fn commit_ready(
        &self,
        staged: StagedProcessOwnership,
    ) -> Result<(), SandboxCoordinatorError> {
        self.require_staged(staged, HostProcessOwnershipSlot::SandboxAttachedRunsc)?;
        let (_, mut ownership) = self.intent_ownership(staged.generation_id())?;
        if ownership.attached_runsc != PersistedHostProcessOwnership::staged(staged.identity())
            || !ownership.runtime_supervisor.is_current_committed()
            || !ownership.host_coordinator.is_current_committed()
        {
            return Err(SandboxCoordinatorError::InvalidState(
                "Sandbox readiness does not match staged ownership",
            ));
        }
        ownership.attached_runsc = PersistedHostProcessOwnership::committed(staged.identity());
        self.domain_store
            .apply(&[DomainStoreMutation::PutParticipant(sandbox_participant(
                staged.generation_id(),
                ParticipantState::Committed,
                ownership,
            )?)])?;
        Ok(())
    }

    /// # Errors
    /// Refuses any Sandbox participant that is not exactly Committed.
    pub fn begin_release(
        &self,
        generation_id: GenerationId,
    ) -> Result<(), SandboxCoordinatorError> {
        let current = self.participant(generation_id)?;
        if current.state() != ParticipantState::Committed {
            return Err(SandboxCoordinatorError::InvalidState(
                "Sandbox is not committed",
            ));
        }
        let ownership =
            current
                .ownership()
                .cloned()
                .ok_or(SandboxCoordinatorError::InvalidState(
                    "Sandbox ownership is absent",
                ))?;
        self.domain_store
            .apply(&[DomainStoreMutation::PutParticipant(
                Participant::snapshot(
                    generation_id,
                    ParticipantRole::Sandbox,
                    ParticipantState::ReleaseIntent,
                    Some(ownership),
                )
                .map_err(|_| SandboxCoordinatorError::InvalidState("Sandbox release is invalid"))?,
            )])?;
        Ok(())
    }

    /// # Errors
    /// Refuses cleanup without exact `ReleaseIntent` and observed runtime absence.
    pub fn complete_clean_after_absence(
        &self,
        generation_id: GenerationId,
    ) -> Result<(), SandboxCoordinatorError> {
        let current = self.participant(generation_id)?;
        if current.state() != ParticipantState::ReleaseIntent {
            return Err(SandboxCoordinatorError::InvalidState(
                "Sandbox release intent is absent",
            ));
        }
        let clean = Participant::snapshot(
            generation_id,
            ParticipantRole::Sandbox,
            ParticipantState::Clean,
            Some(ParticipantOwnership::empty(ParticipantRole::Sandbox)),
        )
        .map_err(|_| SandboxCoordinatorError::InvalidState("Sandbox Clean is invalid"))?;
        self.domain_store
            .apply(&[DomainStoreMutation::CleanParticipantAfterAbsence(clean)])?;
        Ok(())
    }

    fn require_staged(
        &self,
        staged: StagedProcessOwnership,
        slot: HostProcessOwnershipSlot,
    ) -> Result<(), SandboxCoordinatorError> {
        let snapshot = self.domain_store.snapshot()?;
        let generation = snapshot
            .generations
            .iter()
            .find(|generation| generation.generation_id == staged.generation_id())
            .ok_or(SandboxCoordinatorError::InvalidState(
                "staged Sandbox generation is absent",
            ))?;
        if staged.domain_id() != snapshot.domain.domain_id
            || staged.coordinator_id() != generation.coordinator_id
            || staged.slot() != slot
        {
            return Err(SandboxCoordinatorError::InvalidState(
                "staged Sandbox ownership tuple mismatches generation",
            ));
        }
        Ok(())
    }

    fn intent_ownership(
        &self,
        generation_id: GenerationId,
    ) -> Result<(Participant, SandboxOwnership), SandboxCoordinatorError> {
        let current = self.participant(generation_id)?;
        if current.state() != ParticipantState::Intent {
            return Err(SandboxCoordinatorError::InvalidState(
                "Sandbox acquisition is not in Intent",
            ));
        }
        let Some(ParticipantOwnership::Sandbox(ownership)) = current.ownership() else {
            return Err(SandboxCoordinatorError::InvalidState(
                "Sandbox ownership shape is invalid",
            ));
        };
        Ok((current.clone(), ownership.clone()))
    }

    fn put_intent(
        &self,
        generation_id: GenerationId,
        ownership: SandboxOwnership,
    ) -> Result<(), SandboxCoordinatorError> {
        self.domain_store
            .apply(&[DomainStoreMutation::PutParticipant(sandbox_participant(
                generation_id,
                ParticipantState::Intent,
                ownership,
            )?)])?;
        Ok(())
    }

    fn participant(
        &self,
        generation_id: GenerationId,
    ) -> Result<Participant, SandboxCoordinatorError> {
        sandbox_role(&self.domain_store.snapshot()?.participants, generation_id)
            .cloned()
            .ok_or(SandboxCoordinatorError::InvalidState(
                "Sandbox participant is absent",
            ))
    }
}

pub struct AwaitingAttachedRunsc {
    supervisor: Option<RetainedProcess>,
    control: Option<CoordinatorChannel>,
    lifetime: Option<std::os::fd::OwnedFd>,
    staged: Option<StagedProcessOwnership>,
}

impl AwaitingAttachedRunsc {
    /// Launches the RS through the HC broker and commits its staged ownership before execution.
    ///
    /// # Errors
    /// Refuses any pipe, channel, broker, staging, or acknowledgment failure.
    pub fn launch(
        coordinator: &SandboxCoordinator<'_>,
        launch: &RuntimeSupervisorLaunch,
    ) -> Result<Self, SandboxCoordinatorError> {
        let parent = current_process_identity()?;
        let (coordinator_control, supervisor_control) = private_stream_pair()?;
        let (supervisor_lifetime, coordinator_lifetime) =
            pipe_with(PipeFlags::CLOEXEC).map_err(io::Error::from)?;
        let process = child_broker().acquire_with_stdio(
            parent,
            launch.ownership(),
            launch.target(),
            ChildStandardIo::redirected(
                Stdio::from(supervisor_control),
                Stdio::from(supervisor_lifetime),
            ),
            |staged| {
                coordinator
                    .stage_runtime_supervisor(staged)
                    .map_err(io::Error::other)
            },
            |staged| {
                coordinator
                    .acknowledge_runtime_supervisor(staged)
                    .map_err(io::Error::other)
            },
        )?;
        Ok(Self {
            supervisor: Some(process),
            control: Some(CoordinatorChannel::new(coordinator_control)),
            lifetime: Some(coordinator_lifetime),
            staged: None,
        })
    }

    /// Receives and durably stages the exact RS child without releasing its barrier.
    ///
    /// # Errors
    /// Refuses unexpected, duplicate, or mismatched messages and durable-state failures.
    pub fn receive_staged(
        &mut self,
        coordinator: &SandboxCoordinator<'_>,
    ) -> Result<StagedProcessOwnership, SandboxCoordinatorError> {
        if self.staged.is_some() {
            return Err(SandboxCoordinatorError::InvalidState(
                "attached runsc stage was already received",
            ));
        }
        let event = self
            .control
            .as_ref()
            .ok_or(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor control channel is absent",
            ))?
            .receive()?;
        let staged = match event {
            SupervisorEvent::Staged(staged) => staged,
            SupervisorEvent::Refused(detail) => {
                return Err(SandboxCoordinatorError::Io(io::Error::other(detail)));
            }
            _ => {
                return Err(SandboxCoordinatorError::InvalidState(
                    "Runtime Supervisor sent readiness before staged ownership",
                ));
            }
        };
        let acknowledgment = coordinator.stage_attached_runsc(staged)?;
        self.staged = Some(acknowledgment.committed());
        Ok(staged)
    }

    /// Releases the exact pending runsc barrier with its durable typed acknowledgment.
    ///
    /// # Errors
    /// Refuses a different or absent staged tuple or channel failure.
    pub fn acknowledge(
        mut self,
        staged: StagedProcessOwnership,
    ) -> Result<AwaitingSandboxReadiness, SandboxCoordinatorError> {
        if self.staged != Some(staged) {
            return Err(SandboxCoordinatorError::InvalidState(
                "attached runsc acknowledgment does not match pending stage",
            ));
        }
        self.control
            .as_ref()
            .ok_or(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor control channel is absent",
            ))?
            .acknowledge_staged(ProcessOwnershipCommitAck::v1(staged))?;
        Ok(AwaitingSandboxReadiness {
            supervisor: self.supervisor.take(),
            control: self.control.take(),
            lifetime: self.lifetime.take(),
            staged,
        })
    }

    /// # Errors
    /// Returns an error if the exact RS cannot be terminated and reaped.
    pub fn abort(mut self) -> Result<(), SandboxCoordinatorError> {
        self.lifetime.take();
        cleanup_supervisor(&mut self.supervisor)?;
        Ok(())
    }
}

impl Drop for AwaitingAttachedRunsc {
    fn drop(&mut self) {
        self.lifetime.take();
        let _ = cleanup_supervisor(&mut self.supervisor);
    }
}

pub struct AwaitingSandboxReadiness {
    supervisor: Option<RetainedProcess>,
    control: Option<CoordinatorChannel>,
    lifetime: Option<std::os::fd::OwnedFd>,
    staged: StagedProcessOwnership,
}

impl AwaitingSandboxReadiness {
    /// Commits only the exact ready attached runsc after RS capability removal.
    ///
    /// # Errors
    /// Refuses a mismatched readiness tuple, durable commit failure, or channel failure.
    pub fn commit(
        mut self,
        coordinator: &SandboxCoordinator<'_>,
    ) -> Result<CommittedSandboxRuntime, SandboxCoordinatorError> {
        let control = self
            .control
            .as_ref()
            .ok_or(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor control channel is absent",
            ))?;
        let event = control.receive()?;
        let staged = match event {
            SupervisorEvent::Ready { staged, .. } => staged,
            SupervisorEvent::Refused(detail) => {
                return Err(SandboxCoordinatorError::Io(io::Error::other(detail)));
            }
            _ => {
                return Err(SandboxCoordinatorError::InvalidState(
                    "Runtime Supervisor did not report sandbox readiness",
                ));
            }
        };
        if staged != self.staged {
            return Err(SandboxCoordinatorError::InvalidState(
                "sandbox readiness changed attached-runsc ownership",
            ));
        }
        coordinator.commit_ready(staged)?;
        control.acknowledge_sandbox_committed()?;
        Ok(CommittedSandboxRuntime {
            generation_id: staged.generation_id(),
            supervisor: self.supervisor.take(),
            control: self
                .control
                .take()
                .ok_or(SandboxCoordinatorError::InvalidState(
                    "Runtime Supervisor control channel is absent",
                ))?,
            lifetime: self.lifetime.take(),
        })
    }
}

impl Drop for AwaitingSandboxReadiness {
    fn drop(&mut self) {
        self.lifetime.take();
        let _ = cleanup_supervisor(&mut self.supervisor);
    }
}

pub struct CommittedSandboxRuntime {
    generation_id: GenerationId,
    supervisor: Option<RetainedProcess>,
    control: CoordinatorChannel,
    lifetime: Option<std::os::fd::OwnedFd>,
}

impl CommittedSandboxRuntime {
    #[must_use]
    pub fn runtime_supervisor_identity(&self) -> Option<PidfsProcessIdentityV1> {
        self.supervisor.as_ref().map(RetainedProcess::identity)
    }

    /// # Errors
    /// Returns an error if the exact retained RS pidfd action fails.
    pub fn signal_runtime_supervisor(&self, signal: Signal) -> Result<(), SandboxCoordinatorError> {
        self.supervisor
            .as_ref()
            .ok_or(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor handle is absent",
            ))?
            .send_signal(signal)?;
        Ok(())
    }

    /// # Errors
    /// Returns an error if the RS cannot apply the exact retained-runsc pidfd action.
    pub fn signal_attached_runsc(
        &self,
        signal: AttachedRunscSignal,
    ) -> Result<(), SandboxCoordinatorError> {
        self.control.signal_attached_runsc(signal)?;
        if self.control.receive()? == SupervisorEvent::SignalApplied(signal) {
            Ok(())
        } else {
            Err(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor did not confirm the exact runsc action",
            ))
        }
    }

    /// Waits for an exact attached-runsc loss report and reaps the retained RS.
    ///
    /// # Errors
    /// Refuses any other event or an RS that does not then exit.
    pub fn wait_for_attached_runsc_exit(mut self) -> Result<(), SandboxCoordinatorError> {
        if self.control.receive()? != SupervisorEvent::AttachedRunscExited {
            return Err(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor did not report attached-runsc loss",
            ));
        }
        self.lifetime.take();
        wait_and_reap_supervisor(&mut self.supervisor)
    }

    /// Waits for an externally stopped exact RS handle to exit and reaps it.
    ///
    /// # Errors
    /// Refuses an RS that remains present beyond the bounded observation.
    pub fn wait_for_runtime_supervisor_exit(mut self) -> Result<(), SandboxCoordinatorError> {
        self.lifetime.take();
        wait_and_reap_supervisor(&mut self.supervisor)
    }

    /// Performs ordered normal stop and reaps the same retained RS handle.
    ///
    /// # Errors
    /// Refuses missing release intent, unexpected RS response, or exact cleanup failure.
    pub fn stop(
        mut self,
        coordinator: &SandboxCoordinator<'_>,
    ) -> Result<(), SandboxCoordinatorError> {
        coordinator.begin_release(self.generation_id)?;
        self.control.stop()?;
        if self.control.receive()? != SupervisorEvent::Stopped {
            return Err(SandboxCoordinatorError::InvalidState(
                "Runtime Supervisor did not complete sandbox stop",
            ));
        }
        self.lifetime.take();
        wait_and_reap_supervisor(&mut self.supervisor)?;
        coordinator.complete_clean_after_absence(self.generation_id)
    }
}

impl Drop for CommittedSandboxRuntime {
    fn drop(&mut self) {
        self.lifetime.take();
        let _ = cleanup_supervisor(&mut self.supervisor);
    }
}

fn sandbox_participant(
    generation_id: GenerationId,
    state: ParticipantState,
    ownership: SandboxOwnership,
) -> Result<Participant, SandboxCoordinatorError> {
    Participant::snapshot(
        generation_id,
        ParticipantRole::Sandbox,
        state,
        Some(ParticipantOwnership::Sandbox(ownership)),
    )
    .map_err(|_| SandboxCoordinatorError::InvalidState("Sandbox participant shape is invalid"))
}

fn sandbox_role(participants: &[Participant], generation_id: GenerationId) -> Option<&Participant> {
    participants.iter().find(|participant| {
        participant.generation_id() == generation_id
            && participant.role() == ParticipantRole::Sandbox
    })
}

fn cleanup_supervisor(
    supervisor: &mut Option<RetainedProcess>,
) -> Result<(), SandboxCoordinatorError> {
    let Some(process) = supervisor.take() else {
        return Ok(());
    };
    if !process.has_exited()? {
        process.send_signal(Signal::KILL)?;
    }
    child_broker().reap(process)?;
    Ok(())
}

fn wait_and_reap_supervisor(
    supervisor: &mut Option<RetainedProcess>,
) -> Result<(), SandboxCoordinatorError> {
    let Some(process) = supervisor.take() else {
        return Err(SandboxCoordinatorError::InvalidState(
            "Runtime Supervisor handle is absent",
        ));
    };
    let deadline = Instant::now() + SUPERVISOR_EXIT_TIMEOUT;
    while !process.has_exited()? && Instant::now() < deadline {
        std::thread::sleep(Duration::from_millis(5));
    }
    if !process.has_exited()? {
        process.send_signal(Signal::KILL)?;
        child_broker().reap(process)?;
        return Err(SandboxCoordinatorError::InvalidState(
            "Runtime Supervisor did not exit after sandbox stop",
        ));
    }
    child_broker().reap(process)?;
    Ok(())
}

#[cfg(test)]
mod tests;