Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions rivetkit-rust/packages/rivetkit-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions rivetkit-rust/packages/rivetkit-core/src/actor/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
43 changes: 42 additions & 1 deletion rivetkit-rust/packages/rivetkit-core/src/inspector/auth.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
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};

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;
Expand Down Expand Up @@ -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(())
Expand Down
2 changes: 1 addition & 1 deletion rivetkit-rust/packages/rivetkit-core/src/inspector/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn Fn(InspectorSignal) + Send + Sync>;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ impl RegistryDispatcher {
)
})
}
// TODO: Impelment when ready
(http::Method::GET, "/inspector/traces") => json_http_response(
StatusCode::OK,
&json!({
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -280,15 +280,15 @@ mod moved_tests {
assert_eq!(error.code(), "unauthorized");

unsafe {
std::env::remove_var("RIVET_INSPECTOR_TOKEN");
std::env::remove_var("_RIVET_TEST_INSPECTOR_TOKEN");
}
}

#[tokio::test]
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();
Expand Down Expand Up @@ -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(
Expand Down
Loading