lazaret/crates/lazaret-client/src/lib.rs

use std::fmt;
use std::fs::File;
use std::io::{self, Read, Write};
use std::net::Shutdown;
use std::path::PathBuf;
use std::sync::{Mutex, mpsc};
use std::thread::JoinHandle;
use std::time::Duration;

use ed25519_dalek::SigningKey;
use lazaret_control_client::{ControlClient, ControlClientError, PendingClientEnrollment};
use lazaret_core::identity::{ClientId, SessionAnchorId, SessionId};
use lazaret_core::records::{
    ClientEnrollmentResponse, ControlRequest, ControlResponse, DeveloperDecision,
    DeveloperDecisionKind, InteractiveSessionBatchV2, InteractiveSessionEventV2,
    InteractiveSessionMessage, InteractiveSessionRequestV2, NonInteractiveSessionMessage,
    NonInteractiveSessionRequest, ProtocolPublicationIdentity, ProtocolVersion, SessionBinding,
    SessionExitOutcome, SessionRegistrationRequest, SessionTerminalMessage, TlsSubjectId,
};
use sha2::{Digest, Sha256};

use crate::terminal::{HostTerminal, InteractiveSessionTransport, run_interactive_session};

pub mod cli;
pub mod coordinator_backend;
mod domain_observation;
pub mod private_roles;
pub mod state;
pub mod terminal;

pub const DEFAULT_STARTUP_TIMEOUT: Duration = Duration::from_secs(30);

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ClientCommand {
    pub profile: String,
    pub workspace: PathBuf,
    pub configuration: Option<PathBuf>,
    pub startup_timeout: Duration,
    pub argv: Vec<Vec<u8>>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DomainStatus<E> {
    Absent,
    Starting,
    Ready(E),
    Stale,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BootstrapAcquisition {
    Leader,
    Waiter,
}

#[derive(Debug)]
pub enum ClientError<E> {
    Configuration(String),
    UnsupportedHost(String),
    StaleCoordinator,
    LostConnection,
    BootstrapOwnerDisconnected,
    WaitingDomainFailed,
    EmptyCommand,
    Backend(E),
    Io(io::Error),
    Control(ControlClientError),
    Protocol(&'static str),
}

impl<E: fmt::Display> fmt::Display for ClientError<E> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Configuration(message) => write!(formatter, "configuration is invalid: {message}"),
            Self::UnsupportedHost(message) => write!(formatter, "host is unsupported: {message}"),
            Self::StaleCoordinator => formatter.write_str(
                "the Domain coordinator is stale; remove no files manually and run Lazaret recovery",
            ),
            Self::LostConnection => formatter.write_str(
                "the Domain coordinator connection was lost before the command completed",
            ),
            Self::BootstrapOwnerDisconnected => formatter.write_str(
                "the Domain bootstrap owner disconnected before publication",
            ),
            Self::WaitingDomainFailed => formatter.write_str(
                "the Domain failed while this Client was waiting; recovery is required before retry",
            ),
            Self::EmptyCommand => formatter.write_str("a command is required after `--`"),
            Self::Backend(error) => error.fmt(formatter),
            Self::Io(error) => write!(formatter, "standard stream forwarding failed: {error}"),
            Self::Control(error) => error.fmt(formatter),
            Self::Protocol(message) => write!(formatter, "control protocol failed: {message}"),
        }
    }
}

impl<E: fmt::Debug + fmt::Display> std::error::Error for ClientError<E> {}

pub trait ClientBackend {
    type Endpoint;
    type Error;

    /// # Errors
    /// Returns actionable host or configuration validation failures.
    fn validate(&mut self, command: &ClientCommand) -> Result<(), ClientError<Self::Error>>;
    /// # Errors
    /// Returns failures while inspecting current Domain coordination state.
    fn observe(&mut self) -> Result<DomainStatus<Self::Endpoint>, Self::Error>;
    /// # Errors
    /// Returns failures while attempting exclusive bootstrap leadership.
    fn acquire_bootstrap(&mut self) -> Result<BootstrapAcquisition, Self::Error>;
    /// # Errors
    /// Returns startup failure or loss of the bootstrap owner.
    fn bootstrap(&mut self) -> Result<Self::Endpoint, ClientError<Self::Error>>;
    /// # Errors
    /// Returns wait failure, stale publication, or bootstrap-owner loss.
    fn wait_until_ready(&mut self) -> Result<Self::Endpoint, ClientError<Self::Error>>;
    /// # Errors
    /// Returns session registration, stream forwarding, runtime, or connection failure.
    fn execute(
        &mut self,
        endpoint: Self::Endpoint,
        argv: &[Vec<u8>],
        stdout: &mut dyn Write,
        stderr: &mut dyn Write,
    ) -> Result<i32, ClientError<Self::Error>>;
}

/// # Errors
/// Returns validation, coordination, bootstrap, connection, or command execution failures.
pub fn run<B: ClientBackend>(
    backend: &mut B,
    command: &ClientCommand,
    stdout: &mut dyn Write,
    stderr: &mut dyn Write,
) -> Result<i32, ClientError<B::Error>> {
    if command.argv.is_empty() {
        return Err(ClientError::EmptyCommand);
    }
    backend.validate(command)?;
    let endpoint = match backend.observe().map_err(ClientError::Backend)? {
        DomainStatus::Ready(endpoint) => endpoint,
        DomainStatus::Starting => backend.wait_until_ready()?,
        DomainStatus::Stale | DomainStatus::Absent => {
            match backend.acquire_bootstrap().map_err(ClientError::Backend)? {
                BootstrapAcquisition::Leader => match backend.bootstrap() {
                    Ok(endpoint) => endpoint,
                    Err(ClientError::BootstrapOwnerDisconnected) => {
                        return Err(ClientError::BootstrapOwnerDisconnected);
                    }
                    Err(error) => return Err(error),
                },
                BootstrapAcquisition::Waiter => backend.wait_until_ready()?,
            }
        }
    };
    backend.execute(endpoint, &command.argv, stdout, stderr)
}

pub struct ControlEndpoint {
    pub socket_path: PathBuf,
    pub publication: ProtocolPublicationIdentity,
    socket_authority: Option<File>,
    waiter_transfer: Mutex<Option<std::os::unix::net::UnixStream>>,
}

impl ControlEndpoint {
    #[must_use]
    pub fn new(socket_path: PathBuf, publication: ProtocolPublicationIdentity) -> Self {
        Self {
            socket_path,
            publication,
            socket_authority: None,
            waiter_transfer: Mutex::new(None),
        }
    }

    pub(crate) fn confined(
        socket_path: PathBuf,
        publication: ProtocolPublicationIdentity,
        authority: File,
    ) -> Self {
        Self {
            socket_path,
            publication,
            socket_authority: Some(authority),
            waiter_transfer: Mutex::new(None),
        }
    }

    pub(crate) fn retain_waiter_transfer(mut self, waiter: std::os::unix::net::UnixStream) -> Self {
        self.waiter_transfer = Mutex::new(Some(waiter));
        self
    }

    fn acknowledge_waiter_transfer<E>(&self) -> Result<(), ClientError<E>> {
        self.waiter_transfer
            .lock()
            .map_err(|_| ClientError::Protocol("waiter transfer lock is poisoned"))?
            .take();
        let _ = &self.socket_authority;
        Ok(())
    }
}

pub struct ClientSessionIdentity {
    cli_id: ClientId,
    session_id: SessionId,
    anchor: SigningKey,
}

impl ClientSessionIdentity {
    /// Creates one fresh CLI, Session, and host-private Session Anchor identity.
    ///
    /// # Errors
    /// Returns an error when operating-system randomness is unavailable.
    pub fn generate<E>() -> Result<Self, ClientError<E>> {
        Ok(Self {
            cli_id: ClientId::new(random_bytes()?),
            session_id: SessionId::new(random_bytes()?),
            anchor: SigningKey::from_bytes(&random_bytes()?),
        })
    }

    #[must_use]
    pub const fn cli_id(&self) -> ClientId {
        self.cli_id
    }

    #[must_use]
    pub const fn session_id(&self) -> SessionId {
        self.session_id
    }

    #[must_use]
    pub fn session_anchor_id(&self) -> SessionAnchorId {
        SessionAnchorId::new(Sha256::digest(self.anchor.verifying_key().as_bytes()).into())
    }

    #[must_use]
    pub fn session_anchor_public_key(&self) -> [u8; 32] {
        self.anchor.verifying_key().to_bytes()
    }

    /// Signs one exact prompt decision with this Client Session's host-only Session Anchor.
    ///
    /// # Errors
    /// Returns a protocol encoding failure.
    pub fn sign_developer_decision<E>(
        &self,
        session_binding: SessionBinding,
        prompt: &lazaret_core::records::TerminalPrompt,
        decision: DeveloperDecisionKind,
    ) -> Result<Vec<u8>, ClientError<E>> {
        if session_binding.session_id != self.session_id
            || session_binding.cli_id != self.cli_id
            || prompt.session_id != self.session_id
        {
            return Err(ClientError::Protocol("prompt identity differs"));
        }
        lazaret_core::protocol::encode_signed(
            &DeveloperDecision {
                protocol_version: ProtocolVersion::V1,
                session_binding,
                request_id: prompt.request_id,
                request_sha256: prompt.request_sha256,
                decision,
                approved_command_argv: if decision == DeveloperDecisionKind::Approve {
                    prompt.command_argv.clone()
                } else {
                    Vec::new()
                },
            },
            &self.anchor,
        )
        .map_err(|_| ClientError::Protocol("developer decision encoding failed"))
    }
}

/// # Errors
/// Returns enrollment, registration, framing, stream forwarding, or runtime-loss failures.
pub fn execute_over_control_face(
    endpoint: &ControlEndpoint,
    argv: &[Vec<u8>],
    stdout: &mut dyn Write,
    stderr: &mut dyn Write,
) -> Result<i32, ClientError<std::convert::Infallible>> {
    let identity = ClientSessionIdentity::generate()?;
    execute_over_control_face_as(endpoint, argv, stdout, stderr, &identity)
}

/// Executes through the exact identity that owns cold-start leadership.
///
/// # Errors
/// Returns enrollment, registration, framing, stream forwarding, or runtime-loss failures.
pub fn execute_over_control_face_as<E>(
    endpoint: &ControlEndpoint,
    argv: &[Vec<u8>],
    stdout: &mut dyn Write,
    stderr: &mut dyn Write,
    identity: &ClientSessionIdentity,
) -> Result<i32, ClientError<E>> {
    if argv.is_empty() {
        return Err(ClientError::EmptyCommand);
    }
    let generation = endpoint.publication.generation_identity.clone();
    let cli_id = identity.cli_id;
    let session_id = identity.session_id;
    let pending = PendingClientEnrollment::create(generation.clone(), cli_id)?;
    let anonymous = ControlClient::config(&endpoint.publication, None)?;
    let mut enrollment =
        ControlClient::connect(&endpoint.socket_path, anonymous, &endpoint.publication)?;
    ControlClient::send(
        &mut enrollment,
        &ControlRequest::Enroll(pending.request().clone()),
    )?;
    let ControlResponse::Enrolled(ClientEnrollmentResponse { credential, .. }) =
        ControlClient::receive(&mut enrollment)?
    else {
        return Err(ClientError::Protocol(
            "enrollment response was not Enrolled",
        ));
    };
    if credential.generation_identity != generation
        || credential.subject_id() != TlsSubjectId::ClientCli(cli_id)
    {
        return Err(ClientError::Protocol(
            "enrollment credential identity differs",
        ));
    }
    drop(enrollment);

    let binding = SessionBinding {
        generation_identity: generation,
        cli_id,
        session_id,
        session_anchor_public_key: identity.anchor.verifying_key().to_bytes().to_vec(),
    };
    let authenticated = ControlClient::config(
        &endpoint.publication,
        Some((&credential, pending.private_key())),
    )?;
    let mut stream =
        ControlClient::connect(&endpoint.socket_path, authenticated, &endpoint.publication)?;
    ControlClient::send(
        &mut stream,
        &ControlRequest::Register(SessionRegistrationRequest {
            protocol_version: ProtocolVersion::V1,
            session_binding: binding.clone(),
        }),
    )?;
    let ControlResponse::Registered(registered) = ControlClient::receive(&mut stream)? else {
        return Err(ClientError::Protocol(
            "registration response was not Registered",
        ));
    };
    if registered.session_capability.session_binding != binding {
        return Err(ClientError::Protocol(
            "registration capability identity differs",
        ));
    }
    endpoint.acknowledge_waiter_transfer()?;
    ControlClient::send(
        &mut stream,
        &ControlRequest::Execute(NonInteractiveSessionRequest {
            protocol_version: ProtocolVersion::V1,
            session_id,
            argv: argv.to_vec(),
        }),
    )?;
    loop {
        let message = ControlClient::receive(&mut stream).map_err(|error| match error {
            ControlClientError::Io(_) => ClientError::LostConnection,
            other => ClientError::Control(other),
        })?;
        match message {
            NonInteractiveSessionMessage::Stdout(output) => {
                require_session(output.session_id, session_id)?;
                stdout.write_all(&output.bytes)?;
                stdout.flush()?;
            }
            NonInteractiveSessionMessage::Stderr(output) => {
                require_session(output.session_id, session_id)?;
                stderr.write_all(&output.bytes)?;
                stderr.flush()?;
            }
            NonInteractiveSessionMessage::Exit(exit) => {
                require_session(exit.session_id(), session_id)?;
                let result = match exit.outcome() {
                    SessionExitOutcome::Exited => exit
                        .exit_code()
                        .ok_or(ClientError::Protocol("Exited omitted its exit code")),
                    SessionExitOutcome::Terminated | SessionExitOutcome::RuntimeLost => {
                        Err(ClientError::LostConnection)
                    }
                };
                await_control_close(&mut stream)?;
                return result;
            }
        }
    }
}

/// Executes one interactive Session through the exact identity that owns its Session Anchor.
///
/// # Errors
/// Returns enrollment, registration, terminal, framing, or runtime-loss failures.
pub fn execute_interactive_over_control_face_as<E, H: HostTerminal>(
    endpoint: &ControlEndpoint,
    argv: &[Vec<u8>],
    terminal: &mut H,
    identity: &ClientSessionIdentity,
) -> Result<i32, ClientError<E>> {
    if argv.is_empty() {
        return Err(ClientError::EmptyCommand);
    }
    let generation = endpoint.publication.generation_identity.clone();
    let cli_id = identity.cli_id;
    let session_id = identity.session_id;
    let pending = PendingClientEnrollment::create(generation.clone(), cli_id)?;
    let anonymous = ControlClient::config_v2(&endpoint.publication, None)?;
    let mut enrollment =
        ControlClient::connect_v2(&endpoint.socket_path, anonymous, &endpoint.publication)?;
    ControlClient::send(
        &mut enrollment,
        &ControlRequest::Enroll(pending.request().clone()),
    )?;
    let ControlResponse::Enrolled(ClientEnrollmentResponse { credential, .. }) =
        ControlClient::receive(&mut enrollment)?
    else {
        return Err(ClientError::Protocol(
            "enrollment response was not Enrolled",
        ));
    };
    if credential.generation_identity != generation
        || credential.subject_id() != TlsSubjectId::ClientCli(cli_id)
    {
        return Err(ClientError::Protocol(
            "enrollment credential identity differs",
        ));
    }
    drop(enrollment);
    let binding = SessionBinding {
        generation_identity: generation,
        cli_id,
        session_id,
        session_anchor_public_key: identity.anchor.verifying_key().to_bytes().to_vec(),
    };
    let authenticated = ControlClient::config_v2(
        &endpoint.publication,
        Some((&credential, pending.private_key())),
    )?;
    let mut stream =
        ControlClient::connect_v2(&endpoint.socket_path, authenticated, &endpoint.publication)?;
    ControlClient::send(
        &mut stream,
        &ControlRequest::Register(SessionRegistrationRequest {
            protocol_version: ProtocolVersion::V1,
            session_binding: binding.clone(),
        }),
    )?;
    let ControlResponse::Registered(registered) = ControlClient::receive(&mut stream)? else {
        return Err(ClientError::Protocol(
            "registration response was not Registered",
        ));
    };
    if registered.session_capability.session_binding != binding {
        return Err(ClientError::Protocol(
            "registration capability identity differs",
        ));
    }
    endpoint.acknowledge_waiter_transfer()?;
    let mut transport = ControlFaceInteractiveTransport::new(stream).map_err(widen_client_error)?;
    run_interactive_session(&mut transport, terminal, identity, &binding, argv)
        .map_err(widen_client_error)
}

struct ControlFaceInteractiveTransport {
    commands: Option<mpsc::Sender<InteractiveTransportCommand>>,
    batches:
        mpsc::Receiver<Result<InteractiveSessionBatchV2, ClientError<std::convert::Infallible>>>,
    cancellation: std::os::unix::net::UnixStream,
    worker: Option<JoinHandle<()>>,
}

enum InteractiveTransportCommand {
    Start(InteractiveSessionRequestV2),
    Event(InteractiveSessionEventV2),
}

impl ControlFaceInteractiveTransport {
    fn new(
        mut stream: rustls::StreamOwned<rustls::ClientConnection, std::os::unix::net::UnixStream>,
    ) -> Result<Self, ClientError<std::convert::Infallible>> {
        let cancellation = stream.sock.try_clone()?;
        let (commands, command_rx) = mpsc::channel();
        let (batch_tx, batches) = mpsc::channel();
        let worker = std::thread::spawn(move || {
            while let Ok(command) = command_rx.recv() {
                let result = match command {
                    InteractiveTransportCommand::Start(request) => ControlClient::send(
                        &mut stream,
                        &ControlRequest::ExecuteInteractiveV2(request),
                    )
                    .map_err(map_control_connection)
                    .and_then(|()| receive_interactive_batch(&mut stream)),
                    InteractiveTransportCommand::Event(event) => {
                        ControlClient::send(&mut stream, &event)
                            .map_err(map_control_connection)
                            .and_then(|()| receive_interactive_batch(&mut stream))
                    }
                };
                let failed = result.is_err();
                if batch_tx.send(result).is_err() || failed {
                    break;
                }
            }
        });
        Ok(Self {
            commands: Some(commands),
            batches,
            cancellation,
            worker: Some(worker),
        })
    }

    fn submit(
        &self,
        command: InteractiveTransportCommand,
    ) -> Result<(), ClientError<std::convert::Infallible>> {
        self.commands
            .as_ref()
            .ok_or(ClientError::LostConnection)?
            .send(command)
            .map_err(|_| ClientError::LostConnection)
    }
}

impl InteractiveSessionTransport for ControlFaceInteractiveTransport {
    type Error = std::convert::Infallible;

    fn start(
        &mut self,
        request: InteractiveSessionRequestV2,
    ) -> Result<(), ClientError<Self::Error>> {
        self.submit(InteractiveTransportCommand::Start(request))
    }

    fn send(&mut self, event: InteractiveSessionEventV2) -> Result<(), ClientError<Self::Error>> {
        self.submit(InteractiveTransportCommand::Event(event))
    }

    fn try_receive(
        &mut self,
    ) -> Result<Option<InteractiveSessionBatchV2>, ClientError<Self::Error>> {
        match self.batches.try_recv() {
            Ok(batch) => batch.map(Some),
            Err(mpsc::TryRecvError::Empty) => Ok(None),
            Err(mpsc::TryRecvError::Disconnected) => Err(ClientError::LostConnection),
        }
    }

    fn cancel(&mut self) {
        self.commands.take();
        let _ = self.cancellation.shutdown(Shutdown::Both);
    }
}

impl Drop for ControlFaceInteractiveTransport {
    fn drop(&mut self) {
        self.cancel();
        if let Some(worker) = self.worker.take() {
            let _ = worker.join();
        }
    }
}

fn receive_interactive_batch(
    stream: &mut impl std::io::Read,
) -> Result<InteractiveSessionBatchV2, ClientError<std::convert::Infallible>> {
    let batch: InteractiveSessionBatchV2 =
        ControlClient::receive(stream).map_err(map_control_connection)?;
    if batch.messages.iter().any(|message| {
        matches!(
            message,
            InteractiveSessionMessage::Terminal(SessionTerminalMessage::Exit(_))
        )
    }) {
        await_control_close(stream)?;
    }
    Ok(batch)
}

fn await_control_close<E>(stream: &mut impl Read) -> Result<(), ClientError<E>> {
    let mut trailing = [0_u8; 1];
    match stream.read(&mut trailing) {
        Ok(0) => Ok(()),
        Ok(_) => Err(ClientError::Protocol("control data followed Session exit")),
        Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => Ok(()),
        Err(_) => Err(ClientError::LostConnection),
    }
}

fn map_control_connection(error: ControlClientError) -> ClientError<std::convert::Infallible> {
    match error {
        ControlClientError::Io(_) => ClientError::LostConnection,
        other => ClientError::Control(other),
    }
}

fn widen_client_error<E>(error: ClientError<std::convert::Infallible>) -> ClientError<E> {
    match error {
        ClientError::Configuration(message) => ClientError::Configuration(message),
        ClientError::UnsupportedHost(message) => ClientError::UnsupportedHost(message),
        ClientError::StaleCoordinator => ClientError::StaleCoordinator,
        ClientError::LostConnection => ClientError::LostConnection,
        ClientError::BootstrapOwnerDisconnected => ClientError::BootstrapOwnerDisconnected,
        ClientError::WaitingDomainFailed => ClientError::WaitingDomainFailed,
        ClientError::EmptyCommand => ClientError::EmptyCommand,
        ClientError::Backend(error) => match error {},
        ClientError::Io(error) => ClientError::Io(error),
        ClientError::Control(error) => ClientError::Control(error),
        ClientError::Protocol(message) => ClientError::Protocol(message),
    }
}

fn random_bytes<E>() -> Result<[u8; 32], ClientError<E>> {
    let mut bytes = [0_u8; 32];
    getrandom::fill(&mut bytes)
        .map_err(|_| ClientError::Protocol("operating-system randomness is unavailable"))?;
    Ok(bytes)
}

fn require_session<E>(actual: SessionId, expected: SessionId) -> Result<(), ClientError<E>> {
    if actual == expected {
        Ok(())
    } else {
        Err(ClientError::Protocol("response Session ID differs"))
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    struct UnexpectedEof;

    impl Read for UnexpectedEof {
        fn read(&mut self, _buffer: &mut [u8]) -> io::Result<usize> {
            Err(io::Error::from(io::ErrorKind::UnexpectedEof))
        }
    }

    #[test]
    fn clean_control_close_is_required_after_session_exit() {
        let mut closed = std::io::Cursor::new(Vec::<u8>::new());
        assert!(await_control_close::<std::convert::Infallible>(&mut closed).is_ok());

        assert!(await_control_close::<std::convert::Infallible>(&mut UnexpectedEof).is_ok());

        let mut trailing = std::io::Cursor::new(vec![1_u8]);
        assert!(matches!(
            await_control_close::<std::convert::Infallible>(&mut trailing),
            Err(ClientError::Protocol("control data followed Session exit"))
        ));
    }
}

impl<E> From<ControlClientError> for ClientError<E> {
    fn from(error: ControlClientError) -> Self {
        Self::Control(error)
    }
}