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

use std::collections::{HashMap, VecDeque};
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::fd::AsRawFd;
use std::os::unix::ffi::OsStringExt as _;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use lazaret_core::identity::{
    ContainerId, CoordinatorId, DeviceId, DomainId, FileDescriptor, GenerationId,
    GuestProcessGroupId, GuestProcessId, GuestProcessStartTimeTicks, InodeId,
    PidfsProcessIdentityV1, SessionId,
};
use lazaret_core::protocol::VerifiedSession;
use lazaret_core::records::{
    DeveloperDecision, DomainStore, DurableDecision, GuestProcessIdentity, InteractiveSessionBatch,
    InteractiveSessionEvent, InteractiveSessionMessage, InteractiveSessionRequest,
    ParticipantOwnership, ParticipantRole, ParticipantState, PersistedHostProcessOwnership,
    ProcessOwnershipCommitAck, ProtocolVersion, PtyIdentity, SessionExit, SessionPhase,
    SessionTerminal, SessionTerminalEvent, SessionTerminalFields, SessionTerminalMessage,
    SessionTerminalState, StagedProcessOwnership, TerminalOutput, TerminalPrompt, TerminalResize,
    TerminalSignalKind,
};
use lazaret_linux::process::{
    ChildStandardIo, ProcessExit, ProcessOwnershipContext, RetainedProcess, TargetCommand,
    child_broker, current_process_identity,
};
use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl, fstat};
use rustix::pty::{OpenptFlags, grantpt, ioctl_tiocgptpeer, openpt, unlockpt};
use rustix::termios::{Winsize, tcsetwinsize};
use sha2::{Digest, Sha256};

use crate::control::{ControlError, InteractiveControlSession, InteractiveSessionExecutor};
use crate::domain_store::{DomainStoreError, DomainStoreMutation, HostCoordinatorDomainStore};
use crate::sandbox_init::{PRIVATE_SESSION_EXEC_ROLE, PRIVATE_SESSION_IDENTITY_ROLE};

const MAX_EVENT_BYTES: usize = 16 * 1024;
const MAX_INPUT_QUEUE_BYTES: usize = 64 * 1024;
const MAX_OUTPUT_BYTES: usize = 16 * 1024;
const MAX_DRAIN_BYTES: usize = 64 * 1024;
const POLL_INTERVAL: Duration = Duration::from_millis(5);
const IDENTITY_DEADLINE: Duration = Duration::from_secs(5);
const SIGNAL_GRACE: Duration = Duration::from_millis(250);
const DRAIN_DEADLINE: Duration = Duration::from_millis(250);
const SANDBOX_INIT_PATH: &str = "/run/lazaret/bin/lazaret-sandbox-init";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TerminalRoute {
    Detached,
    Sandbox,
    Prompt,
    HostProcess,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TerminalLoss {
    ClientDisconnected,
    PtyLost,
    ExecClientLost,
    HostCoordinatorLost,
}

#[derive(Clone, Copy)]
enum GuestReleaseAuthority {
    Validate,
    HostOnly,
}

enum GuestWaitOutcome {
    Absent,
    TimedOut,
    Revoked,
}

#[derive(Debug)]
pub enum TerminalError {
    Io(io::Error),
    Store(DomainStoreError),
    UnknownSession,
    DuplicateSession,
    StaleGeneration,
    InactiveRoute,
    Releasing,
    OversizedInput,
    Backpressure,
    GuestIdentityMismatch,
    GuestGroupStillPresent,
    InvalidExit,
    AuthorityMismatch,
    DuplicatePrompt,
    PromptValidation,
    AcquisitionCleanup {
        original: Box<TerminalError>,
        cleanup: Option<Box<TerminalError>>,
    },
}

impl std::fmt::Display for TerminalError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(error) => write!(formatter, "Session Terminal I/O failed: {error}"),
            Self::Store(error) => {
                write!(formatter, "Session Terminal store update failed: {error}")
            }
            Self::UnknownSession => formatter.write_str("Session Terminal is unknown"),
            Self::DuplicateSession => formatter.write_str("Session Terminal already exists"),
            Self::StaleGeneration => formatter.write_str("Session Terminal generation is stale"),
            Self::InactiveRoute => formatter.write_str("Session Terminal route is not sandbox"),
            Self::Releasing => formatter.write_str("Session Terminal is releasing"),
            Self::OversizedInput => formatter.write_str("terminal input exceeds the event bound"),
            Self::Backpressure => formatter.write_str("terminal input is backpressured"),
            Self::GuestIdentityMismatch => {
                formatter.write_str("guest process-group identity changed")
            }
            Self::GuestGroupStillPresent => {
                formatter.write_str("guest process group remains present")
            }
            Self::InvalidExit => {
                formatter.write_str("exec client returned an invalid exit outcome")
            }
            Self::AuthorityMismatch => {
                formatter.write_str("Session Terminal runtime authority does not match")
            }
            Self::DuplicatePrompt => formatter.write_str("Session already has an active prompt"),
            Self::PromptValidation => formatter.write_str("prompt authority validation failed"),
            Self::AcquisitionCleanup { original, cleanup } => {
                write!(formatter, "Session Terminal acquisition failed: {original}")?;
                if let Some(cleanup) = cleanup {
                    write!(formatter, "; partial terminal cleanup failed: {cleanup}")?;
                }
                Ok(())
            }
        }
    }
}

impl std::error::Error for TerminalError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Io(error) => Some(error),
            Self::Store(error) => Some(error),
            Self::AcquisitionCleanup { original, .. } => Some(original.as_ref()),
            Self::UnknownSession
            | Self::DuplicateSession
            | Self::StaleGeneration
            | Self::InactiveRoute
            | Self::Releasing
            | Self::OversizedInput
            | Self::Backpressure
            | Self::GuestIdentityMismatch
            | Self::GuestGroupStillPresent
            | Self::InvalidExit
            | Self::AuthorityMismatch
            | Self::DuplicatePrompt
            | Self::PromptValidation => None,
        }
    }
}

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

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

#[allow(clippy::missing_errors_doc)]
pub trait PromptRouteCancellation: Send + Sync {
    fn cancel_for_close(&self, session_id: SessionId) -> Result<(), TerminalError>;
}

pub struct NoActivePrompt;

impl PromptRouteCancellation for NoActivePrompt {
    fn cancel_for_close(&self, _session_id: SessionId) -> Result<(), TerminalError> {
        Ok(())
    }
}

struct PendingPrompt {
    authority: VerifiedSession,
    prompt: TerminalPrompt,
    emitted: bool,
    resumed_size: Option<TerminalResize>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RoutedPromptDecision {
    pub signed_decision: Vec<u8>,
    pub decision: DeveloperDecision,
    pub resumed_size: Option<TerminalResize>,
}

pub struct SessionPromptTransport {
    store: Arc<HostCoordinatorDomainStore>,
    pending: Mutex<HashMap<SessionId, PendingPrompt>>,
    routed: Mutex<HashMap<SessionId, RoutedPromptDecision>>,
}

impl SessionPromptTransport {
    #[must_use]
    pub fn new(store: Arc<HostCoordinatorDomainStore>) -> Self {
        Self {
            store,
            pending: Mutex::new(HashMap::new()),
            routed: Mutex::new(HashMap::new()),
        }
    }

    /// Queues one sandbox-signed request for its exact registered Session.
    ///
    /// # Errors
    /// Refuses invalid authority, stale Session state, and concurrent prompts.
    pub fn queue_signed_request(
        &self,
        authority: VerifiedSession,
        signed_request: &[u8],
    ) -> Result<TerminalPrompt, TerminalError> {
        let (request, digest) = authority
            .verify_request(signed_request)
            .map_err(|_| TerminalError::PromptValidation)?;
        let binding = authority.session_binding();
        let anchor_id = lazaret_core::identity::SessionAnchorId::new(
            Sha256::digest(&binding.session_anchor_public_key).into(),
        );
        let snapshot = self.store.snapshot()?;
        let mut session = snapshot
            .sessions
            .into_iter()
            .find(|session| {
                session.generation_id == binding.generation_identity.generation_id
                    && session.session_id == binding.session_id
                    && session.cli_id == binding.cli_id
                    && session.session_anchor_id == anchor_id
                    && session.session_phase == SessionPhase::Attached
            })
            .ok_or(TerminalError::AuthorityMismatch)?;
        let prompt = TerminalPrompt {
            session_id: binding.session_id,
            request_id: request.request_id(),
            request_sha256: digest,
            agent_profile: binding.generation_identity.agent_profile.clone(),
            command_argv: request.command_argv().to_vec(),
            justification: request.justification().to_owned(),
        };
        let mut pending = self
            .pending
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?;
        if pending.contains_key(&binding.session_id) {
            return Err(TerminalError::DuplicatePrompt);
        }
        session.session_phase = SessionPhase::Prompting;
        self.store
            .apply(&[DomainStoreMutation::PutSession(session)])?;
        pending.insert(
            binding.session_id,
            PendingPrompt {
                authority,
                prompt: prompt.clone(),
                emitted: false,
                resumed_size: None,
            },
        );
        Ok(prompt)
    }

    fn prompt_for(&self, session_id: SessionId) -> Result<Option<TerminalPrompt>, TerminalError> {
        let mut pending = self
            .pending
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?;
        let Some(pending) = pending.get_mut(&session_id) else {
            return Ok(None);
        };
        if pending.emitted {
            return Ok(None);
        }
        pending.emitted = true;
        Ok(Some(pending.prompt.clone()))
    }

    fn route_decision(
        &self,
        session_id: SessionId,
        signed_decision: Vec<u8>,
    ) -> Result<RoutedPromptDecision, TerminalError> {
        let mut prompts = self
            .pending
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?;
        let pending = prompts
            .get(&session_id)
            .ok_or(TerminalError::InactiveRoute)?;
        let decision = pending
            .authority
            .verify_developer_decision(
                &signed_decision,
                pending.prompt.request_id,
                pending.prompt.request_sha256,
            )
            .map_err(|_| TerminalError::PromptValidation)?;
        let resumed_size = pending.resumed_size;
        prompts.remove(&session_id);
        Ok(RoutedPromptDecision {
            signed_decision,
            decision,
            resumed_size,
        })
    }

    fn retain_prompt_resize(&self, resize: TerminalResize) -> Result<bool, TerminalError> {
        let mut prompts = self
            .pending
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?;
        let Some(pending) = prompts.get_mut(&resize.session_id) else {
            return Ok(false);
        };
        if !pending.emitted {
            return Err(TerminalError::InactiveRoute);
        }
        pending.resumed_size = Some(resize);
        Ok(true)
    }

    fn publish_decision(
        &self,
        session_id: SessionId,
        decision: RoutedPromptDecision,
    ) -> Result<(), TerminalError> {
        self.routed
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?
            .insert(session_id, decision);
        Ok(())
    }

    fn restore_sandbox_route(&self, session_id: SessionId) -> Result<(), TerminalError> {
        let snapshot = self.store.snapshot()?;
        let mut session = snapshot
            .sessions
            .into_iter()
            .find(|session| session.session_id == session_id)
            .ok_or(TerminalError::UnknownSession)?;
        if session.session_phase != SessionPhase::Prompting {
            return Err(TerminalError::InactiveRoute);
        }
        session.session_phase = SessionPhase::Attached;
        self.store
            .apply(&[DomainStoreMutation::PutSession(session)])?;
        Ok(())
    }

    /// # Errors
    /// Refuses access after internal synchronization failure.
    pub fn take_decision(
        &self,
        session_id: SessionId,
    ) -> Result<Option<RoutedPromptDecision>, TerminalError> {
        Ok(self
            .routed
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?
            .remove(&session_id))
    }

    pub(crate) fn begin_host_execution(&self, session_id: SessionId) -> Result<(), TerminalError> {
        self.transition_session(session_id, SessionPhase::HostExecuting)
    }

    pub(crate) fn finish_escalation(&self, session_id: SessionId) -> Result<(), TerminalError> {
        self.pending
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?
            .remove(&session_id);
        self.routed
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?
            .remove(&session_id);
        let snapshot = self.store.snapshot()?;
        let Some(session) = snapshot
            .sessions
            .iter()
            .find(|session| session.session_id == session_id)
        else {
            return Err(TerminalError::UnknownSession);
        };
        match session.session_phase {
            SessionPhase::Attached => return Ok(()),
            SessionPhase::Prompting | SessionPhase::HostExecuting => {}
            SessionPhase::New
            | SessionPhase::Waiting
            | SessionPhase::Registering
            | SessionPhase::Closed
            | SessionPhase::Failed => return Err(TerminalError::InactiveRoute),
        }
        self.transition_session(session_id, SessionPhase::Attached)
    }

    fn transition_session(
        &self,
        session_id: SessionId,
        session_phase: SessionPhase,
    ) -> Result<(), TerminalError> {
        let snapshot = self.store.snapshot()?;
        let mut session = snapshot
            .sessions
            .into_iter()
            .find(|session| session.session_id == session_id)
            .ok_or(TerminalError::UnknownSession)?;
        session.session_phase = session_phase;
        self.store
            .apply(&[DomainStoreMutation::PutSession(session)])?;
        Ok(())
    }
}

impl PromptRouteCancellation for SessionPromptTransport {
    fn cancel_for_close(&self, session_id: SessionId) -> Result<(), TerminalError> {
        self.pending
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?
            .remove(&session_id);
        self.routed
            .lock()
            .map_err(|_| TerminalError::PromptValidation)?
            .remove(&session_id);
        Ok(())
    }
}

#[allow(clippy::missing_errors_doc)]
pub trait SessionTerminalProcess: Send {
    fn write_input(&mut self, bytes: &[u8]) -> io::Result<usize>;
    fn resize(&self, resize: TerminalResize) -> io::Result<()>;
    fn read_output(&mut self, bytes: &mut [u8]) -> io::Result<usize>;
    fn exec_exited(&self) -> io::Result<bool>;
    fn signal_group(&self, signal: &str) -> Result<(), TerminalError>;
    fn group_absent(&self) -> Result<bool, TerminalError>;
    fn terminate_exec(&self) -> io::Result<()>;
    fn reap_exec(&mut self) -> io::Result<ProcessExit>;
    fn close_pty(&mut self);
}

#[allow(clippy::missing_errors_doc)]
pub trait SessionTerminalFactory: Send + Sync {
    fn acquire(
        &self,
        request: &TerminalAcquisition,
        recorder: &HostCoordinatorTerminalRecorder,
    ) -> Result<(SessionTerminal, Box<dyn SessionTerminalProcess>), TerminalError>;
}

#[derive(Clone, Debug)]
pub struct TerminalAcquisition {
    pub domain_id: DomainId,
    pub generation_id: GenerationId,
    pub coordinator_id: CoordinatorId,
    pub session_id: SessionId,
    pub container_id: ContainerId,
    pub host_coordinator: PidfsProcessIdentityV1,
    pub runsc_program: OsString,
    pub runsc_global_arguments: Vec<OsString>,
    pub internal_pid_file: PathBuf,
    pub argv: Vec<Vec<u8>>,
}

pub struct HostCoordinatorTerminalRecorder {
    store: Arc<dyn TerminalDomainStore>,
}

impl HostCoordinatorTerminalRecorder {
    #[must_use]
    pub fn new(store: Arc<HostCoordinatorDomainStore>) -> Self {
        Self { store }
    }

    #[cfg(test)]
    fn with_store(store: Arc<dyn TerminalDomainStore>) -> Self {
        Self { store }
    }

    fn put(&self, terminal: SessionTerminal) -> Result<(), TerminalError> {
        self.store
            .apply(&[DomainStoreMutation::PutSessionTerminal(terminal)])?;
        Ok(())
    }

    fn require_authority(
        &self,
        request: &TerminalAcquisition,
    ) -> Result<TerminalRoute, TerminalError> {
        let snapshot = self.store.snapshot()?;
        let generation = snapshot
            .generations
            .iter()
            .find(|generation| generation.generation_id == request.generation_id);
        let current = snapshot
            .generations
            .iter()
            .all(|candidate| candidate.generation_id <= request.generation_id);
        if snapshot.domain.domain_id != request.domain_id
            || !current
            || generation.is_none_or(|generation| {
                generation.coordinator_id != request.coordinator_id
                    || generation.durable_decision != DurableDecision::Published
            })
        {
            return Err(TerminalError::AuthorityMismatch);
        }
        let route = snapshot
            .sessions
            .iter()
            .find(|session| {
                session.generation_id == request.generation_id
                    && session.session_id == request.session_id
            })
            .map(|session| match session.session_phase {
                SessionPhase::Attached => TerminalRoute::Sandbox,
                SessionPhase::Prompting => TerminalRoute::Prompt,
                SessionPhase::HostExecuting => TerminalRoute::HostProcess,
                SessionPhase::New
                | SessionPhase::Waiting
                | SessionPhase::Registering
                | SessionPhase::Closed
                | SessionPhase::Failed => TerminalRoute::Detached,
            })
            .ok_or(TerminalError::AuthorityMismatch)?;
        let sandbox = snapshot
            .participants
            .iter()
            .find(|participant| {
                participant.generation_id() == request.generation_id
                    && participant.role() == ParticipantRole::Sandbox
                    && participant.state() == ParticipantState::Committed
            })
            .and_then(|participant| match participant.ownership() {
                Some(ParticipantOwnership::Sandbox(sandbox)) => Some(sandbox),
                _ => None,
            })
            .ok_or(TerminalError::AuthorityMismatch)?;
        if sandbox.container_id != Some(request.container_id)
            || sandbox.host_coordinator.serviceable_pidfs().ok() != Some(&request.host_coordinator)
        {
            return Err(TerminalError::AuthorityMismatch);
        }
        Ok(route)
    }

    fn require_guest_signal_authority(
        &self,
        request: &TerminalAcquisition,
        terminal: &SessionTerminal,
    ) -> Result<(), TerminalError> {
        let snapshot = self.store.snapshot()?;
        let current = snapshot
            .generations
            .iter()
            .all(|candidate| candidate.generation_id <= request.generation_id);
        let generation = snapshot.generations.iter().any(|generation| {
            generation.generation_id == request.generation_id
                && generation.coordinator_id == request.coordinator_id
                && generation.durable_decision == DurableDecision::Published
        });
        let session = snapshot.sessions.iter().any(|session| {
            session.generation_id == request.generation_id
                && session.session_id == request.session_id
                && matches!(
                    session.session_phase,
                    SessionPhase::Attached | SessionPhase::Closed
                )
        });
        let sandbox = snapshot.participants.iter().any(|participant| {
            participant.generation_id() == request.generation_id
                && participant.role() == ParticipantRole::Sandbox
                && participant.state() == ParticipantState::Committed
                && matches!(
                    participant.ownership(),
                    Some(ParticipantOwnership::Sandbox(sandbox))
                        if sandbox.container_id == Some(request.container_id)
                            && sandbox.host_coordinator.serviceable_pidfs().ok()
                                == Some(&request.host_coordinator)
                )
        });
        let terminal = snapshot.session_terminals.iter().any(|record| {
            record == terminal && record.state() == SessionTerminalState::ReleaseIntent
        });
        if snapshot.domain.domain_id == request.domain_id
            && current
            && generation
            && session
            && sandbox
            && terminal
        {
            Ok(())
        } else {
            Err(TerminalError::AuthorityMismatch)
        }
    }

    fn begin_release(&self, terminal: SessionTerminal) -> Result<(), TerminalError> {
        let snapshot = self.store.snapshot()?;
        let mut session = snapshot
            .sessions
            .into_iter()
            .find(|session| {
                session.generation_id == terminal.generation_id()
                    && session.session_id == terminal.session_id()
            })
            .ok_or(TerminalError::UnknownSession)?;
        let mut mutations = Vec::with_capacity(2);
        if !matches!(
            session.session_phase,
            SessionPhase::Closed | SessionPhase::Failed
        ) {
            session.session_phase = SessionPhase::Closed;
            mutations.push(DomainStoreMutation::PutSession(session));
        }
        mutations.push(DomainStoreMutation::PutSessionTerminal(terminal));
        self.store.apply(&mutations)?;
        Ok(())
    }

    fn clean(&self, terminal: SessionTerminal) -> Result<(), TerminalError> {
        self.store
            .apply(&[DomainStoreMutation::CleanSessionTerminalAfterAbsence(
                terminal,
            )])?;
        Ok(())
    }
}

trait TerminalDomainStore: Send + Sync {
    fn snapshot(&self) -> Result<DomainStore, DomainStoreError>;
    fn apply(&self, mutations: &[DomainStoreMutation]) -> Result<(), DomainStoreError>;
}

impl TerminalDomainStore for HostCoordinatorDomainStore {
    fn snapshot(&self) -> Result<DomainStore, DomainStoreError> {
        self.snapshot()
    }

    fn apply(&self, mutations: &[DomainStoreMutation]) -> Result<(), DomainStoreError> {
        self.apply(mutations)
    }
}

include!("terminal/interactive.rs");
include!("terminal/engine.rs");
include!("terminal/runsc.rs");