rivetkit_core/
metrics_endpoint.rs

1use std::{collections::HashMap, sync::LazyLock};
2
3use anyhow::{Context, Result};
4use parking_lot::Mutex;
5use rivet_metrics::prometheus::{
6	Encoder, IntGaugeVec, TextEncoder, register_int_gauge_vec_with_registry,
7};
8use subtle::ConstantTimeEq;
9
10const METRICS_ENABLED_ENV: &str = "RIVETKIT_METRICS_ENABLED";
11const METRICS_TOKEN_ENV: &str = "RIVETKIT_METRICS_TOKEN";
12
13static RIVETKIT_INFO: LazyLock<IntGaugeVec> = LazyLock::new(|| {
14	register_int_gauge_vec_with_registry!(
15		"rivetkit_info",
16		"Static RivetKit build information.",
17		&[
18			"runtime",
19			"version",
20			"type",
21			"envoy_version",
22			"envoy_kind",
23			"pool_name",
24		],
25		*rivet_metrics::REGISTRY
26	)
27	.unwrap()
28});
29
30static CURRENT_RIVETKIT_INFO: LazyLock<Mutex<Option<RivetKitInfo>>> =
31	LazyLock::new(|| Mutex::new(None));
32
33#[derive(Clone, Debug, Eq, PartialEq)]
34struct RivetKitInfo {
35	runtime: String,
36	version: String,
37	runtime_type: String,
38	envoy_version: String,
39	envoy_kind: String,
40	pool_name: String,
41}
42
43impl RivetKitInfo {
44	fn labels(&self) -> [&str; 6] {
45		[
46			&self.runtime,
47			&self.version,
48			&self.runtime_type,
49			&self.envoy_version,
50			&self.envoy_kind,
51			&self.pool_name,
52		]
53	}
54}
55
56pub struct RenderedMetrics {
57	pub content_type: String,
58	pub body: Vec<u8>,
59}
60
61pub enum MetricsAccessError {
62	NotEnabled,
63	Unauthorized,
64}
65
66pub fn runtime_type() -> &'static str {
67	if std::env::var("NODE_ENV").as_deref() == Ok("production") {
68		"deployed"
69	} else {
70		"local"
71	}
72}
73
74pub fn record_rivetkit_info(
75	version: impl Into<String>,
76	envoy_version: u32,
77	envoy_kind: impl Into<String>,
78	pool_name: impl Into<String>,
79) {
80	record_rivetkit_info_inner(RivetKitInfo {
81		runtime: "rivetkit".to_owned(),
82		version: version.into(),
83		runtime_type: runtime_type().to_owned(),
84		envoy_version: envoy_version.to_string(),
85		envoy_kind: envoy_kind.into(),
86		pool_name: pool_name.into(),
87	});
88}
89
90pub fn authorize_metrics_request(
91	bearer_token: Option<&str>,
92) -> std::result::Result<(), MetricsAccessError> {
93	let Some(configured_token) = configured_metrics_token() else {
94		return Err(MetricsAccessError::NotEnabled);
95	};
96
97	let Some(bearer_token) = bearer_token.filter(|token| !token.is_empty()) else {
98		return Err(MetricsAccessError::Unauthorized);
99	};
100
101	if bearer_token
102		.as_bytes()
103		.ct_eq(configured_token.as_bytes())
104		.into()
105	{
106		Ok(())
107	} else {
108		Err(MetricsAccessError::Unauthorized)
109	}
110}
111
112pub fn render_prometheus_metrics() -> Result<RenderedMetrics> {
113	ensure_rivetkit_info_recorded();
114
115	let encoder = TextEncoder::new();
116	let metric_families = rivet_metrics::REGISTRY.gather();
117	let mut body = Vec::new();
118	encoder
119		.encode(&metric_families, &mut body)
120		.context("encode prometheus metrics")?;
121
122	Ok(RenderedMetrics {
123		content_type: encoder.format_type().to_owned(),
124		body,
125	})
126}
127
128pub fn authorization_bearer_token(headers: &http::HeaderMap) -> Option<&str> {
129	headers
130		.get(http::header::AUTHORIZATION)
131		.and_then(|value| value.to_str().ok())
132		.and_then(bearer_token_from_authorization)
133}
134
135pub fn authorization_bearer_token_map(headers: &HashMap<String, String>) -> Option<&str> {
136	headers
137		.iter()
138		.find(|(name, _)| name.eq_ignore_ascii_case(http::header::AUTHORIZATION.as_str()))
139		.and_then(|(_, value)| bearer_token_from_authorization(value))
140}
141
142fn configured_metrics_token() -> Option<String> {
143	let enabled = std::env::var(METRICS_ENABLED_ENV).ok()?;
144	if enabled != "1" {
145		return None;
146	}
147
148	std::env::var(METRICS_TOKEN_ENV)
149		.ok()
150		.filter(|token| !token.is_empty())
151}
152
153fn ensure_rivetkit_info_recorded() {
154	if CURRENT_RIVETKIT_INFO.lock().is_some() {
155		return;
156	}
157
158	let envoy_version = std::env::var("RIVET_ENVOY_VERSION")
159		.ok()
160		.and_then(|value| value.parse().ok())
161		.unwrap_or(1);
162	let pool_name = std::env::var("RIVET_POOL_NAME").unwrap_or_else(|_| "rivetkit-rust".to_owned());
163	record_rivetkit_info(
164		env!("CARGO_PKG_VERSION"),
165		envoy_version,
166		"unknown",
167		pool_name,
168	);
169}
170
171fn record_rivetkit_info_inner(info: RivetKitInfo) {
172	let mut current = CURRENT_RIVETKIT_INFO.lock();
173	if let Some(previous) = current.as_ref()
174		&& previous != &info
175	{
176		RIVETKIT_INFO.with_label_values(&previous.labels()).set(0);
177	}
178
179	RIVETKIT_INFO.with_label_values(&info.labels()).set(1);
180	*current = Some(info);
181}
182
183fn bearer_token_from_authorization(value: &str) -> Option<&str> {
184	let value = value.trim_start();
185	let scheme = value.get(..6)?;
186	if !scheme.eq_ignore_ascii_case("bearer") {
187		return None;
188	}
189
190	let rest = value.get(6..)?;
191	if !rest.chars().next().is_some_and(char::is_whitespace) {
192		return None;
193	}
194
195	let token = rest.trim_start();
196	if token.is_empty() { None } else { Some(token) }
197}
rivetkit_core/metrics_endpoint.rs

rivetkit_core/
metrics_endpoint.rs
