From af41630520f4d0d00edf163a5e2b0f1235347b8f Mon Sep 17 00:00:00 2001 From: andrew Date: Sat, 18 Apr 2026 03:31:43 +0300 Subject: [PATCH] Extract TokenSource trait for bearer token loading Pure refactor. No behavior change. Introduces a TokenSource trait so additional backends (AWS Secrets Manager, Vault, etc.) can plug in behind feature flags without touching the server wiring. - New module crates/omnigraph-server/src/auth.rs with the TokenSource trait and a single EnvOrFileTokenSource implementation that delegates to the existing server_bearer_tokens_from_env() function. - serve() now constructs EnvOrFileTokenSource and calls load() instead of calling the free function directly. - The trait has a supports_refresh() hook (false for env/file) for future implementations that can rotate without restart. - async-trait added to omnigraph-server deps; it's already in the workspace. Tests: - Unit tests in auth.rs covering load paths and the default supports_refresh / name values. - Existing 128 tests (lib + integration + openapi) pass unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.lock | 1 + crates/omnigraph-server/Cargo.toml | 1 + crates/omnigraph-server/src/auth.rs | 106 ++++++++++++++++++++++++++++ crates/omnigraph-server/src/lib.rs | 5 +- 4 files changed, 112 insertions(+), 1 deletion(-) create mode 100644 crates/omnigraph-server/src/auth.rs diff --git a/Cargo.lock b/Cargo.lock index db7d525..0998210 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4585,6 +4585,7 @@ dependencies = [ name = "omnigraph-server" version = "0.2.2" dependencies = [ + "async-trait", "axum", "cedar-policy", "clap", diff --git a/crates/omnigraph-server/Cargo.toml b/crates/omnigraph-server/Cargo.toml index d4724dc..a659a69 100644 --- a/crates/omnigraph-server/Cargo.toml +++ b/crates/omnigraph-server/Cargo.toml @@ -30,6 +30,7 @@ cedar-policy = { workspace = true } futures = { workspace = true } sha2 = { workspace = true } subtle = { workspace = true } +async-trait = { workspace = true } [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/omnigraph-server/src/auth.rs b/crates/omnigraph-server/src/auth.rs new file mode 100644 index 0000000..1466c2e --- /dev/null +++ b/crates/omnigraph-server/src/auth.rs @@ -0,0 +1,106 @@ +//! Bearer token sources. +//! +//! A `TokenSource` loads `(actor_id, token)` pairs that the server uses to +//! authenticate incoming bearer tokens. Plaintext tokens returned here are +//! hashed immediately by `AppState` on ingest — see `hash_bearer_token` — +//! and never persist past startup/refresh. +//! +//! The trait exists so that additional backends (AWS Secrets Manager, +//! HashiCorp Vault, etc.) can plug in behind feature flags without +//! touching the server wiring. + +use async_trait::async_trait; +use color_eyre::eyre::Result; + +use crate::server_bearer_tokens_from_env; + +/// A source of bearer tokens, returned as `(actor_id, token)` pairs in +/// plaintext. The caller is expected to hash tokens before storing them. +#[async_trait] +pub trait TokenSource: Send + Sync { + /// Fetch the current set of actor → token pairs. + /// + /// Called once at startup. Implementations that support rotation may + /// also be polled periodically. + async fn load(&self) -> Result>; + + /// Whether this source can be re-fetched for rotation without restart. + /// Default: false (one-shot sources). + fn supports_refresh(&self) -> bool { + false + } + + /// Human-readable name for logs and error messages. + fn name(&self) -> &'static str; +} + +/// Reads bearer tokens from environment variables and / or files, matching +/// the long-standing server configuration: +/// +/// - `OMNIGRAPH_SERVER_BEARER_TOKEN` — a single token assigned to the +/// implicit actor `default`. +/// - `OMNIGRAPH_SERVER_BEARER_TOKENS_JSON` — a JSON object of +/// `{"actor_id": "token", …}`. +/// - `OMNIGRAPH_SERVER_BEARER_TOKENS_FILE` — a path to a JSON file of the +/// same shape. +/// +/// Does not support refresh — reloading means restarting the process. +#[derive(Debug, Default, Clone)] +pub struct EnvOrFileTokenSource; + +#[async_trait] +impl TokenSource for EnvOrFileTokenSource { + async fn load(&self) -> Result> { + server_bearer_tokens_from_env() + } + + fn name(&self) -> &'static str { + "env-or-file" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::env; + use serial_test::serial; + + fn clear_env() { + unsafe { + env::remove_var("OMNIGRAPH_SERVER_BEARER_TOKEN"); + env::remove_var("OMNIGRAPH_SERVER_BEARER_TOKENS_JSON"); + env::remove_var("OMNIGRAPH_SERVER_BEARER_TOKENS_FILE"); + } + } + + #[tokio::test] + #[serial] + async fn env_or_file_source_returns_empty_when_nothing_configured() { + clear_env(); + let source = EnvOrFileTokenSource; + let tokens = source.load().await.unwrap(); + assert!(tokens.is_empty()); + } + + #[tokio::test] + #[serial] + async fn env_or_file_source_reads_single_token_as_default_actor() { + clear_env(); + unsafe { + env::set_var("OMNIGRAPH_SERVER_BEARER_TOKEN", "some-token"); + } + let source = EnvOrFileTokenSource; + let tokens = source.load().await.unwrap(); + unsafe { + env::remove_var("OMNIGRAPH_SERVER_BEARER_TOKEN"); + } + assert_eq!(tokens, vec![("default".to_string(), "some-token".to_string())]); + } + + #[tokio::test] + async fn env_or_file_source_does_not_support_refresh() { + let source = EnvOrFileTokenSource; + assert!(!source.supports_refresh()); + assert_eq!(source.name(), "env-or-file"); + } +} diff --git a/crates/omnigraph-server/src/lib.rs b/crates/omnigraph-server/src/lib.rs index 3fa8787..dd373b8 100644 --- a/crates/omnigraph-server/src/lib.rs +++ b/crates/omnigraph-server/src/lib.rs @@ -1,4 +1,5 @@ pub mod api; +pub mod auth; pub mod config; pub mod policy; @@ -37,6 +38,7 @@ use omnigraph::error::{ManifestErrorKind, OmniError}; use omnigraph_compiler::json_params_to_param_map; use omnigraph_compiler::query::parser::parse_query; use omnigraph_compiler::{JsonParamMode, ParamMap}; +pub use auth::{EnvOrFileTokenSource, TokenSource}; pub use policy::{ PolicyAction, PolicyCompiler, PolicyConfig, PolicyDecision, PolicyEngine, PolicyExpectation, PolicyRequest, PolicyTestConfig, @@ -462,9 +464,10 @@ pub fn build_app(state: AppState) -> Router { } pub async fn serve(config: ServerConfig) -> Result<()> { + let token_source = EnvOrFileTokenSource; let state = AppState::open_with_bearer_tokens_and_policy( config.uri.clone(), - server_bearer_tokens_from_env()?, + token_source.load().await?, config.policy_file.as_ref(), ) .await?;