rivetkit_core/inspector/
auth.rs1use 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
16const 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
72pub 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}