fix(guard, runner): enforce http body size and ws msg size

Author: MasterPtato

engine/artifacts/errors/guard.invalid_request_body.json

@@ -1,6 +1,24 @@
 use rivet_error::*;
 use serde::{Deserialize, Serialize};
 
+#[derive(RivetError, Serialize, Deserialize)]
+#[error(
+	"guard",
+	"invalid_request_body",
+	"Unable to parse request body.",
+	"Unable to parse request body: {0}."
+)]
+pub struct InvalidRequestBody(pub String);
+
+#[derive(RivetError, Serialize, Deserialize)]
+#[error(
+	"guard",
+	"invalid_response_body",
+	"Unable to parse response body.",
+	"Unable to parse response body: {0}."
+)]
+pub struct InvalidResponseBody(pub String);
+
 #[derive(RivetError, Serialize, Deserialize)]
 #[error(
 	"guard",
@@ -42,7 +60,7 @@ pub struct UriParseError(pub String);
 pub struct RequestBuildError(pub String);
 
 #[derive(RivetError)]
-#[error("guard", "upstream_error", "Upstream error.", "Upstream error: {0}")]
+#[error("guard", "upstream_error", "Upstream error.", "Upstream error: {0}.")]
 pub struct UpstreamError(pub String);
 
 #[derive(RivetError, Serialize, Deserialize)]

engine/artifacts/errors/guard.invalid_response_body.json

@@ -1,7 +1,7 @@
 use anyhow::{Context, Result, bail, ensure};
 use bytes::Bytes;
 use futures_util::{SinkExt, StreamExt};
-use http_body_util::{BodyExt, Full};
+use http_body_util::{BodyExt, Full, Limited};
 use hyper::{
 	Request, Response, StatusCode,
 	body::Incoming as BodyIncoming,
@@ -40,6 +40,7 @@ use crate::{
 
 pub const X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for");
 pub const X_RIVET_ERROR: HeaderName = HeaderName::from_static("x-rivet-error");
+pub const MAX_BODY_SIZE: usize = rivet_util::size::mebibytes(20) as usize;
 
 const PROXY_STATE_CACHE_TTL: Duration = Duration::from_secs(60 * 60); // 1 hour
 const WEBSOCKET_CLOSE_LINGER: Duration = Duration::from_millis(100); // Keep TCP connection open briefly after WebSocket close
@@ -723,13 +724,11 @@ impl ProxyService {
 			ResolveRouteOutput::Target(mut target) => {
 				// Read the request body before proceeding with retries
 				let (req_parts, body) = req.into_parts();
-				let req_body = match http_body_util::BodyExt::collect(body).await {
-					Ok(collected) => collected.to_bytes(),
-					Err(err) => {
-						tracing::debug!(?err, "Failed to read request body");
-						Bytes::new()
-					}
-				};
+				let req_body = Limited::new(body, MAX_BODY_SIZE)
+					.collect()
+					.await
+					.map_err(|err| errors::InvalidRequestBody(err.to_string()).build())?
+					.to_bytes();
 
 				// Use a value-returning loop to handle both errors and successful responses
 				let mut attempts = 0;
@@ -742,7 +741,8 @@ impl ProxyService {
 
 					// Create the final request with body
 					let proxied_req = builder
-						.body(Full::<Bytes>::new(req_body.clone()))
+						// NOTE: the `Bytes` type is cheaply cloneable, this is not resource intensive
+						.body(Full::new(req_body.clone()))
 						.map_err(|err| errors::RequestBuildError(err.to_string()).build())?;
 
 					// Send the request with timeout
@@ -800,10 +800,13 @@ impl ProxyService {
 								return Ok(Response::from_parts(parts, streaming_body));
 							} else {
 								// For non-streaming responses, buffer as before
-								let body_bytes = match BodyExt::collect(body).await {
-									Ok(collected) => collected.to_bytes(),
-									Err(_) => Bytes::new(),
-								};
+								let body_bytes = Limited::new(body, MAX_BODY_SIZE)
+									.collect()
+									.await
+									.map_err(|err| {
+										errors::InvalidResponseBody(err.to_string()).build()
+									})?
+									.to_bytes();
 
 								let full_body = ResponseBody::Full(Full::new(body_bytes));
 								return Ok(Response::from_parts(parts, full_body));
@@ -857,15 +860,13 @@ impl ProxyService {
 			ResolveRouteOutput::CustomServe(mut handler) => {
 				// Collect request body
 				let (req_parts, body) = req.into_parts();
-				let collected_body = match http_body_util::BodyExt::collect(body).await {
-					Ok(collected) => collected.to_bytes(),
-					Err(err) => {
-						tracing::debug!(?err, "Failed to read request body");
-						Bytes::new()
-					}
-				};
+				let req_body = Limited::new(body, MAX_BODY_SIZE)
+					.collect()
+					.await
+					.map_err(|err| errors::InvalidRequestBody(err.to_string()).build())?
+					.to_bytes();
 				let req_collected =
-					hyper::Request::from_parts(req_parts, Full::<Bytes>::new(collected_body));
+					hyper::Request::from_parts(req_parts, Full::<Bytes>::new(req_body));
 
 				// Attempt request
 				let mut attempts = 0;

engine/packages/guard-core/src/errors.rs

@@ -181,6 +181,8 @@ pub(crate) fn err_into_response(err: anyhow::Error) -> Result<Response<ResponseB
 				("guard", "service_unavailable") => StatusCode::SERVICE_UNAVAILABLE,
 				("guard", "actor_ready_timeout") => StatusCode::SERVICE_UNAVAILABLE,
 				("guard", "no_route") => StatusCode::NOT_FOUND,
+				("guard", "invalid_request_body") => StatusCode::PAYLOAD_TOO_LARGE,
+				("guard", "invalid_response_body") => StatusCode::PAYLOAD_TOO_LARGE,
 				_ => StatusCode::BAD_REQUEST,
 			};
 

engine/packages/guard-core/src/proxy_service.rs

@@ -157,6 +157,7 @@ impl PegboardGateway {
 			max_age: None,
 		});
 
+		// NOTE: Size constraints have already been applied by guard
 		let body_bytes = req
 			.into_body()
 			.collect()

engine/packages/guard-core/src/utils.rs

@@ -6,6 +6,7 @@ use gas::prelude::*;
 use hyper_tungstenite::tungstenite::Message;
 use pegboard::actor_kv;
 use pegboard::pubsub_subjects::GatewayReceiverSubject;
+use rivet_guard_core::proxy_service::MAX_BODY_SIZE;
 use rivet_guard_core::websocket_handle::WebSocketReceiver;
 use rivet_runner_protocol::{self as protocol, PROTOCOL_MK2_VERSION, versioned};
 use std::sync::{Arc, atomic::Ordering};
@@ -783,12 +784,15 @@ async fn handle_tunnel_message_mk2(
 	ctx: &StandaloneCtx,
 	msg: protocol::mk2::ToServerTunnelMessage,
 ) -> Result<()> {
-	// Publish message to UPS
-	let gateway_reply_to = GatewayReceiverSubject::new(msg.message_id.gateway_id).to_string();
-
 	// Extract inner data length before consuming msg
-	let inner_data_len = tunnel_message_inner_data_len(&msg.message_kind);
+	let inner_data_len = tunnel_message_inner_data_len_mk2(&msg.message_kind);
 
+	// Enforce incoming payload size
+	if inner_data_len > MAX_BODY_SIZE {
+		return Err(errors::WsError::InvalidPacket(format!("payload too large")).build());
+	}
+
+	let gateway_reply_to = GatewayReceiverSubject::new(msg.message_id.gateway_id).to_string();
 	let msg_serialized =
 		versioned::ToGateway::wrap_latest(protocol::mk2::ToGateway::ToServerTunnelMessage(msg))
 			.serialize_with_embedded_version(PROTOCOL_MK2_VERSION)
@@ -800,6 +804,7 @@ async fn handle_tunnel_message_mk2(
 		"publishing tunnel message to gateway"
 	);
 
+	// Publish message to UPS
 	ctx.ups()
 		.context("failed to get UPS instance for tunnel message")?
 		.publish(&gateway_reply_to, &msg_serialized, PublishOpts::one())
@@ -814,22 +819,6 @@ async fn handle_tunnel_message_mk2(
 	Ok(())
 }
 
-/// Returns the length of the inner data payload for a tunnel message kind.
-fn tunnel_message_inner_data_len(kind: &protocol::mk2::ToServerTunnelMessageKind) -> usize {
-	use protocol::mk2::ToServerTunnelMessageKind;
-	match kind {
-		ToServerTunnelMessageKind::ToServerResponseStart(resp) => {
-			resp.body.as_ref().map_or(0, |b| b.len())
-		}
-		ToServerTunnelMessageKind::ToServerResponseChunk(chunk) => chunk.body.len(),
-		ToServerTunnelMessageKind::ToServerWebSocketMessage(msg) => msg.data.len(),
-		ToServerTunnelMessageKind::ToServerResponseAbort
-		| ToServerTunnelMessageKind::ToServerWebSocketOpen(_)
-		| ToServerTunnelMessageKind::ToServerWebSocketMessageAck(_)
-		| ToServerTunnelMessageKind::ToServerWebSocketClose(_) => 0,
-	}
-}
-
 #[tracing::instrument(skip_all)]
 async fn handle_tunnel_message_mk1(
 	ctx: &StandaloneCtx,
@@ -843,6 +832,14 @@ async fn handle_tunnel_message_mk1(
 		return Ok(());
 	}
 
+	// Extract inner data length before consuming msg
+	let inner_data_len = tunnel_message_inner_data_len_mk1(&msg.message_kind);
+
+	// Enforce incoming payload size
+	if inner_data_len > MAX_PAYLOAD_SIZE {
+		return Err(errors::WsError::InvalidPacket(format!("payload too large")).build());
+	}
+
 	// Publish message to UPS
 	let gateway_reply_to = GatewayReceiverSubject::new(msg.message_id.gateway_id).to_string();
 	let msg_serialized = versioned::ToGateway::v3_to_v6(versioned::ToGateway::V3(
@@ -864,6 +861,39 @@ async fn handle_tunnel_message_mk1(
 	Ok(())
 }
 
+/// Returns the length of the inner data payload for a tunnel message kind.
+fn tunnel_message_inner_data_len_mk2(kind: &protocol::mk2::ToServerTunnelMessageKind) -> usize {
+	use protocol::mk2::ToServerTunnelMessageKind;
+	match kind {
+		ToServerTunnelMessageKind::ToServerResponseStart(resp) => {
+			resp.body.as_ref().map_or(0, |b| b.len())
+		}
+		ToServerTunnelMessageKind::ToServerResponseChunk(chunk) => chunk.body.len(),
+		ToServerTunnelMessageKind::ToServerWebSocketMessage(msg) => msg.data.len(),
+		ToServerTunnelMessageKind::ToServerResponseAbort
+		| ToServerTunnelMessageKind::ToServerWebSocketOpen(_)
+		| ToServerTunnelMessageKind::ToServerWebSocketMessageAck(_)
+		| ToServerTunnelMessageKind::ToServerWebSocketClose(_) => 0,
+	}
+}
+
+/// Returns the length of the inner data payload for a tunnel message kind.
+fn tunnel_message_inner_data_len_mk1(kind: &protocol::ToServerTunnelMessageKind) -> usize {
+	use protocol::ToServerTunnelMessageKind;
+	match kind {
+		ToServerTunnelMessageKind::ToServerResponseStart(resp) => {
+			resp.body.as_ref().map_or(0, |b| b.len())
+		}
+		ToServerTunnelMessageKind::ToServerResponseChunk(chunk) => chunk.body.len(),
+		ToServerTunnelMessageKind::ToServerWebSocketMessage(msg) => msg.data.len(),
+		ToServerTunnelMessageKind::ToServerResponseAbort
+		| ToServerTunnelMessageKind::ToServerWebSocketOpen(_)
+		| ToServerTunnelMessageKind::ToServerWebSocketMessageAck(_)
+		| ToServerTunnelMessageKind::ToServerWebSocketClose(_)
+		| ToServerTunnelMessageKind::DeprecatedTunnelAck => 0,
+	}
+}
+
 /// Send ack message for deprecated tunnel versions.
 ///
 /// We have to parse as specifically a v2 message since we need the exact request & message ID

engine/packages/pegboard-gateway/src/lib.rs

@@ -15,7 +15,7 @@ import {
 	stringifyToClientTunnelMessageKind,
 	stringifyToServerTunnelMessageKind,
 } from "./stringify";
-import { arraysEqual, idToStr, stringifyError, unreachable } from "./utils";
+import { arraysEqual, idToStr, MAX_BODY_SIZE, stringifyError, unreachable } from "./utils";
 import {
 	HIBERNATABLE_SYMBOL,
 	WebSocketTunnelAdapter,
@@ -855,6 +855,10 @@ export class Tunnel {
 		// Read the body first to get the actual content
 		const body = response.body ? await response.arrayBuffer() : null;
 
+		if (body && body.byteLength > MAX_BODY_SIZE) {
+			throw new Error("Response body too large");
+		}
+
 		// Convert headers to map and add Content-Length if not present
 		const headers = new Map<string, string>();
 		response.headers.forEach((value, key) => {
@@ -1079,7 +1083,7 @@ export class Tunnel {
 		});
 
 		if (clientMessageIndex < 0 || clientMessageIndex > 65535)
-			throw new Error("invalid websocket ack index");
+			throw new Error("Invalid websocket ack index");
 
 		// Get the actor to find the gatewayId
 		//
@@ -1157,7 +1161,7 @@ function buildRequestForWebSocket(
 	};
 
 	if (!path.startsWith("/")) {
-		throw new Error("path must start with leading slash");
+		throw new Error("Path must start with leading slash");
 	}
 
 	const request = new Request(`http://actor${path}`, {

engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs

@@ -1,5 +1,8 @@
 import { logger } from "./log";
 
+// 20MiB. Keep in sync with MAX_BODY_SIZE from engine/packages/guard-core/src/proxy_service.rs
+export const MAX_BODY_SIZE = 20 * 1024 * 1024;
+
 export function unreachable(x: never): never {
 	throw `Unreachable: ${x}`;
 }

engine/sdks/typescript/runner/src/tunnel.ts

@@ -1,7 +1,7 @@
 import type { Logger } from "pino";
 import { VirtualWebSocket, type UniversalWebSocket, type RivetMessageEvent } from "@rivetkit/virtual-websocket";
 import type { Tunnel } from "./tunnel";
-import { wrappingAddU16, wrappingLteU16, wrappingSubU16 } from "./utils";
+import { MAX_BODY_SIZE, wrappingAddU16, wrappingLteU16, wrappingSubU16 } from "./utils";
 
 export const HIBERNATABLE_SYMBOL = Symbol("hibernatable");
 
@@ -70,11 +70,20 @@ export class WebSocketTunnelAdapter {
 		let messageData: string | ArrayBuffer;
 
 		if (typeof data === "string") {
+			let encoder = new TextEncoder();
+			if (encoder.encode(data).byteLength > MAX_BODY_SIZE) {
+				throw new Error("WebSocket message too large");
+			}
+
 			messageData = data;
 		} else if (data instanceof ArrayBuffer) {
+			if (data.byteLength > MAX_BODY_SIZE) throw new Error("WebSocket message too large");
+
 			isBinary = true;
 			messageData = data;
 		} else if (ArrayBuffer.isView(data)) {
+			if (data.byteLength > MAX_BODY_SIZE) throw new Error("WebSocket message too large");
+
 			isBinary = true;
 			const view = data;
 			const buffer = view.buffer instanceof SharedArrayBuffer