Skip to main content

rivetkit_core/inspector/
auth.rs

1use anyhow::{Context, Result};
2use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
3use parking_lot::RwLock;
4use rand::RngCore;
5use rivet_error::RivetError as RivetErrorDerive;
6use serde::{Deserialize, Serialize};
7#[cfg(test)]
8use std::sync::Mutex;
9use std::sync::OnceLock;
10use subtle::ConstantTimeEq;
11
12use crate::ActorContext;
13use crate::actor::keys::INSPECTOR_TOKEN_KEY;
14use crate::actor::preload::PreloadedKv;
15
16/// Test-only override. Not a public/production auth mechanism; production
17/// inspector auth goes through the per-actor KV token.
18const INSPECTOR_TOKEN_ENV: &str = "_RIVET_TEST_INSPECTOR_TOKEN";
19const INSPECTOR_TOKEN_BYTES: usize = 32;
20
21static INSPECTOR_TEST_TOKEN_OVERRIDE: OnceLock<RwLock<Option<String>>> = OnceLock::new();
22
23#[cfg(test)]
24pub(crate) fn test_inspector_env_lock() -> &'static Mutex<()> {
25	static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
26	LOCK.get_or_init(|| Mutex::new(()))
27}
28
29pub fn set_test_inspector_token_override(token: Option<String>) {
30	*INSPECTOR_TEST_TOKEN_OVERRIDE
31		.get_or_init(|| RwLock::new(None))
32		.write() = token.filter(|token| !token.is_empty());
33}
34
35#[derive(Clone, Copy, Debug, Default)]
36pub struct InspectorAuth;
37
38#[derive(RivetErrorDerive, Clone, Debug, Deserialize, Serialize)]
39#[error(
40	"inspector",
41	"unauthorized",
42	"Inspector request requires a valid bearer token"
43)]
44struct InspectorUnauthorized;
45
46impl InspectorAuth {
47	pub fn new() -> Self {
48		Self
49	}
50
51	pub async fn verify(&self, ctx: &ActorContext, bearer_token: Option<&str>) -> Result<()> {
52		let Some(bearer_token) = bearer_token.filter(|token| !token.is_empty()) else {
53			return Err(InspectorUnauthorized.build());
54		};
55
56		if let Some(configured_token) = configured_test_token() {
57			return verify_token_bytes(bearer_token.as_bytes(), configured_token.as_bytes());
58		}
59
60		let stored_token = ctx
61			.kv_internal()
62			.get(&INSPECTOR_TOKEN_KEY)
63			.await
64			.ok()
65			.flatten()
66			.ok_or_else(|| InspectorUnauthorized.build())?;
67
68		verify_token_bytes(bearer_token.as_bytes(), &stored_token)
69	}
70}
71
72/// Ensures the actor has an inspector token persisted in KV so the
73/// engine-facing KV API can serve the token to the dashboard inspector.
74/// Skips the write when the token already exists. No-ops when the
75/// `_RIVET_TEST_INSPECTOR_TOKEN` env override is set, since that takes
76/// precedence over any KV-stored token and we do not want to pin a per-actor
77/// token that will never be consulted.
78pub async fn init_inspector_token(ctx: &ActorContext) -> Result<()> {
79	init_inspector_token_with_preload(ctx, None).await
80}
81
82pub(crate) async fn init_inspector_token_with_preload(
83	ctx: &ActorContext,
84	preloaded_kv: Option<&PreloadedKv>,
85) -> Result<()> {
86	if configured_test_token().is_some() {
87		return Ok(());
88	}
89
90	let existing =
91		match preloaded_kv.and_then(|preloaded| preloaded.key_entry(&INSPECTOR_TOKEN_KEY)) {
92			Some(existing) => existing,
93			None => ctx
94				.kv_internal()
95				.get(&INSPECTOR_TOKEN_KEY)
96				.await
97				.context("load inspector token")?,
98		};
99	if existing.is_some() {
100		return Ok(());
101	}
102
103	let token = generate_inspector_token();
104	ctx.kv_internal()
105		.put(&INSPECTOR_TOKEN_KEY, token.as_bytes())
106		.await
107		.context("persist inspector token")?;
108	tracing::debug!(actor_id = %ctx.actor_id(), "generated new inspector token");
109	Ok(())
110}
111
112fn generate_inspector_token() -> String {
113	let mut bytes = [0u8; INSPECTOR_TOKEN_BYTES];
114	rand::thread_rng().fill_bytes(&mut bytes);
115	URL_SAFE_NO_PAD.encode(bytes)
116}
117
118fn configured_test_token() -> Option<String> {
119	std::env::var(INSPECTOR_TOKEN_ENV)
120		.ok()
121		.filter(|token| !token.is_empty())
122		.or_else(|| {
123			INSPECTOR_TEST_TOKEN_OVERRIDE
124				.get()
125				.and_then(|token| token.read().clone())
126		})
127}
128
129fn verify_token_bytes(candidate: &[u8], expected: &[u8]) -> Result<()> {
130	if candidate.ct_eq(expected).into() {
131		Ok(())
132	} else {
133		Err(InspectorUnauthorized.build())
134	}
135}