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

use std::cell::RefCell;
use std::collections::BTreeMap;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::fs::File;
use std::io::Write;
use std::os::fd::AsFd;
use std::os::unix::fs::MetadataExt;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use lazaret_core::configuration::{
    ConfigurationLoadRequest, ResolvedDomainConfiguration, load_trusted_configuration,
};
use lazaret_core::identity::{
    ClientId, ConfigurationFingerprint, DomainId, SessionAnchorId, SessionId,
};
use lazaret_core::records::{Domain, Session, SessionPhase};
use lazaret_linux::filesystem::try_exclusive_flock;
use lazaret_runtime_supervisor::sandbox::production_sandbox_layout;
use rustix::fs::{Mode, OFlags, openat};

use crate::bootstrap::{ProductionHostCoordinatorLifecycle, serve_production_control};
use crate::domain_store::HostCoordinatorDomainStore;
use crate::production::{BootstrapOwnerGuard, ProductionColdStart, ProductionColdStartError};
use crate::production_selector::{InstalledProductPaths, select_production_cold_start};
use crate::waiter::BootstrapWaiterFace;

pub const PRIVATE_CLIENT_BOOTSTRAP_ROLE: &str = "--lazaret-private-client-bootstrap-v1";

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

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

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

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

struct BootstrapArguments {
    profile: String,
    workspace: PathBuf,
    configuration: Option<PathBuf>,
    cli_id: ClientId,
    session_id: SessionId,
    session_anchor_id: SessionAnchorId,
    expected_domain_id: DomainId,
    expected_configuration_fingerprint: ConfigurationFingerprint,
}

/// Dispatches the exact private Client cold-start entry.
#[must_use]
pub fn dispatch_private_client_bootstrap(
    arguments: &[OsString],
) -> Option<Result<(), PrivateClientBootstrapError>> {
    if arguments
        .first()
        .is_none_or(|argument| argument != PRIVATE_CLIENT_BOOTSTRAP_ROLE)
    {
        return None;
    }
    Some(run_private_client_bootstrap(arguments))
}

fn run_private_client_bootstrap(arguments: &[OsString]) -> Result<(), PrivateClientBootstrapError> {
    let parsed = parse_arguments(arguments).inspect_err(|_| signal(b'E'))?;
    let owner = BootstrapOwnerGuard::from_stdin()
        .map_err(failure)
        .inspect_err(|_| signal(b'E'))?;
    let configuration = load_configuration(&parsed)
        .map_err(failure)
        .inspect_err(|_| signal(b'E'))?;
    if !matches_client_expectation(
        configuration.domain_id(),
        configuration.configuration_fingerprint(),
        &parsed,
    ) {
        signal(b'E');
        return Err(failure(
            "independently resolved Domain identity differs from the Client expectation",
        ));
    }
    let persistent_domain = configuration
        .persistent_domain_root()
        .open_authority()
        .map_err(failure)
        .inspect_err(|_| signal(b'E'))?;
    let lock_file = open_private_domain_file(&persistent_domain, "domain.lock")
        .inspect_err(|_| signal(b'E'))?;
    let domain_lock = match try_exclusive_flock(lock_file) {
        Ok(lock) => lock,
        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
            signal(b'W');
            return Ok(());
        }
        Err(error) => {
            signal(b'E');
            return Err(failure(error));
        }
    };
    let domain = Domain {
        domain_id: configuration.domain_id(),
        profile_name: configuration.selected_profile().to_owned(),
        canonical_workspace: configuration.workspace().canonical_label().clone(),
    };
    let store = Arc::new(
        HostCoordinatorDomainStore::open_file(
            domain_lock,
            open_private_domain_file(&persistent_domain, "domain.redb")
                .inspect_err(|_| signal(b'E'))?,
            &domain,
        )
        .map_err(failure)
        .inspect_err(|_| signal(b'E'))?,
    );
    signal(b'L');

    let executable = std::env::current_exe().map_err(failure)?;
    let installed = executable
        .parent()
        .ok_or_else(|| failure("Host Coordinator executable has no installation directory"))?;
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .map_err(failure)?;
    let selection = runtime
        .block_on(select_production_cold_start(
            &configuration,
            Arc::clone(&store),
            InstalledProductPaths {
                escalation_client: &installed.join("lazaret-escalation-client"),
                sandbox_init: &executable,
            },
        ))
        .map_err(failure)?;
    let inputs = selection
        .cold_start_inputs(&configuration)
        .map_err(failure)?;
    let bootstrap_owner = Session {
        generation_id: inputs.generation.generation_id,
        session_id: parsed.session_id,
        cli_id: parsed.cli_id,
        session_anchor_id: parsed.session_anchor_id,
        session_phase: SessionPhase::New,
    };
    ProductionColdStart::stage(inputs)
        .map_err(failure)?
        .run(bootstrap_owner, &owner, serve_published)
        .map_err(failure)
}

fn serve_published(
    lifecycle: &Mutex<ProductionHostCoordinatorLifecycle<'_, '_>>,
    waiters: &RefCell<BootstrapWaiterFace>,
    bootstrap_owner: &BootstrapOwnerGuard,
    bootstrap_owner_session_id: SessionId,
) -> Result<(), ProductionColdStartError> {
    lifecycle
        .lock()
        .map_err(|_| ProductionColdStartError::service("production lifecycle lock is poisoned"))?
        .adapter()
        .published_control()
        .ok_or_else(|| ProductionColdStartError::service("control face is absent"))?
        .set_nonblocking(true)
        .map_err(ProductionColdStartError::service)?;
    loop {
        {
            let mut lifecycle = lifecycle.lock().map_err(|_| {
                ProductionColdStartError::service("production lifecycle lock is poisoned")
            })?;
            lifecycle
                .reconcile_published_bootstrap_owner(
                    bootstrap_owner_session_id,
                    bootstrap_owner
                        .is_live()
                        .map_err(ProductionColdStartError::service)?,
                )
                .map_err(ProductionColdStartError::service)?;
            if lifecycle.adapter().published_control().is_none() {
                return Ok(());
            }
            waiters
                .borrow_mut()
                .reconcile_published(&mut lifecycle)
                .map_err(ProductionColdStartError::service)?;
        }
        match serve_production_control(lifecycle) {
            Ok(()) => {}
            Err(crate::control::ControlError::Io(error))
                if error.kind() == std::io::ErrorKind::WouldBlock =>
            {
                std::thread::sleep(std::time::Duration::from_millis(10));
            }
            Err(error) => return Err(ProductionColdStartError::service(error)),
        }
        if lifecycle
            .lock()
            .map_err(|_| {
                ProductionColdStartError::service("production lifecycle lock is poisoned")
            })?
            .adapter()
            .published_control()
            .is_none()
        {
            return Ok(());
        }
    }
}

fn matches_client_expectation(
    domain_id: DomainId,
    configuration_fingerprint: ConfigurationFingerprint,
    arguments: &BootstrapArguments,
) -> bool {
    domain_id == arguments.expected_domain_id
        && configuration_fingerprint == arguments.expected_configuration_fingerprint
}

fn parse_arguments(
    arguments: &[OsString],
) -> Result<BootstrapArguments, PrivateClientBootstrapError> {
    let mut arguments = arguments.iter().skip(1);
    require_option(&mut arguments, "--profile")?;
    let profile = required(&mut arguments, "profile")?
        .to_str()
        .ok_or_else(|| failure("profile is not UTF-8"))?
        .to_owned();
    require_option(&mut arguments, "--workspace")?;
    let workspace = PathBuf::from(required(&mut arguments, "workspace")?);
    let mut configuration = None;
    let next = required(&mut arguments, "identity option")?;
    let client_option = if next == OsStr::new("--config") {
        configuration = Some(PathBuf::from(required(&mut arguments, "configuration")?));
        required(&mut arguments, "Client ID option")?
    } else {
        next
    };
    if client_option != OsStr::new("--client-id") {
        return Err(failure("expected --client-id"));
    }
    let cli_id = ClientId::new(decode_id(required(&mut arguments, "Client ID")?)?);
    require_option(&mut arguments, "--session-id")?;
    let session_id = SessionId::new(decode_id(required(&mut arguments, "Session ID")?)?);
    require_option(&mut arguments, "--anchor-id")?;
    let session_anchor_id =
        SessionAnchorId::new(decode_id(required(&mut arguments, "Session Anchor ID")?)?);
    require_option(&mut arguments, "--domain-id")?;
    let expected_domain_id = DomainId::new(decode_id(required(&mut arguments, "Domain ID")?)?);
    require_option(&mut arguments, "--configuration-fingerprint")?;
    let expected_configuration_fingerprint = ConfigurationFingerprint::new(decode_id(required(
        &mut arguments,
        "Configuration Fingerprint",
    )?)?);
    if arguments.next().is_some() {
        return Err(failure("unexpected private bootstrap argument"));
    }
    Ok(BootstrapArguments {
        profile,
        workspace,
        configuration,
        cli_id,
        session_id,
        session_anchor_id,
        expected_domain_id,
        expected_configuration_fingerprint,
    })
}

fn open_private_domain_file(
    root: &impl AsFd,
    name: &str,
) -> Result<File, PrivateClientBootstrapError> {
    let descriptor = openat(
        root,
        name,
        OFlags::RDWR | OFlags::CREATE | OFlags::CLOEXEC | OFlags::NOFOLLOW,
        Mode::RUSR | Mode::WUSR,
    )
    .map_err(failure)?;
    let file = File::from(descriptor);
    let metadata = file.metadata().map_err(failure)?;
    if !metadata.file_type().is_file()
        || metadata.uid() != rustix::process::getuid().as_raw()
        || metadata.mode() & 0o7777 != 0o600
        || metadata.nlink() != 1
    {
        return Err(failure(
            "private Domain file is not exact invoking-user mode-0600 authority",
        ));
    }
    Ok(file)
}

fn required<'a>(
    arguments: &mut impl Iterator<Item = &'a OsString>,
    label: &str,
) -> Result<&'a OsStr, PrivateClientBootstrapError> {
    arguments
        .next()
        .map(OsString::as_os_str)
        .ok_or_else(|| failure(format!("missing {label}")))
}

fn require_option<'a>(
    arguments: &mut impl Iterator<Item = &'a OsString>,
    expected: &str,
) -> Result<(), PrivateClientBootstrapError> {
    if required(arguments, expected)? == OsStr::new(expected) {
        Ok(())
    } else {
        Err(failure(format!("expected {expected}")))
    }
}

fn decode_id(value: &OsStr) -> Result<[u8; 32], PrivateClientBootstrapError> {
    let value = value
        .to_str()
        .ok_or_else(|| failure("identity is not UTF-8"))?;
    if value.len() != 64 {
        return Err(failure("identity is not 32-byte lowercase hexadecimal"));
    }
    let mut decoded = [0_u8; 32];
    for (output, pair) in decoded.iter_mut().zip(value.as_bytes().chunks_exact(2)) {
        *output = match (hex_nibble(pair[0]), hex_nibble(pair[1])) {
            (Some(high), Some(low)) => high << 4 | low,
            _ => return Err(failure("identity is not 32-byte lowercase hexadecimal")),
        };
    }
    Ok(decoded)
}

const fn hex_nibble(byte: u8) -> Option<u8> {
    match byte {
        b'0'..=b'9' => Some(byte - b'0'),
        b'a'..=b'f' => Some(byte - b'a' + 10),
        _ => None,
    }
}

fn load_configuration(
    arguments: &BootstrapArguments,
) -> Result<ResolvedDomainConfiguration, Box<dyn std::error::Error>> {
    let home = required_environment_path("HOME")?;
    let runtime = required_environment_path("XDG_RUNTIME_DIR")?;
    Ok(load_trusted_configuration(ConfigurationLoadRequest {
        explicit_configuration_path: arguments.configuration.clone(),
        selected_profile: arguments.profile.clone(),
        workspace: arguments.workspace.clone(),
        home,
        xdg_config_home: optional_environment_path("XDG_CONFIG_HOME"),
        xdg_data_home: optional_environment_path("XDG_DATA_HOME"),
        xdg_runtime_dir: runtime,
        environment: environment(),
        sandbox_layout: production_sandbox_layout()?,
    })?)
}

fn required_environment_path(name: &str) -> Result<PathBuf, PrivateClientBootstrapError> {
    std::env::var_os(name)
        .filter(|value| !value.is_empty())
        .map(PathBuf::from)
        .ok_or_else(|| failure(format!("{name} is not set")))
}

fn optional_environment_path(name: &str) -> Option<PathBuf> {
    std::env::var_os(name)
        .filter(|value| !value.is_empty())
        .map(PathBuf::from)
}

fn environment() -> BTreeMap<String, OsString> {
    std::env::vars_os()
        .filter_map(|(name, value)| name.into_string().ok().map(|name| (name, value)))
        .collect()
}

fn signal(byte: u8) {
    let _ = std::io::stdout().write_all(&[byte]);
    let _ = std::io::stdout().flush();
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use super::*;

    #[test]
    fn private_bootstrap_arguments_are_exact_and_public() {
        let arguments = [
            PRIVATE_CLIENT_BOOTSTRAP_ROLE,
            "--profile",
            "agent",
            "--workspace",
            "/workspace/with spaces",
            "--config",
            "/host/config.toml",
            "--client-id",
            &"11".repeat(32),
            "--session-id",
            &"22".repeat(32),
            "--anchor-id",
            &"33".repeat(32),
            "--domain-id",
            &"44".repeat(32),
            "--configuration-fingerprint",
            &"55".repeat(32),
        ]
        .map(OsString::from);
        let parsed = parse_arguments(&arguments).unwrap();
        assert_eq!(parsed.profile, "agent");
        assert_eq!(parsed.workspace, Path::new("/workspace/with spaces"));
        assert_eq!(
            parsed.configuration.as_deref(),
            Some(Path::new("/host/config.toml"))
        );
        assert_eq!(parsed.cli_id, ClientId::new([0x11; 32]));
        assert_eq!(parsed.session_id, SessionId::new([0x22; 32]));
        assert_eq!(parsed.session_anchor_id, SessionAnchorId::new([0x33; 32]));
        assert_eq!(parsed.expected_domain_id, DomainId::new([0x44; 32]));
        assert_eq!(
            parsed.expected_configuration_fingerprint,
            ConfigurationFingerprint::new([0x55; 32])
        );
        assert!(matches_client_expectation(
            DomainId::new([0x44; 32]),
            ConfigurationFingerprint::new([0x55; 32]),
            &parsed
        ));
        assert!(!matches_client_expectation(
            DomainId::new([0x44; 32]),
            ConfigurationFingerprint::new([0x56; 32]),
            &parsed
        ));
    }

    #[test]
    fn private_bootstrap_dispatch_and_identity_grammar_are_closed() {
        assert!(dispatch_private_client_bootstrap(&[]).is_none());
        let extra = [
            PRIVATE_CLIENT_BOOTSTRAP_ROLE,
            "--profile",
            "agent",
            "--workspace",
            "/workspace",
            "--client-id",
            &"11".repeat(32),
            "--session-id",
            &"22".repeat(32),
            "--anchor-id",
            &"AA".repeat(32),
            "--domain-id",
            &"44".repeat(32),
            "--configuration-fingerprint",
            &"55".repeat(32),
        ]
        .map(OsString::from);
        assert!(parse_arguments(&extra).is_err());
    }
}