chore: tunnel auth

Author: NathanFlurry

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,7 @@ async fn handle_message(
 			}
 		}
 		protocol::ToRivet::ToRivetTunnelMessage(tunnel_msg) => {
-			handle_tunnel_message(&ctx, tunnel_msg)
+			handle_tunnel_message(&ctx, conn, tunnel_msg)
 				.await
 				.context("failed to handle tunnel message")?;
 		}
@@ -447,16 +448,41 @@ async fn ack_commands(
 #[tracing::instrument(skip_all)]
 async fn handle_tunnel_message(
 	ctx: &StandaloneCtx,
+	conn: &Conn,
+	msg: protocol::ToRivetTunnelMessage,
+) -> Result<()> {
+	forward_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,
+		msg,
+	)
+	.await
+}
+
+async fn forward_tunnel_message(
+	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 +496,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 +524,86 @@ fn tunnel_message_inner_data_len(kind: &protocol::ToRivetTunnelMessageKind) -> u
 	}
 }
 
+#[cfg(test)]
+mod tests {
+	use std::sync::Arc;
+	use std::time::Duration;
+
+	use super::*;
+	use universalpubsub::driver::memory::MemoryDriver;
+	use universalpubsub::NextOutput;
+
+	fn test_pubsub(channel: &str) -> PubSub {
+		PubSub::new(Arc::new(MemoryDriver::new(channel.to_string())))
+	}
+
+	fn test_message(gateway_id: [u8; 4], request_id: [u8; 4]) -> 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 = test_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 = forward_tunnel_message(
+			&pubsub,
+			1024,
+			&authorized_tunnel_routes,
+			test_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 = test_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;
+
+		forward_tunnel_message(
+			&pubsub,
+			1024,
+			&authorized_tunnel_routes,
+			test_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(_)));
+	}
+}
+
 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

@@ -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),

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

@@ -158,6 +158,10 @@ async fn handle_message_mk2(
 			protocol::mk2::ToClient::ToClientAckEvents(x)
 		}
 		protocol::mk2::ToRunner::ToClientTunnelMessage(x) => {
+			let _ = conn
+				.authorized_tunnel_routes
+				.insert_async((x.message_id.gateway_id, x.message_id.request_id), ())
+				.await;
 			protocol::mk2::ToClient::ToClientTunnelMessage(x)
 		}
 	};
@@ -250,6 +254,10 @@ async fn handle_message_mk1(
 		protocol::ToRunner::ToClientAckEvents(x) => protocol::ToClient::ToClientAckEvents(x),
 		protocol::ToRunner::ToClientKvResponse(x) => protocol::ToClient::ToClientKvResponse(x),
 		protocol::ToRunner::ToClientTunnelMessage(x) => {
+			let _ = conn
+				.authorized_tunnel_routes
+				.insert_async((x.message_id.gateway_id, x.message_id.request_id), ())
+				.await;
 			protocol::ToClient::ToClientTunnelMessage(x)
 		}
 	};