diff --git a/Cargo.lock b/Cargo.lock index 0774d60274..ec71594979 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5233,12 +5233,14 @@ name = "rivetkit-core" version = "2.3.0-rc.4" dependencies = [ "anyhow", + "base64 0.22.1", "ciborium", "futures", "http 1.3.1", "nix 0.30.1", "parking_lot", "prometheus", + "rand 0.8.5", "reqwest", "rivet-envoy-client", "rivet-error", diff --git a/rivetkit-rust/packages/rivetkit-core/Cargo.toml b/rivetkit-rust/packages/rivetkit-core/Cargo.toml index 7b2bd4bb6d..ffef48d6ed 100644 --- a/rivetkit-rust/packages/rivetkit-core/Cargo.toml +++ b/rivetkit-rust/packages/rivetkit-core/Cargo.toml @@ -12,12 +12,14 @@ sqlite = ["dep:rivetkit-sqlite"] [dependencies] anyhow.workspace = true +base64.workspace = true ciborium.workspace = true futures.workspace = true http.workspace = true nix.workspace = true parking_lot.workspace = true prometheus.workspace = true +rand.workspace = true reqwest.workspace = true rivet-pools.workspace = true rivet-util.workspace = true diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index 5018f44947..752d9246f7 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -1119,6 +1119,9 @@ impl ActorTask { .persist_state(SaveStateOpts { immediate: true }) .await .context("persist actor initialization")?; + crate::inspector::init_inspector_token(&self.ctx) + .await + .context("initialize inspector token")?; self.ctx .restore_hibernatable_connections_with_preload(self.preloaded_kv.as_ref()) .await diff --git a/rivetkit-rust/packages/rivetkit-core/src/inspector/auth.rs b/rivetkit-rust/packages/rivetkit-core/src/inspector/auth.rs index 39dd4b35f2..b21dbf54b9 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/inspector/auth.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/inspector/auth.rs @@ -1,4 +1,6 @@ -use anyhow::Result; +use anyhow::{Context, Result}; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use rand::RngCore; use rivet_error::RivetError as RivetErrorDerive; use serde::{Deserialize, Serialize}; @@ -6,6 +8,7 @@ use crate::ActorContext; const INSPECTOR_TOKEN_KEY: [u8; 1] = [3]; const INSPECTOR_TOKEN_ENV: &str = "RIVET_INSPECTOR_TOKEN"; +const INSPECTOR_TOKEN_BYTES: usize = 32; #[derive(Clone, Copy, Debug, Default)] pub struct InspectorAuth; @@ -47,6 +50,44 @@ impl InspectorAuth { } } +/// Ensures the actor has an inspector token persisted in KV at `[3]` so the +/// engine-facing KV API can serve the token to the dashboard inspector. +/// Skips the write when the token already exists. No-ops when the +/// `RIVET_INSPECTOR_TOKEN` env override is set, since that takes precedence +/// over any KV-stored token and we do not want to pin a per-actor token that +/// will never be consulted. +pub async fn init_inspector_token(ctx: &ActorContext) -> Result<()> { + if std::env::var(INSPECTOR_TOKEN_ENV) + .ok() + .is_some_and(|token| !token.is_empty()) + { + return Ok(()); + } + + let existing = ctx + .kv() + .get(&INSPECTOR_TOKEN_KEY) + .await + .context("load inspector token")?; + if existing.is_some() { + return Ok(()); + } + + let token = generate_inspector_token(); + ctx.kv() + .put(&INSPECTOR_TOKEN_KEY, token.as_bytes()) + .await + .context("persist inspector token")?; + tracing::debug!(actor_id = %ctx.actor_id(), "generated new inspector token"); + Ok(()) +} + +fn generate_inspector_token() -> String { + let mut bytes = [0u8; INSPECTOR_TOKEN_BYTES]; + rand::thread_rng().fill_bytes(&mut bytes); + URL_SAFE_NO_PAD.encode(bytes) +} + fn verify_token_bytes(candidate: &[u8], expected: &[u8]) -> Result<()> { if timing_safe_equal(candidate, expected) { Ok(()) diff --git a/rivetkit-rust/packages/rivetkit-core/src/inspector/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/inspector/mod.rs index 88aa6a0eed..8704ee18be 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/inspector/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/inspector/mod.rs @@ -7,7 +7,7 @@ use parking_lot::RwLock; pub mod auth; pub(crate) mod protocol; -pub use auth::InspectorAuth; +pub use auth::{InspectorAuth, init_inspector_token}; type InspectorListener = Arc; diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs index 28e30fa33b..9c2b5966bf 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/inspector.rs @@ -156,6 +156,7 @@ impl RegistryDispatcher { ) }) } + // TODO: Impelment when ready (http::Method::GET, "/inspector/traces") => json_http_response( StatusCode::OK, &json!({ diff --git a/rivetkit-rust/packages/rivetkit-core/tests/modules/inspector.rs b/rivetkit-rust/packages/rivetkit-core/tests/modules/inspector.rs index ffcbdd21c3..60546b237f 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/modules/inspector.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/modules/inspector.rs @@ -251,7 +251,7 @@ mod moved_tests { async fn inspector_auth_uses_env_token_before_kv_fallback() { let _env_guard = INSPECTOR_ENV_LOCK.lock().expect("env lock poisoned"); unsafe { - std::env::set_var("RIVET_INSPECTOR_TOKEN", "env-token"); + std::env::set_var("_RIVET_TEST_INSPECTOR_TOKEN", "env-token"); } let kv = crate::kv::tests::new_in_memory(); @@ -280,7 +280,7 @@ mod moved_tests { assert_eq!(error.code(), "unauthorized"); unsafe { - std::env::remove_var("RIVET_INSPECTOR_TOKEN"); + std::env::remove_var("_RIVET_TEST_INSPECTOR_TOKEN"); } } @@ -288,7 +288,7 @@ mod moved_tests { async fn inspector_auth_falls_back_to_actor_kv_token() { let _env_guard = INSPECTOR_ENV_LOCK.lock().expect("env lock poisoned"); unsafe { - std::env::remove_var("RIVET_INSPECTOR_TOKEN"); + std::env::remove_var("_RIVET_TEST_INSPECTOR_TOKEN"); } let kv = crate::kv::tests::new_in_memory(); @@ -321,7 +321,7 @@ mod moved_tests { async fn inspector_auth_rejects_missing_token() { let _env_guard = INSPECTOR_ENV_LOCK.lock().expect("env lock poisoned"); unsafe { - std::env::remove_var("RIVET_INSPECTOR_TOKEN"); + std::env::remove_var("_RIVET_TEST_INSPECTOR_TOKEN"); } let ctx = new_with_kv(