rivetkit-core 3.0.0-preview.3

use anyhow::{Context, Result};
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use parking_lot::RwLock;
use rand::RngCore;
use rivet_error::RivetError as RivetErrorDerive;
use serde::{Deserialize, Serialize};
use std::sync::OnceLock;
use subtle::ConstantTimeEq;

use crate::ActorContext;
use crate::actor::internal_storage;
use crate::actor::keys::INSPECTOR_TOKEN_KEY;

/// Test-only override. Not a public/production auth mechanism. Production
/// inspector auth verifies the bearer token against the token stored in
/// internal SQLite. The legacy KV entry is a write-only mirror kept for
/// dashboard compatibility and is not consulted for auth.
const INSPECTOR_TOKEN_ENV: &str = "_RIVET_TEST_INSPECTOR_TOKEN";
const INSPECTOR_TOKEN_BYTES: usize = 32;

static INSPECTOR_TEST_TOKEN_OVERRIDE: OnceLock<RwLock<Option<String>>> = OnceLock::new();



pub fn set_test_inspector_token_override(token: Option<String>) {
	*INSPECTOR_TEST_TOKEN_OVERRIDE
		.get_or_init(|| RwLock::new(None))
		.write() = token.filter(|token| !token.is_empty());
}

#[derive(Clone, Copy, Debug, Default)]
pub struct InspectorAuth;

#[derive(RivetErrorDerive, Clone, Debug, Deserialize, Serialize)]
#[error(
	"inspector",
	"unauthorized",
	"Inspector request requires a valid bearer token"
)]
struct InspectorUnauthorized;

impl InspectorAuth {
	pub fn new() -> Self {
		Self
	}

	pub async fn verify(&self, ctx: &ActorContext, bearer_token: Option<&str>) -> Result<()> {
		let Some(bearer_token) = bearer_token.filter(|token| !token.is_empty()) else {
			return Err(InspectorUnauthorized.build());
		};

		if let Some(configured_token) = configured_test_token() {
			return verify_token_bytes(bearer_token.as_bytes(), configured_token.as_bytes());
		}

		let stored_token = internal_storage::load_inspector_token(ctx.sql())
			.await
			.ok()
			.flatten()
			.map(String::into_bytes)
			.ok_or_else(|| InspectorUnauthorized.build())?;

		verify_token_bytes(bearer_token.as_bytes(), &stored_token)
	}
}

/// Ensures the actor has an inspector token persisted in SQLite and, for now,
/// mirrored in KV so the engine-facing KV API can serve it to the dashboard.
/// Skips writes when the token and its mirror match. No-ops when the
/// `_RIVET_TEST_INSPECTOR_TOKEN` env override is set, since that takes
/// precedence over any 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 configured_test_token().is_some() {
		return Ok(());
	}

	if let Some(token) = internal_storage::load_inspector_token(ctx.sql())
		.await
		.context("load inspector token from sqlite")?
	{
		return ensure_inspector_token_mirror(ctx, &token).await;
	}

	if let Some(token) = new_inspector_token() {
		let (sqlite, mirror) = futures::join!(
			internal_storage::persist_inspector_token(ctx.sql(), &token),
			persist_inspector_token_mirror(ctx, &token),
		);
		sqlite.context("persist inspector token to sqlite")?;
		mirror?;
		tracing::debug!(actor_id = %ctx.actor_id(), "generated new inspector token");
	}
	Ok(())
}

/// Rewrites the KV mirror of the stored `token` when it is missing or stale. SQLite is
/// authoritative, and a bootstrap writes the mirror only after its SQLite commit, so a crash
/// between them leaves the mirror missing.
pub(crate) async fn ensure_inspector_token_mirror(ctx: &ActorContext, token: &str) -> Result<()> {
	if configured_test_token().is_some() {
		return Ok(());
	}
	let mirror = ctx
		.legacy_kv()
		.batch_get(&[&INSPECTOR_TOKEN_KEY])
		.await?
		.pop()
		.flatten();
	if mirror.as_deref() != Some(token.as_bytes()) {
		persist_inspector_token_mirror(ctx, token).await?;
	}
	Ok(())
}

/// Produces bootstrap token material unless the process uses the test override.
pub(crate) fn new_inspector_token() -> Option<String> {
	configured_test_token()
		.is_none()
		.then(generate_inspector_token)
}

/// Durably writes the dashboard mirror. SQLite remains authoritative for auth.
pub(crate) async fn persist_inspector_token_mirror(ctx: &ActorContext, token: &str) -> Result<()> {
	ctx.legacy_kv()
		.put(&INSPECTOR_TOKEN_KEY, token.as_bytes())
		.await
		.context("persist inspector token mirror to kv")
}

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 configured_test_token() -> Option<String> {
	std::env::var(INSPECTOR_TOKEN_ENV)
		.ok()
		.filter(|token| !token.is_empty())
		.or_else(|| {
			INSPECTOR_TEST_TOKEN_OVERRIDE
				.get()
				.and_then(|token| token.read().clone())
		})
}

fn verify_token_bytes(candidate: &[u8], expected: &[u8]) -> Result<()> {
	if candidate.ct_eq(expected).into() {
		Ok(())
	} else {
		Err(InspectorUnauthorized.build())
	}
}