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

use std::ffi::OsString;
use std::fmt;
use std::fs::File;
use std::io::{self, Write};
use std::os::unix::ffi::OsStringExt;
use std::path::PathBuf;
use std::sync::Mutex;

use ed25519_dalek::SigningKey;
use lazaret_core::identity::{ClientId, SessionAnchorId, SessionId};
use lazaret_core::records::{
    ClientEnrollmentResponse, ControlRequest, ControlResponse, DeveloperDecision,
    DeveloperDecisionKind, InteractiveSessionBatch, InteractiveSessionEvent,
    InteractiveSessionRequest, NonInteractiveSessionMessage, NonInteractiveSessionRequest,
    ProtocolPublicationIdentity, ProtocolVersion, SessionBinding, SessionExitOutcome,
    SessionRegistrationRequest, TlsSubjectId,
};
use lazaret_host_coordinator::control::{
    ControlError, PendingClientEnrollment, PublishedControlFace,
};
use sha2::{Digest, Sha256};

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

pub mod production;
pub mod state;
pub mod terminal;

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

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ArgumentError(&'static str);

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

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

/// Parses `--profile`, optional host paths, and the exact command following `--`.
///
/// # Errors
/// Returns a concise usage error for missing values, unknown options, or a missing command.
pub fn parse_command_line(
    arguments: impl IntoIterator<Item = OsString>,
    current_workspace: PathBuf,
) -> Result<ClientCommand, ArgumentError> {
    let mut arguments = arguments.into_iter();
    let mut profile = None;
    let mut workspace = None;
    let mut configuration = None;
    let mut argv = Vec::new();
    while let Some(argument) = arguments.next() {
        match argument.to_str() {
            Some("--profile") => profile = Some(required_utf8(&mut arguments, "--profile")?),
            Some("--workspace") => {
                workspace = Some(PathBuf::from(required_value(
                    &mut arguments,
                    "--workspace",
                )?));
            }
            Some("--config") => {
                configuration = Some(PathBuf::from(required_value(&mut arguments, "--config")?));
            }
            Some("--") => {
                argv.extend(arguments.map(OsString::into_vec));
                break;
            }
            Some(_) | None => return Err(ArgumentError("unknown option before `--`")),
        }
    }
    let profile = profile.ok_or(ArgumentError("`--profile <name>` is required"))?;
    if argv.is_empty() {
        return Err(ArgumentError("a command is required after `--`"));
    }
    Ok(ClientCommand {
        profile,
        workspace: workspace.unwrap_or(current_workspace),
        configuration,
        argv,
    })
}

fn required_value(
    arguments: &mut impl Iterator<Item = OsString>,
    option: &'static str,
) -> Result<OsString, ArgumentError> {
    arguments.next().ok_or(match option {
        "--profile" => ArgumentError("`--profile` requires a value"),
        "--workspace" => ArgumentError("`--workspace` requires a path"),
        "--config" => ArgumentError("`--config` requires a path"),
        _ => ArgumentError("option requires a value"),
    })
}

fn required_utf8(
    arguments: &mut impl Iterator<Item = OsString>,
    option: &'static str,
) -> Result<String, ArgumentError> {
    required_value(arguments, option)?
        .into_string()
        .map_err(|_| ArgumentError("profile name must be UTF-8"))
}

#[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(ControlError),
    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 => return Err(ClientError::StaleCoordinator),
        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 = PublishedControlFace::client_config(&endpoint.publication, None)?;
    let mut enrollment =
        PublishedControlFace::connect(&endpoint.socket_path, anonymous, &endpoint.publication)?;
    PublishedControlFace::send(
        &mut enrollment,
        &ControlRequest::Enroll(pending.request().clone()),
    )?;
    let ControlResponse::Enrolled(ClientEnrollmentResponse { credential, .. }) =
        PublishedControlFace::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 = PublishedControlFace::client_config(
        &endpoint.publication,
        Some((&credential, pending.private_key())),
    )?;
    let mut stream =
        PublishedControlFace::connect(&endpoint.socket_path, authenticated, &endpoint.publication)?;
    PublishedControlFace::send(
        &mut stream,
        &ControlRequest::Register(SessionRegistrationRequest {
            protocol_version: ProtocolVersion::V1,
            session_binding: binding.clone(),
        }),
    )?;
    let ControlResponse::Registered(registered) = PublishedControlFace::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()?;
    PublishedControlFace::send(
        &mut stream,
        &ControlRequest::Execute(NonInteractiveSessionRequest {
            protocol_version: ProtocolVersion::V1,
            session_id,
            argv: argv.to_vec(),
        }),
    )?;
    loop {
        let message = PublishedControlFace::receive(&mut stream).map_err(|error| match error {
            ControlError::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)?;
                return match exit.outcome() {
                    SessionExitOutcome::Exited => exit
                        .exit_code()
                        .ok_or(ClientError::Protocol("Exited omitted its exit code")),
                    SessionExitOutcome::Terminated | SessionExitOutcome::RuntimeLost => {
                        Err(ClientError::LostConnection)
                    }
                };
            }
        }
    }
}

/// 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 = PublishedControlFace::client_config(&endpoint.publication, None)?;
    let mut enrollment =
        PublishedControlFace::connect(&endpoint.socket_path, anonymous, &endpoint.publication)?;
    PublishedControlFace::send(
        &mut enrollment,
        &ControlRequest::Enroll(pending.request().clone()),
    )?;
    let ControlResponse::Enrolled(ClientEnrollmentResponse { credential, .. }) =
        PublishedControlFace::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 = PublishedControlFace::client_config(
        &endpoint.publication,
        Some((&credential, pending.private_key())),
    )?;
    let mut stream =
        PublishedControlFace::connect(&endpoint.socket_path, authenticated, &endpoint.publication)?;
    PublishedControlFace::send(
        &mut stream,
        &ControlRequest::Register(SessionRegistrationRequest {
            protocol_version: ProtocolVersion::V1,
            session_binding: binding.clone(),
        }),
    )?;
    let ControlResponse::Registered(registered) = PublishedControlFace::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 { stream };
    run_interactive_session(&mut transport, terminal, identity, &binding, argv)
        .map_err(widen_client_error)
}

struct ControlFaceInteractiveTransport {
    stream: rustls::StreamOwned<rustls::ClientConnection, std::os::unix::net::UnixStream>,
}

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

    fn start(
        &mut self,
        request: InteractiveSessionRequest,
    ) -> Result<InteractiveSessionBatch, ClientError<Self::Error>> {
        PublishedControlFace::send(
            &mut self.stream,
            &ControlRequest::ExecuteInteractive(request),
        )?;
        receive_interactive_batch(&mut self.stream)
    }

    fn exchange(
        &mut self,
        event: InteractiveSessionEvent,
    ) -> Result<InteractiveSessionBatch, ClientError<Self::Error>> {
        PublishedControlFace::send(&mut self.stream, &event).map_err(map_control_connection)?;
        receive_interactive_batch(&mut self.stream)
    }
}

fn receive_interactive_batch(
    stream: &mut impl std::io::Read,
) -> Result<InteractiveSessionBatch, ClientError<std::convert::Infallible>> {
    PublishedControlFace::receive(stream).map_err(map_control_connection)
}

fn map_control_connection(error: ControlError) -> ClientError<std::convert::Infallible> {
    match error {
        ControlError::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)
    }
}

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