Add interactive client terminal transport
This commit is contained in:
parent
d597b05484
commit
6816f0632e
1 changed files with 1454 additions and 59 deletions
316
crates/lazaret-client/tests/interactive_terminal.rs
Normal file
316
crates/lazaret-client/tests/interactive_terminal.rs
Normal file
|
|
@ -0,0 +1,316 @@
|
|||
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,
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue