From d9ee6aa49b68d23a161f97be0002925f8202af3a Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 10:11:45 -0700 Subject: [PATCH 1/9] fix: shut down JSONRPC cleanly on EOF (Fixes #529) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet-conda/src/conda_info.rs | 8 + crates/pet-jsonrpc/src/lib.rs | 67 +- crates/pet-jsonrpc/src/output.rs | 1133 +++++++++++++++++ crates/pet-jsonrpc/src/server.rs | 310 ++++- .../src/environment_locations_spawn.rs | 8 + crates/pet-python-utils/src/env.rs | 22 + crates/pet-python-utils/src/process.rs | 347 ++++- .../src/process/supervisor.rs | 163 +++ crates/pet/src/jsonrpc.rs | 13 +- crates/pet/src/main.rs | 12 +- crates/pet/tests/e2e_performance.rs | 17 +- crates/pet/tests/fixtures/shutdown_probe.py | 61 + crates/pet/tests/jsonrpc_client.rs | 71 +- crates/pet/tests/jsonrpc_server_test.rs | 237 +++- crates/pet/tests/process_utils.rs | 119 ++ docs/JSONRPC.md | 44 +- 16 files changed, 2519 insertions(+), 113 deletions(-) create mode 100644 crates/pet-jsonrpc/src/output.rs create mode 100644 crates/pet-python-utils/src/process/supervisor.rs create mode 100644 crates/pet/tests/fixtures/shutdown_probe.py create mode 100644 crates/pet/tests/process_utils.rs diff --git a/crates/pet-conda/src/conda_info.rs b/crates/pet-conda/src/conda_info.rs index a8789d6b..c2d80d69 100644 --- a/crates/pet-conda/src/conda_info.rs +++ b/crates/pet-conda/src/conda_info.rs @@ -107,6 +107,13 @@ impl CondaInfo { None } } + Err(ProcessError::Cancelled) => { + trace!( + "Cancelled Conda probe during process shutdown: {:?}", + executable + ); + None + } Err(err) => { if !is_missing_default_conda(using_default, &err) { warn!( @@ -191,6 +198,7 @@ mod tests { ); } for error in [ + ProcessError::Cancelled, ProcessError::Spawn(io::Error::from(io::ErrorKind::NotFound)), ProcessError::Io(io::Error::from(io::ErrorKind::BrokenPipe)), ProcessError::Timeout(Duration::from_secs(15)), diff --git a/crates/pet-jsonrpc/src/lib.rs b/crates/pet-jsonrpc/src/lib.rs index bd2b3325..8a3341e8 100644 --- a/crates/pet-jsonrpc/src/lib.rs +++ b/crates/pet-jsonrpc/src/lib.rs @@ -2,10 +2,12 @@ // Licensed under the MIT License. use serde::{Deserialize, Serialize}; -use std::io::{self, Write}; +mod output; pub mod server; +pub use output::{initialize_output, output_error, shutdown_output}; + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(untagged)] pub enum RequestId { @@ -24,51 +26,52 @@ impl From for RequestId { #[serde(rename_all = "camelCase")] #[derive(Debug)] struct AnyMethodMessage { - pub jsonrpc: String, + pub jsonrpc: &'static str, pub method: &'static str, pub params: Option, } pub fn send_message(method: &'static str, params: Option) { let payload = AnyMethodMessage { - jsonrpc: "2.0".to_string(), + jsonrpc: "2.0", method, params, }; - let message = serde_json::to_string(&payload).unwrap(); - print!( - "Content-Length: {}\r\nContent-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n{}", - message.len(), - message - ); - let _ = io::stdout().flush(); + output::send(&payload); } + pub fn send_reply(id: &RequestId, payload: Option) { - let payload = serde_json::json!({ - "jsonrpc": "2.0", - "result": payload, - "id": id + #[derive(Serialize)] + struct Reply<'a, T> { + jsonrpc: &'static str, + result: Option, + id: &'a RequestId, + } + + output::send(&Reply { + jsonrpc: "2.0", + result: payload, + id, }); - let message = serde_json::to_string(&payload).unwrap(); - print!( - "Content-Length: {}\r\nContent-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n{}", - message.len(), - message - ); - let _ = io::stdout().flush(); } pub fn send_error(id: Option<&RequestId>, code: i32, message: String) { - let payload = serde_json::json!({ - "jsonrpc": "2.0", - "error": { "code": code, "message": message }, - "id": id + #[derive(Serialize)] + struct ErrorBody { + code: i32, + message: String, + } + + #[derive(Serialize)] + struct ErrorReply<'a> { + jsonrpc: &'static str, + error: ErrorBody, + id: Option<&'a RequestId>, + } + + output::send(&ErrorReply { + jsonrpc: "2.0", + error: ErrorBody { code, message }, + id, }); - let message = serde_json::to_string(&payload).unwrap(); - print!( - "Content-Length: {}\r\nContent-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n{}", - message.len(), - message - ); - let _ = io::stdout().flush(); } diff --git a/crates/pet-jsonrpc/src/output.rs b/crates/pet-jsonrpc/src/output.rs new file mode 100644 index 00000000..d572027f --- /dev/null +++ b/crates/pet-jsonrpc/src/output.rs @@ -0,0 +1,1133 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use serde::Serialize; +use std::collections::VecDeque; +use std::fs::File; +use std::io::{self, Write}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread; + +const MAX_QUEUED_BYTES: usize = 32 * 1024 * 1024; +const MAX_QUEUED_FRAMES: usize = 1024; +const MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024; +const HEADER_PREFIX: &[u8] = b"Content-Length: "; +const HEADER_SUFFIX: &[u8] = b"\r\nContent-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n"; + +#[derive(Clone, Debug)] +struct ErrorSnapshot { + kind: io::ErrorKind, + message: Arc, +} + +impl ErrorSnapshot { + fn new(error: io::Error) -> Self { + Self { + kind: error.kind(), + message: Arc::from(error.to_string()), + } + } + + fn to_error(&self) -> io::Error { + io::Error::new(self.kind, self.message.to_string()) + } +} + +struct State { + accepting: bool, + queue: VecDeque>, + queued_bytes: usize, + error: Option, +} + +impl State { + fn open() -> Self { + Self { + accepting: true, + queue: VecDeque::new(), + queued_bytes: 0, + error: None, + } + } + + fn fail(&mut self, error: io::Error) { + if !self.accepting { + return; + } + self.error = Some(ErrorSnapshot::new(error)); + self.accepting = false; + self.queue.clear(); + self.queued_bytes = 0; + } +} + +struct Shared { + state: Mutex, + ready: Condvar, +} + +#[derive(Clone, Copy)] +struct Limits { + queued_bytes: usize, + queued_frames: usize, + payload_bytes: usize, +} + +impl Default for Limits { + fn default() -> Self { + Self { + queued_bytes: MAX_QUEUED_BYTES, + queued_frames: MAX_QUEUED_FRAMES, + payload_bytes: MAX_PAYLOAD_BYTES, + } + } +} + +#[derive(Clone)] +struct Output { + shared: Arc, + serialization: Arc>, + limits: Limits, +} + +impl Output { + fn new(writer: impl Write + Send + 'static) -> io::Result { + Self::with_limits(writer, Limits::default()) + } + + fn with_limits(writer: impl Write + Send + 'static, limits: Limits) -> io::Result { + let shared = Arc::new(Shared { + state: Mutex::new(State::open()), + ready: Condvar::new(), + }); + let writer_shared = Arc::clone(&shared); + thread::Builder::new() + .name("pet-jsonrpc-output".to_string()) + .spawn(move || writer_loop(writer, writer_shared))?; + + Ok(Self { + shared, + serialization: Arc::new(Mutex::new(())), + limits, + }) + } + + fn send(&self, value: &T) { + if !self.is_accepting() { + return; + } + + let _serialization = self + .serialization + .lock() + .expect("JSONRPC output serialization lock poisoned"); + if !self.is_accepting() { + return; + } + + let frame = match encode_frame(value, self.limits.payload_bytes) { + Ok(frame) => frame, + Err(error) => { + self.fail(error); + return; + } + }; + + let mut state = self + .shared + .state + .lock() + .expect("JSONRPC output state lock poisoned"); + if !state.accepting { + return; + } + + let Some(queued_bytes) = state.queued_bytes.checked_add(frame.capacity()) else { + state.fail(io::Error::new( + io::ErrorKind::WouldBlock, + "JSONRPC output queue byte count overflowed", + )); + self.shared.ready.notify_all(); + return; + }; + if state.queue.len() >= self.limits.queued_frames || queued_bytes > self.limits.queued_bytes + { + state.fail(io::Error::new( + io::ErrorKind::WouldBlock, + format!( + "JSONRPC output queue is full (limit: {} frames and {} bytes)", + self.limits.queued_frames, self.limits.queued_bytes + ), + )); + self.shared.ready.notify_all(); + return; + } + + state.queued_bytes = queued_bytes; + state.queue.push_back(frame); + self.shared.ready.notify_one(); + } + + fn is_accepting(&self) -> bool { + self.shared + .state + .lock() + .expect("JSONRPC output state lock poisoned") + .accepting + } + + fn fail(&self, error: io::Error) { + let mut state = self + .shared + .state + .lock() + .expect("JSONRPC output state lock poisoned"); + state.fail(error); + self.shared.ready.notify_all(); + } + + fn error(&self) -> Option { + self.shared + .state + .lock() + .expect("JSONRPC output state lock poisoned") + .error + .as_ref() + .map(ErrorSnapshot::to_error) + } + + fn close(&self) { + let mut state = self + .shared + .state + .lock() + .expect("JSONRPC output state lock poisoned"); + state.accepting = false; + state.queue.clear(); + state.queued_bytes = 0; + self.shared.ready.notify_all(); + } +} + +fn writer_loop(mut writer: impl Write, shared: Arc) { + loop { + let frame = { + let mut state = shared + .state + .lock() + .expect("JSONRPC output state lock poisoned"); + loop { + if let Some(frame) = state.queue.pop_front() { + state.queued_bytes -= frame.capacity(); + break frame; + } + if !state.accepting { + return; + } + state = shared + .ready + .wait(state) + .expect("JSONRPC output state lock poisoned while waiting"); + } + }; + + if let Err(error) = writer.write_all(&frame) { + let mut state = shared + .state + .lock() + .expect("JSONRPC output state lock poisoned"); + state.fail(io::Error::new( + error.kind(), + format!("failed to write JSONRPC output: {error}"), + )); + shared.ready.notify_all(); + return; + } + if let Err(error) = writer.flush() { + let mut state = shared + .state + .lock() + .expect("JSONRPC output state lock poisoned"); + state.fail(io::Error::new( + error.kind(), + format!("failed to flush JSONRPC output: {error}"), + )); + shared.ready.notify_all(); + return; + } + } +} + +struct LimitedBuffer { + bytes: Vec, + limit: usize, +} + +impl LimitedBuffer { + fn new(limit: usize) -> Self { + Self { + bytes: Vec::new(), + limit, + } + } + + fn reserve_for(&mut self, new_len: usize) -> io::Result<()> { + if new_len <= self.bytes.capacity() { + return Ok(()); + } + + let doubled_capacity = self.bytes.capacity().checked_mul(2).unwrap_or(self.limit); + let new_capacity = doubled_capacity.max(new_len).min(self.limit); + let additional = new_capacity.checked_sub(self.bytes.len()).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "JSONRPC payload capacity calculation underflowed", + ) + })?; + self.bytes.try_reserve_exact(additional).map_err(|error| { + io::Error::new( + io::ErrorKind::OutOfMemory, + format!("failed to allocate JSONRPC payload: {error}"), + ) + })?; + if self.bytes.capacity() > self.limit { + self.bytes = Vec::new(); + return Err(io::Error::new( + io::ErrorKind::OutOfMemory, + "JSONRPC payload allocation exceeded its maximum capacity", + )); + } + Ok(()) + } +} + +impl Write for LimitedBuffer { + fn write(&mut self, bytes: &[u8]) -> io::Result { + let new_len = + self.bytes.len().checked_add(bytes.len()).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "JSON payload too large") + })?; + if new_len > self.limit { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "JSONRPC payload exceeds the maximum size of {} bytes", + self.limit + ), + )); + } + self.reserve_for(new_len)?; + self.bytes.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +fn encode_frame(value: &T, payload_limit: usize) -> io::Result> { + let mut payload = LimitedBuffer::new(payload_limit); + serde_json::to_writer(&mut payload, value).map_err(|error| { + let kind = error.io_error_kind().unwrap_or(io::ErrorKind::InvalidData); + io::Error::new(kind, format!("failed to serialize JSONRPC output: {error}")) + })?; + + let length = payload.bytes.len().to_string(); + let frame_len = HEADER_PREFIX + .len() + .checked_add(length.len()) + .and_then(|len| len.checked_add(HEADER_SUFFIX.len())) + .and_then(|len| len.checked_add(payload.bytes.len())) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "JSONRPC frame too large"))?; + let mut frame = Vec::new(); + frame.try_reserve_exact(frame_len).map_err(|error| { + io::Error::new( + io::ErrorKind::OutOfMemory, + format!("failed to allocate JSONRPC frame: {error}"), + ) + })?; + frame.extend_from_slice(HEADER_PREFIX); + frame.extend_from_slice(length.as_bytes()); + frame.extend_from_slice(HEADER_SUFFIX); + frame.extend_from_slice(&payload.bytes); + Ok(frame) +} + +struct GlobalState { + output: Option, + shutdown: bool, + initialization_error: Option, +} + +static GLOBAL: Mutex = Mutex::new(GlobalState { + output: None, + shutdown: false, + initialization_error: None, +}); + +/// Starts the process-wide JSONRPC writer using an owned duplicate of stdout. +/// +/// The writer is intentionally process-lifetime infrastructure for PET's standalone +/// server. Shutdown never joins it: an OS write already in progress may remain +/// blocked until the process exits. +pub fn initialize_output() -> io::Result<()> { + let mut global = GLOBAL + .lock() + .expect("global JSONRPC output state lock poisoned"); + initialize_output_state(&mut global, || duplicate_stdout().and_then(Output::new)) +} + +fn initialize_output_state( + global: &mut GlobalState, + create_output: impl FnOnce() -> io::Result, +) -> io::Result<()> { + if global.shutdown { + return Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "JSONRPC output has been shut down", + )); + } + if let Some(error) = &global.initialization_error { + return Err(error.to_error()); + } + if global.output.is_some() { + return Ok(()); + } + + let output = create_output(); + match output { + Ok(output) => { + global.output = Some(output); + Ok(()) + } + Err(error) => { + let error = io::Error::new( + error.kind(), + format!("failed to initialize JSONRPC output: {error}"), + ); + global.initialization_error = Some(ErrorSnapshot::new(error)); + Err(global + .initialization_error + .as_ref() + .expect("initialization error was just stored") + .to_error()) + } + } +} + +/// Returns a copy of the first fatal output error, if one has occurred. +pub fn output_error() -> Option { + let global = GLOBAL + .lock() + .expect("global JSONRPC output state lock poisoned"); + global + .initialization_error + .as_ref() + .map(ErrorSnapshot::to_error) + .or_else(|| global.output.as_ref().and_then(Output::error)) +} + +/// Irreversibly closes process-wide output admission and discards queued frames. +/// +/// This returns without joining the writer. A write already blocked in the OS is +/// left for process termination to clean up. +pub fn shutdown_output() { + let output = { + let mut global = GLOBAL + .lock() + .expect("global JSONRPC output state lock poisoned"); + shutdown_output_state(&mut global) + }; + if let Some(output) = output { + output.close(); + } +} + +fn shutdown_output_state(global: &mut GlobalState) -> Option { + global.shutdown = true; + global.output.clone() +} + +pub(crate) fn send(value: &T) { + let output = { + let mut global = GLOBAL + .lock() + .expect("global JSONRPC output state lock poisoned"); + if global.shutdown || global.initialization_error.is_some() { + return; + } + if global.output.is_none() { + match duplicate_stdout().and_then(Output::new) { + Ok(output) => global.output = Some(output), + Err(error) => { + let error = io::Error::new( + error.kind(), + format!("failed to initialize JSONRPC output: {error}"), + ); + global.initialization_error = Some(ErrorSnapshot::new(error)); + return; + } + } + } + global.output.clone() + }; + + if let Some(output) = output { + output.send(value); + } +} + +#[cfg(unix)] +fn duplicate_stdout() -> io::Result { + use std::os::fd::AsFd; + + io::stdout().as_fd().try_clone_to_owned().map(File::from) +} + +#[cfg(windows)] +fn duplicate_stdout() -> io::Result { + use std::os::windows::io::AsHandle; + + io::stdout() + .as_handle() + .try_clone_to_owned() + .map(File::from) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::Serializer; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::mpsc::{self, Receiver, Sender}; + use std::time::{Duration, Instant}; + + const TIMEOUT: Duration = Duration::from_secs(5); + + #[derive(Clone, Default)] + struct CapturedWriter { + bytes: Arc>>, + } + + impl Write for CapturedWriter { + fn write(&mut self, bytes: &[u8]) -> io::Result { + self.bytes + .lock() + .expect("captured output lock poisoned") + .extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + struct BlockingWriter { + entered: Option>, + release: Arc<(Mutex, Condvar)>, + bytes: Arc>>, + finished: Sender<()>, + } + + impl Write for BlockingWriter { + fn write(&mut self, bytes: &[u8]) -> io::Result { + if let Some(entered) = self.entered.take() { + entered.send(()).expect("blocked writer observer dropped"); + } + let (released, ready) = &*self.release; + let mut released = released.lock().expect("release lock poisoned"); + while !*released { + released = ready.wait(released).expect("release lock poisoned"); + } + self.bytes + .lock() + .expect("captured output lock poisoned") + .extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.finished + .send(()) + .expect("writer completion observer dropped"); + Ok(()) + } + } + + struct FailingWriter { + write_error: Option, + flush_error: Option, + } + + impl Write for FailingWriter { + fn write(&mut self, bytes: &[u8]) -> io::Result { + if let Some(kind) = self.write_error { + Err(io::Error::new(kind, "deliberate write failure")) + } else { + Ok(bytes.len()) + } + } + + fn flush(&mut self) -> io::Result<()> { + if let Some(kind) = self.flush_error { + Err(io::Error::new(kind, "deliberate flush failure")) + } else { + Ok(()) + } + } + } + + struct DelayedFailingWriter { + entered: Option>, + release: Arc<(Mutex, Condvar)>, + completed: Option>, + } + + impl Write for DelayedFailingWriter { + fn write(&mut self, _bytes: &[u8]) -> io::Result { + self.entered + .take() + .expect("writer entered more than once") + .send(()) + .expect("writer observer dropped"); + let (released, ready) = &*self.release; + let mut released = released.lock().expect("release lock poisoned"); + while !*released { + released = ready.wait(released).expect("release lock poisoned"); + } + Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "delayed write failure", + )) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + impl Drop for DelayedFailingWriter { + fn drop(&mut self) { + self.completed + .take() + .expect("writer completion already reported") + .send(()) + .expect("writer completion observer dropped"); + } + } + + fn wait_for_error(output: &Output) -> io::Error { + let deadline = Instant::now() + TIMEOUT; + loop { + if let Some(error) = output.error() { + return error; + } + assert!(Instant::now() < deadline, "timed out waiting for error"); + thread::yield_now(); + } + } + + fn wait_for_len(bytes: &Arc>>, expected: usize) { + let deadline = Instant::now() + TIMEOUT; + loop { + if bytes.lock().expect("captured output lock poisoned").len() >= expected { + return; + } + assert!(Instant::now() < deadline, "timed out waiting for output"); + thread::yield_now(); + } + } + + fn frames(bytes: &[u8]) -> Vec<&[u8]> { + let mut remaining = bytes; + let mut result = Vec::new(); + while !remaining.is_empty() { + let separator = b"\r\n\r\n"; + let header_end = remaining + .windows(separator.len()) + .position(|window| window == separator) + .expect("frame header terminator missing"); + let header = std::str::from_utf8(&remaining[..header_end]).expect("invalid header"); + let length: usize = header + .strip_prefix("Content-Length: ") + .and_then(|header| header.lines().next()) + .expect("content length missing") + .parse() + .expect("invalid content length"); + let payload_start = header_end + separator.len(); + let payload_end = payload_start + length; + result.push(&remaining[payload_start..payload_end]); + remaining = &remaining[payload_end..]; + } + result + } + + struct BlockingHarness { + output: Output, + entered: Receiver<()>, + release: Arc<(Mutex, Condvar)>, + bytes: Arc>>, + finished: Receiver<()>, + } + + fn blocking_output(limits: Limits) -> BlockingHarness { + let (entered_tx, entered_rx) = mpsc::channel(); + let (finished_tx, finished_rx) = mpsc::channel(); + let release = Arc::new((Mutex::new(false), Condvar::new())); + let bytes = Arc::new(Mutex::new(Vec::new())); + let output = Output::with_limits( + BlockingWriter { + entered: Some(entered_tx), + release: Arc::clone(&release), + bytes: Arc::clone(&bytes), + finished: finished_tx, + }, + limits, + ) + .expect("failed to create output"); + BlockingHarness { + output, + entered: entered_rx, + release, + bytes, + finished: finished_rx, + } + } + + fn release_writer(release: &Arc<(Mutex, Condvar)>) { + let (released, ready) = &**release; + *released.lock().expect("release lock poisoned") = true; + ready.notify_all(); + } + + #[test] + fn frames_unicode_by_utf8_length_and_preserves_fifo() { + let writer = CapturedWriter::default(); + let bytes = Arc::clone(&writer.bytes); + let output = Output::new(writer).expect("failed to create output"); + + output.send(&"snowman \u{2603}"); + output.send(&serde_json::json!({"sequence": 2})); + + let expected_first = serde_json::to_vec(&"snowman \u{2603}").expect("serialization failed"); + let expected_second = + serde_json::to_vec(&serde_json::json!({"sequence": 2})).expect("serialization failed"); + let expected_len = expected_first.len() + + expected_second.len() + + 2 * (HEADER_PREFIX.len() + HEADER_SUFFIX.len()) + + expected_first.len().to_string().len() + + expected_second.len().to_string().len(); + wait_for_len(&bytes, expected_len); + output.close(); + + let bytes = bytes.lock().expect("captured output lock poisoned"); + assert_eq!(frames(&bytes), vec![expected_first, expected_second]); + } + + #[test] + fn blocked_sink_does_not_hold_control_or_queue_locks() { + let BlockingHarness { + output, + entered, + release, + finished, + .. + } = blocking_output(Limits { + queued_bytes: 1024, + queued_frames: 1, + payload_bytes: 512, + }); + output.send(&"first"); + entered.recv_timeout(TIMEOUT).expect("writer did not block"); + + let start = Instant::now(); + output.send(&"pending"); + assert!(output.error().is_none()); + output.close(); + assert!(start.elapsed() < Duration::from_secs(1)); + + release_writer(&release); + finished + .recv_timeout(TIMEOUT) + .expect("writer did not finish"); + } + + #[test] + fn saturation_is_fatal_and_clears_pending_output() { + let BlockingHarness { + output, + entered, + release, + bytes, + finished, + } = blocking_output(Limits { + queued_bytes: 1024, + queued_frames: 1, + payload_bytes: 512, + }); + output.send(&"first"); + entered.recv_timeout(TIMEOUT).expect("writer did not block"); + output.send(&"pending"); + output.send(&"overflow"); + + let error = output.error().expect("queue saturation was not recorded"); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + assert!(error.to_string().contains("queue is full")); + + release_writer(&release); + finished + .recv_timeout(TIMEOUT) + .expect("writer did not finish"); + let bytes = bytes.lock().expect("captured output lock poisoned"); + let emitted = frames(&bytes); + assert_eq!(emitted, vec![serde_json::to_vec(&"first").unwrap()]); + } + + #[test] + fn oversized_serialization_is_bounded_and_fatal() { + let writer = CapturedWriter::default(); + let bytes = Arc::clone(&writer.bytes); + let output = Output::with_limits( + writer, + Limits { + queued_bytes: 128, + queued_frames: 2, + payload_bytes: 16, + }, + ) + .expect("failed to create output"); + + output.send(&"this payload is much too large"); + + let error = output.error().expect("oversized payload was not recorded"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert!(error.to_string().contains("maximum size")); + assert!(bytes + .lock() + .expect("captured output lock poisoned") + .is_empty()); + } + + #[test] + fn limited_buffer_caps_capacity_during_multi_chunk_growth() { + let mut buffer = LimitedBuffer::new(10); + + buffer.write_all(b"123456").expect("first write failed"); + assert_eq!(buffer.bytes.len(), 6); + assert!(buffer.bytes.capacity() <= buffer.limit); + + buffer.write_all(b"7890").expect("second write failed"); + assert_eq!(buffer.bytes, b"1234567890"); + assert_eq!(buffer.bytes.capacity(), buffer.limit); + } + + #[test] + fn limited_buffer_accepts_exact_limit_and_rejects_over_limit() { + let mut buffer = LimitedBuffer::new(8); + buffer + .write_all(b"12345678") + .expect("exact-limit write failed"); + assert_eq!(buffer.bytes.len(), buffer.limit); + assert!(buffer.bytes.capacity() <= buffer.limit); + + let capacity = buffer.bytes.capacity(); + let error = buffer + .write_all(b"9") + .expect_err("over-limit write succeeded"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(buffer.bytes, b"12345678"); + assert_eq!(buffer.bytes.capacity(), capacity); + } + + #[test] + fn frame_allocation_is_exact_and_payload_limit_is_enforced() { + let payload_len = serde_json::to_vec("1234") + .expect("failed to encode expected payload") + .len(); + let frame = encode_frame(&"1234", payload_len).expect("exact-limit frame failed"); + assert_eq!(frame.capacity(), frame.len()); + + let error = encode_frame(&"1234", payload_len - 1).expect_err("over-limit frame succeeded"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn queued_byte_limit_is_enforced_independently_of_frame_count() { + let queued_frame_bytes = encode_frame(&"pending", 512) + .expect("failed to encode test frame") + .capacity(); + let BlockingHarness { + output, + entered, + release, + bytes, + finished, + } = blocking_output(Limits { + queued_bytes: queued_frame_bytes, + queued_frames: 4, + payload_bytes: 512, + }); + output.send(&"x"); + entered.recv_timeout(TIMEOUT).expect("writer did not block"); + output.send(&"pending"); + output.send(&"one byte too many"); + + let error = output + .error() + .expect("queue byte saturation was not recorded"); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + + release_writer(&release); + finished + .recv_timeout(TIMEOUT) + .expect("writer did not finish"); + let bytes = bytes.lock().expect("captured output lock poisoned"); + assert_eq!(frames(&bytes), vec![serde_json::to_vec(&"x").unwrap()]); + } + + #[test] + fn write_failure_is_explicit_and_first_error_wins() { + let output = Output::new(FailingWriter { + write_error: Some(io::ErrorKind::BrokenPipe), + flush_error: None, + }) + .expect("failed to create output"); + + output.send(&"message"); + let first = wait_for_error(&output); + assert_eq!(first.kind(), io::ErrorKind::BrokenPipe); + assert!(first.to_string().contains("deliberate write failure")); + + output.fail(io::Error::new( + io::ErrorKind::ConnectionReset, + "later failure", + )); + output.close(); + let preserved = output.error().expect("first error was lost"); + assert_eq!(preserved.kind(), io::ErrorKind::BrokenPipe); + assert!(!preserved.to_string().contains("later failure")); + } + + #[test] + fn flush_failure_is_explicit() { + let output = Output::new(FailingWriter { + write_error: None, + flush_error: Some(io::ErrorKind::ConnectionAborted), + }) + .expect("failed to create output"); + + output.send(&"message"); + let error = wait_for_error(&output); + assert_eq!(error.kind(), io::ErrorKind::ConnectionAborted); + assert!(error.to_string().contains("failed to flush JSONRPC output")); + assert!(error.to_string().contains("deliberate flush failure")); + } + + #[test] + fn in_flight_write_failure_after_close_is_discarded() { + let (entered_tx, entered_rx) = mpsc::channel(); + let (completed_tx, completed_rx) = mpsc::channel(); + let release = Arc::new((Mutex::new(false), Condvar::new())); + let output = Output::new(DelayedFailingWriter { + entered: Some(entered_tx), + release: Arc::clone(&release), + completed: Some(completed_tx), + }) + .expect("failed to create output"); + + output.send(&"message"); + entered_rx + .recv_timeout(TIMEOUT) + .expect("writer did not start"); + output.close(); + release_writer(&release); + completed_rx + .recv_timeout(TIMEOUT) + .expect("writer did not finish"); + + assert!(output.error().is_none()); + } + + #[test] + fn close_discards_pending_frames_without_joining_writer() { + let BlockingHarness { + output, + entered, + release, + bytes, + finished, + } = blocking_output(Limits { + queued_bytes: 1024, + queued_frames: 1, + payload_bytes: 512, + }); + output.send(&"in progress"); + entered.recv_timeout(TIMEOUT).expect("writer did not block"); + output.send(&"pending"); + + let start = Instant::now(); + output.close(); + assert!(start.elapsed() < Duration::from_secs(1)); + + release_writer(&release); + finished + .recv_timeout(TIMEOUT) + .expect("writer did not finish"); + let bytes = bytes.lock().expect("captured output lock poisoned"); + assert_eq!( + frames(&bytes), + vec![serde_json::to_vec(&"in progress").unwrap()] + ); + } + + struct CountedSerialization(Arc); + + impl Serialize for CountedSerialization { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.0.fetch_add(1, Ordering::SeqCst); + serializer.serialize_str("unexpected") + } + } + + struct DelayedFailingSerialization { + entered: Sender<()>, + release: Arc<(Mutex, Condvar)>, + } + + impl Serialize for DelayedFailingSerialization { + fn serialize(&self, _serializer: S) -> Result + where + S: Serializer, + { + self.entered + .send(()) + .expect("serialization observer dropped"); + let (released, ready) = &*self.release; + let mut released = released.lock().expect("release lock poisoned"); + while !*released { + released = ready.wait(released).expect("release lock poisoned"); + } + Err(::custom( + "delayed serialization failure", + )) + } + } + + #[test] + fn in_flight_serialization_failure_after_close_is_discarded() { + let output = Output::new(CapturedWriter::default()).expect("failed to create output"); + let (entered_tx, entered_rx) = mpsc::channel(); + let release = Arc::new((Mutex::new(false), Condvar::new())); + let serialization = DelayedFailingSerialization { + entered: entered_tx, + release: Arc::clone(&release), + }; + let publisher = { + let output = output.clone(); + thread::spawn(move || output.send(&serialization)) + }; + + entered_rx + .recv_timeout(TIMEOUT) + .expect("serialization did not start"); + output.close(); + release_writer(&release); + publisher.join().expect("publisher panicked"); + + assert!(output.error().is_none()); + } + + #[test] + fn post_close_send_does_not_serialize_or_write() { + let writer = CapturedWriter::default(); + let bytes = Arc::clone(&writer.bytes); + let output = Output::new(writer).expect("failed to create output"); + output.close(); + let serializations = Arc::new(AtomicUsize::new(0)); + + output.send(&CountedSerialization(Arc::clone(&serializations))); + + assert_eq!(serializations.load(Ordering::SeqCst), 0); + assert!(bytes + .lock() + .expect("captured output lock poisoned") + .is_empty()); + assert!(output.error().is_none()); + } + + #[test] + fn concurrent_messages_are_complete_and_never_interleaved() { + let writer = CapturedWriter::default(); + let bytes = Arc::clone(&writer.bytes); + let output = Output::new(writer).expect("failed to create output"); + let messages = 32; + let mut publishers = Vec::new(); + for sequence in 0..messages { + let output = output.clone(); + publishers.push(thread::spawn(move || { + output.send(&serde_json::json!({ + "sequence": sequence, + "text": "\u{2603}".repeat(32), + })); + })); + } + for publisher in publishers { + publisher.join().expect("publisher panicked"); + } + + let expected_payload_len: usize = (0..messages) + .map(|sequence| { + serde_json::to_vec(&serde_json::json!({ + "sequence": sequence, + "text": "\u{2603}".repeat(32), + })) + .unwrap() + .len() + }) + .sum(); + let minimum_len = + expected_payload_len + messages * (HEADER_PREFIX.len() + HEADER_SUFFIX.len() + 1); + wait_for_len(&bytes, minimum_len); + output.close(); + + let bytes = bytes.lock().expect("captured output lock poisoned"); + let mut sequences = frames(&bytes) + .into_iter() + .map(|payload| { + serde_json::from_slice::(payload) + .expect("interleaved or invalid JSON payload")["sequence"] + .as_u64() + .expect("sequence missing") + }) + .collect::>(); + sequences.sort_unstable(); + assert_eq!(sequences, (0..messages as u64).collect::>()); + } + + #[test] + fn repeated_initialization_after_shutdown_returns_broken_pipe() { + let mut global = GlobalState { + output: None, + shutdown: false, + initialization_error: None, + }; + initialize_output_state(&mut global, || Output::new(CapturedWriter::default())) + .expect("initial initialization failed"); + + let output = shutdown_output_state(&mut global).expect("initialized output missing"); + output.close(); + let error = initialize_output_state(&mut global, || { + panic!("shutdown output must not be reinitialized") + }) + .expect_err("initialization after shutdown succeeded"); + + assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); + assert!(global.initialization_error.is_none()); + assert!(output.error().is_none()); + } +} diff --git a/crates/pet-jsonrpc/src/server.rs b/crates/pet-jsonrpc/src/server.rs index 72cde98b..a29fa5e0 100644 --- a/crates/pet-jsonrpc/src/server.rs +++ b/crates/pet-jsonrpc/src/server.rs @@ -5,8 +5,10 @@ use crate::{send_error, RequestId}; use serde_json::{self, Value}; use std::{ collections::HashMap, - io::{self, Read}, - sync::Arc, + io::{self, BufRead, BufReader, Read}, + sync::{mpsc, Arc}, + thread, + time::Duration, }; type RequestHandler = Arc, RequestId, Value)>; @@ -109,47 +111,136 @@ impl HandlersKeyedByMethodName { } } -/// Starts the jsonrpc server that listens for requests on stdin. -/// This function will block forever. -pub fn start_server(handlers: &HandlersKeyedByMethodName) -> ! { - let mut stdin = io::stdin(); +const INPUT_POLL_INTERVAL: Duration = Duration::from_millis(20); +const MAX_HEADER_BYTES: usize = 8 * 1024; +const MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024; + +/// Runs the standalone process transport until EOF or a fatal I/O error. +/// Pending output is discarded at shutdown; callers must finish subprocess cleanup +/// and exit the process rather than join workers blocked in external I/O. +pub fn start_server(handlers: &HandlersKeyedByMethodName) -> io::Result<()> { + crate::initialize_output()?; + let result = (|| { + let (sender, receiver) = mpsc::sync_channel(1); + thread::Builder::new() + .name("pet-jsonrpc-input".to_string()) + .spawn(move || read_input(BufReader::new(io::stdin()), sender))?; + dispatch_input(handlers, &receiver, crate::output_error) + })(); + close_transport(result, crate::shutdown_output, crate::output_error) +} + +fn close_transport( + result: io::Result<()>, + shutdown_output: impl FnOnce(), + output_error: impl FnOnce() -> Option, +) -> io::Result<()> { + shutdown_output(); + result.and(output_error().map_or(Ok(()), Err)) +} + +fn read_input(mut reader: impl BufRead, sender: mpsc::SyncSender>>>) { loop { - let mut input = String::new(); - match stdin.read_line(&mut input) { - Ok(_) => { - let mut empty_line = String::new(); - match get_content_length(&input) { - Ok(content_length) => { - let _ = stdin.read_line(&mut empty_line); - let mut buffer = vec![0; content_length]; - - match stdin.read_exact(&mut buffer) { - Ok(_) => { - let request = - String::from_utf8_lossy(&buffer[..content_length]).to_string(); - if let Err(err) = handle_payload(handlers, &request) { - eprint!("Failed to parse LINE: {request}, {err:?}") - } - continue; - } - Err(err) => { - eprint!("Failed to read exactly {content_length} bytes, {err:?}") - } - } - } - Err(err) => eprint!("Failed to get content length from {input}, {err:?}"), - }; + let frame = read_payload(&mut reader); + let terminal = !matches!(&frame, Ok(Some(_))); + if sender.send(frame).is_err() || terminal { + return; + } + } +} + +fn dispatch_input( + handlers: &HandlersKeyedByMethodName, + receiver: &mpsc::Receiver>>>, + output_error: impl Fn() -> Option, +) -> io::Result<()> { + loop { + if let Some(error) = output_error() { + return Err(error); + } + match receiver.recv_timeout(INPUT_POLL_INTERVAL) { + Ok(Ok(Some(payload))) => { + if handle_payload(handlers, &payload).is_err() { + (handlers.send_error)(None, -32700, "Invalid JSONRPC JSON payload".to_string()); + } + } + Ok(Ok(None)) => return Ok(()), + Ok(Err(error)) => return Err(error), + Err(mpsc::RecvTimeoutError::Timeout) => {} + Err(mpsc::RecvTimeoutError::Disconnected) => { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "JSONRPC input reader stopped without a terminal result", + )); } - Err(error) => eprint!("Error in reading a line from stdin: {error}"), } } } +fn read_payload(reader: &mut impl BufRead) -> io::Result>> { + let mut header = String::new(); + let count = (&mut *reader) + .take(MAX_HEADER_BYTES as u64 + 1) + .read_line(&mut header)?; + if count == 0 { + return Ok(None); + } + if count > MAX_HEADER_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSONRPC header exceeds limit", + )); + } + if !header.ends_with('\n') { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "Truncated JSONRPC header", + )); + } + let length = get_content_length(&header) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + if length > MAX_PAYLOAD_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSONRPC payload exceeds limit", + )); + } + let mut separator = String::new(); + (&mut *reader) + .take((MAX_HEADER_BYTES - count) as u64 + 1) + .read_line(&mut separator)?; + if count + separator.len() > MAX_HEADER_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSONRPC header exceeds limit", + )); + } + if !separator.ends_with('\n') { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "Truncated JSONRPC header separator", + )); + } + if separator != "\r\n" && separator != "\n" { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Expected a blank JSONRPC header separator", + )); + } + let mut payload = Vec::new(); + payload + .try_reserve_exact(length) + .map_err(io::Error::other)?; + payload.resize(length, 0); + reader.read_exact(&mut payload)?; + Ok(Some(payload)) +} + fn handle_payload( handlers: &HandlersKeyedByMethodName, - payload: &str, + payload: impl AsRef<[u8]>, ) -> Result<(), serde_json::Error> { - let request = serde_json::from_str(payload)?; + let request = serde_json::from_slice(payload.as_ref())?; handlers.handle_request(request); Ok(()) } @@ -181,6 +272,36 @@ mod tests { use serde_json::json; use std::sync::Mutex; + #[test] + fn output_failure_between_eof_and_close_is_not_reported_as_success() { + let closed = std::cell::Cell::new(false); + let result = close_transport( + Ok(()), + || closed.set(true), + || { + assert!( + closed.get(), + "output must be closed before its final error is read" + ); + Some(io::Error::from(io::ErrorKind::BrokenPipe)) + }, + ); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::BrokenPipe); + } + + #[test] + fn transport_error_remains_primary_and_still_closes_output() { + let closed = std::cell::Cell::new(false); + let result = close_transport( + Err(io::Error::from(io::ErrorKind::InvalidData)), + || closed.set(true), + || Some(io::Error::from(io::ErrorKind::BrokenPipe)), + ); + assert!(closed.get()); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidData); + assert!(close_transport(Ok(()), || {}, || None).is_ok()); + } + #[derive(Default)] struct TestContext { request: Mutex>, @@ -444,4 +565,125 @@ mod tests { )] ); } + #[test] + fn input_distinguishes_clean_eof_from_truncated_frames() { + assert!(read_payload(&mut io::Cursor::new(b"")).unwrap().is_none()); + for bytes in [ + b"Content-Length: 2".as_slice(), + b"Content-Length: 2\r\n".as_slice(), + b"Content-Length: 2\r\n\r".as_slice(), + b"Content-Length: 2\r\n\r\n{".as_slice(), + ] { + let error = read_payload(&mut io::Cursor::new(bytes)).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof); + } + } + + #[test] + fn input_preserves_fragmented_consecutive_payload_bytes() { + let first = "{\"text\":\"\u{03c0}\"}".as_bytes(); + let second = b"{}"; + let mut bytes = format!("Content-Length: {}\r\n\r\n", first.len()).into_bytes(); + bytes.extend_from_slice(first); + bytes.extend_from_slice(b"Content-Length: 2\n\n"); + bytes.extend_from_slice(second); + let mut reader = BufReader::with_capacity(1, io::Cursor::new(bytes)); + assert_eq!(read_payload(&mut reader).unwrap().unwrap(), first); + assert_eq!(read_payload(&mut reader).unwrap().unwrap(), second); + assert!(read_payload(&mut reader).unwrap().is_none()); + } + + #[test] + fn input_rejects_oversized_headers_and_payloads_before_body_reads() { + let bytes = vec![b'x'; MAX_HEADER_BYTES + 1]; + assert_eq!( + read_payload(&mut io::Cursor::new(bytes)) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + let bytes = format!("Content-Length: {}\r\n\r\n", MAX_PAYLOAD_BYTES + 1); + assert_eq!( + read_payload(&mut io::Cursor::new(bytes)) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } + + #[test] + fn input_reader_reports_terminal_error_once() { + struct FailedReader; + impl Read for FailedReader { + fn read(&mut self, _: &mut [u8]) -> io::Result { + Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "injected read failure", + )) + } + } + impl BufRead for FailedReader { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "injected read failure", + )) + } + fn consume(&mut self, _: usize) {} + } + let (sender, receiver) = mpsc::sync_channel(1); + read_input(FailedReader, sender); + assert_eq!( + receiver.recv().unwrap().unwrap_err().kind(), + io::ErrorKind::PermissionDenied + ); + assert!(matches!( + receiver.try_recv(), + Err(mpsc::TryRecvError::Disconnected) + )); + let (sender, receiver) = mpsc::sync_channel(1); + read_input(io::Cursor::new(b""), sender); + assert!(receiver.recv().unwrap().unwrap().is_none()); + assert!(matches!( + receiver.try_recv(), + Err(mpsc::TryRecvError::Disconnected) + )); + } + + #[test] + fn dispatch_recovers_after_malformed_json_and_stops_on_eof() { + let context = Arc::new(TestContext::default()); + let mut handlers = create_handlers_with_recorded_errors(context.clone()); + handlers.add_request_handler("info", |context, id, params| { + *context.request.lock().unwrap() = Some((id, params)); + }); + let (sender, receiver) = mpsc::sync_channel(3); + sender.send(Ok(Some(b"{invalid".to_vec()))).unwrap(); + sender + .send(Ok(Some( + br#"{"jsonrpc":"2.0","id":"next","method":"info","params":{}}"#.to_vec(), + ))) + .unwrap(); + sender.send(Ok(None)).unwrap(); + dispatch_input(&handlers, &receiver, || None).unwrap(); + assert_eq!(context.errors.lock().unwrap()[0].1, -32700); + assert_eq!( + context.request.lock().unwrap().as_ref().unwrap().0, + RequestId::String("next".to_string()) + ); + } + + #[test] + fn dispatch_surfaces_output_failure_without_waiting_for_stdin() { + let handlers = create_handlers_with_recorded_errors(Arc::new(TestContext::default())); + let (_sender, receiver) = mpsc::sync_channel(1); + let error = dispatch_input(&handlers, &receiver, || { + Some(io::Error::new( + io::ErrorKind::BrokenPipe, + "injected output failure", + )) + }) + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); + } } diff --git a/crates/pet-poetry/src/environment_locations_spawn.rs b/crates/pet-poetry/src/environment_locations_spawn.rs index 2cc06cf5..05eafa43 100644 --- a/crates/pet-poetry/src/environment_locations_spawn.rs +++ b/crates/pet-poetry/src/environment_locations_spawn.rs @@ -159,6 +159,13 @@ fn run_poetry( ); None } + Err(ProcessError::Cancelled) => { + trace!( + "Cancelled Poetry probe during process shutdown: {:?}", + executable + ); + None + } Err(error) => { error!( "Failed to execute Poetry {:?} using {:?} in {:?}: {}", @@ -227,6 +234,7 @@ mod tests { Err(ProcessError::Spawn(io::Error::from( io::ErrorKind::NotFound, ))), + Err(ProcessError::Cancelled), Err(ProcessError::Io(io::Error::from(io::ErrorKind::BrokenPipe))), Err(ProcessError::Timeout(Duration::from_secs(15))), Err(ProcessError::OutputLimit(4 * 1024 * 1024)), diff --git a/crates/pet-python-utils/src/env.rs b/crates/pet-python-utils/src/env.rs index 1008fdda..61da3efb 100644 --- a/crates/pet-python-utils/src/env.rs +++ b/crates/pet-python-utils/src/env.rs @@ -121,6 +121,13 @@ fn get_interpreter_details_with_runner( ); match result { Ok(output) => parse_interpreter_result(executable, &output, start), + Err(ProcessError::Cancelled) => { + trace!( + "Cancelled interpreter probe during process shutdown: {:?}", + executable + ); + None + } Err(ProcessError::Timeout(timeout)) => { warn!("Timed out after {:?} resolving Python via spawn for {:?}; terminated direct child.", timeout, executable); None @@ -321,6 +328,21 @@ exit {exit_code} mod parser_tests { use super::*; + #[test] + fn cancelled_interpreter_has_no_result() { + let mut called = false; + let result = get_interpreter_details_with_runner( + Path::new("python"), + Duration::from_secs(5), + |_, _| { + called = true; + Err(ProcessError::Cancelled) + }, + ); + assert!(called); + assert!(result.is_none()); + } + #[test] fn preserves_non_utf8_preamble_and_unicode_json() { let mut bytes = vec![0xff, 0xfe, b'\n']; diff --git a/crates/pet-python-utils/src/process.rs b/crates/pet-python-utils/src/process.rs index 4857f5d6..3244da24 100644 --- a/crates/pet-python-utils/src/process.rs +++ b/crates/pet-python-utils/src/process.rs @@ -5,6 +5,8 @@ mod tree; use tree::ProcessTree; +mod supervisor; +use supervisor::{Admission, Supervisor}; use std::{ fmt, @@ -19,12 +21,14 @@ const CAPTURE_LIMIT: usize = 4 * 1024 * 1024; pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(15); const POLL_INTERVAL: Duration = Duration::from_millis(5); const CLEANUP_TIMEOUT: Duration = Duration::from_secs(2); +static SUPERVISOR: Supervisor = Supervisor::new(); #[derive(Debug)] pub enum ProcessError { Spawn(io::Error), Ownership(io::Error), Io(io::Error), + Cancelled, Timeout(Duration), OutputLimit(usize), IncompleteOutput(ExitStatus), @@ -52,6 +56,7 @@ impl fmt::Display for ProcessError { Self::Spawn(error) => write!(f, "failed to spawn subprocess: {error}"), Self::Ownership(error) => write!(f, "subprocess ownership failed: {error}"), Self::Io(error) => write!(f, "subprocess I/O failed: {error}"), + Self::Cancelled => write!(f, "subprocess cancelled for process shutdown"), Self::Timeout(timeout) => write!(f, "subprocess timed out after {timeout:?}"), Self::OutputLimit(limit) => { write!(f, "subprocess output exceeded {limit} captured bytes") @@ -72,7 +77,10 @@ impl std::error::Error for ProcessError { match self { Self::Spawn(error) | Self::Ownership(error) | Self::Io(error) => Some(error), Self::Cleanup { primary, .. } => Some(primary.as_ref()), - Self::Timeout(_) | Self::OutputLimit(_) | Self::IncompleteOutput(_) => None, + Self::Cancelled + | Self::Timeout(_) + | Self::OutputLimit(_) + | Self::IncompleteOutput(_) => None, } } } @@ -95,23 +103,66 @@ impl From for ProcessError { /// depend on additional creation flags. Arguments, cwd, and environment are kept. /// Retained bytes are bounded; an extra byte detects overflow and is not retained. pub fn output(command: &mut Command, timeout: Duration) -> Result { - output_with_limit(command, timeout, CAPTURE_LIMIT) + output_with_supervisor(command, timeout, CAPTURE_LIMIT, &SUPERVISOR) } +/// Permanently closes discovery-probe admission for this process, cancels all +/// admitted probes, and waits up to `timeout` for their ownership cleanup. +/// +/// A successful return guarantees that every probe accepted before closure has +/// completed cleanup successfully. The first cleanup failure is retained and +/// returned after all admitted probes finish, including from repeated calls. A +/// timeout leaves admission closed and cancellation active; callers may invoke +/// this again to continue waiting. Synchronous OS process creation and OS I/O +/// calls cannot be interrupted by this API, so shutdown can time out while an +/// admitted thread is inside one of those calls. +pub fn shutdown(timeout: Duration) -> io::Result<()> { + SUPERVISOR.shutdown(timeout) +} + +/// Returns whether process-wide probe shutdown has begun. +pub fn is_shutting_down() -> bool { + SUPERVISOR.is_shutting_down() +} + +#[cfg(test)] fn output_with_limit( command: &mut Command, timeout: Duration, limit: usize, ) -> Result { - output_with_setup(command, timeout, limit, |tree, child| tree.attach(child)) + output_with_supervisor(command, timeout, limit, &SUPERVISOR) } +fn output_with_supervisor( + command: &mut Command, + timeout: Duration, + limit: usize, + supervisor: &Supervisor, +) -> Result { + output_with_supervisor_setup(command, timeout, limit, supervisor, |tree, child| { + tree.attach(child) + }) +} + +#[cfg(test)] fn output_with_setup( command: &mut Command, timeout: Duration, limit: usize, setup: impl FnOnce(&mut ProcessTree, &Child) -> io::Result<()>, ) -> Result { + output_with_supervisor_setup(command, timeout, limit, &SUPERVISOR, setup) +} + +fn output_with_supervisor_setup( + command: &mut Command, + timeout: Duration, + limit: usize, + supervisor: &Supervisor, + setup: impl FnOnce(&mut ProcessTree, &Child) -> io::Result<()>, +) -> Result { + let admission = supervisor.admit()?; let mut tree = ProcessTree::prepare(command).map_err(ProcessError::Ownership)?; let mut child = command .stdin(Stdio::null()) @@ -139,14 +190,17 @@ fn output_with_setup( started, timeout, limit, + &admission, ) })(); drop(stdout); drop(stderr); let cleanup = cleanup(child, &mut tree); + let shutdown_won = admission.finish(&cleanup); match result { - Ok(output) => cleanup.map(|()| output).map_err(ProcessError::Ownership), Err(primary) => Err(primary.with_cleanup(cleanup)), + Ok(_) if shutdown_won => Err(ProcessError::Cancelled.with_cleanup(cleanup)), + Ok(output) => cleanup.map(|()| output).map_err(ProcessError::Ownership), } } @@ -157,6 +211,7 @@ fn capture( started: Instant, timeout: Duration, limit: usize, + admission: &Admission<'_>, ) -> Result { let mut captured_stdout = Vec::new(); let mut captured_stderr = Vec::new(); @@ -165,6 +220,9 @@ fn capture( let mut status = None; let mut scratch = [0; 8192]; loop { + if admission.is_cancelled() { + return Err(ProcessError::Cancelled); + } if status.is_none() { status = poll()?; } @@ -438,13 +496,20 @@ mod tests { } std::process::exit(if mode == "nonzero" { 23 } else { 0 }); } - "hang" => { + "hang" | "noisy-hang" => { fs::write( std::env::var_os("PET_CAPTURE_PID").unwrap(), std::process::id().to_string(), ) .unwrap(); - thread::sleep(Duration::from_secs(30)); + let started = Instant::now(); + while started.elapsed() < Duration::from_secs(30) { + if mode == "noisy-hang" { + io::stdout().write_all(&[0xfe; 1024]).unwrap(); + io::stdout().flush().unwrap(); + } + thread::sleep(Duration::from_millis(10)); + } } "inherit" | "tree-hang" | "tree-limit" | "tree-nonzero" | "tree-no-pipes" | "escape" => { @@ -879,6 +944,40 @@ mod tests { assert_child_reaped(pid); } + #[test] + fn setup_failure_remains_primary_when_shutdown_wins_completion() { + let _guard = FIXTURES.lock().unwrap(); + let supervisor = Supervisor::new(); + let mut pid = 0; + let error = output_with_supervisor_setup( + &mut fixture_command("quiet"), + Duration::from_secs(2), + 100, + &supervisor, + |tree, child| { + #[cfg(unix)] + tree.attach(child)?; + #[cfg(windows)] + let _ = tree; + pid = child.id(); + let shutdown = supervisor.shutdown(Duration::ZERO).unwrap_err(); + assert_eq!(shutdown.kind(), io::ErrorKind::TimedOut); + Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "injected ownership failure after shutdown", + )) + }, + ) + .unwrap_err(); + + assert!(matches!(error, ProcessError::Ownership(_)), "{error}"); + assert!(error + .to_string() + .contains("injected ownership failure after shutdown")); + assert_child_reaped(pid); + supervisor.shutdown(Duration::ZERO).unwrap(); + } + struct BytesPipe(io::Cursor>); impl Pipe for BytesPipe { fn read_available(&mut self, buffer: &mut [u8]) -> io::Result { @@ -973,6 +1072,8 @@ mod tests { bytes: BytesPipe(io::Cursor::new(vec![0xff, 1, 2])), }; let mut stderr = BytesPipe(io::Cursor::new(vec![])); + let supervisor = Supervisor::new(); + let admission = supervisor.admit().unwrap(); let result = capture( &mut || tree.poll(&mut child).map_err(ProcessError::Ownership), &mut stdout, @@ -980,12 +1081,239 @@ mod tests { Instant::now(), Duration::from_secs(1), 3, + &admission, ) .unwrap(); assert_eq!(result.stdout, vec![0xff, 1, 2]); assert!(result.stderr.is_empty()); assert!(result.status.success()); cleanup(child, &mut tree).unwrap(); + assert!(!admission.finish(&Ok(()))); + } + + #[test] + fn closed_supervisor_refuses_probe_before_spawn() { + let supervisor = Supervisor::new(); + supervisor.shutdown(Duration::ZERO).unwrap(); + let directory = tempfile::tempdir().unwrap(); + let pid = directory.path().join("pid"); + let mut command = fixture_command("hang"); + command.env("PET_CAPTURE_PID", &pid); + + let error = output_with_supervisor( + &mut command, + Duration::from_secs(10), + CAPTURE_LIMIT, + &supervisor, + ) + .unwrap_err(); + + assert!(matches!(error, ProcessError::Cancelled), "{error}"); + assert!(!pid.exists(), "cancelled probe was spawned"); + } + + #[test] + fn shutdown_cancels_running_quiet_and_noisy_probes() { + let _guard = FIXTURES.lock().unwrap(); + for mode in ["hang", "noisy-hang"] { + let supervisor = Supervisor::new(); + let directory = tempfile::tempdir().unwrap(); + let pid_file = directory.path().join("pid"); + let mut command = fixture_command(mode); + command.env("PET_CAPTURE_PID", &pid_file); + + thread::scope(|scope| { + let probe = scope.spawn(|| { + output_with_supervisor( + &mut command, + Duration::from_secs(15), + CAPTURE_LIMIT, + &supervisor, + ) + }); + let pid = wait_for_pid(&pid_file, "probe did not reach fixture"); + let started = Instant::now(); + supervisor.shutdown(Duration::from_secs(5)).unwrap(); + assert!( + started.elapsed() < Duration::from_secs(2), + "shutdown did not cancel {mode} probe promptly" + ); + let error = probe.join().unwrap().unwrap_err(); + assert!(matches!(error, ProcessError::Cancelled), "{error}"); + assert_child_reaped(pid); + }); + } + } + + #[test] + fn shutdown_releases_descendants_and_preserves_unrelated_processes() { + let _guard = FIXTURES.lock().unwrap(); + let unrelated_directory = tempfile::tempdir().unwrap(); + let mut unrelated = FixtureChild( + fixture_command("descendant") + .env("PET_CAPTURE_DIRECTORY", unrelated_directory.path()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(), + ); + wait_for_path( + &unrelated_directory.path().join("ready"), + "unrelated fixture did not start", + ); + + let supervisor = Supervisor::new(); + let owned = ReleaseDescendant(tempfile::tempdir().unwrap()); + let root_pid = owned.0.path().join("root-pid"); + let mut command = fixture_command("tree-hang"); + command + .env("PET_CAPTURE_DIRECTORY", owned.0.path()) + .env("PET_CAPTURE_PID", &root_pid); + thread::scope(|scope| { + let probe = scope.spawn(|| { + output_with_supervisor( + &mut command, + Duration::from_secs(15), + CAPTURE_LIMIT, + &supervisor, + ) + }); + wait_for_path( + &owned.0.path().join("ready"), + "owned descendant did not start", + ); + supervisor.shutdown(Duration::from_secs(5)).unwrap(); + let error = probe.join().unwrap().unwrap_err(); + assert!(matches!(error, ProcessError::Cancelled), "{error}"); + }); + + owned.assert_terminated(); + assert_child_reaped(fs::read_to_string(root_pid).unwrap().parse().unwrap()); + assert!(unrelated.0.try_wait().unwrap().is_none()); + let lease = fs::File::open(unrelated_directory.path().join("lease")).unwrap(); + assert!(matches!( + lease.try_lock(), + Err(fs::TryLockError::WouldBlock) + )); + } + + #[test] + fn shutdown_closes_admission_and_waits_for_accepted_cleanup() { + use std::sync::mpsc; + + let supervisor = Supervisor::new(); + let admission = supervisor.admit().unwrap(); + let (completed_tx, completed_rx) = mpsc::channel(); + thread::scope(|scope| { + scope.spawn(|| { + completed_tx + .send(supervisor.shutdown(Duration::from_secs(2))) + .unwrap(); + }); + let started = Instant::now(); + while !supervisor.is_shutting_down() { + assert!( + started.elapsed() < Duration::from_secs(1), + "shutdown did not close admission" + ); + thread::yield_now(); + } + assert!(matches!(supervisor.admit(), Err(ProcessError::Cancelled))); + assert!( + completed_rx.try_recv().is_err(), + "shutdown completed before accepted cleanup" + ); + assert!(admission.finish(&Ok(()))); + completed_rx + .recv_timeout(Duration::from_secs(1)) + .unwrap() + .unwrap(); + }); + } + + #[test] + fn shutdown_timeout_keeps_admission_closed_and_can_be_retried() { + let supervisor = Supervisor::new(); + let admission = supervisor.admit().unwrap(); + let error = supervisor.shutdown(Duration::ZERO).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); + assert!(matches!(supervisor.admit(), Err(ProcessError::Cancelled))); + assert!(admission.finish(&Ok(()))); + supervisor.shutdown(Duration::ZERO).unwrap(); + } + + #[test] + fn shutdown_preserves_first_cleanup_failure() { + use std::sync::mpsc; + + let supervisor = Supervisor::new(); + let first = supervisor.admit().unwrap(); + let second = supervisor.admit().unwrap(); + let (completed_tx, completed_rx) = mpsc::channel(); + thread::scope(|scope| { + scope.spawn(|| { + completed_tx + .send(supervisor.shutdown(Duration::from_secs(2))) + .unwrap(); + }); + let started = Instant::now(); + while !supervisor.is_shutting_down() { + assert!( + started.elapsed() < Duration::from_secs(1), + "shutdown did not close admission" + ); + thread::yield_now(); + } + assert!(first.finish(&Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "injected first cleanup failure", + )))); + assert!(second.finish(&Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "injected later cleanup failure", + )))); + let error = completed_rx + .recv_timeout(Duration::from_secs(1)) + .unwrap() + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied); + assert!(error.to_string().contains("probe cleanup failed")); + assert!(error.to_string().contains("injected first cleanup failure")); + }); + + let repeated = supervisor.shutdown(Duration::ZERO).unwrap_err(); + assert_eq!(repeated.kind(), io::ErrorKind::PermissionDenied); + assert!(repeated + .to_string() + .contains("injected first cleanup failure")); + + let successful = Supervisor::new(); + let admission = successful.admit().unwrap(); + assert!(!admission.finish(&Ok(()))); + successful.shutdown(Duration::ZERO).unwrap(); + } + + fn wait_for_path(path: &Path, failure: &str) { + let started = Instant::now(); + while !path.exists() { + assert!(started.elapsed() < Duration::from_secs(10), "{failure}"); + thread::sleep(POLL_INTERVAL); + } + } + + fn wait_for_pid(path: &Path, failure: &str) -> u32 { + let started = Instant::now(); + loop { + if let Ok(pid) = fs::read_to_string(path).and_then(|value| { + value + .parse() + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + }) { + return pid; + } + assert!(started.elapsed() < Duration::from_secs(10), "{failure}"); + thread::sleep(POLL_INTERVAL); + } } struct FailingPipe; @@ -1028,6 +1356,13 @@ mod tests { ProcessError::OutputLimit(1).with_cleanup(Ok(())), ProcessError::OutputLimit(1) )); + let cancellation = ProcessError::Cancelled.with_cleanup(Err(io::Error::new( + io::ErrorKind::TimedOut, + "injected cancellation cleanup failure", + ))); + assert!( + matches!(&cancellation, ProcessError::Cleanup { primary, .. } if matches!(primary.as_ref(), ProcessError::Cancelled)) + ); assert!(error.to_string().contains("timed out")); assert!(error.to_string().contains("kill failure")); let source = std::error::Error::source(&error) diff --git a/crates/pet-python-utils/src/process/supervisor.rs b/crates/pet-python-utils/src/process/supervisor.rs new file mode 100644 index 00000000..536e5acd --- /dev/null +++ b/crates/pet-python-utils/src/process/supervisor.rs @@ -0,0 +1,163 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::ProcessError; +use std::{ + io, + sync::{ + atomic::{AtomicBool, Ordering}, + Condvar, Mutex, + }, + time::{Duration, Instant}, +}; + +struct State { + closed: bool, + active: usize, + cleanup_failure: Option, +} + +struct CleanupFailure { + kind: io::ErrorKind, + message: String, +} + +impl CleanupFailure { + fn from_error(error: &io::Error) -> Self { + Self { + kind: error.kind(), + message: format!("probe cleanup failed: {error}"), + } + } + + fn to_error(&self) -> io::Error { + io::Error::new(self.kind, self.message.clone()) + } +} + +pub(super) struct Supervisor { + cancelled: AtomicBool, + state: Mutex, + completed: Condvar, +} + +impl Supervisor { + pub(super) const fn new() -> Self { + Self { + cancelled: AtomicBool::new(false), + state: Mutex::new(State { + closed: false, + active: 0, + cleanup_failure: None, + }), + completed: Condvar::new(), + } + } + + pub(super) fn admit(&self) -> Result, ProcessError> { + let mut state = self + .state + .lock() + .expect("probe supervisor state lock poisoned"); + if state.closed { + return Err(ProcessError::Cancelled); + } + state.active += 1; + Ok(Admission { + supervisor: self, + finished: false, + }) + } + + pub(super) fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::Acquire) + } + + pub(super) fn is_shutting_down(&self) -> bool { + self.is_cancelled() + } + + pub(super) fn shutdown(&self, timeout: Duration) -> io::Result<()> { + let started = Instant::now(); + let mut state = self + .state + .lock() + .expect("probe supervisor state lock poisoned"); + state.closed = true; + self.cancelled.store(true, Ordering::Release); + + while state.active != 0 { + let remaining = timeout.saturating_sub(started.elapsed()); + if remaining.is_zero() { + return Err(shutdown_timeout(state.active, timeout)); + } + let (next, wait) = self + .completed + .wait_timeout(state, remaining) + .expect("probe supervisor state lock poisoned while waiting"); + state = next; + if wait.timed_out() && state.active != 0 { + return Err(shutdown_timeout(state.active, timeout)); + } + } + match &state.cleanup_failure { + Some(failure) => Err(failure.to_error()), + None => Ok(()), + } + } +} + +fn shutdown_timeout(active: usize, timeout: Duration) -> io::Error { + io::Error::new( + io::ErrorKind::TimedOut, + format!("{active} probe(s) did not finish cleanup within {timeout:?}"), + ) +} + +pub(super) struct Admission<'a> { + supervisor: &'a Supervisor, + finished: bool, +} + +impl Admission<'_> { + pub(super) fn is_cancelled(&self) -> bool { + self.supervisor.is_cancelled() + } + + /// Deregisters after cleanup and reports whether shutdown won the + /// completion race while this probe was still admitted. + pub(super) fn finish(mut self, cleanup: &io::Result<()>) -> bool { + let mut state = self + .supervisor + .state + .lock() + .expect("probe supervisor state lock poisoned"); + let cancelled = state.closed; + if state.cleanup_failure.is_none() { + state.cleanup_failure = cleanup.as_ref().err().map(CleanupFailure::from_error); + } + state.active -= 1; + self.finished = true; + if state.active == 0 { + self.supervisor.completed.notify_all(); + } + cancelled + } +} + +impl Drop for Admission<'_> { + fn drop(&mut self) { + if self.finished { + return; + } + let mut state = self + .supervisor + .state + .lock() + .expect("probe supervisor state lock poisoned"); + state.active -= 1; + if state.active == 0 { + self.supervisor.completed.notify_all(); + } + } +} diff --git a/crates/pet/src/jsonrpc.rs b/crates/pet/src/jsonrpc.rs index 3ffaf637..6dc5f2bd 100644 --- a/crates/pet/src/jsonrpc.rs +++ b/crates/pet/src/jsonrpc.rs @@ -559,7 +559,7 @@ const MISSING_ENVS_COMPLETED: u64 = u64::MAX - 1; static MISSING_ENVS_REPORTING_STATE: AtomicU64 = AtomicU64::new(MISSING_ENVS_AVAILABLE); static NEXT_REFRESH_ID: AtomicU64 = AtomicU64::new(1); -pub fn start_jsonrpc_server() { +pub fn start_jsonrpc_server() -> std::io::Result<()> { // Initialize tracing for performance profiling (controlled by RUST_LOG env var) // Note: This includes log compatibility, so we don't call jsonrpc::initialize_logger initialize_tracing(false); @@ -587,7 +587,16 @@ pub fn start_jsonrpc_server() { handlers.add_request_handler("find", handle_find); handlers.add_request_handler("condaInfo", handle_conda_telemetry); handlers.add_request_handler("clear", handle_clear_cache); - start_server(&handlers) + let transport_result = start_server(&handlers); + let cleanup_result = pet_python_utils::process::shutdown(Duration::from_secs(3)); + match (transport_result, cleanup_result) { + (Ok(()), cleanup) => cleanup, + (Err(error), Ok(())) => Err(error), + (Err(error), Err(cleanup)) => Err(std::io::Error::new( + error.kind(), + format!("{error}; subprocess shutdown failed: {cleanup}"), + )), + } } #[derive(Debug, Clone, Deserialize, Serialize)] diff --git a/crates/pet/src/main.rs b/crates/pet/src/main.rs index d0a494bc..36c9591e 100644 --- a/crates/pet/src/main.rs +++ b/crates/pet/src/main.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -use std::path::PathBuf; +use std::{path::PathBuf, process::ExitCode}; use clap::{Parser, Subcommand}; use jsonrpc::start_jsonrpc_server; @@ -95,7 +95,7 @@ enum Commands { Server, } -fn main() { +fn main() -> ExitCode { let cli = Cli::parse(); match cli.command.unwrap_or(Commands::Find { @@ -159,6 +159,12 @@ fn main() { cache_directory, json, } => resolve_report_stdio(executable, verbose, cache_directory, json), - Commands::Server => start_jsonrpc_server(), + Commands::Server => { + if let Err(error) = start_jsonrpc_server() { + eprintln!("PET server stopped with an error: {error}"); + return ExitCode::FAILURE; + } + } } + ExitCode::SUCCESS } diff --git a/crates/pet/tests/e2e_performance.rs b/crates/pet/tests/e2e_performance.rs index e1d653e5..6e5dc380 100644 --- a/crates/pet/tests/e2e_performance.rs +++ b/crates/pet/tests/e2e_performance.rs @@ -557,14 +557,25 @@ impl PetClient { } } +#[path = "process_utils.rs"] +mod process_utils; + impl Drop for PetClient { fn drop(&mut self) { + drop(std::mem::replace( + &mut self.stdin, + Box::new(std::io::sink()), + )); if let Some(process) = self.process.as_mut() { - let _ = process.kill(); - let _ = process.wait(); + if let Err(error) = process_utils::shutdown_fixture(process, Duration::from_secs(4)) { + eprintln!("Failed to stop performance fixture; detaching its reader: {error}"); + return; + } } if let Some(stderr_handle) = self.stderr_handle.take() { - let _ = stderr_handle.join(); + if let Err(error) = process_utils::join_reader(stderr_handle, Duration::from_secs(4)) { + eprintln!("Failed to finish performance fixture stderr reader: {error}"); + } } } } diff --git a/crates/pet/tests/fixtures/shutdown_probe.py b/crates/pet/tests/fixtures/shutdown_probe.py new file mode 100644 index 00000000..7e829048 --- /dev/null +++ b/crates/pet/tests/fixtures/shutdown_probe.py @@ -0,0 +1,61 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import errno +import os +from pathlib import Path +import subprocess +import sys +import time + + +def try_lock(stream): + try: + if os.name == "nt": + import msvcrt + stream.seek(0) + msvcrt.locking(stream.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + fcntl.flock(stream.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + return True + except OSError as error: + if error.errno not in (errno.EACCES, errno.EAGAIN, errno.EDEADLK): + raise + return False + + +def parent_test_is_alive(): + try: + stream = open(os.environ["PET_SHUTDOWN_CONTROL"], "r+b") + except FileNotFoundError: + return False + with stream: + return not try_lock(stream) + + +def wait_for_shutdown(): + deadline = time.monotonic() + 30 + while parent_test_is_alive() and time.monotonic() < deadline: + time.sleep(0.01) + + +if sys.argv[-1] == "child": + with open(os.environ["PET_SHUTDOWN_LEASE"], "w+b") as lease: + lease.write(b"x") + lease.flush() + if not try_lock(lease): + raise RuntimeError("fixture descendant could not acquire its lease") + Path(os.environ["PET_SHUTDOWN_READY"]).write_text("ready", encoding="ascii") + wait_for_shutdown() +else: + child = subprocess.Popen( + [sys.executable, "-S", __file__, "child"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + wait_for_shutdown() + finally: + child.terminate() + child.wait(timeout=2) diff --git a/crates/pet/tests/jsonrpc_client.rs b/crates/pet/tests/jsonrpc_client.rs index 31a7881e..be70aa9c 100644 --- a/crates/pet/tests/jsonrpc_client.rs +++ b/crates/pet/tests/jsonrpc_client.rs @@ -6,12 +6,16 @@ use serde::Deserialize; use serde_json::{json, Value}; use std::collections::HashMap; use std::io::{self, BufRead, BufReader, Read, Write}; -use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; +use std::process::{Child, ChildStdin, ChildStdout, Command, ExitStatus, Stdio}; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{mpsc, Arc, Mutex}; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; +#[path = "process_utils.rs"] +mod process_utils; +pub(crate) use process_utils::{join_reader, shutdown_fixture, wait_for_exit}; + static REQUEST_ID: AtomicU32 = AtomicU32::new(1); const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); @@ -81,19 +85,21 @@ impl Drop for ClientInner { fn drop(&mut self) { let _ = self.stdin.lock().unwrap().take(); + if let Err(error) = + shutdown_fixture(&mut self.child.lock().unwrap(), Duration::from_secs(4)) { - let mut child = self.child.lock().unwrap(); - if child.try_wait().ok().flatten().is_none() { - let _ = child.kill(); - let _ = child.wait(); - } + eprintln!("Failed to stop PET fixture; detaching its readers: {error}"); + return; } - if let Some(handle) = self.reader_handle.lock().unwrap().take() { - let _ = handle.join(); + if let Err(error) = join_reader(handle, Duration::from_secs(4)) { + eprintln!("Failed to finish PET stdout reader: {error}"); + } } if let Some(handle) = self.stderr_handle.lock().unwrap().take() { - let _ = handle.join(); + if let Err(error) = join_reader(handle, Duration::from_secs(4)) { + eprintln!("Failed to finish PET stderr reader: {error}"); + } } } } @@ -105,6 +111,12 @@ pub struct PetJsonRpcClient { impl PetJsonRpcClient { pub fn spawn() -> Result { + Self::spawn_with_environment(&[]) + } + + pub fn spawn_with_environment( + environment: &[(&str, &std::ffi::OsStr)], + ) -> Result { let mut cmd = Command::new(env!("CARGO_BIN_EXE_pet")); cmd.arg("server") .stdin(Stdio::piped()) @@ -114,12 +126,13 @@ impl PetJsonRpcClient { // configuration from leaking into the test environment, then // restore only the minimum required for the OS to function. .env_clear() - .env("PATH", ""); - // On Windows, SYSTEMROOT is required for basic OS functionality - // (crypto, networking, etc.). Only set it when present. + .env("PATH", "") + .envs(environment.iter().copied()); #[cfg(windows)] - if let Ok(val) = std::env::var("SYSTEMROOT") { - cmd.env("SYSTEMROOT", val); + for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { + if let Some(value) = std::env::var_os(name) { + cmd.env(name, value); + } } let mut process = cmd .spawn() @@ -153,6 +166,22 @@ impl PetJsonRpcClient { }) } + pub fn shutdown(&self, timeout: Duration) -> Result { + let started = Instant::now(); + self.inner.stdin.lock().unwrap().take(); + let status = wait_for_exit(&mut self.inner.child.lock().unwrap(), timeout) + .map_err(|error| format!("PET did not shut down normally: {error}"))?; + if let Some(reader) = self.inner.reader_handle.lock().unwrap().take() { + join_reader(reader, timeout.saturating_sub(started.elapsed())) + .map_err(|error| format!("PET stdout reader did not finish: {error}"))?; + } + if let Some(reader) = self.inner.stderr_handle.lock().unwrap().take() { + join_reader(reader, timeout.saturating_sub(started.elapsed())) + .map_err(|error| format!("PET stderr reader did not finish: {error}"))?; + } + Ok(status) + } + pub fn configure(&self, config: Value) -> Result<(), String> { self.send_request_value("configure", config, DEFAULT_REQUEST_TIMEOUT) .map(|_| ()) @@ -311,13 +340,15 @@ impl PetJsonRpcClient { let write_result = { let mut stdin_guard = self.inner.stdin.lock().unwrap(); - let stdin = stdin_guard + stdin_guard .as_mut() - .ok_or_else(|| "PET stdin is no longer available".to_string())?; - stdin - .write_all(wire_message.as_bytes()) - .and_then(|_| stdin.flush()) - .map_err(|e| format!("Failed to send {method} request: {e}")) + .ok_or_else(|| "PET stdin is no longer available".to_string()) + .and_then(|stdin| { + stdin + .write_all(wire_message.as_bytes()) + .and_then(|_| stdin.flush()) + .map_err(|error| format!("Failed to send {method} request: {error}")) + }) }; if let Err(err) = write_result { diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index 77e91fb2..cd7febaf 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -4,13 +4,13 @@ use pet_fs::path::norm_case; use serde_json::{json, Value}; use std::fs; -use std::io::{BufReader, Write}; +use std::io::{BufReader, Read, Write}; use std::path::{Path, PathBuf}; use std::process::{Child, Command, Stdio}; use std::sync::mpsc; use std::thread; use std::thread::JoinHandle; -use std::time::Duration; +use std::time::{Duration, Instant}; use tempfile::TempDir; mod jsonrpc_client; @@ -34,8 +34,10 @@ impl RawRpcClient { .env_clear() .env("PATH", ""); #[cfg(windows)] - if let Some(system_root) = std::env::var_os("SYSTEMROOT") { - command.env("SYSTEMROOT", system_root); + for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { + if let Some(value) = std::env::var_os(name) { + command.env(name, value); + } } let mut child = command.spawn().expect("raw fixture must spawn PET"); let stdout = child.stdout.take().expect("PET stdout must be piped"); @@ -82,17 +84,16 @@ impl RawRpcClient { impl Drop for RawRpcClient { fn drop(&mut self) { - // EOF shutdown is tracked separately; kill only this fixture's child before closing stdin. - if let Err(error) = self.child.kill() { + self.child.stdin.take(); + if let Err(error) = + jsonrpc_client::shutdown_fixture(&mut self.child, Duration::from_secs(4)) + { eprintln!("Failed to stop raw RPC fixture: {error}"); + return; } - if let Err(error) = self.child.wait() { - eprintln!("Failed to reap raw RPC fixture: {error}"); - } - self.child.stdin.take(); if let Some(reader) = self.reader.take() { - if reader.join().is_err() { - eprintln!("Raw RPC fixture reader panicked"); + if let Err(error) = jsonrpc_client::join_reader(reader, Duration::from_secs(4)) { + eprintln!("Failed to finish raw RPC fixture reader: {error}"); } } } @@ -503,3 +504,215 @@ fn concurrent_distinct_refresh_requests_run_separately() { "distinct refresh requests should emit separate performance events" ); } + +struct ShutdownFixture { + child: Child, +} + +impl ShutdownFixture { + fn spawn() -> Self { + let mut command = Command::new(env!("CARGO_BIN_EXE_pet")); + command + .arg("server") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .env_clear() + .env("PATH", ""); + #[cfg(windows)] + for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { + if let Some(value) = std::env::var_os(name) { + command.env(name, value); + } + } + Self { + child: command.spawn().expect("shutdown fixture must spawn PET"), + } + } + + fn send(&mut self, body: &[u8]) { + let stdin = self.child.stdin.as_mut().unwrap(); + write!(stdin, "Content-Length: {}\r\n\r\n", body.len()).unwrap(); + stdin.write_all(body).unwrap(); + stdin.flush().unwrap(); + } +} + +impl Drop for ShutdownFixture { + fn drop(&mut self) { + self.child.stdin.take(); + if let Err(error) = + jsonrpc_client::shutdown_fixture(&mut self.child, Duration::from_secs(4)) + { + eprintln!("Failed to stop shutdown fixture: {error}"); + } + } +} + +#[test] +fn stdin_eof_after_exchange_exits_cleanly_within_one_second() { + let client = PetJsonRpcClient::spawn().unwrap(); + client.info().unwrap(); + let started = Instant::now(); + let status = client.shutdown(Duration::from_secs(1)).unwrap(); + assert!(status.success(), "normal EOF shutdown failed: {status}"); + assert!(started.elapsed() < Duration::from_secs(1)); + assert!( + client.stderr_output().len() < 4096, + "EOF must not produce an error flood" + ); +} + +#[test] +fn truncated_input_exits_unsuccessfully_without_an_error_flood() { + for bytes in [ + b"Content-Length: 2".as_slice(), + b"Content-Length: 2\r\n\r\n{".as_slice(), + ] { + let mut fixture = ShutdownFixture::spawn(); + fixture + .child + .stdin + .as_mut() + .unwrap() + .write_all(bytes) + .unwrap(); + let started = Instant::now(); + fixture.child.stdin.take(); + let status = + jsonrpc_client::wait_for_exit(&mut fixture.child, Duration::from_secs(1)).unwrap(); + assert!(!status.success()); + assert!(started.elapsed() < Duration::from_secs(1)); + let mut stderr = Vec::new(); + fixture + .child + .stderr + .take() + .unwrap() + .read_to_end(&mut stderr) + .unwrap(); + assert!(!stderr.is_empty()); + assert!(stderr.len() < 4096, "truncated input must be reported once"); + } +} + +#[test] +fn closed_output_exits_without_waiting_for_stdin_eof() { + let mut fixture = ShutdownFixture::spawn(); + drop(fixture.child.stdout.take()); + let started = Instant::now(); + fixture.send(br#"{"jsonrpc":"2.0","id":1,"method":"info"}"#); + let status = jsonrpc_client::wait_for_exit(&mut fixture.child, Duration::from_secs(1)).unwrap(); + assert!( + !status.success(), + "broken output must produce a nonzero exit" + ); + assert!(started.elapsed() < Duration::from_secs(1)); + assert!( + fixture.child.stdin.is_some(), + "input remains open throughout this check" + ); +} + +#[test] +fn stdin_eof_exits_while_output_is_not_drained() { + let mut fixture = ShutdownFixture::spawn(); + let mut stdout = fixture.child.stdout.take().unwrap(); + let (sender, receiver) = mpsc::sync_channel(1); + let reader = thread::spawn(move || { + let mut byte = [0]; + let result = stdout.read_exact(&mut byte); + sender + .send((stdout, result, byte)) + .expect("fixture must wait for output to start"); + }); + let body = + serde_json::to_vec(&json!({"jsonrpc":"2.0","id":"x".repeat(128 * 1024),"method":"info"})) + .unwrap(); + fixture.send(&body); + let (_unread_output, result, first_byte) = + receiver.recv_timeout(Duration::from_secs(5)).unwrap(); + result.unwrap(); + assert_eq!(first_byte, [b'C']); + let started = Instant::now(); + fixture.child.stdin.take(); + let status = jsonrpc_client::wait_for_exit(&mut fixture.child, Duration::from_secs(1)).unwrap(); + assert!(status.success()); + assert!(started.elapsed() < Duration::from_secs(1)); + reader.join().unwrap(); +} + +#[cfg(feature = "ci")] +#[test] +fn stdin_eof_cancels_an_active_interpreter_and_its_descendant() { + let output = Command::new(if cfg!(windows) { "python" } else { "python3" }) + .args([ + "-S", + "-c", + "import sys; sys.stdout.buffer.write(sys.executable.encode('utf-8'))", + ]) + .output() + .expect("CI must provide Python for the active-probe fixture"); + assert!(output.status.success()); + let python = String::from_utf8(output.stdout).unwrap(); + let directory = tempfile::tempdir().unwrap(); + fs::write( + directory.path().join("sitecustomize.py"), + include_str!("fixtures/shutdown_probe.py"), + ) + .unwrap(); + let control_path = directory.path().join("control"); + let lease_path = directory.path().join("lease"); + let ready_path = directory.path().join("ready"); + let mut control = fs::OpenOptions::new() + .read(true) + .write(true) + .create_new(true) + .open(&control_path) + .unwrap(); + control.write_all(b"x").unwrap(); + control.try_lock().unwrap(); + let client = PetJsonRpcClient::spawn_with_environment(&[ + ("PYTHONPATH", directory.path().as_os_str()), + ("PET_SHUTDOWN_CONTROL", control_path.as_os_str()), + ("PET_SHUTDOWN_LEASE", lease_path.as_os_str()), + ("PET_SHUTDOWN_READY", ready_path.as_os_str()), + ]) + .unwrap(); + client.info().unwrap(); + let worker = client.clone(); + let request = thread::spawn(move || worker.resolve(&python)); + let started = Instant::now(); + while !ready_path.is_file() { + assert!( + started.elapsed() < Duration::from_secs(10), + "probe must establish its descendant lease; stderr: {}", + client.stderr_output() + ); + thread::sleep(Duration::from_millis(10)); + } + let lease = fs::OpenOptions::new() + .read(true) + .write(true) + .open(&lease_path) + .unwrap(); + assert!(matches!( + lease.try_lock(), + Err(fs::TryLockError::WouldBlock) + )); + let started = Instant::now(); + let status = client.shutdown(Duration::from_secs(4)).unwrap(); + assert!( + status.success(), + "active-probe shutdown failed: {status}; stderr: {}", + client.stderr_output() + ); + assert!(started.elapsed() < Duration::from_secs(4)); + lease + .try_lock() + .expect("shutdown must release the actual descendant's lease"); + assert!( + request.join().unwrap().is_err(), + "an active request must be cancelled, not reported as successful" + ); +} diff --git a/crates/pet/tests/process_utils.rs b/crates/pet/tests/process_utils.rs new file mode 100644 index 00000000..4292809d --- /dev/null +++ b/crates/pet/tests/process_utils.rs @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::{ + io, + process::{Child, ExitStatus}, + thread, + time::{Duration, Instant}, +}; + +pub fn wait_for_exit(child: &mut Child, timeout: Duration) -> io::Result { + let started = Instant::now(); + loop { + if let Some(status) = child.try_wait()? { + return Ok(status); + } + let remaining = timeout.saturating_sub(started.elapsed()); + if remaining.is_zero() { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("process {} did not exit within {timeout:?}", child.id()), + )); + } + thread::sleep(Duration::from_millis(10).min(remaining)); + } +} + +pub fn shutdown_fixture(child: &mut Child, timeout: Duration) -> io::Result { + match wait_for_exit(child, timeout) { + Ok(status) => Ok(status), + Err(error) => { + eprintln!("Normal fixture shutdown failed; forcing owned child termination: {error}"); + child.kill()?; + wait_for_exit(child, timeout) + } + } +} + +pub fn join_reader(handle: thread::JoinHandle<()>, timeout: Duration) -> io::Result<()> { + let started = Instant::now(); + while !handle.is_finished() { + let remaining = timeout.saturating_sub(started.elapsed()); + if remaining.is_zero() { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "fixture reader did not finish before its deadline", + )); + } + thread::sleep(Duration::from_millis(10).min(remaining)); + } + handle + .join() + .map_err(|_| io::Error::other("fixture reader panicked")) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::process::{Command, Stdio}; + use std::sync::mpsc; + + #[test] + fn fixture_ignores_input_eof() { + if std::env::var_os("PET_TEST_SHUTDOWN_CHILD").is_some() { + thread::sleep(Duration::from_secs(30)); + } + } + + #[test] + fn fixture_shutdown_waits_normally_and_forces_only_after_timeout() { + let mut normal = Command::new(std::env::current_exe().unwrap()) + .arg("--list") + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + assert!(shutdown_fixture(&mut normal, Duration::from_secs(4)) + .unwrap() + .success()); + + let mut hanging = Command::new(std::env::current_exe().unwrap()) + .arg("fixture_ignores_input_eof") + .env("PET_TEST_SHUTDOWN_CHILD", "1") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + let started = Instant::now(); + let status = shutdown_fixture(&mut hanging, Duration::from_millis(200)).unwrap(); + assert!( + !status.success(), + "the fixture must require forced termination" + ); + assert!(started.elapsed() < Duration::from_secs(2)); + assert!(hanging.try_wait().unwrap().is_some()); + } + + #[test] + fn blocked_reader_times_out_without_joining_and_can_still_finish() { + let (release, wait) = mpsc::channel(); + let (finished, completion) = mpsc::channel(); + let reader = thread::spawn(move || { + wait.recv().unwrap(); + finished.send(()).unwrap(); + }); + let started = Instant::now(); + assert_eq!( + join_reader(reader, Duration::from_millis(20)) + .unwrap_err() + .kind(), + io::ErrorKind::TimedOut, + ); + assert!(started.elapsed() < Duration::from_secs(1)); + release.send(()).unwrap(); + completion.recv_timeout(Duration::from_secs(1)).unwrap(); + join_reader(thread::spawn(|| {}), Duration::from_secs(1)).unwrap(); + } +} diff --git a/docs/JSONRPC.md b/docs/JSONRPC.md index fc82e644..0d371a2b 100644 --- a/docs/JSONRPC.md +++ b/docs/JSONRPC.md @@ -11,6 +11,48 @@ For samples using JSONRPC, please have a look at the [sample.js](./sample.js) fi Any requests/notifications not documented here are not supported. +## Transport lifetime and limits + +Close the server process's stdin after receiving all responses you need. EOF at a +frame boundary is a normal shutdown, not an empty request. PET stops accepting +requests and output, discards pending notifications/replies, cancels admitted +interpreter and manager probes, and exits successfully once their ownership +cleanup finishes. Closing stdin is cancellation, not a request to drain unfinished +requests. Test clients wait for normal exit and forcibly terminate only their own +child as a bounded failure fallback. + +EOF inside a header or payload, malformed framing, and read/write/flush failures +are terminal errors: PET reports the failure on stderr and exits unsuccessfully. +Malformed JSON in a complete frame instead receives a Parse Error (`-32700`, +`id: null`), after which subsequent frames can still be processed. Protocol stdout +contains framed JSONRPC only. + +Input is currently one `Content-Length` header followed by a blank line and the +specified number of UTF-8 payload bytes. Both CRLF and LF line endings are accepted. +Headers including the separator are limited to 8 KiB, and payloads to 16 MiB, before +payload allocation. Multi-header input parsing is tracked separately in +[#532](https://github.com/microsoft/python-environment-tools/issues/532). + +One process-lifetime writer emits accepted frames in FIFO order, so a refresh reply +cannot overtake notifications already admitted before it. Each serialized output +payload is limited to 16 MiB; the queue holds at most 1,024 frames and 32 MiB of +retained frame capacity. The queue limit excludes one in-flight frame and the +single producer's serialization/frame-building buffers, whose payloads are also +limited to 16 MiB each. Queue saturation is a terminal connection error rather +than an unbounded allocation or a wait on a slow consumer. These are output bounds, +not a bound on total request/discovery memory or worker concurrency. + +The writer owns a duplicate OS stdout handle and does not hold Rust's global stdout +lock. Shutdown drops queued output and abandons in-flight output without joining +an OS-blocked writer or stdin reader; those process-lifetime threads end when the +standalone server exits. The first output failure recorded before closure remains +fatal; errors arriving after normal closure are discarded with the cancelled work. +Admitted probes are tracked separately through ownership cleanup, with a shared +three-second server shutdown wait. Cleanup failures or an expired wait are reported +as errors, never as clean shutdown. Synchronous OS process creation and individual +OS I/O calls cannot be interrupted by this mechanism; nor does it extend the Unix +ownership boundary to descendants that deliberately escape their process group. + ## Request identifiers Requests include an `id` that is a string, JSON number, or explicit `null`. PET preserves @@ -280,7 +322,7 @@ Interpreter probes used for resolution and manager probes used during refresh/di - Unix probes start in a new process group. PET signals that group before reaping its leader, avoiding process-ID reuse. On macOS, an otherwise ambiguous permission error is accepted only when bounded, unchanged membership and process-birth identities confirm that all group members are already zombies; incomplete or failed inspection remains an error. Descendants that deliberately start another group/session escape this boundary; PET does not act as a system-wide descendant reaper. - Windows probes start suspended and without a console window. PET assigns an unnamed, non-breakaway, kill-on-close job before resuming the child's sole thread. Stable Rust does not expose the primary-thread handle, so PET uses a per-process thread-metadata snapshot (`PssCaptureSnapshot`, Windows 8.1+) and verifies the selected thread's process identity. It does not enumerate system-wide threads or clone the child's address space. Ambiguous/missing thread ownership or failed job assignment fails the probe rather than running it unsupervised. Command arguments, environment, working directory, and batch-file launch still use Rust's standard process implementation; the discovery-only runner replaces any caller-supplied Windows creation flags with `CREATE_NO_WINDOW | CREATE_SUSPENDED`. All current interpreter/manager callers previously supplied only `CREATE_NO_WINDOW`. - After the direct child exits, draining continues within the same deadline. If output has not reached EOF (for example, an escaped Unix descendant retains a write handle), the probe fails explicitly with incomplete output rather than waiting indefinitely. Output produced only by background helpers after their parent exits is not guaranteed: remaining owned helpers are terminated, including on successful parent exit. -- These are per-subprocess limits, not a total request/workspace budget. They do not bound synchronous OS process creation. Server active-request cancellation and shutdown remain a separate lifecycle concern. +- These are per-subprocess limits, not a total request/workspace budget. They do not bound synchronous OS process creation. During server shutdown, new probes are refused and active probes are cancelled through the same ownership cleanup. Cancellation alone is traced rather than logged as a probe failure; cleanup failures remain errors. Missing default Conda executables are quietly ignored; installed/custom manager failures and all timeouts/output failures are logged. Interpreter and Conda JSON is parsed strictly. Poetry stdout From 94af1dada6097cf2e00807312c9f8a9786e379b9 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 10:37:58 -0700 Subject: [PATCH 2/9] test: preserve native server coverage profiles (Refs #529) Keep LLVM_PROFILE_FILE in isolated PET subprocess environments so normal EOF shutdown writes profiles into the coverage collector instead of untracked default files. Preserve fixture overrides and minimal environment isolation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet/tests/jsonrpc_client.rs | 33 +++++++++++++++---------- crates/pet/tests/jsonrpc_server_test.rs | 22 +++-------------- 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/crates/pet/tests/jsonrpc_client.rs b/crates/pet/tests/jsonrpc_client.rs index be70aa9c..2fb12e41 100644 --- a/crates/pet/tests/jsonrpc_client.rs +++ b/crates/pet/tests/jsonrpc_client.rs @@ -20,6 +20,21 @@ static REQUEST_ID: AtomicU32 = AtomicU32::new(1); const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +pub(crate) fn configure_isolated_pet_environment(command: &mut Command) { + command.env_clear().env("PATH", ""); + + if let Some(value) = std::env::var_os("LLVM_PROFILE_FILE") { + command.env("LLVM_PROFILE_FILE", value); + } + + #[cfg(windows)] + for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { + if let Some(value) = std::env::var_os(name) { + command.env(name, value); + } + } +} + #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub struct RefreshResult { @@ -121,19 +136,11 @@ impl PetJsonRpcClient { cmd.arg("server") .stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - // Clear all inherited env vars to prevent host-specific tool - // configuration from leaking into the test environment, then - // restore only the minimum required for the OS to function. - .env_clear() - .env("PATH", "") - .envs(environment.iter().copied()); - #[cfg(windows)] - for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { - if let Some(value) = std::env::var_os(name) { - cmd.env(name, value); - } - } + .stderr(Stdio::piped()); + configure_isolated_pet_environment(&mut cmd); + // Explicit fixture variables take precedence over the runner values + // retained by the isolated environment. + cmd.envs(environment.iter().copied()); let mut process = cmd .spawn() .map_err(|e| format!("Failed to spawn pet server: {e}"))?; diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index cd7febaf..04b78502 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -30,15 +30,8 @@ impl RawRpcClient { .arg("server") .stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::inherit()) - .env_clear() - .env("PATH", ""); - #[cfg(windows)] - for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { - if let Some(value) = std::env::var_os(name) { - command.env(name, value); - } - } + .stderr(Stdio::inherit()); + jsonrpc_client::configure_isolated_pet_environment(&mut command); let mut child = command.spawn().expect("raw fixture must spawn PET"); let stdout = child.stdout.take().expect("PET stdout must be piped"); let (sender, responses) = mpsc::channel(); @@ -516,15 +509,8 @@ impl ShutdownFixture { .arg("server") .stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .env_clear() - .env("PATH", ""); - #[cfg(windows)] - for name in ["SYSTEMROOT", "SYSTEMDRIVE"] { - if let Some(value) = std::env::var_os(name) { - command.env(name, value); - } - } + .stderr(Stdio::piped()); + jsonrpc_client::configure_isolated_pet_environment(&mut command); Self { child: command.spawn().expect("shutdown fixture must spawn PET"), } From bfbe17e76186986a8137989e1e8cf9a6f72f9815 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 10:46:30 -0700 Subject: [PATCH 3/9] fix: validate bounded JSONRPC frames and envelopes (Fixes #532) Read complete bounded multi-header frames without losing byte boundaries, and reject invalid request envelopes before dispatch while preserving supported IDs and legacy method errors. Exercise exact limits, fragmentation, recovery and native wire compatibility. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet-jsonrpc/src/framing.rs | 507 ++++++++++++++++++++++++ crates/pet-jsonrpc/src/lib.rs | 1 + crates/pet-jsonrpc/src/server.rs | 396 +++++++++++------- crates/pet/tests/jsonrpc_server_test.rs | 288 +++++++++++++- docs/JSONRPC.md | 43 +- 5 files changed, 1092 insertions(+), 143 deletions(-) create mode 100644 crates/pet-jsonrpc/src/framing.rs diff --git a/crates/pet-jsonrpc/src/framing.rs b/crates/pet-jsonrpc/src/framing.rs new file mode 100644 index 00000000..37b9de2b --- /dev/null +++ b/crates/pet-jsonrpc/src/framing.rs @@ -0,0 +1,507 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::io::{self, BufRead, ErrorKind}; + +/// Maximum number of raw bytes in a frame's headers, including line endings +/// and the blank line separating the headers from the payload. +pub const MAX_HEADER_BYTES: usize = 8 * 1024; + +/// Maximum accepted payload size. The limit is checked before allocating. +pub const MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024; + +/// Reads one Content-Length-framed message without decoding its payload. +/// +/// Both CRLF and LF line endings are accepted. EOF is clean only before any +/// bytes of the next frame have been read. +pub fn read_frame(reader: &mut R) -> io::Result>> { + let mut header_bytes = 0; + let mut line = Vec::new(); + let mut content_length = None; + + loop { + line.clear(); + let found_line_ending = read_bounded_line( + reader, + &mut line, + MAX_HEADER_BYTES.saturating_sub(header_bytes), + )?; + + if line.is_empty() && !found_line_ending { + return if header_bytes == 0 { + Ok(None) + } else { + Err(unexpected_eof("EOF while reading frame headers")) + }; + } + if !found_line_ending { + return Err(unexpected_eof("EOF while reading frame headers")); + } + + header_bytes = header_bytes + .checked_add(line.len()) + .ok_or_else(|| invalid_data("frame headers exceed the byte limit"))?; + + let header = strip_line_ending(&line)?; + if header.is_empty() { + break; + } + + let (name, value) = parse_header(header)?; + if name.eq_ignore_ascii_case(b"Content-Length") { + if content_length.is_some() { + return Err(invalid_data("duplicate Content-Length header")); + } + content_length = Some(parse_content_length(trim_ascii_whitespace(value))?); + } + } + + let content_length = + content_length.ok_or_else(|| invalid_data("missing Content-Length header"))?; + if content_length > MAX_PAYLOAD_BYTES { + return Err(invalid_data("frame payload exceeds the byte limit")); + } + + let mut payload = Vec::new(); + payload + .try_reserve_exact(content_length) + .map_err(|error| io::Error::other(format!("failed to allocate frame payload: {error}")))?; + payload.resize(content_length, 0); + reader + .read_exact(&mut payload) + .map_err(|error| match error.kind() { + ErrorKind::UnexpectedEof => unexpected_eof("EOF while reading frame payload"), + _ => error, + })?; + Ok(Some(payload)) +} + +fn read_bounded_line( + reader: &mut R, + line: &mut Vec, + byte_limit: usize, +) -> io::Result { + loop { + let buffer = loop { + match reader.fill_buf() { + Err(error) if error.kind() == ErrorKind::Interrupted => continue, + result => break result?, + } + }; + if buffer.is_empty() { + return Ok(false); + } + + let bytes_to_take = buffer + .iter() + .position(|byte| *byte == b'\n') + .map_or(buffer.len(), |index| index + 1); + let remaining = byte_limit.saturating_sub(line.len()); + if bytes_to_take > remaining { + return Err(invalid_data("frame headers exceed the byte limit")); + } + + line.extend_from_slice(&buffer[..bytes_to_take]); + let found_line_ending = buffer[bytes_to_take - 1] == b'\n'; + reader.consume(bytes_to_take); + if found_line_ending { + return Ok(true); + } + } +} + +fn strip_line_ending(line: &[u8]) -> io::Result<&[u8]> { + let without_lf = line + .strip_suffix(b"\n") + .ok_or_else(|| invalid_data("header line is not terminated"))?; + Ok(without_lf.strip_suffix(b"\r").unwrap_or(without_lf)) +} + +fn parse_header(header: &[u8]) -> io::Result<(&[u8], &[u8])> { + let colon = header + .iter() + .position(|byte| *byte == b':') + .ok_or_else(|| invalid_data("malformed header"))?; + let name = &header[..colon]; + let value = &header[colon + 1..]; + + if name.is_empty() || !name.iter().all(|byte| is_header_name_byte(*byte)) { + return Err(invalid_data("malformed header name")); + } + if !value + .iter() + .all(|byte| *byte == b'\t' || (b' '..=b'~').contains(byte)) + { + return Err(invalid_data("malformed header value")); + } + + Ok((name, value)) +} + +fn is_header_name_byte(byte: u8) -> bool { + byte.is_ascii_alphanumeric() + || matches!( + byte, + b'!' | b'#' + | b'$' + | b'%' + | b'&' + | b'\'' + | b'*' + | b'+' + | b'-' + | b'.' + | b'^' + | b'_' + | b'`' + | b'|' + | b'~' + ) +} + +fn trim_ascii_whitespace(mut value: &[u8]) -> &[u8] { + while matches!(value.first(), Some(b' ' | b'\t')) { + value = &value[1..]; + } + while matches!(value.last(), Some(b' ' | b'\t')) { + value = &value[..value.len() - 1]; + } + value +} + +fn parse_content_length(value: &[u8]) -> io::Result { + if value.is_empty() || !value.iter().all(u8::is_ascii_digit) { + return Err(invalid_data( + "Content-Length must be a non-negative decimal integer", + )); + } + + value.iter().try_fold(0usize, |length, byte| { + length + .checked_mul(10) + .and_then(|length| length.checked_add(usize::from(*byte - b'0'))) + .ok_or_else(|| invalid_data("Content-Length overflows usize")) + }) +} + +fn invalid_data(message: &'static str) -> io::Error { + io::Error::new(ErrorKind::InvalidData, message) +} + +fn unexpected_eof(message: &'static str) -> io::Error { + io::Error::new(ErrorKind::UnexpectedEof, message) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cmp; + use std::io::{BufReader, Cursor, Read}; + + fn frame_with_line_ending(payload: &[u8], line_ending: &[u8]) -> Vec { + let mut frame = format!("Content-Length: {}", payload.len()).into_bytes(); + frame.extend_from_slice(line_ending); + frame.extend_from_slice(line_ending); + frame.extend_from_slice(payload); + frame + } + + fn error_kind(input: &[u8]) -> ErrorKind { + read_frame(&mut Cursor::new(input)) + .expect_err("frame should be rejected") + .kind() + } + + #[test] + fn accepts_multiple_headers_in_any_order_and_case() { + let input = b"Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n\ + X-Extension: value\r\n\ + cOnTeNt-LeNgTh: 7\r\n\r\n\ + {\"x\":1}"; + assert_eq!( + read_frame(&mut Cursor::new(input)).unwrap(), + Some(br#"{"x":1}"#.to_vec()) + ); + + let input = b"X-Before: yes\nContent-Length: 2\nX-After: yes\n\n{}"; + assert_eq!( + read_frame(&mut Cursor::new(input)).unwrap(), + Some(b"{}".to_vec()) + ); + } + + #[test] + fn reads_consecutive_crlf_and_lf_frames() { + let mut input = frame_with_line_ending(b"one", b"\r\n"); + input.extend(frame_with_line_ending(b"two", b"\n")); + let mut reader = Cursor::new(input); + + assert_eq!(read_frame(&mut reader).unwrap(), Some(b"one".to_vec())); + assert_eq!(read_frame(&mut reader).unwrap(), Some(b"two".to_vec())); + assert_eq!(read_frame(&mut reader).unwrap(), None); + } + + #[test] + fn preserves_unicode_payload_bytes() { + let payload = "snowman: \u{2603}; crab: \u{1f980}".as_bytes(); + let input = frame_with_line_ending(payload, b"\r\n"); + assert_eq!( + read_frame(&mut Cursor::new(input)).unwrap(), + Some(payload.to_vec()) + ); + } + + #[test] + fn supports_zero_length_payload() { + let mut reader = Cursor::new(b"Content-Length: 0\r\n\r\nnext"); + assert_eq!(read_frame(&mut reader).unwrap(), Some(Vec::new())); + assert_eq!(reader.position(), 21); + } + + #[test] + fn distinguishes_clean_and_truncated_eof() { + assert_eq!(read_frame(&mut Cursor::new(b"")).unwrap(), None); + assert_eq!(error_kind(b"Content-Length: 1"), ErrorKind::UnexpectedEof); + assert_eq!(error_kind(b"Content-Length: 1\n"), ErrorKind::UnexpectedEof); + assert_eq!( + error_kind(b"Content-Length: 3\n\nab"), + ErrorKind::UnexpectedEof + ); + } + + #[test] + fn rejects_invalid_content_lengths() { + for input in [ + b"Content-Length:\n\n".as_slice(), + b"Content-Length: -1\n\n", + b"Content-Length: +1\n\n", + b"Content-Length: 1.0\n\n", + b"Content-Length: 1 0\n\n", + b"Content-Length: \xff\n\n", + ] { + assert_eq!(error_kind(input), ErrorKind::InvalidData, "{input:?}"); + } + + let overflow = format!("Content-Length: {}0\n\n", usize::MAX); + assert_eq!(error_kind(overflow.as_bytes()), ErrorKind::InvalidData); + let over_limit = format!("Content-Length: {}\n\n", MAX_PAYLOAD_BYTES + 1); + assert_eq!(error_kind(over_limit.as_bytes()), ErrorKind::InvalidData); + } + + #[test] + fn rejects_duplicate_and_missing_content_length() { + assert_eq!( + error_kind(b"Content-Length: 0\ncontent-length: 0\n\n"), + ErrorKind::InvalidData + ); + assert_eq!( + error_kind(b"Content-Type: application/json\n\n"), + ErrorKind::InvalidData + ); + } + + #[test] + fn rejects_malformed_headers() { + for input in [ + b"Content Length: 0\n\n".as_slice(), + b"Content-Length 0\n\n", + b": value\nContent-Length: 0\n\n", + b"X: value\rcontinued\nContent-Length: 0\n\n", + b"X: \x7f\nContent-Length: 0\n\n", + ] { + assert_eq!(error_kind(input), ErrorKind::InvalidData, "{input:?}"); + } + } + + fn header_of_size(size: usize) -> Vec { + let fixed = b"X: \nContent-Length: 0\n\n"; + assert!(size >= fixed.len()); + let mut header = b"X: ".to_vec(); + header.resize(size - (fixed.len() - 3), b'a'); + header.extend_from_slice(b"\nContent-Length: 0\n\n"); + assert_eq!(header.len(), size); + header + } + + #[test] + fn enforces_exact_header_byte_limit() { + let at_limit = header_of_size(MAX_HEADER_BYTES); + assert_eq!( + read_frame(&mut Cursor::new(at_limit)).unwrap(), + Some(Vec::new()) + ); + + let over_limit = header_of_size(MAX_HEADER_BYTES + 1); + assert_eq!(error_kind(&over_limit), ErrorKind::InvalidData); + } + + #[test] + fn accepts_exact_payload_limit_without_consuming_the_next_frame() { + let header = format!("Content-Length: {MAX_PAYLOAD_BYTES}\r\n\r\n"); + let input = Cursor::new(header.into_bytes()) + .chain(io::repeat(b'x').take(MAX_PAYLOAD_BYTES as u64)) + .chain(Cursor::new(b"Content-Length: 0\n\n")); + let mut reader = BufReader::new(input); + let payload = read_frame(&mut reader).unwrap().unwrap(); + assert_eq!(payload.len(), MAX_PAYLOAD_BYTES); + assert!(payload.iter().all(|byte| *byte == b'x')); + assert_eq!(read_frame(&mut reader).unwrap(), Some(Vec::new())); + assert_eq!(read_frame(&mut reader).unwrap(), None); + } + + #[test] + fn bounds_unterminated_header_lines() { + assert_eq!( + error_kind(&vec![b'a'; MAX_HEADER_BYTES]), + ErrorKind::UnexpectedEof + ); + assert_eq!( + error_kind(&vec![b'a'; MAX_HEADER_BYTES + 1]), + ErrorKind::InvalidData + ); + } + + struct FragmentedReader { + data: Vec, + position: usize, + chunks: Vec, + next_chunk: usize, + } + + impl FragmentedReader { + fn new(data: Vec, chunks: Vec) -> Self { + Self { + data, + position: 0, + chunks, + next_chunk: 0, + } + } + } + + impl Read for FragmentedReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if self.position == self.data.len() { + return Ok(0); + } + let chunk = self.chunks[self.next_chunk % self.chunks.len()]; + self.next_chunk += 1; + let length = cmp::min( + chunk, + cmp::min(buffer.len(), self.data.len() - self.position), + ); + buffer[..length].copy_from_slice(&self.data[self.position..self.position + length]); + self.position += length; + Ok(length) + } + } + + #[test] + fn handles_deterministic_fragmentation_patterns() { + let payloads = [b"".as_slice(), b"x", br#"{"unicode":"\u2603","value":42}"#]; + let patterns = [ + vec![1], + vec![2, 1, 3], + vec![7, 4, 1, 9, 2], + (1..=17).collect(), + ]; + + for payload in payloads { + let frame = frame_with_line_ending(payload, b"\r\n"); + for chunks in &patterns { + let fragmented = FragmentedReader::new(frame.clone(), chunks.clone()); + let mut reader = BufReader::with_capacity(5, fragmented); + assert_eq!( + read_frame(&mut reader).unwrap(), + Some(payload.to_vec()), + "payload {payload:?}, chunks {chunks:?}" + ); + assert_eq!(read_frame(&mut reader).unwrap(), None); + } + } + + for seed in 0..32u32 { + let mut state = seed.wrapping_add(1); + let payload_length = (seed as usize * 37) % 257; + let payload: Vec = (0..payload_length) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + state.to_le_bytes()[2] + }) + .collect(); + let chunks: Vec = (0..23) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + usize::from(state.to_le_bytes()[1] % 19) + 1 + }) + .collect(); + let line_ending: &[u8] = if seed % 2 == 0 { b"\r\n" } else { b"\n" }; + let frame = frame_with_line_ending(&payload, line_ending); + let fragmented = FragmentedReader::new(frame, chunks); + let mut reader = BufReader::with_capacity((seed as usize % 11) + 1, fragmented); + + assert_eq!(read_frame(&mut reader).unwrap(), Some(payload)); + assert_eq!(read_frame(&mut reader).unwrap(), None); + } + } + + struct InterruptingReader { + inner: Cursor>, + interrupt_next: bool, + } + + impl Read for InterruptingReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if self.interrupt_next { + self.interrupt_next = false; + return Err(io::Error::new(ErrorKind::Interrupted, "interrupted")); + } + self.interrupt_next = true; + self.inner.read(buffer) + } + } + + #[test] + fn retries_interrupted_header_and_payload_reads() { + let input = frame_with_line_ending(b"payload", b"\r\n"); + let interrupting = InterruptingReader { + inner: Cursor::new(input), + interrupt_next: true, + }; + let mut reader = BufReader::with_capacity(3, interrupting); + assert_eq!(read_frame(&mut reader).unwrap(), Some(b"payload".to_vec())); + } + + struct ErrorAfterReader { + inner: Cursor>, + bytes_before_error: usize, + } + + impl Read for ErrorAfterReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if self.bytes_before_error == 0 { + return Err(io::Error::other("reader failed")); + } + let limit = cmp::min(buffer.len(), self.bytes_before_error); + let read = self.inner.read(&mut buffer[..limit])?; + self.bytes_before_error -= read; + Ok(read) + } + } + + #[test] + fn propagates_actual_reader_errors_in_headers_and_payload() { + for bytes_before_error in [5, b"Content-Length: 4\r\n\r\n".len() + 2] { + let input = frame_with_line_ending(b"data", b"\r\n"); + let failing = ErrorAfterReader { + inner: Cursor::new(input), + bytes_before_error, + }; + let mut reader = BufReader::with_capacity(3, failing); + assert_eq!( + read_frame(&mut reader).unwrap_err().kind(), + ErrorKind::Other + ); + } + } +} diff --git a/crates/pet-jsonrpc/src/lib.rs b/crates/pet-jsonrpc/src/lib.rs index 8a3341e8..99eb0705 100644 --- a/crates/pet-jsonrpc/src/lib.rs +++ b/crates/pet-jsonrpc/src/lib.rs @@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize}; +mod framing; mod output; pub mod server; diff --git a/crates/pet-jsonrpc/src/server.rs b/crates/pet-jsonrpc/src/server.rs index a29fa5e0..b6f320a4 100644 --- a/crates/pet-jsonrpc/src/server.rs +++ b/crates/pet-jsonrpc/src/server.rs @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -use crate::{send_error, RequestId}; +use crate::{framing::read_frame, send_error, RequestId}; use serde_json::{self, Value}; use std::{ collections::HashMap, - io::{self, BufRead, BufReader, Read}, + io::{self, BufRead, BufReader}, sync::{mpsc, Arc}, thread, time::Duration, @@ -70,6 +70,11 @@ impl HandlersKeyedByMethodName { } fn handle_request(&self, message: Value) { + let Value::Object(message) = message else { + (self.send_error)(None, -32600, "Invalid JSONRPC request".to_string()); + return; + }; + let id = match message.get("id") { None => None, Some(Value::String(id)) => Some(RequestId::String(id.clone())), @@ -80,11 +85,36 @@ impl HandlersKeyedByMethodName { return; } }; - match message["method"].as_str() { + + if !matches!(message.get("jsonrpc"), Some(Value::String(version)) if version == "2.0") { + (self.send_error)(id.as_ref(), -32600, "Invalid JSONRPC request".to_string()); + return; + } + + match message.get("method").and_then(Value::as_str) { Some(method) => { + let params = match message.get("params") { + None | Some(Value::Null) => Value::Null, + Some(params @ (Value::Object(_) | Value::Array(_))) => params.clone(), + Some(_) => { + if let Some(id) = id.as_ref() { + (self.send_error)( + Some(id), + -32602, + "JSONRPC params must be an object or array".to_string(), + ); + } else { + log::error!( + "Ignoring JSONRPC notification with invalid params for method {method}" + ); + } + return; + } + }; + if let Some(id) = id { if let Some(handler) = self.requests.get(method) { - handler(self.context.clone(), id, message["params"].clone()); + handler(self.context.clone(), id, params); } else { eprint!("Failed to find handler for method: {method}"); (self.send_error)( @@ -94,12 +124,13 @@ impl HandlersKeyedByMethodName { ); } } else if let Some(handler) = self.notifications.get(method) { - handler(self.context.clone(), message["params"].clone()); + handler(self.context.clone(), params); } else { eprint!("Failed to find handler for method: {method}"); } } None => { + let message = Value::Object(message); eprint!("Failed to get method from message: {message}"); (self.send_error)( id.as_ref(), @@ -112,8 +143,6 @@ impl HandlersKeyedByMethodName { } const INPUT_POLL_INTERVAL: Duration = Duration::from_millis(20); -const MAX_HEADER_BYTES: usize = 8 * 1024; -const MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024; /// Runs the standalone process transport until EOF or a fatal I/O error. /// Pending output is discarded at shutdown; callers must finish subprocess cleanup @@ -141,7 +170,7 @@ fn close_transport( fn read_input(mut reader: impl BufRead, sender: mpsc::SyncSender>>>) { loop { - let frame = read_payload(&mut reader); + let frame = read_frame(&mut reader); let terminal = !matches!(&frame, Ok(Some(_))); if sender.send(frame).is_err() || terminal { return; @@ -177,65 +206,6 @@ fn dispatch_input( } } -fn read_payload(reader: &mut impl BufRead) -> io::Result>> { - let mut header = String::new(); - let count = (&mut *reader) - .take(MAX_HEADER_BYTES as u64 + 1) - .read_line(&mut header)?; - if count == 0 { - return Ok(None); - } - if count > MAX_HEADER_BYTES { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "JSONRPC header exceeds limit", - )); - } - if !header.ends_with('\n') { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "Truncated JSONRPC header", - )); - } - let length = get_content_length(&header) - .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; - if length > MAX_PAYLOAD_BYTES { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "JSONRPC payload exceeds limit", - )); - } - let mut separator = String::new(); - (&mut *reader) - .take((MAX_HEADER_BYTES - count) as u64 + 1) - .read_line(&mut separator)?; - if count + separator.len() > MAX_HEADER_BYTES { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "JSONRPC header exceeds limit", - )); - } - if !separator.ends_with('\n') { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "Truncated JSONRPC header separator", - )); - } - if separator != "\r\n" && separator != "\n" { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "Expected a blank JSONRPC header separator", - )); - } - let mut payload = Vec::new(); - payload - .try_reserve_exact(length) - .map_err(io::Error::other)?; - payload.resize(length, 0); - reader.read_exact(&mut payload)?; - Ok(Some(payload)) -} - fn handle_payload( handlers: &HandlersKeyedByMethodName, payload: impl AsRef<[u8]>, @@ -245,31 +215,12 @@ fn handle_payload( Ok(()) } -/// Parses the content length from the given line. -fn get_content_length(line: &str) -> Result { - let line = line.trim(); - if let Some(content_length) = line.find("Content-Length: ") { - let start = content_length + "Content-Length: ".len(); - if let Ok(length) = line[start..].parse::() { - Ok(length) - } else { - Err(format!( - "Failed to parse content length from {} for {}", - &line[start..], - line - )) - } - } else { - Err(format!( - "String 'Content-Length' not found in input => {line}" - )) - } -} - #[cfg(test)] mod tests { use super::*; + use crate::framing::{MAX_HEADER_BYTES, MAX_PAYLOAD_BYTES}; use serde_json::json; + use std::io::Read; use std::sync::Mutex; #[test] @@ -322,9 +273,8 @@ mod tests { }) } - #[test] - fn request_ids_preserve_values_for_dispatch_and_errors() { - for value in [ + fn supported_request_id_values() -> [Value; 10] { + [ json!("request-1"), json!(""), json!("\u{03c0}-request"), @@ -335,7 +285,12 @@ mod tests { json!(i64::MIN), json!(1.5), Value::Null, - ] { + ] + } + + #[test] + fn request_ids_preserve_values_for_dispatch_and_errors() { + for value in supported_request_id_values() { let id = serde_json::from_value::(value.clone()).unwrap(); assert_eq!(serde_json::to_value(&id).unwrap(), value); let context = Arc::new(TestContext::default()); @@ -346,16 +301,21 @@ mod tests { handlers.add_notification_handler("method", |context, params| { *context.notification.lock().unwrap() = Some(params); }); - handlers.handle_request(json!({"id": value, "method": "method", "params": 42})); + handlers.handle_request( + json!({"jsonrpc": "2.0", "id": value, "method": "method", "params": {"value": 42}}), + ); assert_eq!( *context.request.lock().unwrap(), - Some((id.clone(), json!(42))) + Some((id.clone(), json!({"value": 42}))) ); assert!(context.notification.lock().unwrap().is_none()); - handlers.handle_request(json!({"method": "method", "params": 7})); - assert_eq!(context.notification.lock().unwrap().take(), Some(json!(7))); - handlers.handle_request(json!({"id": value, "method": "unknown"})); - handlers.handle_request(json!({"id": value})); + handlers.handle_request(json!({"jsonrpc": "2.0", "method": "method", "params": [7]})); + assert_eq!( + context.notification.lock().unwrap().take(), + Some(json!([7])) + ); + handlers.handle_request(json!({"jsonrpc": "2.0", "id": value, "method": "unknown"})); + handlers.handle_request(json!({"jsonrpc": "2.0", "id": value})); let errors = context.errors.lock().unwrap(); assert_eq!(errors.len(), 2); assert_eq!(errors[0].0, Some(id.clone())); @@ -375,13 +335,16 @@ mod tests { handlers.add_notification_handler("method", |context, params| { *context.notification.lock().unwrap() = Some(params); }); - handlers.handle_request(json!({"method": "method", "params": 7})); - assert_eq!(context.notification.lock().unwrap().take(), Some(json!(7))); + handlers.handle_request(json!({"jsonrpc": "2.0", "method": "method", "params": [7]})); + assert_eq!( + context.notification.lock().unwrap().take(), + Some(json!([7])) + ); assert!(context.request.lock().unwrap().is_none()); assert!(context.errors.lock().unwrap().is_empty()); for value in [json!(true), json!(false), json!([]), json!({"id": 1})] { assert!(serde_json::from_value::(value.clone()).is_err()); - handlers.handle_request(json!({"id": value, "method": "method"})); + handlers.handle_request(json!({"jsonrpc": "2.0", "id": value, "method": "method"})); assert!(context.notification.lock().unwrap().is_none()); assert!(context.request.lock().unwrap().is_none()); assert_eq!( @@ -390,32 +353,180 @@ mod tests { ); } assert!(context.errors.lock().unwrap().is_empty()); - handlers.handle_request(json!({"id": "after-invalid", "method": "method", "params": 42})); + handlers.handle_request( + json!({"jsonrpc": "2.0", "id": "after-invalid", "method": "method", "params": [42]}), + ); assert_eq!( context.request.lock().unwrap().take(), - Some((RequestId::String("after-invalid".into()), json!(42))) + Some((RequestId::String("after-invalid".into()), json!([42]))) ); assert!(context.notification.lock().unwrap().is_none()); assert!(context.errors.lock().unwrap().is_empty()); } #[test] - fn get_content_length_parses_valid_header() { - assert_eq!(get_content_length("Content-Length: 42\r\n").unwrap(), 42); + fn invalid_top_level_values_and_batches_do_not_invoke_handlers() { + let context = Arc::new(TestContext::default()); + let mut handlers = create_handlers_with_recorded_errors(context.clone()); + handlers.add_request_handler("method", |context, id, params| { + *context.request.lock().unwrap() = Some((id, params)); + }); + handlers.add_notification_handler("method", |context, params| { + *context.notification.lock().unwrap() = Some(params); + }); + + for message in [ + Value::Null, + json!(false), + json!(42), + json!("request"), + json!([]), + json!([{"jsonrpc": "2.0", "id": 1, "method": "method"}]), + ] { + handlers.handle_request(message); + } + + assert!(context.request.lock().unwrap().is_none()); + assert!(context.notification.lock().unwrap().is_none()); + assert_eq!( + context.errors.lock().unwrap().as_slice(), + vec![(None, -32600, "Invalid JSONRPC request".to_string()); 6] + ); + } + + #[test] + fn invalid_jsonrpc_versions_do_not_invoke_handlers() { + let context = Arc::new(TestContext::default()); + let mut handlers = create_handlers_with_recorded_errors(context.clone()); + handlers.add_request_handler("method", |context, id, params| { + *context.request.lock().unwrap() = Some((id, params)); + }); + handlers.add_notification_handler("method", |context, params| { + *context.notification.lock().unwrap() = Some(params); + }); + + for message in [ + json!({"method": "method"}), + json!({"jsonrpc": null, "method": "method"}), + json!({"jsonrpc": 2.0, "method": "method"}), + json!({"jsonrpc": true, "method": "method"}), + json!({"jsonrpc": "1.0", "method": "method"}), + json!({"jsonrpc": "2.0 ", "method": "method"}), + ] { + handlers.handle_request(message); + } + + assert!(context.request.lock().unwrap().is_none()); + assert!(context.notification.lock().unwrap().is_none()); + assert_eq!( + context.errors.lock().unwrap().as_slice(), + vec![(None, -32600, "Invalid JSONRPC request".to_string()); 6] + ); + } + + #[test] + fn valid_ids_round_trip_in_new_envelope_and_params_errors() { + for value in supported_request_id_values() { + let id = serde_json::from_value::(value.clone()).unwrap(); + let context = Arc::new(TestContext::default()); + let mut handlers = create_handlers_with_recorded_errors(context.clone()); + handlers.add_request_handler("method", |context, id, params| { + *context.request.lock().unwrap() = Some((id, params)); + }); + + handlers.handle_request(json!({"jsonrpc": "1.0", "id": value, "method": "method"})); + handlers.handle_request( + json!({"jsonrpc": "2.0", "id": value, "method": "method", "params": true}), + ); + + assert!(context.request.lock().unwrap().is_none()); + assert_eq!( + context.errors.lock().unwrap().as_slice(), + &[ + ( + Some(id.clone()), + -32600, + "Invalid JSONRPC request".to_string() + ), + ( + Some(id), + -32602, + "JSONRPC params must be an object or array".to_string() + ) + ] + ); + } } #[test] - fn get_content_length_rejects_missing_header() { - let error = get_content_length("Content-Type: application/json\r\n").unwrap_err(); + fn valid_parameter_containers_dispatch_requests_and_notifications() { + for params in [None, Some(Value::Null), Some(json!([])), Some(json!({}))] { + let expected = params.clone().unwrap_or(Value::Null); + let context = Arc::new(TestContext::default()); + let mut handlers = create_handlers_with_recorded_errors(context.clone()); + handlers.add_request_handler("method", |context, id, params| { + *context.request.lock().unwrap() = Some((id, params)); + }); + handlers.add_notification_handler("method", |context, params| { + *context.notification.lock().unwrap() = Some(params); + }); - assert!(error.contains("String 'Content-Length' not found")); + let mut request = json!({"jsonrpc": "2.0", "id": 1, "method": "method"}); + let mut notification = json!({"jsonrpc": "2.0", "method": "method"}); + if let Some(params) = params { + request["params"] = params.clone(); + notification["params"] = params; + } + handlers.handle_request(request); + handlers.handle_request(notification); + + assert_eq!( + context.request.lock().unwrap().take(), + Some((1.into(), expected.clone())) + ); + assert_eq!(context.notification.lock().unwrap().take(), Some(expected)); + assert!(context.errors.lock().unwrap().is_empty()); + } } #[test] - fn get_content_length_rejects_non_numeric_length() { - let error = get_content_length("Content-Length: nope\r\n").unwrap_err(); + fn invalid_parameter_containers_do_not_invoke_handlers() { + let context = Arc::new(TestContext::default()); + let mut handlers = create_handlers_with_recorded_errors(context.clone()); + handlers.add_request_handler("method", |context, id, params| { + *context.request.lock().unwrap() = Some((id, params)); + }); + handlers.add_notification_handler("method", |context, params| { + *context.notification.lock().unwrap() = Some(params); + }); - assert!(error.contains("Failed to parse content length")); + for (index, params) in [json!(false), json!(42), json!("params")] + .into_iter() + .enumerate() + { + handlers.handle_request( + json!({"jsonrpc": "2.0", "id": index, "method": "method", "params": params}), + ); + handlers + .handle_request(json!({"jsonrpc": "2.0", "method": "method", "params": params})); + } + + assert!(context.request.lock().unwrap().is_none()); + assert!(context.notification.lock().unwrap().is_none()); + assert_eq!( + context + .errors + .lock() + .unwrap() + .iter() + .map(|(id, code, _)| (id.clone(), *code)) + .collect::>(), + vec![ + (Some(0.into()), -32602), + (Some(1.into()), -32602), + (Some(2.into()), -32602) + ] + ); } #[test] @@ -526,21 +637,42 @@ mod tests { let context = Arc::new(TestContext::default()); let handlers = create_handlers_with_recorded_errors(context.clone()); - let message = json!({ - "jsonrpc": "2.0", - "id": 1, - "params": { "value": 42 } - }); - - handlers.handle_request(message.clone()); + let messages = [ + json!({ + "jsonrpc": "2.0", + "id": 1, + "params": { "value": 42 } + }), + json!({ + "jsonrpc": "2.0", + "id": 1, + "method": 42 + }), + ]; + for message in &messages { + handlers.handle_request(message.clone()); + } assert_eq!( context.errors.lock().unwrap().as_slice(), - &[( - Some(1.into()), - -3, - format!("Failed to extract method from JSONRPC payload {message:?}") - )] + &[ + ( + Some(1.into()), + -3, + format!( + "Failed to extract method from JSONRPC payload {:?}", + messages[0] + ) + ), + ( + Some(1.into()), + -3, + format!( + "Failed to extract method from JSONRPC payload {:?}", + messages[1] + ) + ) + ] ); } @@ -567,14 +699,14 @@ mod tests { } #[test] fn input_distinguishes_clean_eof_from_truncated_frames() { - assert!(read_payload(&mut io::Cursor::new(b"")).unwrap().is_none()); + assert!(read_frame(&mut io::Cursor::new(b"")).unwrap().is_none()); for bytes in [ b"Content-Length: 2".as_slice(), b"Content-Length: 2\r\n".as_slice(), b"Content-Length: 2\r\n\r".as_slice(), b"Content-Length: 2\r\n\r\n{".as_slice(), ] { - let error = read_payload(&mut io::Cursor::new(bytes)).unwrap_err(); + let error = read_frame(&mut io::Cursor::new(bytes)).unwrap_err(); assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof); } } @@ -588,25 +720,21 @@ mod tests { bytes.extend_from_slice(b"Content-Length: 2\n\n"); bytes.extend_from_slice(second); let mut reader = BufReader::with_capacity(1, io::Cursor::new(bytes)); - assert_eq!(read_payload(&mut reader).unwrap().unwrap(), first); - assert_eq!(read_payload(&mut reader).unwrap().unwrap(), second); - assert!(read_payload(&mut reader).unwrap().is_none()); + assert_eq!(read_frame(&mut reader).unwrap().unwrap(), first); + assert_eq!(read_frame(&mut reader).unwrap().unwrap(), second); + assert!(read_frame(&mut reader).unwrap().is_none()); } #[test] fn input_rejects_oversized_headers_and_payloads_before_body_reads() { let bytes = vec![b'x'; MAX_HEADER_BYTES + 1]; assert_eq!( - read_payload(&mut io::Cursor::new(bytes)) - .unwrap_err() - .kind(), + read_frame(&mut io::Cursor::new(bytes)).unwrap_err().kind(), io::ErrorKind::InvalidData ); let bytes = format!("Content-Length: {}\r\n\r\n", MAX_PAYLOAD_BYTES + 1); assert_eq!( - read_payload(&mut io::Cursor::new(bytes)) - .unwrap_err() - .kind(), + read_frame(&mut io::Cursor::new(bytes)).unwrap_err().kind(), io::ErrorKind::InvalidData ); } diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index 04b78502..0687f6c9 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -17,6 +17,19 @@ mod jsonrpc_client; use jsonrpc_client::{EnvironmentNotification, PetJsonRpcClient}; +fn frame_with_headers(payload: &[u8], headers: &[(&str, &str)], line_ending: &[u8]) -> Vec { + let mut frame = Vec::new(); + for (name, value) in headers { + frame.extend_from_slice(name.as_bytes()); + frame.extend_from_slice(b": "); + frame.extend_from_slice(value.as_bytes()); + frame.extend_from_slice(line_ending); + } + frame.extend_from_slice(line_ending); + frame.extend_from_slice(payload); + frame +} + struct RawRpcClient { child: Child, responses: mpsc::Receiver>, @@ -61,9 +74,22 @@ impl RawRpcClient { fn send(&mut self, message: Value) { let body = serde_json::to_vec(&message).unwrap(); + self.send_payload(&body); + } + + fn send_payload(&mut self, body: &[u8]) { + let content_length = body.len().to_string(); + let frame = frame_with_headers( + body, + &[("Content-Length", content_length.as_str())], + b"\r\n", + ); + self.write_raw(&frame); + } + + fn write_raw(&mut self, bytes: &[u8]) { let stdin = self.child.stdin.as_mut().expect("PET stdin must be piped"); - write!(stdin, "Content-Length: {}\r\n\r\n", body.len()).unwrap(); - stdin.write_all(&body).unwrap(); + stdin.write_all(bytes).unwrap(); stdin.flush().unwrap(); } @@ -92,6 +118,188 @@ impl Drop for RawRpcClient { } } +fn assert_rpc_error(response: &Value, expected_id: &Value, expected_code: i64) { + assert_eq!(response["jsonrpc"], "2.0"); + assert_eq!(response.get("id"), Some(expected_id)); + assert_eq!(response["error"]["code"], expected_code); + assert!(response.get("result").is_none()); +} + +#[test] +fn native_wire_accepts_header_variants_fragmentation_and_consecutive_frames() { + let mut client = RawRpcClient::spawn(); + let crlf_payload = br#"{"jsonrpc":"2.0","id":"crlf-content-type-first","method":"info"}"#; + let lf_payload = + "{\"jsonrpc\":\"2.0\",\"id\":\"lf-snowman-\u{2603}\",\"method\":\"info\"}".as_bytes(); + let no_content_type_payload = br#"{"jsonrpc":"2.0","id":"no-content-type","method":"info"}"#; + assert!( + lf_payload.iter().any(|byte| !byte.is_ascii()), + "fixture must exercise byte lengths rather than character counts" + ); + + let crlf_length = crlf_payload.len().to_string(); + let lf_length = lf_payload.len().to_string(); + let no_content_type_length = no_content_type_payload.len().to_string(); + let mut wire = frame_with_headers( + crlf_payload, + &[ + ("cOnTeNt-TyPe", "application/vscode-jsonrpc; charset=utf-8"), + ("X-Before-Length", "accepted"), + ("cOnTeNt-LeNgTh", crlf_length.as_str()), + ], + b"\r\n", + ); + wire.extend(frame_with_headers( + lf_payload, + &[ + ("CONTENT-LENGTH", lf_length.as_str()), + ("x-after-length", "accepted"), + ("CONTENT-TYPE", "application/vscode-jsonrpc; charset=utf-8"), + ], + b"\n", + )); + wire.extend(frame_with_headers( + no_content_type_payload, + &[ + ("X-Optional-Content-Type", "omitted"), + ("Content-Length", no_content_type_length.as_str()), + ], + b"\r\n", + )); + + for fragment in wire.chunks(3) { + client.write_raw(fragment); + } + + for expected_id in [ + "crlf-content-type-first", + "lf-snowman-\u{2603}", + "no-content-type", + ] { + let response = client.receive(); + assert_eq!(response["jsonrpc"], "2.0"); + assert_eq!(response["id"], expected_id); + assert!(response["result"]["petVersion"].is_string()); + assert!(response.get("error").is_none()); + } +} + +#[test] +fn complete_invalid_json_and_utf8_frames_recover_for_the_next_frame() { + let mut client = RawRpcClient::spawn(); + + client.send_payload(br#"{"jsonrpc":"2.0","id":"malformed","method":"info""#); + assert_rpc_error(&client.receive(), &Value::Null, -32700); + + client.send_payload(&[b'{', b'"', 0xff, b'"', b':', b'1', b'}']); + assert_rpc_error(&client.receive(), &Value::Null, -32700); + + client.send(json!({ + "jsonrpc": "2.0", + "id": "after-parse-errors", + "method": "info" + })); + let response = client.receive(); + assert_eq!(response["id"], "after-parse-errors"); + assert!(response["result"]["petVersion"].is_string()); +} + +#[test] +fn native_wire_validates_envelopes_params_and_legacy_errors() { + let mut client = RawRpcClient::spawn(); + + for invalid in [ + Value::Null, + json!(false), + json!(42), + json!("request"), + json!([]), + json!([{"jsonrpc": "2.0", "id": "batch", "method": "info"}]), + ] { + client.send(invalid); + assert_rpc_error(&client.receive(), &Value::Null, -32600); + } + + for request in [ + json!({"id": "missing-version", "method": "info"}), + json!({"jsonrpc": "1.0", "id": "wrong-version", "method": "info"}), + json!({"jsonrpc": 2.0, "id": 17, "method": "info"}), + ] { + let expected_id = request["id"].clone(); + client.send(request); + assert_rpc_error(&client.receive(), &expected_id, -32600); + } + + for invalid_id in [json!([]), json!({"nested": "id"})] { + client.send(json!({ + "jsonrpc": "2.0", + "id": invalid_id, + "method": "info" + })); + assert_rpc_error(&client.receive(), &Value::Null, -32600); + } + + for (index, params) in [json!(false), json!(7), json!("scalar")] + .into_iter() + .enumerate() + { + let id = json!(format!("invalid-params-{index}")); + client.send(json!({ + "jsonrpc": "2.0", + "id": id, + "method": "info", + "params": params + })); + assert_rpc_error(&client.receive(), &id, -32602); + } + + for (id, params) in [ + ("missing-params", None), + ("null-params", Some(Value::Null)), + ("array-params", Some(json!([]))), + ("object-params", Some(json!({}))), + ] { + let mut request = json!({"jsonrpc": "2.0", "id": id, "method": "info"}); + if let Some(params) = params { + request["params"] = params; + } + client.send(request); + let response = client.receive(); + assert_eq!(response["id"], id); + assert!(response["result"]["petVersion"].is_string()); + } + + client.send(json!({ + "jsonrpc": "2.0", + "method": "info", + "params": "invalid-notification-params" + })); + client.send(json!({ + "jsonrpc": "2.0", + "id": "notification-sentinel", + "method": "info" + })); + let response = client.receive(); + assert_eq!(response["id"], "notification-sentinel"); + assert!(response["result"]["petVersion"].is_string()); + + client.send(json!({"jsonrpc": "2.0", "id": "missing-method"})); + assert_rpc_error(&client.receive(), &json!("missing-method"), -3); + client.send(json!({ + "jsonrpc": "2.0", + "id": "unknown-method", + "method": "unknown" + })); + assert_rpc_error(&client.receive(), &json!("unknown-method"), -1); + client.send(json!({ + "jsonrpc": "2.0", + "id": "handler-invalid-param", + "method": "resolve", + "params": {"executable": 42} + })); + assert_rpc_error(&client.receive(), &json!("handler-invalid-param"), -4); +} + #[test] fn request_ids_round_trip_through_success_and_error_responses() { let mut client = RawRpcClient::spawn(); @@ -522,6 +730,12 @@ impl ShutdownFixture { stdin.write_all(body).unwrap(); stdin.flush().unwrap(); } + + fn write_raw(&mut self, bytes: &[u8]) { + let stdin = self.child.stdin.as_mut().unwrap(); + stdin.write_all(bytes).unwrap(); + stdin.flush().unwrap(); + } } impl Drop for ShutdownFixture { @@ -535,6 +749,76 @@ impl Drop for ShutdownFixture { } } +fn assert_fatal_framing_input(name: &str, input: &[u8]) { + let mut fixture = ShutdownFixture::spawn(); + fixture.write_raw(input); + let started = Instant::now(); + fixture.child.stdin.take(); + let status = jsonrpc_client::wait_for_exit(&mut fixture.child, Duration::from_secs(1)) + .unwrap_or_else(|error| panic!("{name} did not terminate within one second: {error}")); + assert!( + !status.success(), + "{name} must terminate the server unsuccessfully" + ); + assert!( + started.elapsed() < Duration::from_secs(1), + "{name} exceeded the shutdown budget" + ); + let mut stderr = Vec::new(); + fixture + .child + .stderr + .take() + .unwrap() + .read_to_end(&mut stderr) + .unwrap(); + assert!(!stderr.is_empty(), "{name} must be reported"); + assert!( + stderr.len() < 4096, + "{name} produced an error flood of {} bytes", + stderr.len() + ); +} + +#[test] +fn invalid_and_oversize_framing_terminates_with_bounded_diagnostics() { + const MAX_HEADER_BYTES: usize = 8 * 1024; + const MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024; + + let mut oversized_header = b"X-Oversized: ".to_vec(); + oversized_header.resize(MAX_HEADER_BYTES + 1, b'x'); + oversized_header.extend_from_slice(b"\r\n\r\n"); + + let cases = [ + ( + "invalid Content-Length", + b"Content-Length: twelve\r\n\r\n".to_vec(), + ), + ( + "missing Content-Length", + b"Content-Type: application/json\r\n\r\n".to_vec(), + ), + ( + "duplicate Content-Length", + b"Content-Length: 0\r\nContent-Length: 0\r\n\r\n".to_vec(), + ), + ( + "overflowing Content-Length", + b"Content-Length: 184467440737095516160\r\n\r\n".to_vec(), + ), + ( + "oversize Content-Length", + format!("Content-Length: {}\r\n\r\n", MAX_PAYLOAD_BYTES + 1).into_bytes(), + ), + ("malformed header", b"Not-A-Header\r\n\r\n".to_vec()), + ("oversize header", oversized_header), + ]; + + for (name, input) in cases { + assert_fatal_framing_input(name, &input); + } +} + #[test] fn stdin_eof_after_exchange_exits_cleanly_within_one_second() { let client = PetJsonRpcClient::spawn().unwrap(); diff --git a/docs/JSONRPC.md b/docs/JSONRPC.md index 0d371a2b..dde61965 100644 --- a/docs/JSONRPC.md +++ b/docs/JSONRPC.md @@ -27,11 +27,21 @@ Malformed JSON in a complete frame instead receives a Parse Error (`-32700`, `id: null`), after which subsequent frames can still be processed. Protocol stdout contains framed JSONRPC only. -Input is currently one `Content-Length` header followed by a blank line and the -specified number of UTF-8 payload bytes. Both CRLF and LF line endings are accepted. -Headers including the separator are limited to 8 KiB, and payloads to 16 MiB, before -payload allocation. Multi-header input parsing is tracked separately in -[#532](https://github.com/microsoft/python-environment-tools/issues/532). +Each input frame contains ASCII headers through a blank line, followed by exactly +`Content-Length` UTF-8 payload **bytes** (not characters). Both CRLF and LF line +endings are accepted. Header names are ASCII-case-insensitive; optional +`Content-Type` and other well-formed headers may precede or follow `Content-Length`. +PET always decodes payloads as UTF-8; it does not negotiate another encoding from +`Content-Type`. Fragmented reads and consecutive frames preserve byte boundaries. + +Exactly one `Content-Length` is required. Its value is decimal digits, optionally +surrounded by spaces or tabs; signs, fractions, duplicate/missing lengths, overflow, +and malformed header names or values are rejected. Total raw headers (including +line endings and the blank separator) are limited to 8 KiB, and payloads to 16 MiB, +before payload allocation. Both exact limits are accepted. Invalid framing closes +the connection unsuccessfully rather than attempting to guess the next boundary. +A complete frame containing invalid UTF-8 or invalid JSON gets the recoverable +Parse Error described above. One process-lifetime writer emits accepted frames in FIFO order, so a refresh reply cannot overtake notifications already admitted before it. Each serialized output @@ -63,8 +73,27 @@ spelling is not preserved. Prefer string IDs when exact values exceed the 64-bit or the precision of a client's numeric type. JSONRPC recommends avoiding fractional and null IDs. Only an absent `id` denotes a notification. Boolean, array, and object IDs produce an Invalid -Request error (`-32600`) with `id: null` and do not invoke a handler. Other existing method and -parameter error codes are unchanged. Notifications do not receive request replies. +Request error (`-32600`) with `id: null` and do not invoke a handler. Valid notifications +do not receive request replies. + +## Request envelopes and errors + +Each payload must be a JSON object with `"jsonrpc": "2.0"`; batches are not supported. +Non-object payloads and missing/invalid versions produce Invalid Request (`-32600`) +without invoking a handler. Errors preserve a valid parsed ID where available, or +use `id: null` otherwise. An invalid envelope with no ID is still an error, not a +valid notification. + +If present, `params` must be an object or array. Absent params and explicit `null` +remain supported for compatibility; method-specific schemas still apply. Other +parameter containers produce Invalid Params (`-32602`) for requests. Otherwise-valid +notifications with invalid params are logged and not dispatched or replied to. +Existing PET method-level codes are retained: missing/nonstring method (`-3`), +unknown request method (`-1`), and handler-specific parameter errors (`-4`). Unknown +notification methods are logged without a reply. These legacy codes are not the +standard JSONRPC equivalents and clients should retain their existing handling. +Complete-frame envelope/parameter errors do not prevent subsequent frames from +being processed. # Info Request From 4276b382d427a6bd66c63fbd08125db301fd7576 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 11:02:01 -0700 Subject: [PATCH 4/9] test: measure broken-output shutdown after readiness (Refs #529) Establish a bounded info exchange before closing the same stdout pipe, so startup is not charged to the unchanged one-second shutdown deadline. Explicitly terminate and bounded-join the readiness helper on setup failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet/tests/jsonrpc_server_test.rs | 26 ++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index 04b78502..2c9a8478 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -585,7 +585,31 @@ fn truncated_input_exits_unsuccessfully_without_an_error_flood() { #[test] fn closed_output_exits_without_waiting_for_stdin_eof() { let mut fixture = ShutdownFixture::spawn(); - drop(fixture.child.stdout.take()); + let mut stdout = BufReader::new(fixture.child.stdout.take().unwrap()); + let (sender, receiver) = mpsc::sync_channel(1); + let reader = thread::spawn(move || { + let response = jsonrpc_client::read_message(&mut stdout); + let _ = sender.send((stdout, response)); + }); + fixture.send(br#"{"jsonrpc":"2.0","id":"ready","method":"info"}"#); + let (stdout, response) = match receiver.recv_timeout(Duration::from_secs(10)) { + Ok(result) => result, + Err(error) => { + drop(receiver); + fixture.child.stdin.take(); + let shutdown = + jsonrpc_client::shutdown_fixture(&mut fixture.child, Duration::from_secs(4)); + let joined = jsonrpc_client::join_reader(reader, Duration::from_secs(4)); + panic!("server readiness failed: {error}; shutdown: {shutdown:?}; reader: {joined:?}"); + } + }; + jsonrpc_client::join_reader(reader, Duration::from_secs(1)).unwrap(); + let response = response + .unwrap() + .expect("ready server must respond to info"); + assert_eq!(response["id"], "ready"); + drop(stdout); + let started = Instant::now(); fixture.send(br#"{"jsonrpc":"2.0","id":1,"method":"info"}"#); let status = jsonrpc_client::wait_for_exit(&mut fixture.child, Duration::from_secs(1)).unwrap(); From 60d8d5c19a8c7395933e107aa0f02d700d22cb2d Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 11:26:13 -0700 Subject: [PATCH 5/9] test: isolate broken-pipe fixture from concurrent forks (Refs #529) Concurrent Unix fork/exec can retain another fixture pipe reader until exec, allowing the sole response write to succeed before the final reader disappears. Run the same real-pipe scenario in an isolated test subprocess; keep the one-second measured deadline and bounded cleanup unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet/tests/jsonrpc_server_test.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index 2c9a8478..ac31605c 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -584,6 +584,29 @@ fn truncated_input_exits_unsuccessfully_without_an_error_flood() { #[test] fn closed_output_exits_without_waiting_for_stdin_eof() { + // Concurrent fork/exec can temporarily inherit a pipe reader despite CLOEXEC. + // Isolate this scenario so its dropped handle really is the final reader. + if std::env::var_os("PET_TEST_CLOSED_OUTPUT_CHILD").is_none() { + let mut child = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "closed_output_exits_without_waiting_for_stdin_eof", + "--nocapture", + ]) + .env("PET_TEST_CLOSED_OUTPUT_CHILD", "1") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()) + .spawn() + .expect("isolated closed-output fixture must spawn"); + let status = jsonrpc_client::shutdown_fixture(&mut child, Duration::from_secs(20)).unwrap(); + assert!( + status.success(), + "isolated closed-output fixture failed: {status}" + ); + return; + } + let mut fixture = ShutdownFixture::spawn(); let mut stdout = BufReader::new(fixture.child.stdout.take().unwrap()); let (sender, receiver) = mpsc::sync_channel(1); From 83df603d674889b85984d59e8f785a7d74d6b270 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 11:35:05 -0700 Subject: [PATCH 6/9] test: cover composed fixture teardown deadline (Refs #529) Allow the outer test-process guard to cover readiness, both forced-shutdown waits, reader completion, and fallback cleanup without changing the one-second shutdown assertion. Record the verified subprocess testing pitfalls in the existing Rust skill. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/skills/rust-coding-skill/SKILL.md | 4 ++++ crates/pet/tests/jsonrpc_server_test.rs | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/skills/rust-coding-skill/SKILL.md b/.github/skills/rust-coding-skill/SKILL.md index c934d8e5..e89b0e15 100644 --- a/.github/skills/rust-coding-skill/SKILL.md +++ b/.github/skills/rust-coding-skill/SKILL.md @@ -98,3 +98,7 @@ Before every Rust commit, run targeted tests and invoke the `rust-precommit` ski ## Learnings Do not execute freshly written scripts as concurrent Unix subprocess fixtures: spawning can fail with `ETXTBSY` (Text file busy). Prefer an existing interpreter such as `/bin/sh -c` with an inline script, or the existing test executable. Assert the typed runner outcome before checking an optional parsed result, so a spawn failure cannot masquerade as a successful negative parsing or timeout test. + +For real-pipe EOF/EPIPE tests, create the pipe inside an isolated test subprocess when other test threads spawn children. Unix `CLOEXEC` closes descriptors at exec, not fork: a concurrent child can temporarily retain a reader, allowing the only write to succeed before the final reader disappears. A readiness handshake alone does not prevent this race. Keep the operation's measured deadline separate from setup, and make an outer fixture deadline cover readiness, waits both before and after forced termination, reader joins, and fallback `Drop` cleanup. + +Use a per-worktree Cargo target directory when validating stacked changes so native fixtures cannot execute another worktree's stale binary. On WSL, run timing-sensitive Linux binaries from the native Linux filesystem rather than a Windows mount, where page faults can stall in filesystem RPC. When launching instrumented PET with `env_clear()`, retain `LLVM_PROFILE_FILE` exactly so child coverage reaches the collector instead of an uncollected default profile. diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index ac31605c..e510ffe8 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -599,7 +599,8 @@ fn closed_output_exits_without_waiting_for_stdin_eof() { .stderr(Stdio::inherit()) .spawn() .expect("isolated closed-output fixture must spawn"); - let status = jsonrpc_client::shutdown_fixture(&mut child, Duration::from_secs(20)).unwrap(); + // Cover readiness, both forced-shutdown waits, reader joining, and Drop cleanup. + let status = jsonrpc_client::shutdown_fixture(&mut child, Duration::from_secs(40)).unwrap(); assert!( status.success(), "isolated closed-output fixture failed: {status}" From 22a0e9328b0491678bca3179fabf23144a237789 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 14:50:59 -0700 Subject: [PATCH 7/9] test: observe descendant release within shutdown deadline (Refs #529) Poll the actual descendant lease using only the remainder of the existing four-second shutdown budget instead of requiring an instantaneous lock transition at server exit. Test held-lock timeout and success after release; preserve cancellation and elapsed-time assertions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet/tests/jsonrpc_server_test.rs | 51 +++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index e510ffe8..7dd1cd79 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -676,6 +676,48 @@ fn stdin_eof_exits_while_output_is_not_drained() { reader.join().unwrap(); } +fn wait_for_descendant_lease(lease: &fs::File, timeout: Duration) -> std::io::Result<()> { + let started = Instant::now(); + loop { + match lease.try_lock() { + Ok(()) => return Ok(()), + Err(fs::TryLockError::Error(error)) => return Err(error), + Err(fs::TryLockError::WouldBlock) => {} + } + let remaining = timeout.saturating_sub(started.elapsed()); + if remaining.is_zero() { + return Err(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "descendant still holds its lease at the shutdown deadline", + )); + } + thread::sleep(Duration::from_millis(10).min(remaining)); + } +} + +#[test] +fn descendant_lease_wait_is_bounded_and_requires_release() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("lease"); + let holder = fs::File::create(&path).unwrap(); + holder.try_lock().unwrap(); + let lease = fs::OpenOptions::new() + .read(true) + .write(true) + .open(path) + .unwrap(); + let started = Instant::now(); + assert_eq!( + wait_for_descendant_lease(&lease, Duration::from_millis(20)) + .unwrap_err() + .kind(), + std::io::ErrorKind::TimedOut + ); + assert!(started.elapsed() < Duration::from_secs(1)); + drop(holder); + wait_for_descendant_lease(&lease, Duration::ZERO).unwrap(); +} + #[cfg(feature = "ci")] #[test] fn stdin_eof_cancels_an_active_interpreter_and_its_descendant() { @@ -741,10 +783,13 @@ fn stdin_eof_cancels_an_active_interpreter_and_its_descendant() { "active-probe shutdown failed: {status}; stderr: {}", client.stderr_output() ); + // Observe OS lease release within the same budget as server shutdown. + wait_for_descendant_lease( + &lease, + Duration::from_secs(4).saturating_sub(started.elapsed()), + ) + .expect("shutdown must release the actual descendant's lease"); assert!(started.elapsed() < Duration::from_secs(4)); - lease - .try_lock() - .expect("shutdown must release the actual descendant's lease"); assert!( request.join().unwrap().is_err(), "an active request must be cancelled, not reported as successful" From 0fee3760d1edec98547c852892679dd37e1e4806 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Fri, 25 Sep 2026 14:57:55 -0700 Subject: [PATCH 8/9] test: prove lease retry after observed contention (Refs #529) Release the independent holder only after the same polling invocation observes WouldBlock, and assert exactly two real lock attempts. This removes scheduling assumptions and rejects an immediate-timeout mutation without changing shutdown semantics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet/tests/jsonrpc_server_test.rs | 29 ++++++++++++++++++++----- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/crates/pet/tests/jsonrpc_server_test.rs b/crates/pet/tests/jsonrpc_server_test.rs index 7dd1cd79..3c75670e 100644 --- a/crates/pet/tests/jsonrpc_server_test.rs +++ b/crates/pet/tests/jsonrpc_server_test.rs @@ -676,10 +676,13 @@ fn stdin_eof_exits_while_output_is_not_drained() { reader.join().unwrap(); } -fn wait_for_descendant_lease(lease: &fs::File, timeout: Duration) -> std::io::Result<()> { +fn wait_for_descendant_lease( + mut try_lock: impl FnMut() -> Result<(), fs::TryLockError>, + timeout: Duration, +) -> std::io::Result<()> { let started = Instant::now(); loop { - match lease.try_lock() { + match try_lock() { Ok(()) => return Ok(()), Err(fs::TryLockError::Error(error)) => return Err(error), Err(fs::TryLockError::WouldBlock) => {} @@ -708,14 +711,28 @@ fn descendant_lease_wait_is_bounded_and_requires_release() { .unwrap(); let started = Instant::now(); assert_eq!( - wait_for_descendant_lease(&lease, Duration::from_millis(20)) + wait_for_descendant_lease(|| lease.try_lock(), Duration::from_millis(20)) .unwrap_err() .kind(), std::io::ErrorKind::TimedOut ); assert!(started.elapsed() < Duration::from_secs(1)); - drop(holder); - wait_for_descendant_lease(&lease, Duration::ZERO).unwrap(); + let mut holder = Some(holder); + let mut attempts = 0; + wait_for_descendant_lease( + || { + attempts += 1; + let result = lease.try_lock(); + if attempts == 1 { + assert!(matches!(result, Err(fs::TryLockError::WouldBlock))); + drop(holder.take()); + } + result + }, + Duration::from_secs(1), + ) + .expect("lease polling must observe release after initial contention"); + assert_eq!(attempts, 2); } #[cfg(feature = "ci")] @@ -785,7 +802,7 @@ fn stdin_eof_cancels_an_active_interpreter_and_its_descendant() { ); // Observe OS lease release within the same budget as server shutdown. wait_for_descendant_lease( - &lease, + || lease.try_lock(), Duration::from_secs(4).saturating_sub(started.elapsed()), ) .expect("shutdown must release the actual descendant's lease"); From 3b2eec8e52aabcaf3163778ad6d73aa9acd967e0 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Mon, 28 Sep 2026 10:37:14 -0700 Subject: [PATCH 9/9] test: prove dispatch recovery after invalid envelopes (Refs #532) Exercise valid requests and notifications after rejection checks, proving handler wiring and exact ID/parameter preservation without adding errors. Covers previously unexecuted negative-test callbacks without changing production behavior or coverage budgets. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/pet-jsonrpc/src/server.rs | 38 +++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/crates/pet-jsonrpc/src/server.rs b/crates/pet-jsonrpc/src/server.rs index b6f320a4..a88fb6e8 100644 --- a/crates/pet-jsonrpc/src/server.rs +++ b/crates/pet-jsonrpc/src/server.rs @@ -273,6 +273,30 @@ mod tests { }) } + fn assert_valid_dispatch_after_invalid_input( + handlers: &HandlersKeyedByMethodName, + context: &TestContext, + ) { + assert!(context.request.lock().unwrap().is_none()); + assert!(context.notification.lock().unwrap().is_none()); + let error_count = context.errors.lock().unwrap().len(); + handlers.handle_request( + json!({"jsonrpc": "2.0", "id": "after-invalid", "method": "method", "params": [42]}), + ); + assert_eq!( + context.request.lock().unwrap().take(), + Some((RequestId::String("after-invalid".into()), json!([42]))) + ); + assert!(context.notification.lock().unwrap().is_none()); + handlers.handle_request(json!({"jsonrpc": "2.0", "method": "method", "params": [7]})); + assert_eq!( + context.notification.lock().unwrap().take(), + Some(json!([7])) + ); + assert!(context.request.lock().unwrap().is_none()); + assert_eq!(context.errors.lock().unwrap().len(), error_count); + } + fn supported_request_id_values() -> [Value; 10] { [ json!("request-1"), @@ -392,6 +416,7 @@ mod tests { context.errors.lock().unwrap().as_slice(), vec![(None, -32600, "Invalid JSONRPC request".to_string()); 6] ); + assert_valid_dispatch_after_invalid_input(&handlers, &context); } #[test] @@ -422,6 +447,7 @@ mod tests { context.errors.lock().unwrap().as_slice(), vec![(None, -32600, "Invalid JSONRPC request".to_string()); 6] ); + assert_valid_dispatch_after_invalid_input(&handlers, &context); } #[test] @@ -449,12 +475,21 @@ mod tests { "Invalid JSONRPC request".to_string() ), ( - Some(id), + Some(id.clone()), -32602, "JSONRPC params must be an object or array".to_string() ) ] ); + handlers.handle_request( + json!({"jsonrpc": "2.0", "id": value, "method": "method", "params": [42]}), + ); + assert_eq!( + context.request.lock().unwrap().take(), + Some((id, json!([42]))) + ); + assert!(context.notification.lock().unwrap().is_none()); + assert_eq!(context.errors.lock().unwrap().len(), 2); } } @@ -527,6 +562,7 @@ mod tests { (Some(2.into()), -32602) ] ); + assert_valid_dispatch_after_invalid_input(&handlers, &context); } #[test]