Merge pull request #862 from schoeller/main

request_with_poll TCP desync fix

crates/ocs_plugin_api/src/ipc/proxy.rs

@@ -79,6 +79,78 @@ pub fn recv_framed<R: Read, T: serde::de::DeserializeOwned>(reader: &mut R) -> R
     Ok(msg)
 }
 
+/// Error during a polled read: either an IO error on the stream, or an
+/// interrupt reported by the poll callback (e.g. the user pressed Ctrl+C).
+#[derive(Debug)]
+enum PollError {
+    Io(ProxyError),
+    Interrupted(PluginRequestError),
+}
+
+impl From<ProxyError> for PollError {
+    fn from(e: ProxyError) -> Self {
+        PollError::Io(e)
+    }
+}
+
+impl From<PollError> for PluginRequestError {
+    fn from(e: PollError) -> Self {
+        match e {
+            PollError::Io(err) => err.into(),
+            PollError::Interrupted(err) => err,
+        }
+    }
+}
+
+/// 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<R: Read>(
+    reader: &mut R,
+    buf: &mut [u8],
+    poll: &mut dyn FnMut() -> Result<(), PluginRequestError>,
+) -> Result<(), PollError> {
+    use std::io::ErrorKind;
+    let mut off = 0;
+    while off < buf.len() {
+        match reader.read(&mut buf[off..]) {
+            Ok(0) => return Err(PollError::Io(ProxyError::Io(std::io::Error::new(
+                std::io::ErrorKind::UnexpectedEof,
+                "proxy disconnected",
+            )))),
+            Ok(n) => off += n,
+            Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => {
+                if let Err(err) = poll() {
+                    return Err(PollError::Interrupted(err));
+                }
+            }
+            Err(e) => return Err(PollError::Io(ProxyError::Io(e))),
+        }
+    }
+    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<R: Read>(
+    reader: &mut R,
+    poll: &mut dyn FnMut() -> Result<(), PluginRequestError>,
+) -> Result<PluginResponse, PollError> {
+    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 +185,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<PluginResponse, PluginRequestError> {
-        use std::io::ErrorKind;
         let mut stream = self
             .stream
             .lock()
@@ -130,21 +205,19 @@ 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 = match recv_framed_with_poll(&mut *stream, poll) {
+            Ok(resp) => Ok(resp),
+            Err(PollError::Interrupted(e)) => {
+                // The caller (e.g. Python's Ctrl+C handler) asked us to stop
+                // waiting. Shut down the stream immediately so the proxy server
+                // and any spawned runners can observe the disconnect.
+                let _ = stream.shutdown(Shutdown::Both);
+                Err(e)
             }
+            Err(PollError::Io(e)) => Err(e.into()),
         };
+
         let _ = stream.set_read_timeout(original_timeout);
         result
     }
@@ -331,4 +404,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<usize> {
+            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);
+    }
 }