From 4fd94d324a980911bb5d3016845770ec9321a99f Mon Sep 17 00:00:00 2001 From: Sebastian <106036+schoeller@users.noreply.github.com> Date: Sun, 23 Aug 2026 10:40:39 +0200 Subject: [PATCH] fix(ipc): read response byte-by-byte in request_with_poll to avoid TCP desync --- crates/ocs_plugin_api/src/ipc/proxy.rs | 120 +++++++++++++++++++++---- 1 file changed, 103 insertions(+), 17 deletions(-) diff --git a/crates/ocs_plugin_api/src/ipc/proxy.rs b/crates/ocs_plugin_api/src/ipc/proxy.rs index 1300e18d..e4d19825 100644 --- a/crates/ocs_plugin_api/src/ipc/proxy.rs +++ b/crates/ocs_plugin_api/src/ipc/proxy.rs @@ -11,7 +11,7 @@ //! formats or enum variants. use std::io::{Read, Write}; -use std::net::{Shutdown, TcpListener, TcpStream}; +use std::net::{TcpListener, TcpStream}; use std::sync::Arc; use std::sync::Mutex; use std::time::Duration; @@ -79,6 +79,50 @@ pub fn recv_framed(reader: &mut R) -> R Ok(msg) } +/// Read exactly `buf.len()` bytes from `reader`, calling `poll` whenever no +/// data is currently available. Unlike `read_exact`, this preserves partial +/// progress across short read timeouts so the stream can never get out of sync. +fn read_with_poll( + reader: &mut R, + buf: &mut [u8], + poll: &mut dyn FnMut() -> Result<(), PluginRequestError>, +) -> Result<(), PluginRequestError> { + use std::io::ErrorKind; + let mut off = 0; + while off < buf.len() { + match reader.read(&mut buf[off..]) { + Ok(0) => return Err(PluginRequestError("proxy disconnected".to_string())), + Ok(n) => off += n, + Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => { + poll()?; + } + Err(e) => return Err(ProxyError::Io(e).into()), + } + } + Ok(()) +} + +/// Receive a length-framed bincode message with polling. The length prefix and +/// payload are read incrementally so a timeout never splits a message. +fn recv_framed_with_poll( + reader: &mut R, + poll: &mut dyn FnMut() -> Result<(), PluginRequestError>, +) -> Result { + let mut len_buf = [0u8; 8]; + read_with_poll(reader, &mut len_buf, poll)?; + let len = u64::from_le_bytes(len_buf) as usize; + if len == 0 { + return Err(ProxyError::Empty.into()); + } + if len > MAX_MESSAGE_SIZE { + return Err(ProxyError::TooLarge(len).into()); + } + let mut buf = vec![0u8; len]; + read_with_poll(reader, &mut buf, poll)?; + let msg = bincode::deserialize(&buf).map_err(ProxyError::Encode)?; + Ok(msg) +} + /// Thread-safe [`PluginRequestSender`] implementation that forwards requests /// over a single TCP stream to a request proxy server. pub struct ProxyPluginRequestSender { @@ -113,12 +157,15 @@ impl ProxyPluginRequestSender { /// [`POLL_INTERVAL`] while waiting. This lets Python callers run /// `py.check_signals()` so that `Ctrl+C`/`KeyboardInterrupt` is raised /// promptly instead of waiting for the full socket timeout. + /// + /// The response is read byte-by-byte with polling so that a short timeout + /// can never leave the TCP stream in a partially-read state between the + /// length prefix and the payload. pub fn request_with_poll( &self, req: PluginRequest, poll: &mut dyn FnMut() -> Result<(), PluginRequestError>, ) -> Result { - use std::io::ErrorKind; let mut stream = self .stream .lock() @@ -130,21 +177,9 @@ impl ProxyPluginRequestSender { stream .set_read_timeout(Some(POLL_INTERVAL)) .map_err(|e| PluginRequestError(e.to_string()))?; - let result = loop { - match recv_framed(&mut *stream) { - Ok(resp) => break Ok(resp), - Err(ProxyError::Io(e)) - if e.kind() == ErrorKind::WouldBlock - || e.kind() == ErrorKind::TimedOut => - { - if let Err(e) = poll() { - let _ = stream.shutdown(Shutdown::Both); - break Err(e); - } - } - Err(e) => break Err(e.into()), - } - }; + + let result = recv_framed_with_poll(&mut *stream, poll); + let _ = stream.set_read_timeout(original_timeout); result } @@ -331,4 +366,55 @@ mod tests { .unwrap(); assert!(matches!(resp, PluginResponse::Ok)); } + + /// A reader that returns one byte per call and then `WouldBlock` on the + /// next call, exercising both the incremental read and the poll path. + struct OneByteThenBlock<'a> { + data: &'a [u8], + block_next: bool, + } + + impl<'a> Read for OneByteThenBlock<'a> { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + if self.block_next { + self.block_next = false; + return Err(std::io::Error::new( + std::io::ErrorKind::WouldBlock, + "simulated block", + )); + } + if self.data.is_empty() { + return Ok(0); + } + let n = buf.len().min(1); + buf[..n].copy_from_slice(&self.data[..n]); + self.data = &self.data[n..]; + self.block_next = true; + Ok(n) + } + } + + #[test] + fn recv_framed_with_poll_reads_one_byte_at_a_time() { + let msg = PluginResponse::Bool(true); + let bytes = bincode::serialize(&msg).unwrap(); + let len = bytes.len() as u64; + let mut framed = Vec::with_capacity(8 + bytes.len()); + framed.extend_from_slice(&len.to_le_bytes()); + framed.extend_from_slice(&bytes); + + let mut reader = OneByteThenBlock { + data: &framed, + block_next: false, + }; + let mut polls = 0; + let resp = recv_framed_with_poll(&mut reader, &mut || { + polls += 1; + Ok(()) + }) + .unwrap(); + assert!(matches!(resp, PluginResponse::Bool(true))); + // One byte, one block for every byte except the last (no poll after EOF). + assert_eq!(polls, framed.len() - 1); + } }