lazaret/crates/lazaret-host-coordinator/tests/client-control-face.rs

use std::fs::OpenOptions;
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use std::sync::Arc;

use lazaret_client::{ControlEndpoint, execute_over_control_face};
use lazaret_core::identity::{
    ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId, HostPath, NetworkSidecarId,
    SandboxId,
};
use lazaret_core::records::{
    Domain, DurableDecision, Generation, GenerationIdentity, NonInteractiveOutput,
    NonInteractiveSessionMessage, NonInteractiveSessionRequest, SessionExit,
};
use lazaret_host_coordinator::control::{
    ControlError, NonInteractiveSessionExecutor, PublishedControlFace, TrustAnchor,
};
use lazaret_host_coordinator::domain_store::{DomainStoreMutation, HostCoordinatorDomainStore};
use lazaret_linux::filesystem::try_exclusive_flock;

struct TestDirectory(PathBuf);

impl TestDirectory {
    fn new() -> Self {
        let path = std::env::temp_dir().join(format!(
            "lazaret-client-control-{}-{}",
            std::process::id(),
            random_suffix()
        ));
        std::fs::create_dir(&path).unwrap();
        std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700)).unwrap();
        Self(path)
    }
}

impl Drop for TestDirectory {
    fn drop(&mut self) {
        std::fs::remove_dir_all(&self.0).unwrap();
    }
}

struct Executor;

impl NonInteractiveSessionExecutor for Executor {
    fn execute(
        &self,
        request: &NonInteractiveSessionRequest,
        emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
    ) -> Result<SessionExit, ControlError> {
        assert_eq!(request.argv, [b"tool".to_vec(), b"argument".to_vec()]);
        emit(NonInteractiveSessionMessage::Stdout(NonInteractiveOutput {
            session_id: request.session_id,
            bytes: b"first".to_vec(),
        }))?;
        emit(NonInteractiveSessionMessage::Stderr(NonInteractiveOutput {
            session_id: request.session_id,
            bytes: b"second".to_vec(),
        }))?;
        Ok(SessionExit::exited(request.session_id, 19))
    }
}

#[test]
fn client_enrolls_registers_forwards_both_streams_and_returns_guest_status() {
    let directory = TestDirectory::new();
    let identity = generation_identity();
    let store = open_store(&directory, &identity);
    let face = Arc::new(
        PublishedControlFace::publish_with_executor(
            &directory.0,
            rustix::process::getuid().as_raw(),
            TrustAnchor::create(identity).unwrap(),
            store,
            Arc::new(Executor),
        )
        .unwrap(),
    );
    let endpoint = ControlEndpoint::new(
        face.socket_path().to_path_buf(),
        face.publication_identity().clone(),
    );
    let server = {
        let face = Arc::clone(&face);
        std::thread::spawn(move || {
            face.serve_next().unwrap();
            face.serve_next().unwrap();
        })
    };
    let mut stdout = Vec::new();
    let mut stderr = Vec::new();

    let status = execute_over_control_face(
        &endpoint,
        &[b"tool".to_vec(), b"argument".to_vec()],
        &mut stdout,
        &mut stderr,
    )
    .unwrap();

    server.join().unwrap();
    assert_eq!(status, 19);
    assert_eq!(stdout, b"first");
    assert_eq!(stderr, b"second");
}

fn generation_identity() -> GenerationIdentity {
    GenerationIdentity {
        agent_profile: "trusted".into(),
        domain_id: DomainId::new([1; 32]),
        generation_id: GenerationId::new(1),
        configuration_fingerprint: ConfigurationFingerprint::new([2; 32]),
        coordinator_id: CoordinatorId::new([3; 32]),
        sandbox_id: SandboxId::new([4; 32]),
    }
}

fn open_store(
    directory: &TestDirectory,
    identity: &GenerationIdentity,
) -> Arc<HostCoordinatorDomainStore> {
    let lock = OpenOptions::new()
        .read(true)
        .write(true)
        .create_new(true)
        .open(directory.0.join("domain.lock"))
        .unwrap();
    let domain = Domain {
        domain_id: identity.domain_id,
        profile_name: identity.agent_profile.clone(),
        canonical_workspace: HostPath::new(b"/workspace/project".to_vec()),
    };
    let store = Arc::new(
        HostCoordinatorDomainStore::open(
            try_exclusive_flock(lock).unwrap(),
            &directory.0.join("domain.redb"),
            &domain,
        )
        .unwrap(),
    );
    for durable_decision in [
        DurableDecision::None,
        DurableDecision::Preparing,
        DurableDecision::Published,
    ] {
        store
            .apply(&[DomainStoreMutation::PutGeneration(Generation {
                domain_id: identity.domain_id,
                generation_id: identity.generation_id,
                configuration_fingerprint: identity.configuration_fingerprint,
                coordinator_id: identity.coordinator_id,
                network_sidecar_id: NetworkSidecarId::new([5; 32]),
                sandbox_id: identity.sandbox_id,
                durable_decision,
            })])
            .unwrap();
    }
    store
}

fn random_suffix() -> u64 {
    let mut bytes = [0_u8; 8];
    getrandom::fill(&mut bytes).unwrap();
    u64::from_ne_bytes(bytes)
}