From 02d57721ca676c8c144f932393cb79dc8bf24082 Mon Sep 17 00:00:00 2001 From: Dotta Date: Tue, 25 Aug 2026 16:40:41 -0500 Subject: [PATCH] fix(runnerd): preserve partial websocket frames --- .../crates/runner-core/src/durable/runner.rs | 25 +++++++ .../runner-core/src/durable/transport.rs | 74 +++++++++++++++---- 2 files changed, 85 insertions(+), 14 deletions(-) diff --git a/packages/paperclip-runner/runner/crates/runner-core/src/durable/runner.rs b/packages/paperclip-runner/runner/crates/runner-core/src/durable/runner.rs index e01d287443d..9de1c825b69 100644 --- a/packages/paperclip-runner/runner/crates/runner-core/src/durable/runner.rs +++ b/packages/paperclip-runner/runner/crates/runner-core/src/durable/runner.rs @@ -3272,6 +3272,31 @@ mod tests { assert!(error.to_string().contains("control frame exceeds 125 bytes")); } + #[test] + fn runner_listener_preserves_a_partially_received_frame_across_poll_timeouts() { + let (mut client, server) = listener_client( + "GET /api/runner/v1/connect/run-1 HTTP/1.1\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: dGVzdA==\r\nSec-WebSocket-Version: 13\r\n\r\n", + ); + let mut accepted = WsClient::accept( + server, + "/api/runner/v1/connect/run-1", + 1024, + ) + .unwrap(); + let _ = read_http_headers(&mut client, "test response").unwrap(); + let frame = encode_frame(0x1, br#"{"crosses":"timeout"}"#, Some([1, 2, 3, 4]), 1024) + .unwrap(); + let split = frame.len() - 4; + let writer = thread::spawn(move || { + client.write_all(&frame[..split]).unwrap(); + thread::sleep(Duration::from_millis(400)); + client.write_all(&frame[split..]).unwrap(); + }); + let value = accepted.receive_plain_json().unwrap().unwrap(); + writer.join().unwrap(); + assert_eq!(value, json!({ "crosses": "timeout" })); + } + #[test] fn revoked_runner_rejects_turn_with_zero_effects_and_preserves_existing_outbox() { let root = temporary_root("revoked-command"); diff --git a/packages/paperclip-runner/runner/crates/runner-core/src/durable/transport.rs b/packages/paperclip-runner/runner/crates/runner-core/src/durable/transport.rs index 04bd4ba2e39..f0bba11baa4 100644 --- a/packages/paperclip-runner/runner/crates/runner-core/src/durable/transport.rs +++ b/packages/paperclip-runner/runner/crates/runner-core/src/durable/transport.rs @@ -194,6 +194,54 @@ struct WsClient { secure_channel: Option, } +const IN_PROGRESS_FRAME_TIMEOUT: Duration = Duration::from_secs(30); + +fn read_frame_bytes( + stream: &mut dyn WsReadWrite, + bytes: &mut [u8], + label: &str, +) -> Result<(), DurableRunnerError> { + let deadline = Instant::now() + .checked_add(IN_PROGRESS_FRAME_TIMEOUT) + .ok_or_else(|| DurableRunnerError::invalid("WebSocket frame deadline overflowed"))?; + let mut offset = 0; + while offset < bytes.len() { + match stream.read(&mut bytes[offset..]) { + Ok(0) => { + return Err(DurableRunnerError::invalid( + "WebSocket connection closed while receiving a frame", + )) + } + Ok(read) => offset += read, + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) && Instant::now() < deadline => + { + continue; + } + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) => + { + return Err(DurableRunnerError::invalid(format!( + "WebSocket {label} timed out while the frame was in progress", + ))) + } + Err(error) => { + return Err(DurableRunnerError::invalid(format!( + "WebSocket connection closed while receiving {label}: {error}", + ))) + } + } + } + Ok(()) +} + struct SecureChannel { send_cipher: Aes256Gcm, receive_cipher: Aes256Gcm, @@ -625,8 +673,14 @@ impl WsClient { fn receive_plain_json(&mut self) -> Result, DurableRunnerError> { loop { let mut header = [0_u8; 2]; - match self.stream.read_exact(&mut header) { - Ok(()) => {} + match self.stream.read(&mut header[..1]) { + Ok(0) => { + return Err(DurableRunnerError::invalid( + "WebSocket connection closed while receiving a frame", + )) + } + Ok(_) => read_frame_bytes(self.stream.as_mut(), &mut header[1..], "frame header")?, + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, Err(error) if matches!( error.kind(), @@ -656,15 +710,11 @@ impl WsClient { let mut length = u64::from(header[1] & 0x7f); if length == 126 { let mut extended = [0_u8; 2]; - self.stream - .read_exact(&mut extended) - .map_err(|error| DurableRunnerError::invalid(error.to_string()))?; + read_frame_bytes(self.stream.as_mut(), &mut extended, "extended frame length")?; length = u64::from(u16::from_be_bytes(extended)); } else if length == 127 { let mut extended = [0_u8; 8]; - self.stream - .read_exact(&mut extended) - .map_err(|error| DurableRunnerError::invalid(error.to_string()))?; + read_frame_bytes(self.stream.as_mut(), &mut extended, "extended frame length")?; length = u64::from_be_bytes(extended); } let length = checked_inbound_frame_length(length, self.max_frame_bytes)?; @@ -675,14 +725,10 @@ impl WsClient { } let mut mask = [0_u8; 4]; if masked { - self.stream - .read_exact(&mut mask) - .map_err(|error| DurableRunnerError::invalid(error.to_string()))?; + read_frame_bytes(self.stream.as_mut(), &mut mask, "frame mask")?; } let mut payload = vec![0_u8; length]; - self.stream - .read_exact(&mut payload) - .map_err(|error| DurableRunnerError::invalid(error.to_string()))?; + read_frame_bytes(self.stream.as_mut(), &mut payload, "frame payload")?; if masked { for (index, byte) in payload.iter_mut().enumerate() { *byte ^= mask[index % 4];