apunkt/omnigraph
1
0
Fork 0
mirror of https://github.com/ModernRelay/omnigraph.git synced 2026-06-09 01:35:18 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add aws feature + SecretsManagerTokenSource backend

Browse source 
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) <noreply@anthropic.com>
This commit is contained in:
andrew 2026-04-18 03:48:51 +03:00
parent af41630520
commit 7a3bf5c758
5 changed files with 432 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

46
.github/workflows/ci.yml vendored
View file

@ -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:

195
Cargo.lock generated
View file

@ -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",
]

8
crates/omnigraph-server/Cargo.toml
View file

@ -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 }

206
crates/omnigraph-server/src/auth.rs
View file

@ -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<Box<dyn TokenSource>> {
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<Vec<(String, String)>> {
use std::collections::HashMap;
let map: HashMap<String, String> = 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<String>) -> Result<Self> {
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<Vec<(String, String)>> {
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"));
}
}

5
crates/omnigraph-server/src/lib.rs
View file

@ -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?,