From e0aac2049af970ac6ee73ef755bf704ea54ba128 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 10 Jul 2026 06:29:59 -0700 Subject: [PATCH 1/2] fix(envoy-client): remove native websocket input limits --- .../envoy-client/src/connection/native.rs | 64 ++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/engine/sdks/rust/envoy-client/src/connection/native.rs b/engine/sdks/rust/envoy-client/src/connection/native.rs index 69ef21b2ad..372ca1ddc3 100644 --- a/engine/sdks/rust/envoy-client/src/connection/native.rs +++ b/engine/sdks/rust/envoy-client/src/connection/native.rs @@ -21,6 +21,12 @@ pub fn start_connection(shared: Arc) { tokio::spawn(connection_loop(shared).instrument(span)); } +fn websocket_config() -> tungstenite::protocol::WebSocketConfig { + tungstenite::protocol::WebSocketConfig::default() + .max_message_size(None) + .max_frame_size(None) +} + async fn connection_loop(shared: Arc) { let mut attempt = 0u32; @@ -106,7 +112,12 @@ async fn single_connection( .body(()) .map_err(|e| anyhow::anyhow!("failed to build ws request: {e}"))?; - let (ws_stream, _) = tokio_tungstenite::connect_async(request).await?; + let (ws_stream, _) = tokio_tungstenite::connect_async_with_config( + request, + Some(websocket_config()), + false, + ) + .await?; let (mut write, mut read) = ws_stream.split(); let (ws_tx, mut ws_rx) = mpsc::unbounded_channel::(); @@ -239,3 +250,54 @@ fn extract_host(url: &str) -> String { .unwrap_or("localhost") .to_string() } + +#[cfg(test)] +mod tests { + use futures_util::{SinkExt, StreamExt}; + use tokio::net::TcpListener; + + use super::*; + + #[test] + fn websocket_config_has_no_input_size_limits() { + let config = websocket_config(); + + assert_eq!(config.max_frame_size, None); + assert_eq!(config.max_message_size, None); + } + + #[tokio::test] + async fn receives_frame_larger_than_default_limit() { + const FRAME_SIZE: usize = (16 << 20) + 64 * 1024; + const FRAME_BYTE: u8 = 0xa5; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(stream).await.unwrap(); + websocket + .send(tungstenite::Message::Binary( + vec![FRAME_BYTE; FRAME_SIZE].into(), + )) + .await + .unwrap(); + }); + + let (mut client, _) = tokio_tungstenite::connect_async_with_config( + format!("ws://{addr}"), + Some(websocket_config()), + false, + ) + .await + .unwrap(); + let message = client.next().await.unwrap().unwrap(); + let tungstenite::Message::Binary(data) = message else { + panic!("expected a binary websocket message"); + }; + + assert_eq!(data.len(), FRAME_SIZE); + assert!(data.iter().all(|&byte| byte == FRAME_BYTE)); + server.await.unwrap(); + } +} From 26a111c735fb5fa2111c5f4aee160dfb8f560271 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 10 Jul 2026 15:29:08 -0700 Subject: [PATCH 2/2] [SLOP(gpt-5)] test(envoy-client): reproduce oversized guard payload disconnect --- .../engine/tests/common/test_helpers.rs | 9 +- .../engine/tests/envoy_large_payload.rs | 112 ++++++++++++++++++ .../envoy-client/src/connection/native.rs | 9 +- 3 files changed, 122 insertions(+), 8 deletions(-) create mode 100644 engine/packages/engine/tests/envoy_large_payload.rs diff --git a/engine/packages/engine/tests/common/test_helpers.rs b/engine/packages/engine/tests/common/test_helpers.rs index ca19c1719b..4e875981c4 100644 --- a/engine/packages/engine/tests/common/test_helpers.rs +++ b/engine/packages/engine/tests/common/test_helpers.rs @@ -306,9 +306,14 @@ pub async fn upsert_normal_runner_config( datacenters.insert( dc.config.dc_name().unwrap().to_string(), rivet_api_types::namespaces::runner_configs::RunnerConfig { - kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {}, + kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal { + drain_on_version_upgrade: None, + actor_eviction_delay: None, + actor_eviction_period: None, + actor_eviction_rate: None, + }, metadata: None, - drain_on_version_upgrade: true, + drain_on_version_upgrade: Some(true), }, ); diff --git a/engine/packages/engine/tests/envoy_large_payload.rs b/engine/packages/engine/tests/envoy_large_payload.rs new file mode 100644 index 0000000000..975f4bd93c --- /dev/null +++ b/engine/packages/engine/tests/envoy_large_payload.rs @@ -0,0 +1,112 @@ +#[path = "common/mod.rs"] +mod common; + +use futures_util::{SinkExt, StreamExt}; +use tokio_tungstenite::tungstenite::{ + Message, client::IntoClientRequest, protocol::WebSocketConfig, +}; + +const DEFAULT_TUNGSTENITE_MESSAGE_LIMIT: usize = 16 << 20; +const BELOW_DEFAULT_LIMIT: usize = DEFAULT_TUNGSTENITE_MESSAGE_LIMIT - 64 * 1024; +const ABOVE_DEFAULT_LIMIT: usize = DEFAULT_TUNGSTENITE_MESSAGE_LIMIT + 64 * 1024; + +#[test] +fn envoy_websocket_accepts_payloads_across_default_tungstenite_limit() { + common::run( + common::TestOpts::new(1).with_timeout(60), + |ctx| async move { + let (namespace, _) = common::setup_test_namespace(ctx.leader_dc()).await; + let envoy = common::setup_envoy(ctx.leader_dc(), &namespace, |builder| { + builder.with_actor_behavior("test-actor", |_| { + Box::new(common::test_envoy::EchoActor::new()) + }) + }) + .await; + + let res = common::create_actor( + ctx.leader_dc().guard_port(), + &namespace, + "test-actor", + envoy.pool_name(), + rivet_types::actors::CrashPolicy::Sleep, + ) + .await; + let actor_id = res.actor.actor_id.to_string(); + wait_for_envoy_actor(&envoy, &actor_id).await; + + let mut request = format!("ws://127.0.0.1:{}/ws", ctx.leader_dc().guard_port()) + .into_client_request() + .expect("failed to create WebSocket request"); + request.headers_mut().insert( + "Sec-WebSocket-Protocol", + format!( + "rivet, rivet_target.actor, rivet_actor.{}", + urlencoding::encode(&actor_id) + ) + .parse() + .unwrap(), + ); + + let websocket_config = WebSocketConfig::default() + .max_message_size(None) + .max_frame_size(None); + let (ws_stream, response) = tokio_tungstenite::connect_async_with_config( + request, + Some(websocket_config), + false, + ) + .await + .expect("failed to connect WebSocket through guard"); + assert_eq!(response.status(), 101); + let (mut write, mut read) = ws_stream.split(); + + let mut disconnect = envoy.wait_for_next_connection_event( + common::test_envoy::EnvoyConnectionEvent::Disconnected, + ); + disconnect.assert_no_event(); + + for payload_size in [BELOW_DEFAULT_LIMIT, ABOVE_DEFAULT_LIMIT] { + tracing::info!(payload_size, "sending WebSocket payload through guard"); + write + .send(Message::Binary(vec![b'x'; payload_size].into())) + .await + .expect("failed to send WebSocket payload through guard"); + + let response = + tokio::time::timeout(std::time::Duration::from_secs(20), read.next()) + .await + .unwrap_or_else(|_| { + panic!("timed out waiting for {payload_size}-byte WebSocket echo") + }) + .unwrap_or_else(|| { + panic!( + "WebSocket stream ended before {payload_size}-byte payload was echoed" + ) + }) + .expect("failed to receive WebSocket echo"); + + let Message::Text(response) = response else { + panic!("expected text echo, got {response:?}"); + }; + let response = response.as_bytes(); + assert_eq!(response.len(), "Echo: ".len() + payload_size); + assert_eq!(&response[.."Echo: ".len()], b"Echo: "); + assert!(response["Echo: ".len()..].iter().all(|byte| *byte == b'x')); + disconnect.assert_no_event(); + } + }, + ); +} + +async fn wait_for_envoy_actor(envoy: &common::test_envoy::TestEnvoy, actor_id: &str) { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if envoy.has_actor(actor_id).await { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + }) + .await + .expect("envoy should receive actor"); +} diff --git a/engine/sdks/rust/envoy-client/src/connection/native.rs b/engine/sdks/rust/envoy-client/src/connection/native.rs index 372ca1ddc3..36c7fa0c63 100644 --- a/engine/sdks/rust/envoy-client/src/connection/native.rs +++ b/engine/sdks/rust/envoy-client/src/connection/native.rs @@ -112,12 +112,9 @@ async fn single_connection( .body(()) .map_err(|e| anyhow::anyhow!("failed to build ws request: {e}"))?; - let (ws_stream, _) = tokio_tungstenite::connect_async_with_config( - request, - Some(websocket_config()), - false, - ) - .await?; + let (ws_stream, _) = + tokio_tungstenite::connect_async_with_config(request, Some(websocket_config()), false) + .await?; let (mut write, mut read) = ws_stream.split(); let (ws_tx, mut ws_rx) = mpsc::unbounded_channel::();