Skip to main content

rivetkit_core/
serverless.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::time::Duration;
5
6use anyhow::{Context, Result};
7use http::StatusCode;
8use rivet_envoy_client::config::{ActorName as EnvoyActorName, EnvoyConfig};
9use rivet_envoy_client::envoy::start_envoy as start_envoy_client;
10use rivet_envoy_client::handle::EnvoyHandle;
11use rivet_envoy_client::protocol;
12use rivetkit_shared_types::serverless_metadata::{
13	ActorName, ServerlessMetadataEnvoy, ServerlessMetadataEnvoyKind, ServerlessMetadataPayload,
14};
15use serde::Serialize;
16use serde_json::json;
17use tokio::sync::{Mutex as TokioMutex, mpsc};
18use tokio_util::sync::CancellationToken;
19use url::Url;
20
21use crate::actor::factory::ActorFactory;
22#[cfg(feature = "native-runtime")]
23use crate::engine_process::{EngineProcessManager, EngineResolverConfig};
24use crate::registry::{
25	CoreEnvoyHandle, CoreEnvoyStatus, RegistryCallbacks, RegistryDispatcher, ServeConfig,
26	should_manage_engine,
27};
28use crate::runtime::RuntimeSpawner;
29use crate::time::{sleep, timeout};
30
31const DEFAULT_BASE_PATH: &str = "/api/rivet";
32const SSE_PING_INTERVAL: Duration = Duration::from_secs(1);
33const SSE_PING_FRAME: &[u8] = b"event: ping\ndata:\n\n";
34const SSE_STOPPING_FRAME: &[u8] = b"event: stopping\ndata:\n\n";
35/// Bound on `handle.shutdown_and_wait` inside teardown paths. If envoy cannot
36/// reach the engine (reconnect loop stuck), we fall back to immediate `Stop`
37/// rather than hanging indefinitely. Must stay below the outer TS grace ceiling.
38const SHUTDOWN_DRAIN_TIMEOUT: Duration = Duration::from_secs(20);
39
40#[derive(Clone)]
41pub struct CoreServerlessRuntime {
42	settings: Arc<ServerlessSettings>,
43	dispatcher: Arc<RegistryDispatcher>,
44	envoy: Arc<TokioMutex<Option<EnvoyHandle>>>,
45	#[cfg(feature = "native-runtime")]
46	_engine_process: Arc<TokioMutex<Option<EngineProcessManager>>>,
47	shutting_down: Arc<AtomicBool>,
48}
49
50#[derive(Clone, Debug)]
51struct ServerlessSettings {
52	version: u32,
53	configured_endpoint: String,
54	configured_namespace: String,
55	base_path: String,
56	package_version: String,
57	client_endpoint: Option<String>,
58	client_namespace: Option<String>,
59	client_token: Option<String>,
60	validate_endpoint: bool,
61	max_start_payload_bytes: usize,
62	cache_envoy: bool,
63}
64
65#[derive(Debug)]
66pub struct ServerlessRequest {
67	pub method: String,
68	pub url: String,
69	pub headers: HashMap<String, String>,
70	pub body: Vec<u8>,
71	pub cancel_token: CancellationToken,
72}
73
74#[derive(Debug)]
75pub struct ServerlessResponse {
76	pub status: u16,
77	pub headers: HashMap<String, String>,
78	pub body: mpsc::UnboundedReceiver<Result<Vec<u8>, ServerlessStreamError>>,
79}
80
81#[derive(Clone, Debug, Serialize)]
82pub struct ServerlessStreamError {
83	pub group: String,
84	pub code: String,
85	pub message: String,
86}
87
88#[derive(Debug)]
89struct StartHeaders {
90	endpoint: String,
91	token: Option<String>,
92	pool_name: String,
93	namespace: String,
94}
95
96#[derive(Debug, Serialize)]
97struct ServerlessErrorBody<'a> {
98	group: &'a str,
99	code: &'a str,
100	message: String,
101	metadata: serde_json::Value,
102}
103
104#[derive(rivet_error::RivetError, Serialize)]
105#[error("request", "invalid", "Invalid request.", "Invalid request: {reason}")]
106struct InvalidRequest {
107	reason: String,
108}
109
110#[derive(rivet_error::RivetError, Serialize)]
111#[error(
112	"config",
113	"endpoint_mismatch",
114	"Endpoint mismatch.",
115	"Endpoint mismatch: expected \"{expected}\", received \"{received}\""
116)]
117struct EndpointMismatch {
118	expected: String,
119	received: String,
120}
121
122#[derive(rivet_error::RivetError, Serialize)]
123#[error(
124	"config",
125	"namespace_mismatch",
126	"Namespace mismatch.",
127	"Namespace mismatch: expected \"{expected}\", received \"{received}\""
128)]
129struct NamespaceMismatch {
130	expected: String,
131	received: String,
132}
133
134#[derive(rivet_error::RivetError, Serialize)]
135#[error(
136	"message",
137	"incoming_too_long",
138	"Incoming message too long.",
139	"Incoming message too long. Exceeded limit of {limit} bytes."
140)]
141struct IncomingMessageTooLong {
142	limit: usize,
143}
144
145#[derive(rivet_error::RivetError, Serialize)]
146#[error(
147	"registry",
148	"shut_down",
149	"Registry is shut down.",
150	"Registry is shut down; no new requests can be accepted."
151)]
152struct RuntimeShutDown;
153
154impl CoreServerlessRuntime {
155	pub(crate) async fn new(
156		factories: HashMap<String, Arc<ActorFactory>>,
157		config: ServeConfig,
158	) -> Result<Self> {
159		#[cfg(feature = "native-runtime")]
160		let engine_process = if should_manage_engine(&config.endpoint, config.engine_spawn)? {
161			Some(
162				EngineProcessManager::start_or_reuse(EngineResolverConfig::from_parts(
163					&config.endpoint,
164					config.engine_binary_path.clone(),
165					config.engine_host.clone(),
166					config.engine_port,
167					config.engine_auto_download,
168				))
169				.await?,
170			)
171		} else {
172			None
173		};
174		#[cfg(not(feature = "native-runtime"))]
175		if should_manage_engine(&config.endpoint, config.engine_spawn)? {
176			anyhow::bail!("engine process spawning requires the `native-runtime` feature");
177		}
178
179		let dispatcher = Arc::new(RegistryDispatcher::new(
180			factories,
181			config.handle_inspector_http_in_runtime,
182		));
183		let base_path = normalize_base_path(config.serverless_base_path.as_deref());
184		crate::metrics_endpoint::record_rivetkit_info(
185			config.serverless_package_version.clone(),
186			config.version,
187			"serverless",
188			config.pool_name.clone(),
189		);
190
191		Ok(Self {
192			settings: Arc::new(ServerlessSettings {
193				version: config.version,
194				configured_endpoint: config.endpoint,
195				configured_namespace: config.namespace,
196				base_path,
197				package_version: config.serverless_package_version,
198				client_endpoint: config.serverless_client_endpoint,
199				client_namespace: config.serverless_client_namespace,
200				client_token: config.serverless_client_token,
201				validate_endpoint: config.serverless_validate_endpoint,
202				max_start_payload_bytes: config.serverless_max_start_payload_bytes,
203				cache_envoy: config.serverless_cache_envoy,
204			}),
205			dispatcher,
206			envoy: Arc::new(TokioMutex::new(None)),
207			#[cfg(feature = "native-runtime")]
208			_engine_process: Arc::new(TokioMutex::new(engine_process)),
209			shutting_down: Arc::new(AtomicBool::new(false)),
210		})
211	}
212
213	/// Tear down the cached envoy handle. Idempotent.
214	///
215	/// Sets `shutting_down` so concurrent `ensure_envoy` callers short-circuit
216	/// instead of starting a fresh envoy after teardown, and waits (with a
217	/// bounded timeout) for `envoy_loop` to exit. If the drain exceeds the
218	/// timeout (e.g. engine unreachable), falls back to an immediate `Stop`.
219	pub async fn shutdown(&self) {
220		self.shutting_down.store(true, Ordering::Release);
221		let handle = { self.envoy.lock().await.take() };
222		let Some(handle) = handle else { return };
223		match timeout(SHUTDOWN_DRAIN_TIMEOUT, handle.shutdown_and_wait(false)).await {
224			Ok(()) => {}
225			Err(_) => {
226				tracing::warn!(
227					"serverless runtime envoy drain exceeded timeout; forcing immediate stop"
228				);
229				handle.shutdown(true);
230				handle.wait_stopped().await;
231			}
232		}
233	}
234
235	pub async fn active_envoy_actor_count(&self) -> Option<usize> {
236		self.active_envoy_status()
237			.await
238			.map(|status| status.active_actor_count)
239	}
240
241	pub async fn active_envoy_status(&self) -> Option<CoreEnvoyStatus> {
242		self.envoy
243			.lock()
244			.await
245			.as_ref()
246			.map(|handle| CoreEnvoyHandle::new(handle.clone()).status())
247	}
248
249	pub async fn active_envoy_actor_stop_threshold_ms(&self) -> Option<i64> {
250		let handle = self.envoy.lock().await.as_ref().cloned()?;
251		CoreEnvoyHandle::new(handle).actor_stop_threshold_ms().await
252	}
253
254	/// Listener-side body cap; reuses the `/start` payload limit.
255	pub fn max_request_body_bytes(&self) -> usize {
256		self.settings.max_start_payload_bytes
257	}
258
259	/// Returns whether the native listener should reserve this URL for
260	/// RivetKit's framework routes instead of forwarding it to an application
261	/// fallback.
262	pub fn handles_listener_request(&self, url: &str) -> bool {
263		handles_listener_request(&self.settings.base_path, url)
264	}
265
266	/// Canonical 413 response built through the `RivetError` system.
267	pub fn incoming_too_long_response(&self) -> ServerlessResponse {
268		let error = IncomingMessageTooLong {
269			limit: self.settings.max_start_payload_bytes,
270		}
271		.build();
272		error_response(error)
273	}
274
275	/// Canonical 400 response for malformed requests.
276	pub fn invalid_request_response(&self, reason: impl Into<String>) -> ServerlessResponse {
277		let error = InvalidRequest {
278			reason: reason.into(),
279		}
280		.build();
281		error_response(error)
282	}
283
284	pub async fn handle_request(&self, req: ServerlessRequest) -> ServerlessResponse {
285		let cors = cors_headers(&req);
286		match self.handle_request_inner(req).await {
287			Ok(mut response) => {
288				apply_cors(&mut response.headers, cors);
289				response
290			}
291			Err(error) => {
292				let mut response = error_response(error);
293				apply_cors(&mut response.headers, cors);
294				response
295			}
296		}
297	}
298
299	async fn handle_request_inner(&self, req: ServerlessRequest) -> Result<ServerlessResponse> {
300		let path = route_path(&self.settings.base_path, &req.url)?;
301		match (req.method.as_str(), path.as_str()) {
302			("GET", "") | ("GET", "/") => Ok(text_response(
303				StatusCode::OK,
304				"text/plain; charset=utf-8",
305				"This is a RivetKit server.\n\nLearn more at https://rivet.dev",
306			)),
307			("GET", "/health") => {
308				// Healthy if no envoy is connected yet or if the envoy has received a
309				// recent engine ping. Unhealthy only when an envoy exists but has not
310				// received a recent ping. 503 is the conventional "recycle me" signal
311				// for container hosts running behind an HTTP health probe.
312				let runtime_healthy = {
313					let guard = self.envoy.lock().await;
314					guard
315						.as_ref()
316						.map(|handle| handle.is_ping_healthy())
317						.unwrap_or(true)
318				};
319				if runtime_healthy {
320					Ok(json_response(
321						StatusCode::OK,
322						json!({
323							"status": "ok",
324							"runtime": "rivetkit",
325							"version": self.settings.package_version,
326						}),
327					))
328				} else {
329					Ok(json_response(
330						StatusCode::SERVICE_UNAVAILABLE,
331						json!({
332							"status": "engine_ping_stale",
333							"runtime": "rivetkit",
334							"version": self.settings.package_version,
335						}),
336					))
337				}
338			}
339			("GET", "/metadata") => Ok(self.metadata_response()),
340			("GET", "/metrics") => Ok(metrics_response(&req.headers)),
341			("GET", "/start") | ("POST", "/start") => self.start_response(req).await,
342			("OPTIONS", _) => Ok(bytes_response(
343				StatusCode::NO_CONTENT,
344				HashMap::new(),
345				Vec::new(),
346			)),
347			_ => Ok(text_response(
348				StatusCode::NOT_FOUND,
349				"text/plain; charset=utf-8",
350				"Not Found (RivetKit)",
351			)),
352		}
353	}
354
355	async fn start_response(&self, req: ServerlessRequest) -> Result<ServerlessResponse> {
356		let headers = parse_start_headers(&req.headers)?;
357		self.validate_start_headers(&headers)?;
358		crate::metrics_endpoint::record_rivetkit_info(
359			self.settings.package_version.clone(),
360			self.settings.version,
361			"serverless",
362			headers.pool_name.clone(),
363		);
364		if req.body.len() > self.settings.max_start_payload_bytes {
365			return Err(IncomingMessageTooLong {
366				limit: self.settings.max_start_payload_bytes,
367			}
368			.build());
369		}
370
371		let handle = self.ensure_envoy(&headers).await?;
372		let payload = req.body;
373		let actor_start = handle.decode_serverless_actor_start(&payload)?;
374		let cancel_token = req.cancel_token;
375		let cache_envoy = self.settings.cache_envoy;
376		let (tx, rx) = mpsc::unbounded_channel();
377		let _ = tx.send(Ok(SSE_PING_FRAME.to_vec()));
378
379		RuntimeSpawner::spawn(async move {
380			let shutdown_handle = handle.clone();
381			let result = tokio::select! {
382				_ = cancel_token.cancelled() => {
383					if !cache_envoy {
384						shutdown_handle.shutdown_and_wait(false).await;
385					}
386					return;
387				}
388				result = handle.start_serverless_actor(&payload) => result,
389			};
390			if let Err(error) = result {
391				let error = stream_error(error);
392				let _ = tx.send(Err(error));
393				if !cache_envoy {
394					handle.shutdown_and_wait(false).await;
395				}
396				return;
397			}
398
399			loop {
400				tokio::select! {
401					_ = cancel_token.cancelled() => {
402						break;
403					}
404					_ = handle.wait_actor_registered_then_stopped(&actor_start.actor_id, actor_start.generation) => {
405						let _ = tx.send(Ok(SSE_STOPPING_FRAME.to_vec()));
406						break;
407					}
408					_ = sleep(SSE_PING_INTERVAL) => {
409						if tx.send(Ok(SSE_PING_FRAME.to_vec())).is_err() {
410							break;
411						}
412					}
413				}
414			}
415
416			if !cache_envoy {
417				handle.shutdown_and_wait(false).await;
418			}
419		});
420
421		Ok(ServerlessResponse {
422			status: StatusCode::OK.as_u16(),
423			headers: HashMap::from([
424				("content-type".to_owned(), "text/event-stream".to_owned()),
425				("cache-control".to_owned(), "no-cache".to_owned()),
426				("connection".to_owned(), "keep-alive".to_owned()),
427			]),
428			body: rx,
429		})
430	}
431
432	fn metadata_response(&self) -> ServerlessResponse {
433		let actor_names = self
434			.dispatcher
435			.build_actor_metadata_map()
436			.into_iter()
437			.map(|(name, metadata)| {
438				(
439					name,
440					ActorName {
441						metadata: Some(metadata),
442					},
443				)
444			})
445			.collect::<HashMap<_, _>>();
446
447		let payload = ServerlessMetadataPayload {
448			runtime: "rivetkit".to_owned(),
449			version: self.settings.package_version.clone(),
450			envoy_protocol_version: Some(protocol::PROTOCOL_VERSION),
451			actor_names,
452			envoy: Some(ServerlessMetadataEnvoy {
453				kind: Some(ServerlessMetadataEnvoyKind::Serverless {}),
454				version: Some(self.settings.version),
455			}),
456			runner: None,
457			client_endpoint: self.settings.client_endpoint.clone(),
458			client_namespace: self.settings.client_namespace.clone(),
459			client_token: self.settings.client_token.clone(),
460		};
461
462		let response = serde_json::to_value(payload).unwrap_or_else(|_| json!({}));
463
464		json_response(StatusCode::OK, response)
465	}
466
467	fn validate_start_headers(&self, headers: &StartHeaders) -> Result<()> {
468		if self.settings.validate_endpoint {
469			if !endpoints_match(&headers.endpoint, &self.settings.configured_endpoint) {
470				tracing::warn!(
471					configured_endpoint = %self.settings.configured_endpoint,
472					received_endpoint = %headers.endpoint,
473					"serverless start rejected: endpoint mismatch",
474				);
475				return Err(EndpointMismatch {
476					expected: self.settings.configured_endpoint.clone(),
477					received: headers.endpoint.clone(),
478				}
479				.build());
480			}
481
482			if headers.namespace != self.settings.configured_namespace {
483				tracing::warn!(
484					configured_namespace = %self.settings.configured_namespace,
485					received_namespace = %headers.namespace,
486					"serverless start rejected: namespace mismatch",
487				);
488				return Err(NamespaceMismatch {
489					expected: self.settings.configured_namespace.clone(),
490					received: headers.namespace.clone(),
491				}
492				.build());
493			}
494		}
495
496		Ok(())
497	}
498
499	async fn ensure_envoy(&self, headers: &StartHeaders) -> Result<EnvoyHandle> {
500		if self.shutting_down.load(Ordering::Acquire) {
501			return Err(RuntimeShutDown.build());
502		}
503		if !self.settings.cache_envoy {
504			return self.start_envoy(headers).await;
505		}
506		let mut guard = self.envoy.lock().await;
507		if let Some(handle) = guard.as_ref() {
508			// The start request token authenticates the serverless callback. It is not part
509			// of envoy identity, and may differ from the token used for the engine connection.
510			if !endpoints_match(handle.endpoint(), &headers.endpoint)
511				|| handle.namespace() != headers.namespace
512				|| handle.pool_name() != headers.pool_name
513			{
514				anyhow::bail!("serverless start headers do not match active envoy");
515			}
516			return Ok(handle.clone());
517		}
518
519		let handle = self.start_envoy(headers).await?;
520		// Re-check under the lock: shutdown may have run while we were awaiting
521		// `start_envoy`. If so, tear down the freshly-built envoy rather than
522		// installing it into the cache.
523		if self.shutting_down.load(Ordering::Acquire) {
524			drop(guard);
525			match timeout(SHUTDOWN_DRAIN_TIMEOUT, handle.shutdown_and_wait(false)).await {
526				Ok(()) => {}
527				Err(_) => {
528					handle.shutdown(true);
529					handle.wait_stopped().await;
530				}
531			}
532			return Err(RuntimeShutDown.build());
533		}
534		*guard = Some(handle.clone());
535		Ok(handle)
536	}
537
538	async fn start_envoy(&self, headers: &StartHeaders) -> Result<EnvoyHandle> {
539		let callbacks = Arc::new(RegistryCallbacks {
540			dispatcher: self.dispatcher.clone(),
541		});
542		let prepopulate_actor_names = self
543			.dispatcher
544			.build_actor_metadata_map()
545			.into_iter()
546			.map(|(name, metadata)| (name, EnvoyActorName { metadata }))
547			.collect();
548		// not_global: true to avoid caching the handle in the process-wide
549		// `GLOBAL_ENVOY` OnceLock. Without this, a shutdown-during-build race
550		// (spec ยง3 step 7) leaves a dead handle cached for the life of the
551		// process and any subsequent consumer gets it back.
552		Ok(start_envoy_client(EnvoyConfig {
553			version: self.settings.version,
554			endpoint: headers.endpoint.clone(),
555			token: headers.token.clone(),
556			namespace: headers.namespace.clone(),
557			pool_name: headers.pool_name.clone(),
558			prepopulate_actor_names,
559			metadata: Some(json!({
560				"rivetkit": { "version": self.settings.package_version },
561			})),
562			not_global: true,
563			debug_latency_ms: None,
564			callbacks,
565		})
566		.await)
567	}
568}
569
570fn route_path(base_path: &str, url: &str) -> Result<String> {
571	let parsed = Url::parse(url).with_context(|| format!("parse request URL `{url}`"))?;
572	let path = parsed.path();
573	if path == base_path {
574		return Ok(String::new());
575	}
576	let prefix = format!("{base_path}/");
577	if let Some(rest) = path.strip_prefix(&prefix) {
578		return Ok(format!("/{rest}"));
579	}
580	Ok(path.to_owned())
581}
582
583fn handles_listener_request(base_path: &str, url: &str) -> bool {
584	let Ok(parsed) = Url::parse(url) else {
585		// Let the normal framework handler return its structured invalid URL
586		// response rather than passing malformed input to user code.
587		return true;
588	};
589	let request_path = parsed.path();
590	if request_path != base_path && !request_path.starts_with(&format!("{base_path}/")) {
591		return false;
592	}
593	let path = route_path(base_path, url).expect("URL was parsed and the same base path is valid");
594	matches!(
595		path.as_str(),
596		"" | "/" | "/health" | "/metadata" | "/metrics" | "/start"
597	)
598}
599
600fn parse_start_headers(headers: &HashMap<String, String>) -> Result<StartHeaders> {
601	let pool_name = match optional_header(headers, "x-rivet-pool-name") {
602		Some(pool_name) => pool_name,
603		None => optional_header(headers, "x-rivet-runner-name").ok_or_else(|| {
604			InvalidRequest {
605				reason: "x-rivet-pool-name header is required".to_string(),
606			}
607			.build()
608		})?,
609	};
610
611	Ok(StartHeaders {
612		endpoint: required_header(headers, "x-rivet-endpoint")?,
613		token: optional_header(headers, "x-rivet-token"),
614		pool_name,
615		namespace: required_header(headers, "x-rivet-namespace-name")?,
616	})
617}
618
619fn required_header(headers: &HashMap<String, String>, name: &str) -> Result<String> {
620	headers
621		.get(name)
622		.filter(|value| !value.is_empty())
623		.cloned()
624		.ok_or_else(|| {
625			InvalidRequest {
626				reason: format!("{name} header is required"),
627			}
628			.build()
629		})
630}
631
632fn optional_header(headers: &HashMap<String, String>, name: &str) -> Option<String> {
633	headers.get(name).filter(|value| !value.is_empty()).cloned()
634}
635
636fn cors_headers(req: &ServerlessRequest) -> HashMap<String, String> {
637	let origin = req
638		.headers
639		.get("origin")
640		.cloned()
641		.unwrap_or_else(|| "*".to_owned());
642	let mut headers = HashMap::from([
643		("access-control-allow-origin".to_owned(), origin.clone()),
644		(
645			"access-control-allow-credentials".to_owned(),
646			"true".to_owned(),
647		),
648		("access-control-expose-headers".to_owned(), "*".to_owned()),
649	]);
650	if origin != "*" {
651		headers.insert("vary".to_owned(), "Origin".to_owned());
652	}
653
654	if req.method == "OPTIONS" {
655		headers.insert(
656			"access-control-allow-methods".to_owned(),
657			"GET, POST, PUT, DELETE, OPTIONS, PATCH".to_owned(),
658		);
659		headers.insert(
660			"access-control-allow-headers".to_owned(),
661			req.headers
662				.get("access-control-request-headers")
663				.cloned()
664				.unwrap_or_else(|| "*".to_owned()),
665		);
666		headers.insert("access-control-max-age".to_owned(), "86400".to_owned());
667	}
668
669	headers
670}
671
672fn apply_cors(headers: &mut HashMap<String, String>, cors: HashMap<String, String>) {
673	headers.extend(cors);
674}
675
676fn normalize_base_path(base_path: Option<&str>) -> String {
677	let base_path = base_path
678		.filter(|base_path| !base_path.is_empty())
679		.unwrap_or(DEFAULT_BASE_PATH);
680	let prefixed = if base_path.starts_with('/') {
681		base_path.to_owned()
682	} else {
683		format!("/{base_path}")
684	};
685	let trimmed = prefixed.trim_end_matches('/');
686	if trimmed.is_empty() {
687		"/".to_owned()
688	} else {
689		trimmed.to_owned()
690	}
691}
692
693fn text_response(status: StatusCode, content_type: &str, body: &str) -> ServerlessResponse {
694	bytes_response(
695		status,
696		HashMap::from([("content-type".to_owned(), content_type.to_owned())]),
697		body.as_bytes().to_vec(),
698	)
699}
700
701fn json_response(status: StatusCode, body: serde_json::Value) -> ServerlessResponse {
702	bytes_response(
703		status,
704		HashMap::from([("content-type".to_owned(), "application/json".to_owned())]),
705		serde_json::to_vec(&body).unwrap_or_else(|_| b"{}".to_vec()),
706	)
707}
708
709fn metrics_response(headers: &HashMap<String, String>) -> ServerlessResponse {
710	let bearer_token = crate::metrics_endpoint::authorization_bearer_token_map(headers);
711	match crate::metrics_endpoint::authorize_metrics_request(bearer_token) {
712		Ok(()) => match crate::metrics_endpoint::render_prometheus_metrics() {
713			Ok(metrics) => bytes_response(
714				StatusCode::OK,
715				HashMap::from([("content-type".to_owned(), metrics.content_type)]),
716				metrics.body,
717			),
718			Err(error) => error_response(error),
719		},
720		Err(crate::metrics_endpoint::MetricsAccessError::NotEnabled) => text_response(
721			StatusCode::FORBIDDEN,
722			"text/plain; charset=utf-8",
723			"metrics not enabled\n",
724		),
725		Err(crate::metrics_endpoint::MetricsAccessError::Unauthorized) => text_response(
726			StatusCode::UNAUTHORIZED,
727			"text/plain; charset=utf-8",
728			"metrics request requires a valid bearer token\n",
729		),
730	}
731}
732
733fn bytes_response(
734	status: StatusCode,
735	headers: HashMap<String, String>,
736	body: Vec<u8>,
737) -> ServerlessResponse {
738	let (tx, rx) = mpsc::unbounded_channel();
739	let _ = tx.send(Ok(body));
740	ServerlessResponse {
741		status: status.as_u16(),
742		headers,
743		body: rx,
744	}
745}
746
747fn error_response(error: anyhow::Error) -> ServerlessResponse {
748	let extracted = rivet_error::RivetError::extract(&error);
749	let status = serverless_error_status(extracted.group(), extracted.code());
750	let body = ServerlessErrorBody {
751		group: extracted.group(),
752		code: extracted.code(),
753		message: extracted.message().to_owned(),
754		metadata: extracted.metadata().unwrap_or(serde_json::Value::Null),
755	};
756	bytes_response(
757		status,
758		HashMap::from([("content-type".to_owned(), "application/json".to_owned())]),
759		serde_json::to_vec(&body).unwrap_or_else(|_| b"{}".to_vec()),
760	)
761}
762
763fn serverless_error_status(group: &str, code: &str) -> StatusCode {
764	match (group, code) {
765		("auth", "forbidden") => StatusCode::FORBIDDEN,
766		("message", "incoming_too_long") => StatusCode::PAYLOAD_TOO_LARGE,
767		_ => StatusCode::BAD_REQUEST,
768	}
769}
770
771fn stream_error(error: anyhow::Error) -> ServerlessStreamError {
772	let extracted = rivet_error::RivetError::extract(&error);
773	ServerlessStreamError {
774		group: extracted.group().to_owned(),
775		code: extracted.code().to_owned(),
776		message: extracted.message().to_owned(),
777	}
778}
779
780pub fn normalize_endpoint_url(url: &str) -> Option<String> {
781	let parsed = Url::parse(url).ok()?;
782	let pathname = if parsed.path() == "/" {
783		"/".to_owned()
784	} else {
785		parsed.path().trim_end_matches('/').to_owned()
786	};
787	let mut hostname = parsed.host_str()?.to_owned();
788	if is_loopback_address(&hostname) {
789		hostname = "localhost".to_owned();
790	}
791	hostname = normalize_regional_hostname(&hostname);
792	let host = match parsed.port() {
793		Some(port) => format!("{hostname}:{port}"),
794		None => hostname,
795	};
796	Some(format!("{}://{}{}", parsed.scheme(), host, pathname))
797}
798
799fn normalized_endpoint_candidates(value: &str) -> Vec<String> {
800	value
801		.split(',')
802		.map(str::trim)
803		.filter(|candidate| !candidate.is_empty())
804		.map(|candidate| normalize_endpoint_url(candidate).unwrap_or_else(|| candidate.to_owned()))
805		.collect()
806}
807
808pub fn endpoints_match(a: &str, b: &str) -> bool {
809	let a_candidates = normalized_endpoint_candidates(a);
810	let b_candidates = normalized_endpoint_candidates(b);
811	a_candidates.iter().any(|a_candidate| {
812		b_candidates
813			.iter()
814			.any(|b_candidate| a_candidate == b_candidate)
815	})
816}
817
818fn normalize_regional_hostname(hostname: &str) -> String {
819	if !hostname.ends_with(".rivet.dev") || !hostname.starts_with("api-") {
820		return hostname.to_owned();
821	}
822	let without_prefix = &hostname[4..];
823	let Some(first_dot_index) = without_prefix.find('.') else {
824		return hostname.to_owned();
825	};
826	let domain = &without_prefix[first_dot_index + 1..];
827	format!("api.{domain}")
828}
829
830fn is_loopback_address(hostname: &str) -> bool {
831	matches!(hostname, "127.0.0.1" | "0.0.0.0" | "::1" | "[::1]")
832}
833
834// Test shim keeps moved tests in crate-root tests/ with private-module access.
835#[cfg(test)]
836#[path = "../tests/serverless.rs"]
837mod tests;