@rivetkit/engine 1.0.0

package/packages/guard/src/routing/pegboard_gateway.rs

@@ -0,0 +1,260 @@
+use std::time::Duration;
+
+use anyhow::Result;
+use gas::prelude::*;
+use hyper::header::HeaderName;
+use rivet_guard_core::proxy_service::{RouteConfig, RouteTarget, RoutingOutput, RoutingTimeout};
+
+use super::{SEC_WEBSOCKET_PROTOCOL, WS_PROTOCOL_ACTOR, WS_PROTOCOL_TOKEN, X_RIVET_TOKEN};
+use crate::{errors, shared_state::SharedState};
+
+const ACTOR_FORCE_WAKE_PENDING_TIMEOUT: i64 = util::duration::seconds(60);
+const ACTOR_READY_TIMEOUT: Duration = Duration::from_secs(10);
+pub const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor");
+
+/// Route requests to actor services using path-based routing
+#[tracing::instrument(skip_all)]
+pub async fn route_request_path_based(
+	ctx: &StandaloneCtx,
+	shared_state: &SharedState,
+	actor_id_str: &str,
+	token: Option<&str>,
+	path: &str,
+	_headers: &hyper::HeaderMap,
+	_is_websocket: bool,
+) -> Result<Option<RoutingOutput>> {
+	// Parse actor ID
+	let actor_id = Id::parse(actor_id_str).context("invalid actor id in path")?;
+
+	route_request_inner(ctx, shared_state, actor_id, path, token).await
+}
+
+/// Route requests to actor services based on headers
+#[tracing::instrument(skip_all)]
+pub async fn route_request(
+	ctx: &StandaloneCtx,
+	shared_state: &SharedState,
+	target: &str,
+	_host: &str,
+	path: &str,
+	headers: &hyper::HeaderMap,
+	is_websocket: bool,
+) -> Result<Option<RoutingOutput>> {
+	// Check target
+	if target != "actor" {
+		return Ok(None);
+	}
+
+	// Extract actor ID and token from WebSocket protocol or HTTP headers
+	let (actor_id_str, token) = if is_websocket {
+		// For WebSocket, parse the sec-websocket-protocol header
+		let protocols_header = headers
+			.get(SEC_WEBSOCKET_PROTOCOL)
+			.and_then(|protocols| protocols.to_str().ok())
+			.ok_or_else(|| {
+				crate::errors::MissingHeader {
+					header: "sec-websocket-protocol".to_string(),
+				}
+				.build()
+			})?;
+
+		let protocols: Vec<&str> = protocols_header.split(',').map(|p| p.trim()).collect();
+
+		let actor_id_raw = protocols
+			.iter()
+			.find_map(|p| p.strip_prefix(WS_PROTOCOL_ACTOR))
+			.ok_or_else(|| {
+				crate::errors::MissingHeader {
+					header: "`rivet_actor.*` protocol in sec-websocket-protocol".to_string(),
+				}
+				.build()
+			})?;
+
+		let actor_id = urlencoding::decode(actor_id_raw)
+			.context("invalid url encoding in actor id")?
+			.to_string();
+
+		let token = protocols
+			.iter()
+			.find_map(|p| p.strip_prefix(WS_PROTOCOL_TOKEN));
+
+		(actor_id, token)
+	} else {
+		// For HTTP, use headers
+		let actor_id = headers
+			.get(X_RIVET_ACTOR)
+			.map(|x| x.to_str())
+			.transpose()
+			.context("invalid x-rivet-actor header")?
+			.ok_or_else(|| {
+				crate::errors::MissingHeader {
+					header: X_RIVET_ACTOR.to_string(),
+				}
+				.build()
+			})?;
+
+		let token = headers
+			.get(X_RIVET_TOKEN)
+			.map(|x| x.to_str())
+			.transpose()
+			.context("invalid x-rivet-token header")?;
+
+		(actor_id.to_string(), token)
+	};
+
+	// Find actor to route to
+	let actor_id = Id::parse(&actor_id_str).context("invalid x-rivet-actor header")?;
+
+	route_request_inner(ctx, shared_state, actor_id, path, token).await
+}
+
+async fn route_request_inner(
+	ctx: &StandaloneCtx,
+	shared_state: &SharedState,
+	actor_id: Id,
+	path: &str,
+	_token: Option<&str>,
+) -> Result<Option<RoutingOutput>> {
+	// NOTE: Token validation implemented in EE
+
+	// Route to peer dc where the actor lives
+	if actor_id.label() != ctx.config().dc_label() {
+		tracing::debug!(peer_dc_label=?actor_id.label(), "re-routing actor to peer dc");
+
+		let peer_dc = ctx
+			.config()
+			.dc_for_label(actor_id.label())
+			.context("dc with the given label not found")?;
+
+		return Ok(Some(RoutingOutput::Route(RouteConfig {
+			targets: vec![RouteTarget {
+				actor_id: Some(actor_id),
+				host: peer_dc
+					.proxy_url_host()
+					.context("bad peer dc proxy url host")?
+					.to_string(),
+				port: peer_dc
+					.proxy_url_port()
+					.context("bad peer dc proxy url port")?,
+				path: path.to_owned(),
+			}],
+			timeout: RoutingTimeout {
+				routing_timeout: 10,
+			},
+		})));
+	}
+
+	// Create subs before checking if actor exists/is not destroyed
+	let mut ready_sub = ctx
+		.subscribe::<pegboard::workflows::actor::Ready>(("actor_id", actor_id))
+		.await?;
+	let mut stopped_sub = ctx
+		.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", actor_id))
+		.await?;
+	let mut fail_sub = ctx
+		.subscribe::<pegboard::workflows::actor::Failed>(("actor_id", actor_id))
+		.await?;
+	let mut destroy_sub = ctx
+		.subscribe::<pegboard::workflows::actor::DestroyStarted>(("actor_id", actor_id))
+		.await?;
+
+	// Fetch actor info
+	let Some(actor) = ctx
+		.op(pegboard::ops::actor::get_for_gateway::Input { actor_id })
+		.await?
+	else {
+		return Err(pegboard::errors::Actor::NotFound.build());
+	};
+
+	if actor.destroyed {
+		return Err(pegboard::errors::Actor::NotFound.build());
+	}
+
+	// Wake actor if sleeping
+	if actor.sleeping {
+		tracing::debug!(?actor_id, "actor sleeping, waking");
+
+		ctx.signal(pegboard::workflows::actor::Wake {
+			allocation_override: pegboard::workflows::actor::AllocationOverride::DontSleep {
+				pending_timeout: Some(ACTOR_FORCE_WAKE_PENDING_TIMEOUT),
+			},
+		})
+		.to_workflow_id(actor.workflow_id)
+		.send()
+		.await?;
+	}
+
+	let runner_id = if let (Some(runner_id), true) = (actor.runner_id, actor.connectable) {
+		runner_id
+	} else {
+		tracing::debug!(?actor_id, "waiting for actor to become ready");
+
+		let mut wake_retries = 0;
+
+		// Wait for ready, fail, or destroy
+		loop {
+			tokio::select! {
+				res = ready_sub.next() => break res?.runner_id,
+				res = stopped_sub.next() => {
+					res?;
+
+					if wake_retries < 16 {
+						tracing::debug!(?actor_id, ?wake_retries, "actor stopped while we were waiting for it to become ready, attempting rewake");
+						wake_retries += 1;
+
+						let res = ctx.signal(pegboard::workflows::actor::Wake {
+							allocation_override: pegboard::workflows::actor::AllocationOverride::DontSleep {
+								pending_timeout: Some(ACTOR_FORCE_WAKE_PENDING_TIMEOUT),
+							},
+						})
+						.to_workflow_id(actor.workflow_id)
+						.send()
+						.await;
+
+						if let Some(WorkflowError::WorkflowNotFound) = res
+							.as_ref()
+							.err()
+							.and_then(|x| x.chain().find_map(|x| x.downcast_ref::<WorkflowError>()))
+						{
+							tracing::warn!(
+								?actor_id,
+								"actor workflow not found for rewake"
+							);
+						} else {
+							res?;
+						}
+					} else {
+						tracing::warn!("actor retried waking 16 times, has not yet started");
+						return Err(rivet_guard_core::errors::ServiceUnavailable.build());
+					}
+				}
+				res = fail_sub.next() => {
+					let msg = res?;
+					return Err(msg.error.clone().build());
+				}
+				res = destroy_sub.next() => {
+					res?;
+					return Err(pegboard::errors::Actor::DestroyedWhileWaitingForReady.build());
+				}
+				// Ready timeout
+				_ = tokio::time::sleep(ACTOR_READY_TIMEOUT) => {
+					return Err(errors::ActorReadyTimeout { actor_id }.build());
+				}
+			}
+		}
+	};
+
+	tracing::debug!(?actor_id, ?runner_id, "actor ready");
+
+	// Return pegboard-gateway instance with path
+	let gateway = pegboard_gateway::PegboardGateway::new(
+		ctx.clone(),
+		shared_state.pegboard_gateway.clone(),
+		runner_id,
+		actor_id,
+		path.to_string(),
+	);
+	Ok(Some(RoutingOutput::CustomServe(std::sync::Arc::new(
+		gateway,
+	))))
+}

package/packages/guard/src/routing/runner.rs

@@ -0,0 +1,129 @@
+use anyhow::*;
+use gas::prelude::*;
+use rivet_guard_core::proxy_service::RoutingOutput;
+use std::sync::Arc;
+
+use super::{SEC_WEBSOCKET_PROTOCOL, X_RIVET_TOKEN};
+pub(crate) const WS_PROTOCOL_TOKEN: &str = "rivet_token.";
+
+/// Route requests to the runner service using header-based routing
+#[tracing::instrument(skip_all)]
+pub async fn route_request(
+	ctx: &StandaloneCtx,
+	target: &str,
+	host: &str,
+	path: &str,
+	headers: &hyper::HeaderMap,
+) -> Result<Option<RoutingOutput>> {
+	if target != "runner" {
+		return Ok(None);
+	}
+
+	tracing::debug!(?host, path, "routing to runner via header");
+
+	route_runner_internal(ctx, host, headers).await.map(Some)
+}
+
+/// Route requests to the runner service using path-based routing
+/// Matches path: /runners/connect
+#[tracing::instrument(skip_all)]
+pub async fn route_request_path_based(
+	ctx: &StandaloneCtx,
+	host: &str,
+	path: &str,
+	headers: &hyper::HeaderMap,
+) -> Result<Option<RoutingOutput>> {
+	// Check if path matches /runners/connect
+	let path_without_query = path.split('?').next().unwrap_or(path);
+	if path_without_query != "/runners/connect" {
+		return Ok(None);
+	}
+
+	tracing::debug!(?host, path, "routing to runner via path");
+
+	route_runner_internal(ctx, host, headers).await.map(Some)
+}
+
+/// Internal runner routing logic shared by both header-based and path-based routing
+#[tracing::instrument(skip_all)]
+async fn route_runner_internal(
+	ctx: &StandaloneCtx,
+	host: &str,
+	headers: &hyper::HeaderMap,
+) -> Result<RoutingOutput> {
+	// Validate that the host is valid for the current datacenter
+	let current_dc = ctx.config().topology().current_dc()?;
+	if !current_dc.is_valid_regional_host(host) {
+		tracing::warn!(?host, datacenter=?current_dc.name, "invalid host for current datacenter");
+
+		// Determine valid hosts for error message
+		let valid_hosts = if let Some(hosts) = &current_dc.valid_hosts {
+			hosts.join(", ")
+		} else {
+			current_dc
+				.public_url
+				.host_str()
+				.map(|h| h.to_string())
+				.unwrap_or_else(|| "unknown".to_string())
+		};
+
+		return Err(crate::errors::MustUseRegionalHost {
+			host: host.to_string(),
+			datacenter: current_dc.name.clone(),
+			valid_hosts,
+		}
+		.build());
+	}
+
+	tracing::debug!(datacenter = ?current_dc.name, "validated host for datacenter");
+
+	let is_websocket = headers
+		.get("upgrade")
+		.and_then(|v| v.to_str().ok())
+		.map(|v| v.eq_ignore_ascii_case("websocket"))
+		.unwrap_or(false);
+
+	tracing::debug!(is_websocket, "connection type");
+
+	// Check auth (if enabled)
+	if let Some(auth) = &ctx.config().auth {
+		// Extract token from protocol or header
+		let token = if is_websocket {
+			headers
+				.get(SEC_WEBSOCKET_PROTOCOL)
+				.and_then(|protocols| protocols.to_str().ok())
+				.and_then(|protocols| {
+					protocols
+						.split(',')
+						.map(|p| p.trim())
+						.find_map(|p| p.strip_prefix(WS_PROTOCOL_TOKEN))
+				})
+				.ok_or_else(|| {
+					crate::errors::MissingHeader {
+						header: "`rivet_token.*` protocol in sec-websocket-protocol".to_string(),
+					}
+					.build()
+				})?
+		} else {
+			headers
+				.get(X_RIVET_TOKEN)
+				.and_then(|x| x.to_str().ok())
+				.ok_or_else(|| {
+					crate::errors::MissingHeader {
+						header: X_RIVET_TOKEN.to_string(),
+					}
+					.build()
+				})?
+		};
+
+		// Validate token
+		if token != auth.admin_token.read() {
+			return Err(rivet_api_builder::ApiForbidden.build());
+		}
+
+		tracing::debug!("authenticated runner connection");
+	}
+
+	let tunnel = pegboard_runner::PegboardRunnerWsCustomServe::new(ctx.clone());
+	Ok(RoutingOutput::CustomServe(Arc::new(tunnel)))
+}

package/packages/guard/src/shared_state.rs

@@ -0,0 +1,31 @@
+use anyhow::*;
+use std::{ops::Deref, sync::Arc};
+use universalpubsub::PubSub;
+
+#[derive(Clone)]
+pub struct SharedState(Arc<SharedStateInner>);
+
+impl SharedState {
+	pub fn new(config: &rivet_config::Config, pubsub: PubSub) -> SharedState {
+		SharedState(Arc::new(SharedStateInner {
+			pegboard_gateway: pegboard_gateway::shared_state::SharedState::new(config, pubsub),
+		}))
+	}
+
+	pub async fn start(&self) -> Result<()> {
+		self.pegboard_gateway.start().await?;
+		Ok(())
+	}
+}
+
+impl Deref for SharedState {
+	type Target = SharedStateInner;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}
+
+pub struct SharedStateInner {
+	pub pegboard_gateway: pegboard_gateway::shared_state::SharedState,
+}

package/packages/guard/src/tls.rs

@@ -0,0 +1,224 @@
+use anyhow::*;
+use gas::prelude::*;
+use rivet_guard_core::CertResolverFn;
+
+// /// Certificate pair with name for logging
+// struct CertificatePair {
+// 	name: &'static str,
+// 	cert_path: Box<Path>,
+// 	key_path: Box<Path>,
+// }
+
+// /// Helper function to load a certificate and key into a CertifiedKey
+// fn load_certified_key(cert_pair: &CertificatePair) -> GlobalResult<Arc<CertifiedKey>> {
+// 	// Validate that paths exist
+// 	if !cert_pair.cert_path.exists() {
+// 		bail!(
+// 			"{} certificate file not found at {:?}",
+// 			cert_pair.name,
+// 			cert_pair.cert_path
+// 		);
+// 	}
+// 	if !cert_pair.key_path.exists() {
+// 		bail!(
+// 			"{} key file not found at {:?}",
+// 			cert_pair.name,
+// 			cert_pair.key_path
+// 		);
+// 	}
+//
+// 	tracing::debug!("Loading {} certificate from:", cert_pair.name);
+// 	tracing::debug!("  Cert: {:?}", cert_pair.cert_path);
+// 	tracing::debug!("  Key: {:?}", cert_pair.key_path);
+//
+// 	// Load certificate
+// 	let cert_file = match File::open(&cert_pair.cert_path) {
+// 		Ok(file) => file,
+// 		Err(e) => bail!("Failed to open {} certificate file: {}", cert_pair.name, e),
+// 	};
+// 	let cert_reader = &mut BufReader::new(cert_file);
+//
+// 	let cert_chain = match certs(cert_reader).collect::<Result<Vec<_>, _>>() {
+// 		Ok(chain) => chain,
+// 		Err(e) => bail!("Failed to parse {} certificate: {}", cert_pair.name, e),
+// 	};
+//
+// 	if cert_chain.is_empty() {
+// 		bail!(
+// 			"No certificates found in {} certificate file",
+// 			cert_pair.name
+// 		);
+// 	}
+//
+// 	// Load private key
+// 	let key_file = match File::open(&cert_pair.key_path) {
+// 		Ok(file) => file,
+// 		Err(e) => bail!("Failed to open {} key file: {}", cert_pair.name, e),
+// 	};
+// 	let key_reader = &mut BufReader::new(key_file);
+//
+// 	let key_der = match private_key(key_reader) {
+// 		Ok(Some(key)) => key,
+// 		Ok(None) => bail!("No private key found in {} key file", cert_pair.name),
+// 		Err(e) => bail!("Failed to parse {} key: {}", cert_pair.name, e),
+// 	};
+//
+// 	let signing_key = match any_supported_type(&key_der) {
+// 		Ok(key) => key,
+// 		Err(e) => bail!("Failed to load {} signing key: {}", cert_pair.name, e),
+// 	};
+//
+// 	tracing::info!("{} certificate loaded successfully", cert_pair.name);
+// 	Ok(Arc::new(CertifiedKey::new(cert_chain, signing_key)))
+// }
+
+/// Create a certificate resolver function for TLS
+///
+/// This function sets up a certificate resolver that will serve:
+/// - Actor certificate for hostnames that match the actor routing logic
+/// - API certificate for all other hostnames
+///
+/// It follows the same routing logic as the main routing function to ensure
+/// consistent behavior between routing and certificate selection.
+#[tracing::instrument(skip_all)]
+pub async fn create_cert_resolver(
+	_ctx: &gas::prelude::StandaloneCtx,
+) -> Result<Option<CertResolverFn>> {
+	return Ok(None);
+
+	// // Get the Guard configuration
+	// let guard_config = match ctx.config().guard() {
+	// 	Ok(config) => config,
+	// 	Err(e) => {
+	// 		tracing::warn!("Failed to get Guard configuration: {}", e);
+	// 		return Ok(None);
+	// 	}
+	// };
+	//
+	// // If HTTPS is not configured, return None
+	// let https_config = match &guard_config.https {
+	// 	Some(config) => config,
+	// 	None => {
+	// 		tracing::info!("HTTPS configuration not found in Guard config - TLS disabled");
+	// 		return Ok(None);
+	// 	}
+	// };
+	//
+	// // TLS config is directly in the HTTPS config section
+	// let tls_config = &https_config.tls;
+	//
+	// // Load certificates
+	// let api_cert = load_certified_key(&CertificatePair {
+	// 	name: "API",
+	// 	cert_path: Path::new(&tls_config.api_cert_path).into(),
+	// 	key_path: Path::new(&tls_config.api_key_path).into(),
+	// })?;
+	//
+	// // Load actor certificate
+	// let actor_cert = load_certified_key(&CertificatePair {
+	// 	name: "Actor",
+	// 	cert_path: Path::new(&tls_config.actor_cert_path).into(),
+	// 	key_path: Path::new(&tls_config.actor_key_path).into(),
+	// })?;
+	//
+	// // Get the datacenter ID from config
+	// let dc_id = ctx.config().server()?.rivet.edge()?.datacenter_id;
+	//
+	// // Get datacenter information to get the guard public hostname
+	// let dc_res = ctx
+	// 	.op(cluster::ops::datacenter::get::Input {
+	// 		datacenter_ids: vec![dc_id],
+	// 	})
+	// 	.await?;
+	//
+	// let dc = unwrap!(dc_res.datacenters.first());
+	// let guard_hostname = &dc.guard_public_hostname;
+	// let api_hostname = ctx
+	// 	.config()
+	// 	.server()?
+	// 	.rivet
+	// 	.edge_api_routing_host(&dc.name_id)?;
+	//
+	// tracing::info!("Using datacenter guard hostname: {:?}", guard_hostname);
+	// if let Some(api_host) = &api_hostname {
+	// 	tracing::info!("Using datacenter API hostname: {:?}", api_host);
+	// }
+	//
+	// // Get the hostname regexes for actor routing with hostname-based endpoint type
+	// let actor_hostname_regex_dynamic =
+	// 	match build_actor_hostname_and_path_regex(EndpointType::Hostname, guard_hostname) {
+	// 		Ok(Some((x, _))) => {
+	// 			tracing::info!("Successfully built dynamic hostname actor routing regex");
+	// 			Some(x)
+	// 		}
+	// 		Ok(None) => {
+	// 			tracing::warn!(
+	// 				"Could not build dynamic hostname actor routing regex - pattern will be skipped"
+	// 			);
+	// 			None
+	// 		}
+	// 		Err(e) => bail!(
+	// 			"Failed to build dynamic hostname actor routing regex: {}",
+	// 			e
+	// 		),
+	// 	};
+	// let actor_hostname_regex_static =
+	// 	match build_actor_hostname_and_path_regex(EndpointType::Path, guard_hostname) {
+	// 		Ok(Some((x, _))) => {
+	// 			tracing::info!("Successfully built static path actor routing regex");
+	// 			Some(x)
+	// 		}
+	// 		Ok(None) => {
+	// 			tracing::warn!(
+	// 				"Could not build static path actor routing regex - pattern will be skipped"
+	// 			);
+	// 			None
+	// 		}
+	// 		Err(e) => bail!("Failed to build static path actor routing regex: {}", e),
+	// 	};
+	//
+	// // Create resolver function that matches the routing logic
+	// let api_cert_clone = api_cert.clone();
+	// let actor_cert_clone = actor_cert.clone();
+	// let api_hostname_clone = api_hostname.clone();
+	//
+	// let resolver_fn: CertResolverFn = Arc::new(
+	// 	move |hostname: &str| -> Result<Arc<CertifiedKey>, Box<dyn Error + Send + Sync>> {
+	// 		// Extract just the host, stripping the port if present
+	// 		let host = hostname.split(':').next().unwrap_or(hostname);
+	//
+	// 		// First check if hostname matches the actor pattern
+	// 		// This follows the same routing precedence as in routing/mod.rs
+	// 		if let Some(x) = &actor_hostname_regex_dynamic {
+	// 			if x.is_match(host) {
+	// 				tracing::debug!(
+	// 					"Using dynamic hostname actor certificate for hostname: {}",
+	// 					host
+	// 				);
+	// 				return Ok(actor_cert_clone.clone());
+	// 			}
+	// 		}
+	// 		if let Some(x) = &actor_hostname_regex_static {
+	// 			if x.is_match(host) {
+	// 				tracing::debug!(
+	// 					"Using static hostname actor certificate for hostname: {}",
+	// 					host
+	// 				);
+	// 				return Ok(actor_cert_clone.clone());
+	// 			}
+	// 		}
+	//
+	// 		// Then check if it matches the API hostname
+	// 		if let Some(api_host) = &api_hostname_clone {
+	// 			if host == api_host {
+	// 				tracing::debug!("Using API certificate for API hostname: {}", host);
+	// 				return Ok(api_cert_clone.clone());
+	// 			}
+	// 		}
+	//
+	// 		bail!("Did not match any routes")
+	// 	},
+	// );
+	//
+	// Ok(Some(resolver_fn))
+}