@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/pegboard-runner/tests/support/ws_to_tunnel_task.rs

@@ -0,0 +1,150 @@
+use std::{sync::Arc, time::Duration};
+
+use pegboard::pubsub_subjects::GatewayReceiverSubject;
+use rivet_runner_protocol as protocol;
+use scc::HashMap;
+use universalpubsub::{NextOutput, PubSub, driver::memory::MemoryDriver};
+
+use super::{handle_tunnel_message_mk1, handle_tunnel_message_mk2};
+
+fn memory_pubsub(channel: &str) -> PubSub {
+	PubSub::new(Arc::new(MemoryDriver::new(channel.to_string())))
+}
+
+fn response_abort_message_mk2(
+	gateway_id: protocol::mk2::GatewayId,
+	request_id: protocol::mk2::RequestId,
+) -> protocol::mk2::ToServerTunnelMessage {
+	protocol::mk2::ToServerTunnelMessage {
+		message_id: protocol::mk2::MessageId {
+			gateway_id,
+			request_id,
+			message_index: 0,
+		},
+		message_kind: protocol::mk2::ToServerTunnelMessageKind::ToServerResponseAbort,
+	}
+}
+
+fn response_abort_message_mk1(
+	gateway_id: protocol::mk2::GatewayId,
+	request_id: protocol::mk2::RequestId,
+) -> protocol::ToServerTunnelMessage {
+	protocol::ToServerTunnelMessage {
+		message_id: protocol::MessageId {
+			gateway_id,
+			request_id,
+			message_index: 0,
+		},
+		message_kind: protocol::ToServerTunnelMessageKind::ToServerResponseAbort,
+	}
+}
+
+#[tokio::test]
+async fn rejects_unissued_mk2_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-runner-ws-to-tunnel-test-reject-mk2");
+	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_mk2(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message_mk2(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_mk2_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-runner-ws-to-tunnel-test-allow-mk2");
+	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_mk2(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message_mk2(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(_)));
+}
+
+#[tokio::test]
+async fn rejects_unissued_mk1_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-runner-ws-to-tunnel-test-reject-mk1");
+	let gateway_id = [17, 18, 19, 20];
+	let request_id = [21, 22, 23, 24];
+	let mut sub = pubsub
+		.subscribe(&GatewayReceiverSubject::new(gateway_id).to_string())
+		.await
+		.unwrap();
+	let authorized_tunnel_routes = HashMap::new();
+
+	let err = handle_tunnel_message_mk1(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message_mk1(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_mk1_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-runner-ws-to-tunnel-test-allow-mk1");
+	let gateway_id = [25, 26, 27, 28];
+	let request_id = [29, 30, 31, 32];
+	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_mk1(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message_mk1(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(_)));
+}

package/packages/pools/Cargo.toml

@@ -8,12 +8,11 @@ edition.workspace = true
 [dependencies]
 anyhow.workspace = true
 async-nats.workspace = true
-clickhouse-inserter.workspace = true
 clickhouse.workspace = true
 futures-util.workspace = true
 governor.workspace = true
 hyper-tls.workspace = true
-hyper.workspace = true
+hyper-util.workspace = true
 lazy_static.workspace = true
 reqwest.workspace = true
 rivet-config.workspace = true

package/packages/pools/src/db/clickhouse.rs

@@ -6,18 +6,19 @@ use rivet_config::Config;
 pub type ClickHousePool = clickhouse::Client;
 
 #[tracing::instrument(skip(config))]
-pub fn setup(config: Config) -> Result<Option<ClickHousePool>> {
-	if let Some(clickhouse) = &config.clickhouse() {
+pub fn setup(config: &Config) -> Result<Option<ClickHousePool>> {
+	if let Some(clickhouse) = config.clickhouse() {
 		tracing::debug!("clickhouse connecting");
 
 		// Build HTTP client
-		let mut http_connector = hyper::client::connect::HttpConnector::new();
+		let mut http_connector = hyper_util::client::legacy::connect::HttpConnector::new();
 		http_connector.enforce_http(false);
 		http_connector.set_keepalive(Some(Duration::from_secs(15)));
 		let https_connector = hyper_tls::HttpsConnector::new_with_connector(http_connector);
-		let http_client = hyper::Client::builder()
-			.pool_idle_timeout(Duration::from_secs(2))
-			.build(https_connector);
+		let http_client =
+			hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new())
+				.pool_idle_timeout(Duration::from_secs(2))
+				.build(https_connector);
 
 		// Build ClickHouse client
 		let mut client = clickhouse::Client::with_http_client(http_client)

package/packages/pools/src/db/udb.rs

@@ -17,11 +17,24 @@ impl Deref for UdbPool {
 }
 
 #[tracing::instrument(skip(config))]
-pub async fn setup(config: Config) -> Result<Option<UdbPool>> {
+pub async fn setup(config: &Config) -> Result<Option<UdbPool>> {
 	let db_driver = match config.database() {
 		config::Database::Postgres(pg) => {
-			Arc::new(universaldb::driver::PostgresDatabaseDriver::new(pg.url.read().clone()).await?)
-				as universaldb::DatabaseDriverHandle
+			let postgres_config = universaldb::driver::postgres::PostgresConfig {
+				connection_string: pg.url.read().clone(),
+				ssl_config: pg.ssl.as_ref().map(|ssl| {
+					universaldb::driver::postgres::PostgresSslConfig {
+						ssl_root_cert_path: ssl.root_cert_path.clone(),
+						ssl_client_cert_path: ssl.client_cert_path.clone(),
+						ssl_client_key_path: ssl.client_key_path.clone(),
+					}
+				}),
+			};
+
+			Arc::new(
+				universaldb::driver::PostgresDatabaseDriver::new_with_config(postgres_config)
+					.await?,
+			) as universaldb::DatabaseDriverHandle
 		}
 		config::Database::FileSystem(fs) => {
 			Arc::new(universaldb::driver::RocksDbDatabaseDriver::new(fs.path.clone()).await?)

package/packages/pools/src/db/ups.rs

@@ -9,7 +9,7 @@ use crate::Error;
 pub type UpsPool = ups::PubSub;
 
 #[tracing::instrument(skip(config))]
-pub async fn setup(config: Config, client_name: String) -> Result<UpsPool> {
+pub async fn setup(config: &Config, client_name: &str) -> Result<UpsPool> {
 	let driver = match config.pubsub() {
 		config::PubSub::Nats(nats) => {
 			// Parse nodes
@@ -44,13 +44,20 @@ pub async fn setup(config: Config, client_name: String) -> Result<UpsPool> {
 									tracing::debug!(?server_addrs, "nats reconnected");
 								}
 								async_nats::Event::Disconnected => {
-									tracing::error!(?server_addrs, "nats disconnected");
+									tracing::warn!(?server_addrs, "nats disconnected");
 								}
 								async_nats::Event::LameDuckMode => {
 									tracing::warn!(?server_addrs, "nats lame duck mode");
 								}
-								async_nats::Event::SlowConsumer(_) => {
-									tracing::warn!(?server_addrs, "nats slow consumer");
+								async_nats::Event::Draining => {
+									tracing::warn!(?server_addrs, "nats draining");
+								}
+								async_nats::Event::Closed => {
+									// Engine is shutting down, not an error
+									tracing::info!(?server_addrs, "nats closed");
+								}
+								async_nats::Event::SlowConsumer(sid) => {
+									tracing::warn!(?server_addrs, ?sid, "nats slow consumer");
 								}
 								async_nats::Event::ServerError(err) => {
 									tracing::error!(?server_addrs, ?err, "nats server error");
@@ -71,16 +78,31 @@ pub async fn setup(config: Config, client_name: String) -> Result<UpsPool> {
 		}
 		config::PubSub::PostgresNotify(pg) => {
 			tracing::debug!("creating postgres pubsub driver");
+
+			let (ssl_root_cert_path, ssl_client_cert_path, ssl_client_key_path) =
+				if let Some(ssl) = &pg.ssl {
+					(
+						ssl.root_cert_path.clone(),
+						ssl.client_cert_path.clone(),
+						ssl.client_key_path.clone(),
+					)
+				} else {
+					(None, None, None)
+				};
+
 			Arc::new(
 				ups::driver::postgres::PostgresDriver::connect(
 					pg.url.read().clone(),
 					pg.memory_optimization,
+					ssl_root_cert_path,
+					ssl_client_cert_path,
+					ssl_client_key_path,
 				)
 				.await?,
 			) as ups::PubSubDriverHandle
 		}
 		config::PubSub::Memory(memory) => {
-			tracing::debug!(channel = ?memory.channel, "creating memory pubsub driver");
+			tracing::debug!(channel=%memory.channel, "creating memory pubsub driver");
 			Arc::new(ups::driver::memory::MemoryDriver::new(
 				memory.channel.clone(),
 			)) as ups::PubSubDriverHandle

package/packages/pools/src/error.rs

@@ -6,9 +6,6 @@ pub enum Error {
 	#[error("missing clickhouse pool")]
 	MissingClickHousePool,
 
-	#[error("missing clickhouse inserter")]
-	MissingClickHouseInserter,
-
 	#[error("missing udb pool")]
 	MissingUdbPool,
 

package/packages/pools/src/lib.rs

@@ -9,8 +9,6 @@ pub use crate::{
 	db::clickhouse::ClickHousePool, db::udb::UdbPool, db::ups::UpsPool, error::Error, pools::Pools,
 };
 
-pub use clickhouse_inserter::ClickHouseInserterHandle;
-
 // Re-export for macros
 #[doc(hidden)]
 pub use rivet_util as __rivet_util;

package/packages/pools/src/metrics.rs

@@ -1,32 +1,37 @@
-use rivet_metrics::{
-	BUCKETS,
-	otel::{global::*, metrics::*},
-};
+use rivet_metrics::{BUCKETS, REGISTRY, prometheus::*};
 
 lazy_static::lazy_static! {
-	static ref METER: Meter = meter("rivet-pools");
-
 	// MARK: SQL
-	/// Expected attributes: "action", "context_name", "location"
-	pub static ref SQL_QUERY_TOTAL: Counter<u64> = METER.u64_counter("rivet_sql_query_total")
-		.with_description("Total number of queries.")
-		.build();
-	/// Expected attributes: "action", "context_name", "location"
-	pub static ref SQL_QUERY_DURATION: Histogram<f64> = METER.f64_histogram("rivet_sql_query_duration")
-		.with_description("Total duration of sql query.")
-		.with_boundaries(BUCKETS.to_vec())
-		.build();
-	/// Expected attributes: "action", "context_name", "location"
-	pub static ref SQL_ACQUIRE_DURATION: Histogram<f64> = METER.f64_histogram("rivet_sql_acquire_duration")
-		.with_description("Total duration to acquire an sql connection.")
-		.with_boundaries(BUCKETS.to_vec())
-		.build();
-	/// Expected attributes: "action", "context_name", "location", "acquire_result"
-	pub static ref SQL_ACQUIRE_TOTAL: Counter<u64> = METER.u64_counter("rivet_sql_acquire_total")
-		.with_description("Amount times a pool connection was acquired.")
-		.build();
-	/// Expected attributes: "action", "context_name", "location", "acquire_result"
-	pub static ref SQL_ACQUIRE_TRIES: Counter<u64> = METER.u64_counter("rivet_sql_acquire_tries")
-		.with_description("Amount of tries required to get a pool connection.")
-		.build();
+	pub static ref SQL_QUERY_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"sql_query_total",
+		"Total number of queries.",
+		&["action", "context_name", "location"],
+		*REGISTRY
+	).unwrap();
+	pub static ref SQL_QUERY_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"sql_query_duration",
+		"Total duration of sql query.",
+		&["action", "context_name", "location"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+	pub static ref SQL_ACQUIRE_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"sql_acquire_duration",
+		"Total duration to acquire an sql connection.",
+		&["action", "context_name", "location"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+	pub static ref SQL_ACQUIRE_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"sql_acquire_total",
+		"Amount times a pool connection was acquired.",
+		&["action", "context_name", "location", "acquire_result"],
+		*REGISTRY
+	).unwrap();
+	pub static ref SQL_ACQUIRE_TRIES: IntCounterVec = register_int_counter_vec_with_registry!(
+		"sql_acquire_tries",
+		"Amount of tries required to get a pool connection.",
+		&["action", "context_name", "location", "acquire_result"],
+		*REGISTRY
+	).unwrap();
 }

package/packages/pools/src/pools.rs

@@ -1,7 +1,6 @@
 use std::sync::Arc;
 
 use anyhow::*;
-use clickhouse_inserter::ClickHouseInserterHandle;
 use rivet_config::Config;
 use tokio_util::sync::{CancellationToken, DropGuard};
 
@@ -12,9 +11,7 @@ pub(crate) struct PoolsInner {
 	pub(crate) _guard: DropGuard,
 	pub(crate) ups: Option<UpsPool>,
 	pub(crate) clickhouse: Option<clickhouse::Client>,
-	pub(crate) clickhouse_inserter: Option<ClickHouseInserterHandle>,
 	pub(crate) udb: Option<UdbPool>,
-	clickhouse_enabled: bool,
 }
 
 #[derive(Clone)]
@@ -24,33 +21,26 @@ impl Pools {
 	#[tracing::instrument(skip(config))]
 	pub async fn new(config: Config) -> Result<Pools> {
 		// TODO: Choose client name for this service
-		let client_name = "rivet".to_string();
+		let client_name = "rivet";
 		let token = CancellationToken::new();
 
 		let (ups, udb) = tokio::try_join!(
-			crate::db::ups::setup(config.clone(), client_name.clone()),
-			crate::db::udb::setup(config.clone()),
+			crate::db::ups::setup(&config, client_name),
+			crate::db::udb::setup(&config),
 		)?;
-		let clickhouse = crate::db::clickhouse::setup(config.clone())?;
-
-		// Create the ClickHouse inserter if vector is enabled
-		let clickhouse_inserter = if let Some(vector_http) = config.vector_http().as_ref() {
-			let inserter =
-				clickhouse_inserter::create_inserter(&vector_http.host, vector_http.port)?;
-			Some(inserter)
-		} else {
-			None
-		};
+		let clickhouse = crate::db::clickhouse::setup(&config)?;
 
 		let pool = Pools(Arc::new(PoolsInner {
 			_guard: token.clone().drop_guard(),
 			ups: Some(ups),
 			clickhouse,
-			clickhouse_inserter,
 			udb,
-			clickhouse_enabled: config.clickhouse().is_some(),
 		}));
 
+		// Initialize here to avoid cold starts elsewhere
+		crate::reqwest::client().await?;
+		crate::reqwest::client_no_timeout().await?;
+
 		Ok(pool)
 	}
 
@@ -58,30 +48,27 @@ impl Pools {
 	#[tracing::instrument(skip(config))]
 	pub async fn test(config: Config) -> Result<Pools> {
 		// TODO: Choose client name for this service
-		let client_name = "rivet".to_string();
+		let client_name = "rivet";
 		let token = CancellationToken::new();
 
 		let (ups, udb) = tokio::try_join!(
-			crate::db::ups::setup(config.clone(), client_name.clone()),
-			crate::db::udb::setup(config.clone()),
+			crate::db::ups::setup(&config, client_name),
+			crate::db::udb::setup(&config),
 		)?;
 
-		// Test setup doesn't use ClickHouse inserter
 		let pool = Pools(Arc::new(PoolsInner {
 			_guard: token.clone().drop_guard(),
 			ups: Some(ups),
 			clickhouse: None,
-			clickhouse_inserter: None,
 			udb,
-			clickhouse_enabled: config.clickhouse().is_some(),
 		}));
 
 		Ok(pool)
 	}
 
 	// MARK: Getters
-	pub fn ups_option(&self) -> &Option<UpsPool> {
-		&self.0.ups
+	pub fn ups_option(&self) -> Option<&UpsPool> {
+		self.0.ups.as_ref()
 	}
 
 	// MARK: Pool lookups
@@ -89,23 +76,12 @@ impl Pools {
 		self.0.ups.clone().ok_or(Error::MissingUpsPool.into())
 	}
 
-	pub fn clickhouse_enabled(&self) -> bool {
-		self.0.clickhouse_enabled
-	}
-
 	pub fn clickhouse(&self) -> Result<ClickHousePool> {
-		ensure!(self.clickhouse_enabled(), "clickhouse disabled");
-
 		self.0.clickhouse.clone().context("missing clickhouse pool")
 	}
 
-	pub fn clickhouse_inserter(&self) -> Result<ClickHouseInserterHandle> {
-		ensure!(self.clickhouse_enabled(), "clickhouse disabled");
-
-		self.0
-			.clickhouse_inserter
-			.clone()
-			.context("missing clickhouse inserter")
+	pub fn clickhouse_option(&self) -> Option<&ClickHousePool> {
+		self.0.clickhouse.as_ref()
 	}
 
 	pub fn udb(&self) -> Result<UdbPool> {

package/packages/pools/src/prelude.rs

@@ -1,3 +1,3 @@
 pub use clickhouse;
 
-pub use crate::{ClickHouseInserterHandle, ClickHousePool, UdbPool, UpsPool};
+pub use crate::{ClickHousePool, UdbPool, UpsPool};

package/packages/postgres-util/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "rivet-postgres-util"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+license.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+rustls.workspace = true
+rustls-pemfile.workspace = true
+tracing.workspace = true
+webpki-roots.workspace = true

package/packages/postgres-util/src/lib.rs

@@ -0,0 +1,84 @@
+use std::path::PathBuf;
+
+use anyhow::{Context, Result};
+use rustls::ClientConfig;
+
+/// Helper function to build TLS configuration with optional custom certificates
+///
+/// # Arguments
+///
+/// * `ssl_root_cert_path` - Optional path to root CA certificate for server verification
+/// * `ssl_client_cert_path` - Optional path to client certificate for mutual TLS
+/// * `ssl_client_key_path` - Optional path to client private key for mutual TLS
+///
+/// # Returns
+///
+/// A configured `ClientConfig` for use with tokio-postgres-rustls
+pub fn build_tls_config(
+	ssl_root_cert_path: Option<&PathBuf>,
+	ssl_client_cert_path: Option<&PathBuf>,
+	ssl_client_key_path: Option<&PathBuf>,
+) -> Result<ClientConfig> {
+	let mut root_store = rustls::RootCertStore::empty();
+
+	// Add custom root certificate if provided
+	if let Some(root_cert_path) = ssl_root_cert_path {
+		tracing::debug!(?root_cert_path, "loading custom root certificate");
+		let cert_data = std::fs::read(root_cert_path).with_context(|| {
+			format!("failed to read root certificate from {:?}", root_cert_path)
+		})?;
+
+		let certs = rustls_pemfile::certs(&mut cert_data.as_slice())
+			.collect::<std::result::Result<Vec<_>, _>>()
+			.context("failed to parse root certificate")?;
+
+		for cert in certs {
+			root_store
+				.add(cert)
+				.context("failed to add root certificate to store")?;
+		}
+	} else {
+		// Use default webpki root certificates
+		root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
+	}
+
+	// Build TLS config with client certificates if provided
+	let tls_config = if let (Some(client_cert_path), Some(client_key_path)) =
+		(ssl_client_cert_path, ssl_client_key_path)
+	{
+		tracing::debug!(
+			?client_cert_path,
+			?client_key_path,
+			"loading client certificates"
+		);
+
+		// Load client certificate
+		let cert_data = std::fs::read(client_cert_path).with_context(|| {
+			format!(
+				"failed to read client certificate from {:?}",
+				client_cert_path
+			)
+		})?;
+		let certs = rustls_pemfile::certs(&mut cert_data.as_slice())
+			.collect::<std::result::Result<Vec<_>, _>>()
+			.context("failed to parse client certificate")?;
+
+		// Load client key
+		let key_data = std::fs::read(client_key_path)
+			.with_context(|| format!("failed to read client key from {:?}", client_key_path))?;
+		let key = rustls_pemfile::private_key(&mut key_data.as_slice())
+			.context("failed to parse client key")?
+			.context("no private key found in client key file")?;
+
+		ClientConfig::builder()
+			.with_root_certificates(root_store)
+			.with_client_auth_cert(certs, key)
+			.context("failed to configure client authentication")?
+	} else {
+		ClientConfig::builder()
+			.with_root_certificates(root_store)
+			.with_no_client_auth()
+	};
+
+	Ok(tls_config)
+}