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::internal_storage;
14use crate::actor::keys::INSPECTOR_TOKEN_KEY;
15
16/// Test-only override. Not a public/production auth mechanism. Production
17/// inspector auth verifies the bearer token against the token stored in
18/// internal SQLite. The legacy KV entry is a write-only mirror kept for
19/// dashboard compatibility and is not consulted for auth.
20const INSPECTOR_TOKEN_ENV: &str = "_RIVET_TEST_INSPECTOR_TOKEN";
21const INSPECTOR_TOKEN_BYTES: usize = 32;
22
23static INSPECTOR_TEST_TOKEN_OVERRIDE: OnceLock<RwLock<Option<String>>> = OnceLock::new();
24
25#[cfg(test)]
26pub(crate) fn test_inspector_env_lock() -> &'static Mutex<()> {
27	static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
28	LOCK.get_or_init(|| Mutex::new(()))
29}
30
31pub fn set_test_inspector_token_override(token: Option<String>) {
32	*INSPECTOR_TEST_TOKEN_OVERRIDE
33		.get_or_init(|| RwLock::new(None))
34		.write() = token.filter(|token| !token.is_empty());
35}
36
37#[derive(Clone, Copy, Debug, Default)]
38pub struct InspectorAuth;
39
40#[derive(RivetErrorDerive, Clone, Debug, Deserialize, Serialize)]
41#[error(
42	"inspector",
43	"unauthorized",
44	"Inspector request requires a valid bearer token"
45)]
46struct InspectorUnauthorized;
47
48impl InspectorAuth {
49	pub fn new() -> Self {
50		Self
51	}
52
53	pub async fn verify(&self, ctx: &ActorContext, bearer_token: Option<&str>) -> Result<()> {
54		let Some(bearer_token) = bearer_token.filter(|token| !token.is_empty()) else {
55			return Err(InspectorUnauthorized.build());
56		};
57
58		if let Some(configured_token) = configured_test_token() {
59			return verify_token_bytes(bearer_token.as_bytes(), configured_token.as_bytes());
60		}
61
62		let stored_token = internal_storage::load_inspector_token(ctx.sql())
63			.await
64			.ok()
65			.flatten()
66			.map(String::into_bytes)
67			.ok_or_else(|| InspectorUnauthorized.build())?;
68
69		verify_token_bytes(bearer_token.as_bytes(), &stored_token)
70	}
71}
72
73/// Ensures the actor has an inspector token persisted in SQLite and, for now,
74/// mirrored in KV so the engine-facing KV API can serve it to the dashboard.
75/// Skips writes when the token already exists. No-ops when the
76/// `_RIVET_TEST_INSPECTOR_TOKEN` env override is set, since that takes
77/// precedence over any stored token and we do not want to pin a per-actor token
78/// that will never be consulted.
79pub async fn init_inspector_token(ctx: &ActorContext) -> Result<()> {
80	if configured_test_token().is_some() {
81		return Ok(());
82	}
83
84	let existing = internal_storage::load_inspector_token(ctx.sql())
85		.await
86		.context("load inspector token from sqlite")?
87		.map(String::into_bytes);
88	if existing.is_some() {
89		// Token creation and legacy import both leave the compatibility mirror in
90		// KV, so an existing SQLite token does not need another startup write.
91		return Ok(());
92	}
93
94	let token = generate_inspector_token();
95	internal_storage::persist_inspector_token(ctx.sql(), &token)
96		.await
97		.context("persist inspector token to sqlite")?;
98	// The dashboard reads the token through the engine's public actor-KV
99	// endpoint. Keep this sole KV write until the dashboard fetches inspector
100	// tokens another way; see ~/.agents/specs/kv-to-sqlite-internal-storage.md.
101	ctx.legacy_kv()
102		.put(&INSPECTOR_TOKEN_KEY, token.as_bytes())
103		.await
104		.context("persist inspector token mirror to kv")?;
105	tracing::debug!(actor_id = %ctx.actor_id(), "generated new inspector token");
106	Ok(())
107}
108
109fn generate_inspector_token() -> String {
110	let mut bytes = [0u8; INSPECTOR_TOKEN_BYTES];
111	rand::thread_rng().fill_bytes(&mut bytes);
112	URL_SAFE_NO_PAD.encode(bytes)
113}
114
115fn configured_test_token() -> Option<String> {
116	std::env::var(INSPECTOR_TOKEN_ENV)
117		.ok()
118		.filter(|token| !token.is_empty())
119		.or_else(|| {
120			INSPECTOR_TEST_TOKEN_OVERRIDE
121				.get()
122				.and_then(|token| token.read().clone())
123		})
124}
125
126fn verify_token_bytes(candidate: &[u8], expected: &[u8]) -> Result<()> {
127	if candidate.ct_eq(expected).into() {
128		Ok(())
129	} else {
130		Err(InspectorUnauthorized.build())
131	}
132}