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, PromptOutcome,
    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, InteractiveSessionBatchV2,
    InteractiveSessionEventV2, InteractiveSessionMessage, InteractiveSessionRequestV2,
    SessionBinding, SessionExit, SessionTerminalEventV2, SessionTerminalMessage,
    SignedProtocolRecord, TerminalModesV2, TerminalOutput, TerminalPrompt, TerminalResize,
};
use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl};
use rustix::pty::{OpenptFlags, grantpt, ioctl_tiocgptpeer, openpt, unlockpt};
use rustix::termios::{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();
        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 modes(&self) -> TerminalModesV2 {
        test_modes()
    }

    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<PromptOutcome> {
        assert_eq!(prompt.agent_profile, "trusted");
        assert_eq!(prompt.command_argv, [b"/usr/bin/tool".to_vec()]);
        assert_eq!(prompt.justification, "update host metadata");
        let resumed_size = if matches!(self.events.front(), Some(HostTerminalEvent::LocalSuspend)) {
            self.events.pop_front();
            self.suspend_and_resume()?;
            Some(self.size(prompt.session_id)?)
        } else {
            None
        };
        Ok(PromptOutcome {
            decision: self.decision,
            resumed_size,
        })
    }

    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<InteractiveSessionEventV2>,
    prompt_sent: bool,
    failure: Option<TransportFailure>,
    signed_decision: Option<Vec<u8>>,
    prompt_after_events: usize,
    batches: VecDeque<InteractiveSessionBatchV2>,
    cancelled: bool,
    blocked: bool,
}

impl FixtureTransport {
    fn successful(session_id: SessionId) -> Self {
        Self {
            session_id,
            events: Vec::new(),
            prompt_sent: false,
            failure: None,
            signed_decision: None,
            prompt_after_events: 5,
            batches: VecDeque::new(),
            cancelled: false,
            blocked: false,
        }
    }
}

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

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

    fn send(&mut self, event: InteractiveSessionEventV2) -> Result<(), ClientError<Self::Error>> {
        if matches!(self.failure, Some(TransportFailure::Lost)) {
            return Err(ClientError::LostConnection);
        }
        self.events.push(event.clone());
        if let InteractiveSessionEventV2::PromptDecision {
            signed_decision, ..
        } = event
        {
            self.signed_decision = Some(signed_decision);
            self.batches.push_back(batch(
                self.session_id,
                vec![InteractiveSessionMessage::Terminal(
                    SessionTerminalMessage::Exit(SessionExit::exited(self.session_id, 23)),
                )],
            ));
            return Ok(());
        }
        if !self.prompt_sent && self.events.len() >= self.prompt_after_events {
            self.prompt_sent = true;
            self.batches.push_back(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(),
                })],
            ));
            return Ok(());
        }
        self.batches.push_back(batch(self.session_id, Vec::new()));
        Ok(())
    }

    fn try_receive(
        &mut self,
    ) -> Result<Option<InteractiveSessionBatchV2>, ClientError<Self::Error>> {
        if self.blocked {
            return Ok(None);
        }
        Ok(self.batches.pop_front())
    }

    fn cancel(&mut self) {
        self.cancelled = true;
    }
}

#[test]
fn prompt_suspend_sends_resumed_size_before_decision() {
    let identity = ClientSessionIdentity::generate::<std::convert::Infallible>().unwrap();
    let binding = binding(&identity);
    let mut terminal = PtyTerminal::new([HostTerminalEvent::Idle, HostTerminalEvent::LocalSuspend]);
    let mut transport = FixtureTransport::successful(identity.session_id());
    transport.prompt_after_events = 1;

    assert_eq!(
        run_interactive_session(
            &mut transport,
            &mut terminal,
            &identity,
            &binding,
            &[b"/bin/sh".to_vec()],
        )
        .unwrap(),
        23
    );

    assert_eq!(terminal.suspend_count, 1);
    assert!(matches!(
        transport.events.as_slice(),
        [
            InteractiveSessionEventV2::Poll { .. },
            InteractiveSessionEventV2::Terminal(SessionTerminalEventV2::Resize(_)),
            InteractiveSessionEventV2::PromptDecision { .. }
        ]
    ));
}

#[test]
fn terminal_fixture_routes_resize_input_break_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::Input(vec![3]),
        HostTerminalEvent::Break,
        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,
        InteractiveSessionEventV2::Terminal(SessionTerminalEventV2::Resize(_))
    )));
    assert!(transport.events.iter().any(|event| matches!(
        event,
        InteractiveSessionEventV2::Terminal(SessionTerminalEventV2::Input(input))
            if input.bytes == [3]
    )));
    assert!(transport.events.iter().any(|event| matches!(
        event,
        InteractiveSessionEventV2::Terminal(SessionTerminalEventV2::Break(_))
    )));
}

#[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
        ));
    }
}

#[test]
fn local_disconnect_cancels_blocked_transport_and_restores_terminal() {
    let identity = ClientSessionIdentity::generate::<std::convert::Infallible>().unwrap();
    let binding = binding(&identity);
    let mut terminal = PtyTerminal::new([
        HostTerminalEvent::LocalSuspend,
        HostTerminalEvent::Terminate,
    ]);
    let mut transport = FixtureTransport {
        blocked: true,
        ..FixtureTransport::successful(identity.session_id())
    };

    assert!(matches!(
        run_interactive_session(
            &mut transport,
            &mut terminal,
            &identity,
            &binding,
            &[b"agent".to_vec()],
        ),
        Err(ClientError::LostConnection)
    ));
    assert!(transport.cancelled);
    assert_eq!(terminal.suspend_count, 1);
    assert!(terminal.restored());
}

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>,
) -> InteractiveSessionBatchV2 {
    InteractiveSessionBatchV2 {
        session_id,
        messages,
    }
}

fn test_modes() -> TerminalModesV2 {
    TerminalModesV2 {
        input_modes: 0,
        output_modes: 0,
        control_modes: 0,
        local_modes: 0,
        line_discipline: 0,
        input_speed: 38_400,
        output_speed: 38_400,
        interrupt: 3,
        quit: 28,
        erase: 127,
        kill: 21,
        eof: 4,
        time: 0,
        minimum: 1,
        switch: 0,
        start: 17,
        stop: 19,
        suspend: 26,
        end_of_line: 0,
        reprint: 18,
        discard: 15,
        word_erase: 23,
        literal_next: 22,
        end_of_line_2: 0,
    }
}