From 7a3bf5c75851830bc00640e41c51b79b29e9769d Mon Sep 17 00:00:00 2001 From: andrew Date: Sat, 18 Apr 2026 03:48:51 +0300 Subject: [PATCH] Add aws feature + SecretsManagerTokenSource backend Introduces an opt-in AWS Secrets Manager backend for bearer tokens, behind the `aws` Cargo feature. Default builds (on-prem, local dev) don't pull in the AWS SDK and don't pay its compile cost. - New Cargo feature `aws` gates the `aws-config` + `aws-sdk-secretsmanager` optional deps. Default features remain empty. - New `auth::aws::SecretsManagerTokenSource` implements `TokenSource` by fetching a JSON `{"actor_id": "token", ...}` payload from a named Secrets Manager secret. Credentials resolve via the AWS default chain (env, shared config, IMDSv2 instance role, ECS task role) so no explicit plumbing is needed under an IAM role. - New `resolve_token_source()` dispatches based on the `OMNIGRAPH_SERVER_BEARER_TOKENS_AWS_SECRET` env var. If the var is set but the binary was built without `--features aws`, returns a clear rebuild instruction rather than silently falling back. - `serve()` now uses `resolve_token_source()` and logs which source was selected at startup. - `parse_json_secret_payload()` is factored out as a free function so the payload validation (trim whitespace, reject blank actor/token, reject non-object) is unit-testable without the AWS SDK. - New CI job `test_aws_feature` builds + tests with `--features aws`. Not in this PR (follow-ups): - Background refresh loop for rotation. `SecretsManagerTokenSource` advertises `supports_refresh: true` but the AppState-level refresh task isn't wired yet. - Config-YAML dispatch (today the AWS source is selected via env var only; eventually `server.bearer_tokens.source` in `omnigraph.yaml`). Tests: - Default-feature build: 33 lib + 41 integration + 64 openapi. - `--features aws` build: 32 lib (one test is cfg-gated) + 41 + 64. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/ci.yml | 46 +++++++ Cargo.lock | 195 ++++++++++++++++++++++---- crates/omnigraph-server/Cargo.toml | 8 ++ crates/omnigraph-server/src/auth.rs | 206 +++++++++++++++++++++++++++- crates/omnigraph-server/src/lib.rs | 5 +- 5 files changed, 432 insertions(+), 28 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7827df8..40eba61 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -140,6 +140,52 @@ jobs: if: needs.classify_changes.outputs.run_full_ci == 'true' run: cargo test --workspace --locked + test_aws_feature: + name: Test omnigraph-server --features aws + needs: classify_changes + runs-on: ubuntu-latest + timeout-minutes: 30 + permissions: + contents: read + env: + CARGO_TERM_COLOR: always + steps: + - name: Skip for text-only changes + if: needs.classify_changes.outputs.run_full_ci != 'true' + run: echo "Text-only change detected; skipping aws feature build." + + - name: Checkout source + if: needs.classify_changes.outputs.run_full_ci == 'true' + uses: actions/checkout@v5.0.1 + + - name: Install system dependencies + if: needs.classify_changes.outputs.run_full_ci == 'true' + run: | + sudo apt-get update + sudo apt-get install -y protobuf-compiler libprotobuf-dev + + - name: Install Rust stable + if: needs.classify_changes.outputs.run_full_ci == 'true' + uses: dtolnay/rust-toolchain@stable + with: + toolchain: stable + + - name: Cache Rust build data + if: needs.classify_changes.outputs.run_full_ci == 'true' + uses: Swatinem/rust-cache@v2 + with: + workspaces: | + . -> target + key: aws-feature + + - name: Build omnigraph-server with aws feature + if: needs.classify_changes.outputs.run_full_ci == 'true' + run: cargo build --locked -p omnigraph-server --features aws + + - name: Test omnigraph-server with aws feature + if: needs.classify_changes.outputs.run_full_ci == 'true' + run: cargo test --locked -p omnigraph-server --features aws + rustfs_integration: name: RustFS S3 Integration needs: diff --git a/Cargo.lock b/Cargo.lock index 0998210..7332d52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -597,6 +597,30 @@ dependencies = [ "uuid", ] +[[package]] +name = "aws-sdk-secretsmanager" +version = "1.103.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae64963d3d16d8070aaa2fb79c11cd3b13f44d2f13bba3fe8f49dcd2c42f2987" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sso" version = "1.97.0" @@ -733,17 +757,23 @@ dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", "aws-smithy-types", - "h2", + "h2 0.3.27", + "h2 0.4.13", + "http 0.2.12", "http 1.4.0", - "hyper", - "hyper-rustls", + "http-body 0.4.6", + "hyper 0.14.32", + "hyper 1.8.1", + "hyper-rustls 0.24.2", + "hyper-rustls 0.27.7", "hyper-util", "pin-project-lite", - "rustls", + "rustls 0.21.12", + "rustls 0.23.36", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower", "tracing", ] @@ -827,6 +857,7 @@ dependencies = [ "base64-simd", "bytes", "bytes-utils", + "futures-core", "http 0.2.12", "http 1.4.0", "http-body 0.4.6", @@ -839,6 +870,8 @@ dependencies = [ "ryu", "serde", "time", + "tokio", + "tokio-util", ] [[package]] @@ -878,7 +911,7 @@ dependencies = [ "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.8.1", "hyper-util", "itoa", "matchit", @@ -2786,6 +2819,25 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap 2.13.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.13" @@ -2966,6 +3018,30 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.10", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.8.1" @@ -2976,7 +3052,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2", + "h2 0.4.13", "http 1.4.0", "http-body 1.0.1", "httparse", @@ -2989,6 +3065,21 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-rustls" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" +dependencies = [ + "futures-util", + "http 0.2.12", + "hyper 0.14.32", + "log", + "rustls 0.21.12", + "tokio", + "tokio-rustls 0.24.1", +] + [[package]] name = "hyper-rustls" version = "0.27.7" @@ -2996,13 +3087,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ "http 1.4.0", - "hyper", + "hyper 1.8.1", "hyper-util", - "rustls", + "rustls 0.23.36", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower-service", "webpki-roots", ] @@ -3019,12 +3110,12 @@ dependencies = [ "futures-util", "http 1.4.0", "http-body 1.0.1", - "hyper", + "hyper 1.8.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.2", "tokio", "tower-service", "tracing", @@ -4465,7 +4556,7 @@ dependencies = [ "http-body-util", "httparse", "humantime", - "hyper", + "hyper 1.8.1", "itertools 0.14.0", "md-5", "parking_lot", @@ -4586,6 +4677,8 @@ name = "omnigraph-server" version = "0.2.2" dependencies = [ "async-trait", + "aws-config", + "aws-sdk-secretsmanager", "axum", "cedar-policy", "clap", @@ -5169,8 +5262,8 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls", - "socket2", + "rustls 0.23.36", + "socket2 0.6.2", "thiserror", "tokio", "tracing", @@ -5189,7 +5282,7 @@ dependencies = [ "rand 0.9.2", "ring", "rustc-hash", - "rustls", + "rustls 0.23.36", "rustls-pki-types", "slab", "thiserror", @@ -5207,7 +5300,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.2", "tracing", "windows-sys 0.60.2", ] @@ -5491,12 +5584,12 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2", + "h2 0.4.13", "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper", - "hyper-rustls", + "hyper 1.8.1", + "hyper-rustls 0.27.7", "hyper-util", "js-sys", "log", @@ -5505,7 +5598,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls", + "rustls 0.23.36", "rustls-native-certs", "rustls-pki-types", "serde", @@ -5513,7 +5606,7 @@ dependencies = [ "serde_urlencoded", "sync_wrapper", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tokio-util", "tower", "tower-http", @@ -5644,6 +5737,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.21.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" +dependencies = [ + "log", + "ring", + "rustls-webpki 0.101.7", + "sct", +] + [[package]] name = "rustls" version = "0.23.36" @@ -5654,7 +5759,7 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki", + "rustls-webpki 0.103.9", "subtle", "zeroize", ] @@ -5690,6 +5795,16 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "rustls-webpki" version = "0.103.9" @@ -5797,6 +5912,16 @@ dependencies = [ "sha2", ] +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "sdd" version = "3.0.10" @@ -6135,6 +6260,16 @@ dependencies = [ "syn 2.0.115", ] +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.2" @@ -6602,7 +6737,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.2", "tokio-macros", "windows-sys 0.61.2", ] @@ -6618,13 +6753,23 @@ dependencies = [ "syn 2.0.115", ] +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls 0.21.12", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls", + "rustls 0.23.36", "tokio", ] diff --git a/crates/omnigraph-server/Cargo.toml b/crates/omnigraph-server/Cargo.toml index a659a69..1d2029f 100644 --- a/crates/omnigraph-server/Cargo.toml +++ b/crates/omnigraph-server/Cargo.toml @@ -12,6 +12,12 @@ documentation = "https://docs.rs/omnigraph-server" name = "omnigraph-server" path = "src/main.rs" +[features] +default = [] +# Enables the AWS Secrets Manager bearer-token source. Off by default — on-prem +# and local-dev builds don't pay the AWS SDK compile cost. +aws = ["dep:aws-config", "dep:aws-sdk-secretsmanager"] + [dependencies] omnigraph = { package = "omnigraph-engine", path = "../omnigraph", version = "0.2.2" } omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.2.2" } @@ -31,6 +37,8 @@ futures = { workspace = true } sha2 = { workspace = true } subtle = { workspace = true } async-trait = { workspace = true } +aws-config = { version = "1", optional = true, default-features = false, features = ["rustls", "rt-tokio", "credentials-process", "sso"] } +aws-sdk-secretsmanager = { version = "1", optional = true, default-features = false, features = ["rustls", "rt-tokio"] } [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/omnigraph-server/src/auth.rs b/crates/omnigraph-server/src/auth.rs index 1466c2e..80b6ed5 100644 --- a/crates/omnigraph-server/src/auth.rs +++ b/crates/omnigraph-server/src/auth.rs @@ -10,10 +10,15 @@ //! touching the server wiring. use async_trait::async_trait; -use color_eyre::eyre::Result; +use color_eyre::eyre::{Result, bail}; use crate::server_bearer_tokens_from_env; +/// Environment variable that, when set, selects AWS Secrets Manager as the +/// token source. Its value is the secret ID or ARN. Only honored when the +/// binary is compiled with `--features aws`. +pub const AWS_SECRET_ENV: &str = "OMNIGRAPH_SERVER_BEARER_TOKENS_AWS_SECRET"; + /// 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] @@ -59,6 +64,139 @@ impl TokenSource for EnvOrFileTokenSource { } } +/// Pick the token source based on configuration. +/// +/// Preference order: +/// 1. If `OMNIGRAPH_SERVER_BEARER_TOKENS_AWS_SECRET` is set AND the binary was +/// built with `--features aws`, returns an AWS Secrets Manager source. +/// 2. If that env var is set but the binary was built without the feature, +/// errors with a clear rebuild instruction rather than silently falling +/// back to the env/file source (which would hide the misconfiguration). +/// 3. Otherwise, returns `EnvOrFileTokenSource`. +pub async fn resolve_token_source() -> Result> { + if let Ok(secret_id) = std::env::var(AWS_SECRET_ENV) { + let secret_id = secret_id.trim().to_string(); + if !secret_id.is_empty() { + #[cfg(feature = "aws")] + { + let source = aws::SecretsManagerTokenSource::new(secret_id).await?; + return Ok(Box::new(source)); + } + #[cfg(not(feature = "aws"))] + { + bail!( + "{} is set but this binary was not built with --features aws. \ + Rebuild: cargo build --release --features aws", + AWS_SECRET_ENV + ); + } + } + } + Ok(Box::new(EnvOrFileTokenSource)) +} + +/// Parse a JSON secret payload (from AWS Secrets Manager or any equivalent +/// source) into actor → token pairs. +/// +/// Payload shape: `{"actor_id_1": "token_1", "actor_id_2": "token_2", ...}`. +/// Extracted as a free function so it can be unit-tested without the AWS SDK. +#[cfg(any(test, feature = "aws"))] +pub(crate) fn parse_json_secret_payload(payload: &str) -> Result> { + use std::collections::HashMap; + + let map: HashMap = serde_json::from_str(payload).map_err(|err| { + color_eyre::eyre::eyre!( + "bearer-token secret payload is not a JSON object of actor→token: {}", + err + ) + })?; + + let mut pairs: Vec<(String, String)> = Vec::with_capacity(map.len()); + for (actor, token) in map { + let actor = actor.trim().to_string(); + let token = token.trim().to_string(); + if actor.is_empty() { + bail!("bearer-token secret contains a blank actor id"); + } + if token.is_empty() { + bail!("bearer-token secret has a blank token for actor '{}'", actor); + } + pairs.push((actor, token)); + } + pairs.sort_by(|(a, _), (b, _)| a.cmp(b)); + Ok(pairs) +} + +#[cfg(feature = "aws")] +pub mod aws { + //! AWS Secrets Manager bearer-token backend. + //! + //! Fetches a JSON payload from a named secret on startup. Credentials are + //! resolved via the AWS default chain — env vars, shared config, IMDSv2 + //! instance role, or ECS task role — so no explicit credential plumbing + //! is needed when running under an IAM role. + //! + //! Background refresh for rotation is a follow-up. + use super::TokenSource; + use async_trait::async_trait; + use color_eyre::eyre::{Result, WrapErr, eyre}; + + /// Loads bearer tokens from a named AWS Secrets Manager secret. + pub struct SecretsManagerTokenSource { + client: aws_sdk_secretsmanager::Client, + secret_id: String, + } + + impl SecretsManagerTokenSource { + /// Construct a new source. Resolves AWS credentials + region via the + /// default chain — no explicit configuration needed on EC2/ECS/EKS. + pub async fn new(secret_id: impl Into) -> Result { + let config = + aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await; + let client = aws_sdk_secretsmanager::Client::new(&config); + Ok(Self { + client, + secret_id: secret_id.into(), + }) + } + } + + #[async_trait] + impl TokenSource for SecretsManagerTokenSource { + async fn load(&self) -> Result> { + let output = self + .client + .get_secret_value() + .secret_id(&self.secret_id) + .send() + .await + .wrap_err_with(|| { + format!("fetch AWS Secrets Manager secret '{}'", self.secret_id) + })?; + + let payload = output.secret_string().ok_or_else(|| { + eyre!( + "secret '{}' has no SecretString — binary secrets are not supported", + self.secret_id + ) + })?; + + super::parse_json_secret_payload(payload) + } + + fn supports_refresh(&self) -> bool { + true + } + + fn name(&self) -> &'static str { + "aws-secrets-manager" + } + } +} + +#[cfg(feature = "aws")] +pub use aws::SecretsManagerTokenSource; + #[cfg(test)] mod tests { use super::*; @@ -103,4 +241,70 @@ mod tests { assert!(!source.supports_refresh()); assert_eq!(source.name(), "env-or-file"); } + + #[test] + fn parse_json_secret_payload_reads_actor_token_map() { + let pairs = parse_json_secret_payload(r#"{"alice": "tok-a", "bob": "tok-b"}"#).unwrap(); + assert_eq!( + pairs, + vec![ + ("alice".to_string(), "tok-a".to_string()), + ("bob".to_string(), "tok-b".to_string()), + ] + ); + } + + #[test] + fn parse_json_secret_payload_trims_whitespace() { + let pairs = parse_json_secret_payload(r#"{" alice ": " tok-a "}"#).unwrap(); + assert_eq!(pairs, vec![("alice".to_string(), "tok-a".to_string())]); + } + + #[test] + fn parse_json_secret_payload_rejects_blank_actor() { + let err = parse_json_secret_payload(r#"{" ": "tok"}"#).unwrap_err(); + assert!(err.to_string().contains("blank actor")); + } + + #[test] + fn parse_json_secret_payload_rejects_blank_token() { + let err = parse_json_secret_payload(r#"{"alice": " "}"#).unwrap_err(); + assert!(err.to_string().contains("blank token")); + } + + #[test] + fn parse_json_secret_payload_rejects_non_object() { + let err = parse_json_secret_payload("[1, 2, 3]").unwrap_err(); + assert!(err.to_string().contains("not a JSON object")); + } + + #[tokio::test] + #[serial] + async fn resolve_token_source_falls_back_to_env_or_file_when_aws_var_unset() { + clear_env(); + unsafe { + env::remove_var(AWS_SECRET_ENV); + } + let source = resolve_token_source().await.unwrap(); + assert_eq!(source.name(), "env-or-file"); + } + + #[cfg(not(feature = "aws"))] + #[tokio::test] + #[serial] + async fn resolve_token_source_errors_when_aws_var_set_without_feature() { + clear_env(); + unsafe { + env::set_var(AWS_SECRET_ENV, "some-secret-id"); + } + let result = resolve_token_source().await; + unsafe { + env::remove_var(AWS_SECRET_ENV); + } + let err = match result { + Ok(_) => panic!("expected resolve_token_source to error without aws feature"), + Err(err) => err, + }; + assert!(err.to_string().contains("--features aws")); + } } diff --git a/crates/omnigraph-server/src/lib.rs b/crates/omnigraph-server/src/lib.rs index dd373b8..5d08734 100644 --- a/crates/omnigraph-server/src/lib.rs +++ b/crates/omnigraph-server/src/lib.rs @@ -38,7 +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 auth::{AWS_SECRET_ENV, EnvOrFileTokenSource, TokenSource, resolve_token_source}; pub use policy::{ PolicyAction, PolicyCompiler, PolicyConfig, PolicyDecision, PolicyEngine, PolicyExpectation, PolicyRequest, PolicyTestConfig, @@ -464,7 +464,8 @@ pub fn build_app(state: AppState) -> Router { } pub async fn serve(config: ServerConfig) -> Result<()> { - let token_source = EnvOrFileTokenSource; + let token_source = resolve_token_source().await?; + info!(source = token_source.name(), "loaded bearer token source"); let state = AppState::open_with_bearer_tokens_and_policy( config.uri.clone(), token_source.load().await?,