chore: tunnel auth

Author: NathanFlurry

CLAUDE.md

@@ -216,6 +216,14 @@ When the user asks to track something in a note, store it in `.agent/notes/` by
 - Any behavior, protocol handling, or test coverage added to one runner should be mirrored in the other runner in the same change whenever possible.
 - When parity cannot be completed in the same change, explicitly document the gap and add a follow-up task.
 
+### Trust Boundaries
+- Treat `client <-> engine` as untrusted.
+- Treat `envoy <-> pegboard-envoy` as untrusted.
+- Treat traffic inside the engine over `nats`, `fdb`, and other internal backends as trusted.
+- Treat `gateway`, `api`, `pegboard-envoy`, `nats`, `fdb`, and similar engine-internal services as one trusted internal boundary once traffic is inside the engine.
+- Validate and authorize all client-originated data at the engine edge before it reaches trusted internal systems.
+- Validate and authorize all envoy-originated data at `pegboard-envoy` before it reaches trusted internal systems.
+
 ### Important Patterns
 
 **Error Handling**

Cargo.lock

@@ -46,6 +46,7 @@ members = [
   "engine/packages/universaldb",
   "engine/packages/universalpubsub",
   "engine/packages/util",
+  "engine/packages/util-serde",
   "engine/packages/util-id",
   "engine/packages/workflow-worker",
   "engine/sdks/rust/api-full",
@@ -491,6 +492,9 @@ members = [
     [workspace.dependencies.rivet-util-id]
     path = "engine/packages/util-id"
 
+    [workspace.dependencies.rivet-util-serde]
+    path = "engine/packages/util-serde"
+
     [workspace.dependencies.rivet-workflow-worker]
     path = "engine/packages/workflow-worker"
 

Cargo.toml

@@ -27,6 +27,7 @@ rivet-metrics.workspace = true
 rivet-envoy-protocol.workspace = true
 rivet-runtime.workspace = true
 rivet-types.workspace = true
+scc.workspace = true
 serde_bare.workspace = true
 serde_json.workspace = true
 serde.workspace = true

engine/packages/pegboard-envoy/Cargo.toml

@@ -15,6 +15,7 @@ use rivet_data::converted::{ActorNameKeyData, MetadataKeyData};
 use rivet_envoy_protocol::{self as protocol, versioned};
 use rivet_guard_core::WebSocketHandle;
 use rivet_types::runner_configs::RunnerConfigKind;
+use scc::HashMap;
 use universaldb::prelude::*;
 use vbare::OwnedVersionedData;
 
@@ -26,6 +27,7 @@ pub struct Conn {
 	pub envoy_key: String,
 	pub protocol_version: u16,
 	pub ws_handle: WebSocketHandle,
+	pub authorized_tunnel_routes: HashMap<(protocol::GatewayId, protocol::RequestId), ()>,
 	pub is_serverless: bool,
 	pub last_rtt: AtomicU32,
 	/// Timestamp (epoch ms) of the last pong received from the envoy.
@@ -101,6 +103,7 @@ pub async fn init_conn(
 		envoy_key,
 		protocol_version,
 		ws_handle,
+		authorized_tunnel_routes: HashMap::new(),
 		is_serverless: false,
 		last_rtt: AtomicU32::new(0),
 		last_ping_ts: AtomicI64::new(util::timestamp::now()),
@@ -114,7 +117,6 @@ pub async fn init_conn(
 
 	Ok(Arc::new(conn))
 }
-
 #[tracing::instrument(skip_all)]
 pub async fn handle_init(
 	ctx: &StandaloneCtx,

engine/packages/pegboard-envoy/src/conn.rs

@@ -162,6 +162,10 @@ async fn handle_message(
 		}
 		protocol::ToEnvoyConn::ToEnvoyAckEvents(x) => protocol::ToEnvoy::ToEnvoyAckEvents(x),
 		protocol::ToEnvoyConn::ToEnvoyTunnelMessage(x) => {
+			let _ = conn
+				.authorized_tunnel_routes
+				.insert_async((x.message_id.gateway_id, x.message_id.request_id), ())
+				.await;
 			protocol::ToEnvoy::ToEnvoyTunnelMessage(x)
 		}
 	};

engine/packages/pegboard-envoy/src/tunnel_to_ws_task.rs

@@ -8,10 +8,11 @@ use pegboard::actor_kv;
 use pegboard::pubsub_subjects::GatewayReceiverSubject;
 use rivet_envoy_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
 use rivet_guard_core::websocket_handle::WebSocketReceiver;
+use scc::HashMap;
 use std::sync::{Arc, atomic::Ordering};
 use tokio::sync::{Mutex, MutexGuard, watch};
 use universaldb::utils::end_of_key_range;
-use universalpubsub::PublishOpts;
+use universalpubsub::{PubSub, PublishOpts};
 use vbare::OwnedVersionedData;
 
 use crate::{LifecycleResult, actor_event_demuxer::ActorEventDemuxer, conn::Conn, errors};
@@ -366,7 +367,12 @@ async fn handle_message(
 			}
 		}
 		protocol::ToRivet::ToRivetTunnelMessage(tunnel_msg) => {
-			handle_tunnel_message(&ctx, tunnel_msg)
+			handle_tunnel_message(
+				&ctx.ups().context("failed to get UPS instance for tunnel message")?,
+				ctx.config().pegboard().envoy_max_response_payload_size(),
+				&conn.authorized_tunnel_routes,
+				tunnel_msg,
+			)
 				.await
 				.context("failed to handle tunnel message")?;
 		}
@@ -446,17 +452,28 @@ async fn ack_commands(
 
 #[tracing::instrument(skip_all)]
 async fn handle_tunnel_message(
-	ctx: &StandaloneCtx,
+	ups: &PubSub,
+	max_payload_size: usize,
+	authorized_tunnel_routes: &HashMap<(protocol::GatewayId, protocol::RequestId), ()>,
 	msg: protocol::ToRivetTunnelMessage,
 ) -> Result<()> {
 	// Extract inner data length before consuming msg
 	let inner_data_len = tunnel_message_inner_data_len(&msg.message_kind);
 
 	// Enforce incoming payload size
-	if inner_data_len > ctx.config().pegboard().envoy_max_response_payload_size() {
+	if inner_data_len > max_payload_size {
 		return Err(errors::WsError::InvalidPacket("payload too large".to_string()).build());
 	}
 
+	if !authorized_tunnel_routes
+		.contains_async(&(msg.message_id.gateway_id, msg.message_id.request_id))
+		.await
+	{
+		return Err(
+			errors::WsError::InvalidPacket("unauthorized tunnel message".to_string()).build(),
+		);
+	}
+
 	let gateway_reply_to = GatewayReceiverSubject::new(msg.message_id.gateway_id).to_string();
 	let msg_serialized =
 		versioned::ToGateway::wrap_latest(protocol::ToGateway::ToRivetTunnelMessage(msg))
@@ -470,9 +487,7 @@ async fn handle_tunnel_message(
 	);
 
 	// Publish message to UPS
-	ctx.ups()
-		.context("failed to get UPS instance for tunnel message")?
-		.publish(&gateway_reply_to, &msg_serialized, PublishOpts::one())
+	ups.publish(&gateway_reply_to, &msg_serialized, PublishOpts::one())
 		.await
 		.with_context(|| {
 			format!(
@@ -500,6 +515,10 @@ fn tunnel_message_inner_data_len(kind: &protocol::ToRivetTunnelMessageKind) -> u
 	}
 }
 
+#[cfg(test)]
+#[path = "../tests/support/ws_to_tunnel_task.rs"]
+mod tests;
+
 async fn send_actor_kv_error(conn: &Conn, request_id: u32, message: &str) -> Result<()> {
 	let res_msg = versioned::ToEnvoy::wrap_latest(protocol::ToEnvoy::ToEnvoyKvResponse(
 		protocol::ToEnvoyKvResponse {

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

@@ -0,0 +1,81 @@
+use std::{sync::Arc, time::Duration};
+
+use pegboard::pubsub_subjects::GatewayReceiverSubject;
+use rivet_envoy_protocol as protocol;
+use scc::HashMap;
+use universalpubsub::{NextOutput, PubSub, driver::memory::MemoryDriver};
+
+use super::handle_tunnel_message;
+
+fn memory_pubsub(channel: &str) -> PubSub {
+	PubSub::new(Arc::new(MemoryDriver::new(channel.to_string())))
+}
+
+fn response_abort_message(
+	gateway_id: protocol::GatewayId,
+	request_id: protocol::RequestId,
+) -> protocol::ToRivetTunnelMessage {
+	protocol::ToRivetTunnelMessage {
+		message_id: protocol::MessageId {
+			gateway_id,
+			request_id,
+			message_index: 0,
+		},
+		message_kind: protocol::ToRivetTunnelMessageKind::ToRivetResponseAbort,
+	}
+}
+
+#[tokio::test]
+async fn rejects_unissued_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-envoy-ws-to-tunnel-test-reject");
+	let gateway_id = [1, 2, 3, 4];
+	let request_id = [5, 6, 7, 8];
+	let mut sub = pubsub
+		.subscribe(&GatewayReceiverSubject::new(gateway_id).to_string())
+		.await
+		.unwrap();
+	let authorized_tunnel_routes = HashMap::new();
+
+	let err = handle_tunnel_message(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message(gateway_id, request_id),
+	)
+	.await
+	.unwrap_err();
+	assert!(err.to_string().contains("unauthorized tunnel message"));
+
+	let recv = tokio::time::timeout(Duration::from_millis(50), sub.next()).await;
+	assert!(recv.is_err());
+}
+
+#[tokio::test]
+async fn republishes_issued_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-envoy-ws-to-tunnel-test-allow");
+	let gateway_id = [9, 10, 11, 12];
+	let request_id = [13, 14, 15, 16];
+	let mut sub = pubsub
+		.subscribe(&GatewayReceiverSubject::new(gateway_id).to_string())
+		.await
+		.unwrap();
+	let authorized_tunnel_routes = HashMap::new();
+	let _ = authorized_tunnel_routes
+		.insert_async((gateway_id, request_id), ())
+		.await;
+
+	handle_tunnel_message(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message(gateway_id, request_id),
+	)
+	.await
+	.unwrap();
+
+	let msg = tokio::time::timeout(Duration::from_secs(1), sub.next())
+		.await
+		.unwrap()
+		.unwrap();
+	assert!(matches!(msg, NextOutput::Message(_)));
+}

engine/packages/pegboard-envoy/tests/support/ws_to_tunnel_task.rs

@@ -28,6 +28,7 @@ rivet-metrics.workspace = true
 rivet-runner-protocol.workspace = true
 rivet-runtime.workspace = true
 rivet-types.workspace = true
+scc.workspace = true
 serde_bare.workspace = true
 serde_json.workspace = true
 serde.workspace = true

engine/packages/pegboard-runner/Cargo.toml

@@ -16,6 +16,7 @@ use rivet_data::converted::{ActorNameKeyData, MetadataKeyData};
 use rivet_guard_core::WebSocketHandle;
 use rivet_runner_protocol::{self as protocol, versioned};
 use rivet_types::runner_configs::RunnerConfigKind;
+use scc::HashMap;
 use universaldb::prelude::*;
 use vbare::OwnedVersionedData;
 
@@ -29,6 +30,7 @@ pub struct Conn {
 	pub workflow_id: Id,
 	pub protocol_version: u16,
 	pub ws_handle: WebSocketHandle,
+	pub authorized_tunnel_routes: HashMap<(protocol::mk2::GatewayId, protocol::mk2::RequestId), ()>,
 	pub last_rtt: AtomicU32,
 	/// Timestamp (epoch ms) of the last pong received from the runner.
 	pub last_ping_ts: AtomicI64,
@@ -188,6 +190,7 @@ pub async fn init_conn(
 		workflow_id,
 		protocol_version,
 		ws_handle,
+		authorized_tunnel_routes: HashMap::new(),
 		last_rtt: AtomicU32::new(0),
 		last_ping_ts: AtomicI64::new(util::timestamp::now()),
 	});
@@ -213,7 +216,6 @@ pub async fn init_conn(
 
 	Ok(conn)
 }
-
 enum Init {
 	Mk2(protocol::mk2::ToServerInit),
 	Mk1(protocol::ToServerInit),