lazaret/crates/lazaret-client/tests/interactive_terminal.rs

use std::collections::VecDeque;
use std::fs::File;
use std::io;

use lazaret_client::terminal::{
    HostTerminal, HostTerminalEvent, InteractiveSessionTransport, run_interactive_session,
};
use lazaret_client::{ClientError, ClientSessionIdentity};
use lazaret_core::identity::{
    ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId, RequestDigest, RequestId,
    SandboxId, SessionId,
};
use lazaret_core::records::{
    DeveloperDecision, DeveloperDecisionKind, GenerationIdentity, InteractiveSessionBatch,
    InteractiveSessionEvent, InteractiveSessionMessage, InteractiveSessionRequest, SessionBinding,
    SessionExit, SessionTerminalMessage, SignedProtocolRecord, TerminalOutput, TerminalPrompt,
    TerminalResize, TerminalSignalKind,
};
use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl};
use rustix::pty::{OpenptFlags, grantpt, ioctl_tiocgptpeer, openpt, unlockpt};
use rustix::termios::{LocalModes, OptionalActions, Termios, tcgetattr, tcsetattr};

struct PtyTerminal {
    _master: File,
    slave: File,
    original: Termios,
    original_flags: OFlags,
    raw: bool,
    events: VecDeque<HostTerminalEvent>,
    output: Vec<u8>,
    decision: DeveloperDecisionKind,
    suspend_count: usize,
}

impl PtyTerminal {
    fn new(events: impl IntoIterator<Item = HostTerminalEvent>) -> Self {
        let master = openpt(OpenptFlags::CLOEXEC | OpenptFlags::NOCTTY).unwrap();
        grantpt(&master).unwrap();
        unlockpt(&master).unwrap();
        let slave = File::from(
            ioctl_tiocgptpeer(
                &master,
                OpenptFlags::RDWR | OpenptFlags::CLOEXEC | OpenptFlags::NOCTTY,
            )
            .unwrap(),
        );
        let original = tcgetattr(&slave).unwrap();
        let original_flags = fcntl_getfl(&slave).unwrap();
        Self {
            _master: File::from(master),
            slave,
            original,
            original_flags,
            raw: false,
            events: events.into_iter().collect(),
            output: Vec::new(),
            decision: DeveloperDecisionKind::Approve,
            suspend_count: 0,
        }
    }

    fn restored(&self) -> bool {
        let current = tcgetattr(&self.slave).unwrap();
        current.input_modes == self.original.input_modes
            && current.output_modes == self.original.output_modes
            && current.control_modes == self.original.control_modes
            && current.local_modes == self.original.local_modes
            && !self.raw
    }
}

impl HostTerminal for PtyTerminal {
    fn enter_raw(&mut self) -> io::Result<()> {
        let mut raw = self.original.clone();
        raw.make_raw();
        raw.local_modes.insert(LocalModes::ISIG);
        tcsetattr(&self.slave, OptionalActions::Now, &raw)?;
        fcntl_setfl(&self.slave, self.original_flags | OFlags::NONBLOCK)?;
        self.raw = true;
        Ok(())
    }

    fn restore(&mut self) -> io::Result<()> {
        fcntl_setfl(&self.slave, self.original_flags)?;
        tcsetattr(&self.slave, OptionalActions::Now, &self.original)?;
        self.raw = false;
        Ok(())
    }

    fn size(&self, session_id: SessionId) -> io::Result<TerminalResize> {
        Ok(TerminalResize {
            session_id,
            rows: std::num::NonZeroU16::new(42).unwrap(),
            columns: std::num::NonZeroU16::new(132).unwrap(),
            pixel_width: 0,
            pixel_height: 0,
        })
    }

    fn next_event(&mut self) -> io::Result<HostTerminalEvent> {
        Ok(self.events.pop_front().unwrap_or(HostTerminalEvent::Idle))
    }

    fn write_output(&mut self, bytes: &[u8]) -> io::Result<()> {
        self.output.extend_from_slice(bytes);
        Ok(())
    }

    fn prompt(&mut self, prompt: &TerminalPrompt) -> io::Result<DeveloperDecisionKind> {
        assert_eq!(prompt.agent_profile, "trusted");
        assert_eq!(prompt.command_argv, [b"/usr/bin/tool".to_vec()]);
        assert_eq!(prompt.justification, "update host metadata");
        Ok(self.decision)
    }

    fn suspend_and_resume(&mut self) -> io::Result<()> {
        self.restore()?;
        assert!(self.restored());
        self.suspend_count += 1;
        self.enter_raw()
    }
}

enum TransportFailure {
    Stale,
    Lost,
}

struct FixtureTransport {
    session_id: SessionId,
    events: Vec<InteractiveSessionEvent>,
    prompt_sent: bool,
    failure: Option<TransportFailure>,
    signed_decision: Option<Vec<u8>>,
}

impl FixtureTransport {
    fn successful(session_id: SessionId) -> Self {
        Self {
            session_id,
            events: Vec::new(),
            prompt_sent: false,
            failure: None,
            signed_decision: None,
        }
    }
}

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

    fn start(
        &mut self,
        request: InteractiveSessionRequest,
    ) -> Result<InteractiveSessionBatch, ClientError<Self::Error>> {
        if matches!(self.failure, Some(TransportFailure::Stale)) {
            return Err(ClientError::StaleCoordinator);
        }
        assert_eq!(request.session_id, self.session_id);
        Ok(batch(
            self.session_id,
            vec![InteractiveSessionMessage::Terminal(
                SessionTerminalMessage::Output(TerminalOutput {
                    session_id: self.session_id,
                    bytes: b"guest-output".to_vec(),
                }),
            )],
        ))
    }

    fn exchange(
        &mut self,
        event: InteractiveSessionEvent,
    ) -> Result<InteractiveSessionBatch, ClientError<Self::Error>> {
        if matches!(self.failure, Some(TransportFailure::Lost)) {
            return Err(ClientError::LostConnection);
        }
        self.events.push(event.clone());
        if let InteractiveSessionEvent::PromptDecision {
            signed_decision, ..
        } = event
        {
            self.signed_decision = Some(signed_decision);
            return Ok(batch(
                self.session_id,
                vec![InteractiveSessionMessage::Terminal(
                    SessionTerminalMessage::Exit(SessionExit::exited(self.session_id, 23)),
                )],
            ));
        }
        if !self.prompt_sent && self.events.len() >= 5 {
            self.prompt_sent = true;
            return Ok(batch(
                self.session_id,
                vec![InteractiveSessionMessage::Prompt(TerminalPrompt {
                    session_id: self.session_id,
                    request_id: RequestId::new([8; 32]),
                    request_sha256: RequestDigest::new([9; 32]),
                    agent_profile: "trusted".into(),
                    command_argv: vec![b"/usr/bin/tool".to_vec()],
                    justification: "update host metadata".into(),
                })],
            ));
        }
        Ok(batch(self.session_id, Vec::new()))
    }
}

#[test]
fn terminal_fixture_routes_resize_signals_suspend_prompt_and_nonzero_exit() {
    let identity = ClientSessionIdentity::generate::<std::convert::Infallible>().unwrap();
    let binding = binding(&identity);
    let mut terminal = PtyTerminal::new([
        HostTerminalEvent::Resize,
        HostTerminalEvent::Signal(TerminalSignalKind::Interrupt),
        HostTerminalEvent::Signal(TerminalSignalKind::Quit),
        HostTerminalEvent::LocalSuspend,
        HostTerminalEvent::Input(b"input".to_vec()),
    ]);
    let mut transport = FixtureTransport::successful(identity.session_id());

    let status = run_interactive_session(
        &mut transport,
        &mut terminal,
        &identity,
        &binding,
        &[b"agent".to_vec()],
    )
    .unwrap();

    assert_eq!(status, 23);
    assert_eq!(terminal.output, b"guest-output");
    assert_eq!(terminal.suspend_count, 1);
    assert!(terminal.restored());
    let signed = lazaret_core::codec::decode::<SignedProtocolRecord>(
        transport.signed_decision.as_ref().unwrap(),
    )
    .unwrap();
    let decision = lazaret_core::codec::decode::<DeveloperDecision>(&signed.payload_bytes).unwrap();
    assert_eq!(signed.signature.len(), 64);
    assert_eq!(decision.session_binding, binding);
    assert_eq!(decision.decision, DeveloperDecisionKind::Approve);
    assert_eq!(decision.approved_command_argv, [b"/usr/bin/tool".to_vec()]);
    assert!(transport.events.iter().any(|event| matches!(
        event,
        InteractiveSessionEvent::Terminal(lazaret_core::records::SessionTerminalEvent::Resize(_))
    )));
    for signal in [TerminalSignalKind::Interrupt, TerminalSignalKind::Quit] {
        assert!(transport.events.iter().any(|event| matches!(
            event,
            InteractiveSessionEvent::Terminal(
                lazaret_core::records::SessionTerminalEvent::Signal(sent)
            ) if sent.signal == signal
        )));
    }
    assert!(!transport.events.iter().any(|event| matches!(
        event,
        InteractiveSessionEvent::Terminal(
            lazaret_core::records::SessionTerminalEvent::Signal(sent)
        ) if sent.signal == TerminalSignalKind::Suspend
    )));
}

#[test]
fn terminal_fixture_restores_raw_mode_for_stale_and_lost_coordinator() {
    for failure in [TransportFailure::Stale, TransportFailure::Lost] {
        let identity = ClientSessionIdentity::generate::<std::convert::Infallible>().unwrap();
        let binding = binding(&identity);
        let mut terminal = PtyTerminal::new([HostTerminalEvent::Idle]);
        let mut transport = FixtureTransport {
            failure: Some(failure),
            ..FixtureTransport::successful(identity.session_id())
        };

        let error = run_interactive_session(
            &mut transport,
            &mut terminal,
            &identity,
            &binding,
            &[b"agent".to_vec()],
        )
        .unwrap_err();

        assert!(terminal.restored());
        assert!(matches!(
            error,
            ClientError::StaleCoordinator | ClientError::LostConnection
        ));
    }
}

fn binding(identity: &ClientSessionIdentity) -> SessionBinding {
    SessionBinding {
        generation_identity: 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]),
        },
        cli_id: identity.cli_id(),
        session_id: identity.session_id(),
        session_anchor_public_key: identity.session_anchor_public_key().to_vec(),
    }
}

fn batch(
    session_id: SessionId,
    messages: Vec<InteractiveSessionMessage>,
) -> InteractiveSessionBatch {
    InteractiveSessionBatch {
        session_id,
        messages,
    }
}