use crate::connection::{FilesystemAccess, SshOptions};
use crate::deadline::{clamp_unix_ms, duration_until_unix_ms};
use crate::error::ErrorKind;
use crate::frame::{FrameDecoder, encode};
use crate::jsonrpc::{Envelope, Request, RequestId, Response};
use crate::launch::{Environment, LaunchOptions};
use crate::protocol::resolver::{
self as resolver_protocol, DeleteParams, DeleteResult, GetParams, GetResult,
InitializeApplication, InitializedApplication, ReleaseParams, ReleaseResult, SetParams,
SetResult,
};
use crate::protocol::{
InitializeParams, InitializeResult, Limits, PROTOCOL_VERSION, Product, RESOLVER_PROTOCOL, rpc,
};
use crate::{ABSOLUTE_MAX_FRAME_BYTES, Error, Result};
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::{Value, json};
use std::collections::{HashSet, VecDeque};
use std::io::{Read, Write};
use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{Receiver, RecvTimeoutError, SyncSender, sync_channel};
use std::sync::{Arc, Condvar, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use zeroize::Zeroizing;
const REAP_GRACE: Duration = Duration::from_secs(2);
const WAIT_POLL: Duration = Duration::from_millis(10);
const READ_CHUNK: usize = 8192;
pub struct ResolverSession {
transport: Transport,
filesystem: FilesystemAccess,
capabilities: HashSet<String>,
initialized: InitializedApplication,
}
impl ResolverSession {
pub fn launch(
options: LaunchOptions,
client: Product,
limits: Limits,
application: InitializeApplication,
startup_deadline_unix_ms: u64,
) -> Result<Self> {
application.validate_for_connection()?;
let initialize = InitializeParams {
protocol: RESOLVER_PROTOCOL.to_string(),
versions: vec![PROTOCOL_VERSION],
client,
limits,
client_methods: Vec::new(),
application,
};
initialize.validate_common(RESOLVER_PROTOCOL)?;
let mut transport = Transport::spawn(&options)?;
let handshake = transport.initialize(&initialize, startup_deadline_unix_ms);
let initialized: InitializeResult<InitializedApplication> = match handshake {
Ok(initialized) => initialized,
Err(error) => {
transport.terminate();
return Err(error);
}
};
let mut session = Self {
transport,
filesystem: FilesystemAccess::Shared,
capabilities: initialized.methods.into_iter().collect(),
initialized: initialized.application,
};
if let Err(error) = session.validate_endpoint() {
session.transport.terminate();
return Err(error);
}
Ok(session)
}
pub fn launch_ssh(
options: SshOptions,
client: Product,
limits: Limits,
application: InitializeApplication,
startup_deadline_unix_ms: u64,
) -> Result<Self> {
let mut session = Self::launch(
options.launch_options()?,
client,
limits,
application,
startup_deadline_unix_ms,
)?;
session.filesystem = options.filesystem;
Ok(session)
}
fn validate_endpoint(&self) -> Result<()> {
self.initialized.validate()?;
if !resolver_protocol::CAPABILITIES
.iter()
.all(|method| self.capabilities.contains(*method))
{
return Err(Error::Protocol(
"resolution endpoint did not advertise all required methods",
));
}
Ok(())
}
pub fn get(&mut self, params: &GetParams, deadline_unix_ms: u64) -> Result<GetResult> {
let params = self.filesystem.prepare_get(params)?;
let result = self.call(
resolver_protocol::method::GET,
params.as_ref(),
deadline_unix_ms,
)?;
if !self.filesystem.accepts(&result) {
self.transport.terminate();
return Err(Error::Protocol(
"resolver returned a path on a remote filesystem",
));
}
Ok(result)
}
pub fn set(&mut self, params: &SetParams, deadline_unix_ms: u64) -> Result<SetResult> {
params.validate()?;
self.call(resolver_protocol::method::SET, params, deadline_unix_ms)
}
pub fn delete(&mut self, params: &DeleteParams, deadline_unix_ms: u64) -> Result<DeleteResult> {
params.validate()?;
self.call(resolver_protocol::method::DELETE, params, deadline_unix_ms)
}
pub fn supports(&self, method: &str) -> bool {
self.capabilities.contains(method)
}
pub fn release(
&mut self,
params: &ReleaseParams,
deadline_unix_ms: u64,
) -> Result<ReleaseResult> {
params.validate()?;
self.call(resolver_protocol::method::RELEASE, params, deadline_unix_ms)
}
pub fn capabilities(&self) -> &HashSet<String> {
&self.capabilities
}
pub fn initialized(&self) -> &InitializedApplication {
&self.initialized
}
pub fn is_closed(&self) -> bool {
self.transport.closed
}
pub fn close(&mut self, deadline_unix_ms: u64) -> Result<()> {
self.transport.close(deadline_unix_ms)
}
fn call<P: Serialize, R: DeserializeOwned>(
&mut self,
method: &str,
params: &P,
deadline_unix_ms: u64,
) -> Result<R> {
if !self.capabilities.contains(method) {
return Err(Error::Protocol("method was not advertised"));
}
self.transport.call(method, params, deadline_unix_ms)
}
}
struct Transport {
child: Arc<Mutex<Child>>,
stdin: Option<ChildStdin>,
stdout: Receiver<StdoutEvent>,
decoder: FrameDecoder,
frames: VecDeque<Zeroizing<Vec<u8>>>,
next_id: u64,
max_frame_bytes: usize,
closed: bool,
}
impl Transport {
fn spawn(options: &LaunchOptions) -> Result<Self> {
options.validate()?;
let decoder = FrameDecoder::new(ABSOLUTE_MAX_FRAME_BYTES)?;
let mut command = Command::new(&options.executable);
command
.args(&options.arguments)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
match &options.environment {
Environment::Inherit(overrides) => {
command.envs(overrides);
}
Environment::Replace(environment) => {
command.env_clear().envs(environment);
}
}
let mut child = command.spawn()?;
let Some(((stdin, stdout), stderr)) = child
.stdin
.take()
.zip(child.stdout.take())
.zip(child.stderr.take())
else {
let _ = child.kill();
let _ = child.wait();
return Err(Error::Protocol("child pipes were not created"));
};
drain_stderr(stderr, options.max_stderr_bytes);
let stdout = read_stdout(stdout);
Ok(Self {
child: Arc::new(Mutex::new(child)),
stdin: Some(stdin),
stdout,
decoder,
frames: VecDeque::new(),
next_id: 1,
max_frame_bytes: ABSOLUTE_MAX_FRAME_BYTES,
closed: false,
})
}
fn initialize<A: Serialize, B: DeserializeOwned>(
&mut self,
initialize: &InitializeParams<A>,
deadline_unix_ms: u64,
) -> Result<InitializeResult<B>> {
let params = serde_json::to_value(initialize)
.map_err(|_| Error::Protocol("failed to serialize initialization"))?;
let value = self.exchange(rpc::INITIALIZE, params, deadline_unix_ms)?;
let initialized: InitializeResult<B> = serde_json::from_value(value)
.map_err(|error| Error::ProtocolOwned(error.to_string()))?;
initialized.validate_common(
&initialize.protocol,
&initialize.versions,
initialize.limits,
)?;
self.max_frame_bytes = initialized.limits.max_frame_bytes;
self.decoder.set_limit(self.max_frame_bytes)?;
Ok(initialized)
}
fn call<P: Serialize, R: DeserializeOwned>(
&mut self,
method: &str,
params: &P,
deadline_unix_ms: u64,
) -> Result<R> {
let params = serde_json::to_value(params)
.map_err(|_| Error::Protocol("failed to serialize call params"))?;
let value = self.exchange(method, params, deadline_unix_ms)?;
serde_json::from_value(value).map_err(|error| Error::ProtocolOwned(error.to_string()))
}
fn exchange(&mut self, method: &str, params: Value, deadline_unix_ms: u64) -> Result<Value> {
if self.closed {
return Err(Error::Closed);
}
let deadline_unix_ms = clamp_unix_ms(deadline_unix_ms);
let remaining = duration_until_unix_ms(deadline_unix_ms);
if remaining.is_zero() {
return Err(Error::DeadlineExceeded);
}
let id = self.next_id()?;
let request = Request::new(id, method, deadline_unix_ms, params)?;
let deadline = Instant::now() + remaining;
let watchdog = Watchdog::arm(&self.child, remaining);
let outcome = self.attempt(request, id, deadline);
let io_timed_out = self.closed && matches!(&outcome, Err(Error::DeadlineExceeded));
let watchdog_fired = watchdog.disarm();
if io_timed_out && !watchdog_fired {
self.stdin = None;
self.kill_and_reap();
}
if watchdog_fired {
self.closed = true;
if outcome.is_err() {
return Err(Error::DeadlineExceeded);
}
}
outcome
}
fn attempt(&mut self, request: Request, id: RequestId, deadline: Instant) -> Result<Value> {
let limit = self.max_frame_bytes;
self.write_envelope(&Envelope::Request(request), limit, deadline)?;
let response = self.read_response(deadline)?;
if response.id() != Some(id) {
self.closed = true;
return Err(Error::Protocol("response ID does not match the request"));
}
response_value(response)
}
fn write_envelope(
&mut self,
envelope: &Envelope,
limit: usize,
deadline: Instant,
) -> Result<()> {
let payload = Zeroizing::new(envelope.to_vec()?);
let frame = Zeroizing::new(encode(&payload, limit)?);
let Some(mut stdin) = self.stdin.take() else {
return Err(Error::Closed);
};
let (sender, receiver) = sync_channel(1);
std::thread::spawn(move || {
let result = stdin.write_all(&frame).and_then(|()| stdin.flush());
let _ = sender.send((stdin, result));
});
match receiver.recv_timeout(deadline.saturating_duration_since(Instant::now())) {
Ok((stdin, result)) => {
self.stdin = Some(stdin);
if let Err(error) = result {
self.closed = true;
return Err(Error::Io(error));
}
Ok(())
}
Err(error) => {
self.closed = true;
Err(match error {
RecvTimeoutError::Timeout => Error::DeadlineExceeded,
RecvTimeoutError::Disconnected => Error::Closed,
})
}
}
}
fn read_response(&mut self, deadline: Instant) -> Result<Response> {
loop {
if let Some(frame) = self.frames.pop_front() {
return match Envelope::parse(&frame) {
Ok(Envelope::Response(response)) => Ok(response),
Ok(Envelope::Notification(_)) => continue,
Ok(Envelope::Request(_)) => {
self.closed = true;
Err(Error::Protocol("peer sent an unexpected request"))
}
Err(error) => {
self.closed = true;
Err(error)
}
};
}
let event = self
.stdout
.recv_timeout(deadline.saturating_duration_since(Instant::now()));
let buffer = match event {
Ok(StdoutEvent::Chunk(buffer)) => buffer,
Ok(StdoutEvent::Error(error)) => {
self.closed = true;
return Err(Error::Io(error));
}
Ok(StdoutEvent::Eof) | Err(RecvTimeoutError::Disconnected) => {
self.closed = true;
self.decoder.finish_eof()?;
return Err(Error::Closed);
}
Err(RecvTimeoutError::Timeout) => {
self.closed = true;
return Err(Error::DeadlineExceeded);
}
};
if buffer.is_empty() {
self.closed = true;
return Err(Error::Protocol("stdout reader returned an empty chunk"));
}
match self.decoder.push(&buffer) {
Ok(frames) => self.frames.extend(frames),
Err(error) => {
self.closed = true;
return Err(error);
}
}
}
}
fn next_id(&mut self) -> Result<RequestId> {
let id = RequestId::new(self.next_id)?;
self.next_id += 1;
Ok(id)
}
fn close(&mut self, deadline_unix_ms: u64) -> Result<()> {
let outcome = if self.closed {
Ok(())
} else {
self.exchange(rpc::SHUTDOWN, json!({}), deadline_unix_ms)
.and_then(|value| {
if value == json!({}) {
Ok(())
} else {
Err(Error::Protocol("shutdown result is not empty"))
}
})
};
self.closed = true;
self.stdin = None;
let graceful = wait_until(&self.child, Instant::now() + REAP_GRACE);
if !matches!(graceful, Ok(true)) {
self.kill_and_reap();
}
graceful?;
outcome
}
fn terminate(&mut self) {
self.closed = true;
self.stdin = None;
self.kill_and_reap();
}
fn kill_and_reap(&mut self) {
let mut child = lock_unpoisoned(&self.child);
let _ = child.kill();
drop(child);
let _ = wait_until(&self.child, Instant::now() + REAP_GRACE);
}
}
impl Drop for Transport {
fn drop(&mut self) {
self.terminate();
}
}
enum StdoutEvent {
Chunk(Zeroizing<Vec<u8>>),
Eof,
Error(std::io::Error),
}
fn read_stdout(mut stdout: ChildStdout) -> Receiver<StdoutEvent> {
let (sender, receiver) = sync_channel(1);
std::thread::spawn(move || {
loop {
let mut buffer = Zeroizing::new(vec![0_u8; READ_CHUNK]);
match stdout.read(&mut buffer) {
Ok(0) => {
send_stdout_event(&sender, StdoutEvent::Eof);
break;
}
Ok(read) => {
buffer.truncate(read);
if sender.send(StdoutEvent::Chunk(buffer)).is_err() {
break;
}
}
Err(error) => {
send_stdout_event(&sender, StdoutEvent::Error(error));
break;
}
}
}
});
receiver
}
fn send_stdout_event(sender: &SyncSender<StdoutEvent>, event: StdoutEvent) {
let _ = sender.send(event);
}
struct Watchdog {
finished: Arc<(Mutex<bool>, Condvar)>,
fired: Arc<AtomicBool>,
thread: Option<std::thread::JoinHandle<()>>,
}
impl Watchdog {
fn arm(child: &Arc<Mutex<Child>>, timeout: Duration) -> Self {
let finished = Arc::new((Mutex::new(false), Condvar::new()));
let fired = Arc::new(AtomicBool::new(false));
let thread = std::thread::spawn({
let finished = Arc::clone(&finished);
let fired = Arc::clone(&fired);
let child = Arc::clone(child);
move || {
let (lock, condvar) = &*finished;
let guard = lock_unpoisoned(lock);
let (_guard, timeout_result) = condvar
.wait_timeout_while(guard, timeout, |finished| !*finished)
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !timeout_result.timed_out() {
return;
}
fired.store(true, Ordering::Release);
let _ = lock_unpoisoned(&child).kill();
}
});
Self {
finished,
fired,
thread: Some(thread),
}
}
fn disarm(mut self) -> bool {
let (lock, condvar) = &*self.finished;
*lock_unpoisoned(lock) = true;
condvar.notify_all();
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
self.fired.load(Ordering::Acquire)
}
}
fn drain_stderr(mut stderr: ChildStderr, max_stderr_bytes: usize) {
std::thread::spawn(move || {
let mut retained = Zeroizing::new(Vec::with_capacity(max_stderr_bytes));
let mut buffer = Zeroizing::new(vec![0_u8; READ_CHUNK]);
loop {
let read = match stderr.read(&mut buffer) {
Ok(0) | Err(_) => break,
Ok(read) => read,
};
let available = max_stderr_bytes.saturating_sub(retained.len());
retained.extend_from_slice(&buffer[..read.min(available)]);
}
});
}
fn wait_until(child: &Arc<Mutex<Child>>, deadline: Instant) -> Result<bool> {
loop {
if lock_unpoisoned(child).try_wait()?.is_some() {
return Ok(true);
}
if Instant::now() >= deadline {
return Ok(false);
}
std::thread::sleep(WAIT_POLL);
}
}
fn response_value(response: Response) -> Result<Value> {
match response {
Response::Success(response) => Ok(response.result),
Response::Error(response) => match response.error.data.kind {
ErrorKind::Cancelled => Err(Error::Cancelled),
ErrorKind::DeadlineExceeded => Err(Error::DeadlineExceeded),
ErrorKind::Unavailable => Err(Error::Unavailable),
_ => Err(Error::Remote(response.error)),
},
}
}
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
#[test]
fn inherited_stdin_does_not_block_write_deadline() {
let options = LaunchOptions {
executable: "sh".into(),
arguments: vec!["-c".into(), "(sleep 4) <&0 & printf ready; wait".into()],
environment: Environment::Inherit(Default::default()),
allow_path_discovery: true,
max_stderr_bytes: 0,
};
let mut transport = Transport::spawn(&options).unwrap();
assert!(matches!(
transport.stdout.recv_timeout(Duration::from_secs(2)),
Ok(StdoutEvent::Chunk(_))
));
let started = Instant::now();
let result = transport.exchange(
rpc::INITIALIZE,
json!({"padding": "x".repeat(512 * 1024)}),
crate::deadline_unix_ms_after(Duration::from_millis(150)),
);
assert!(matches!(result, Err(Error::DeadlineExceeded)), "{result:?}");
assert!(transport.closed);
drop(transport);
assert!(started.elapsed() < Duration::from_secs(2));
}
}