Compare commits

...

2 Commits

Author SHA1 Message Date
Sharang Parnerkar 324b137862 feat(m7.1): wire compliance-agent to compliance-core auth + status gate
CI / Check (pull_request) Successful in 8m14s
CI / Detect Changes (pull_request) Has been skipped
CI / Deploy Agent (pull_request) Has been skipped
CI / Deploy Dashboard (pull_request) Has been skipped
CI / Deploy Docs (pull_request) Has been skipped
CI / Deploy MCP (pull_request) Has been skipped
Supersedes #82.

Deletes the inline JWT middleware in compliance-agent (now stale —
missing JWKS refresh from #84 and tenant extraction from #83) and
imports require_jwt_auth, require_tenant_status, JwksState from
compliance_core::auth. Wires the status gate into the server's layer
stack: Extension(jwks_state) → require_jwt_auth → require_tenant_status
→ handler.

Adds the integration test from #82, retargeted to
compliance_core::auth::require_tenant_status.

Test plan
- cargo fmt --all clean
- cargo clippy --workspace --exclude compliance-dashboard -- -D warnings
  clean (matches baseline)
- cargo test -p compliance-core --lib — 7 tests pass
- cargo test -p compliance-agent --lib — 228 tests pass
- cargo test -p compliance-agent --test tenant_status_middleware —
  6 tests pass
- scripts/smoke.sh against live certifai KC — 15/15 cells pass
  (anon, bogus, active×2, trial, frozen, archived × {GET/health,
  GET/echo, POST/echo})

Caveats
- M7.1 only — status gate + claim extraction. Per-collection
  tenant_id scoping (M7.2) still pending; agent will still serve any
  Active/Trial tenant's data to any caller until the ~38 query
  call-sites use compliance_core::db::tenant_filter.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-17 11:07:56 +02:00
sharang dbadff0aac fix(m7.1): JWKS refresh-on-failure in auth middleware (#84)
CI / Check (push) Has been skipped
CI / Detect Changes (push) Successful in 3s
CI / Deploy Agent (push) Successful in 11m44s
CI / Deploy Dashboard (push) Successful in 13m1s
CI / Deploy Docs (push) Has been skipped
CI / Deploy MCP (push) Successful in 1m53s
2026-06-04 14:46:14 +00:00
7 changed files with 230 additions and 131 deletions
Generated
+1
View File
@@ -687,6 +687,7 @@ dependencies = [
"tokio-cron-scheduler", "tokio-cron-scheduler",
"tokio-stream", "tokio-stream",
"tokio-tungstenite 0.26.2", "tokio-tungstenite 0.26.2",
"tower",
"tower-http", "tower-http",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
+3 -2
View File
@@ -7,7 +7,7 @@ edition = "2021"
workspace = true workspace = true
[dependencies] [dependencies]
compliance-core = { workspace = true, features = ["mongodb", "telemetry"] } compliance-core = { workspace = true, features = ["mongodb", "telemetry", "axum"] }
compliance-graph = { path = "../compliance-graph" } compliance-graph = { path = "../compliance-graph" }
compliance-dast = { path = "../compliance-dast" } compliance-dast = { path = "../compliance-dast" }
serde = { workspace = true } serde = { workspace = true }
@@ -44,7 +44,8 @@ dashmap = { workspace = true }
tokio-stream = { workspace = true } tokio-stream = { workspace = true }
[dev-dependencies] [dev-dependencies]
compliance-core = { workspace = true, features = ["mongodb"] } compliance-core = { workspace = true, features = ["mongodb", "axum"] }
tower = { version = "0.5", features = ["util"] }
reqwest = { workspace = true } reqwest = { workspace = true }
serde_json = { workspace = true } serde_json = { workspace = true }
tokio = { workspace = true } tokio = { workspace = true }
-113
View File
@@ -1,113 +0,0 @@
use std::sync::Arc;
use axum::{
extract::Request,
middleware::Next,
response::{IntoResponse, Response},
};
use jsonwebtoken::{decode, decode_header, jwk::JwkSet, DecodingKey, Validation};
use reqwest::StatusCode;
use serde::Deserialize;
use tokio::sync::RwLock;
/// Cached JWKS from Keycloak for token validation.
#[derive(Clone)]
pub struct JwksState {
pub jwks: Arc<RwLock<Option<JwkSet>>>,
pub jwks_url: String,
}
#[derive(Debug, Deserialize)]
struct Claims {
#[allow(dead_code)]
sub: String,
}
const PUBLIC_ENDPOINTS: &[&str] = &["/api/v1/health"];
/// Middleware that validates Bearer JWT tokens against Keycloak's JWKS.
///
/// Skips validation for health check endpoints.
/// If `JwksState` is not present as an extension (keycloak not configured),
/// all requests pass through.
pub async fn require_jwt_auth(request: Request, next: Next) -> Response {
let path = request.uri().path();
if PUBLIC_ENDPOINTS.contains(&path) {
return next.run(request).await;
}
let jwks_state = match request.extensions().get::<JwksState>() {
Some(s) => s.clone(),
None => return next.run(request).await,
};
let auth_header = match request.headers().get("authorization") {
Some(h) => h,
None => return (StatusCode::UNAUTHORIZED, "Missing authorization header").into_response(),
};
let token = match auth_header.to_str() {
Ok(s) if s.starts_with("Bearer ") => &s[7..],
_ => return (StatusCode::UNAUTHORIZED, "Invalid authorization header").into_response(),
};
match validate_token(token, &jwks_state).await {
Ok(()) => next.run(request).await,
Err(e) => {
tracing::warn!("JWT validation failed: {e}");
(StatusCode::UNAUTHORIZED, "Invalid token").into_response()
}
}
}
async fn validate_token(token: &str, state: &JwksState) -> Result<(), String> {
let header = decode_header(token).map_err(|e| format!("failed to decode JWT header: {e}"))?;
let kid = header
.kid
.ok_or_else(|| "JWT missing kid header".to_string())?;
let jwks = fetch_or_get_jwks(state).await?;
let jwk = jwks
.keys
.iter()
.find(|k| k.common.key_id.as_deref() == Some(&kid))
.ok_or_else(|| "no matching key found in JWKS".to_string())?;
let decoding_key =
DecodingKey::from_jwk(jwk).map_err(|e| format!("failed to create decoding key: {e}"))?;
let mut validation = Validation::new(header.alg);
validation.validate_exp = true;
validation.validate_aud = false;
decode::<Claims>(token, &decoding_key, &validation)
.map_err(|e| format!("token validation failed: {e}"))?;
Ok(())
}
async fn fetch_or_get_jwks(state: &JwksState) -> Result<JwkSet, String> {
{
let cached = state.jwks.read().await;
if let Some(ref jwks) = *cached {
return Ok(jwks.clone());
}
}
let resp = reqwest::get(&state.jwks_url)
.await
.map_err(|e| format!("failed to fetch JWKS: {e}"))?;
let jwks: JwkSet = resp
.json()
.await
.map_err(|e| format!("failed to parse JWKS: {e}"))?;
let mut cached = state.jwks.write().await;
*cached = Some(jwks.clone());
Ok(jwks)
}
-1
View File
@@ -1,4 +1,3 @@
pub mod auth_middleware;
pub mod handlers; pub mod handlers;
pub mod routes; pub mod routes;
pub mod server; pub mod server;
+8 -3
View File
@@ -7,8 +7,9 @@ use tower_http::cors::CorsLayer;
use tower_http::set_header::SetResponseHeaderLayer; use tower_http::set_header::SetResponseHeaderLayer;
use tower_http::trace::TraceLayer; use tower_http::trace::TraceLayer;
use compliance_core::auth::{require_jwt_auth, require_tenant_status, JwksState};
use crate::agent::ComplianceAgent; use crate::agent::ComplianceAgent;
use crate::api::auth_middleware::{require_jwt_auth, JwksState};
use crate::api::routes; use crate::api::routes;
use crate::error::AgentError; use crate::error::AgentError;
@@ -44,9 +45,13 @@ pub async fn start_api_server(agent: ComplianceAgent, port: u16) -> Result<(), A
jwks_url, jwks_url,
}; };
tracing::info!("Keycloak JWT auth enabled for realm '{kc_realm}'"); tracing::info!("Keycloak JWT auth enabled for realm '{kc_realm}'");
// Layers execute outermost-first. Extension(jwks_state) must run
// before require_jwt_auth so the middleware can read it; the
// status gate runs after JWT so TenantContext is in extensions.
app = app app = app
.layer(Extension(jwks_state)) .layer(middleware::from_fn(require_tenant_status))
.layer(middleware::from_fn(require_jwt_auth)); .layer(middleware::from_fn(require_jwt_auth))
.layer(Extension(jwks_state));
} else { } else {
tracing::warn!("Keycloak not configured - API endpoints are unprotected"); tracing::warn!("Keycloak not configured - API endpoints are unprotected");
} }
@@ -0,0 +1,122 @@
//! M7.1 — integration tests for `compliance_core::auth::require_tenant_status`.
//!
//! Exercises the middleware end-to-end through an Axum router so we
//! catch wiring bugs (extension propagation, method matching) that pure
//! unit tests would miss.
#![allow(clippy::expect_used, clippy::unwrap_used)]
use axum::{
body::Body,
extract::Request,
http::{Method, StatusCode},
middleware::{from_fn, Next},
response::Response,
routing::{get, post},
Router,
};
use compliance_core::{auth::require_tenant_status, TenantContext, TenantStatus};
use tower::ServiceExt;
fn ctx_with(status: TenantStatus) -> TenantContext {
TenantContext {
tenant_id: "t-1".to_string(),
tenant_slug: "acme".to_string(),
org_roles: vec![],
products: vec![],
plan: "starter".to_string(),
status,
user_id: "u-1".to_string(),
user_name: None,
}
}
fn router_with_ctx(ctx: Option<TenantContext>) -> Router {
let injector = move |mut req: Request, next: Next| {
let ctx = ctx.clone();
async move {
if let Some(c) = ctx {
req.extensions_mut().insert(c);
}
next.run(req).await
}
};
Router::new()
.route("/r", get(|| async { "read" }))
.route("/w", post(|| async { "write" }))
.layer(from_fn(require_tenant_status))
.layer(from_fn(injector))
}
async fn call(router: Router, method: Method, path: &str) -> Response {
let req = Request::builder()
.method(method)
.uri(path)
.body(Body::empty())
.expect("request build");
router.oneshot(req).await.expect("oneshot")
}
#[tokio::test]
async fn active_tenant_can_read_and_write() {
let r = router_with_ctx(Some(ctx_with(TenantStatus::Active)));
assert_eq!(
call(r.clone(), Method::GET, "/r").await.status(),
StatusCode::OK
);
assert_eq!(call(r, Method::POST, "/w").await.status(), StatusCode::OK);
}
#[tokio::test]
async fn trial_tenant_can_read_and_write() {
let r = router_with_ctx(Some(ctx_with(TenantStatus::Trial)));
assert_eq!(
call(r.clone(), Method::GET, "/r").await.status(),
StatusCode::OK
);
assert_eq!(call(r, Method::POST, "/w").await.status(), StatusCode::OK);
}
#[tokio::test]
async fn demo_tenant_can_read_and_write() {
let r = router_with_ctx(Some(ctx_with(TenantStatus::Demo)));
assert_eq!(
call(r.clone(), Method::GET, "/r").await.status(),
StatusCode::OK
);
assert_eq!(call(r, Method::POST, "/w").await.status(), StatusCode::OK);
}
#[tokio::test]
async fn frozen_tenant_can_read_but_not_write() {
let r = router_with_ctx(Some(ctx_with(TenantStatus::Frozen)));
assert_eq!(
call(r.clone(), Method::GET, "/r").await.status(),
StatusCode::OK
);
assert_eq!(
call(r, Method::POST, "/w").await.status(),
StatusCode::PAYMENT_REQUIRED
);
}
#[tokio::test]
async fn archived_tenant_is_gone_on_every_method() {
let r = router_with_ctx(Some(ctx_with(TenantStatus::Archived)));
assert_eq!(
call(r.clone(), Method::GET, "/r").await.status(),
StatusCode::GONE
);
assert_eq!(call(r, Method::POST, "/w").await.status(), StatusCode::GONE);
}
#[tokio::test]
async fn no_context_passes_through() {
let r = router_with_ctx(None);
assert_eq!(
call(r.clone(), Method::GET, "/r").await.status(),
StatusCode::OK
);
assert_eq!(call(r, Method::POST, "/w").await.status(), StatusCode::OK);
}
+96 -12
View File
@@ -148,27 +148,83 @@ async fn validate_token(token: &str, state: &JwksState) -> Result<TenantContext,
let kid = header let kid = header
.kid .kid
.clone()
.ok_or_else(|| "JWT missing kid header".to_string())?; .ok_or_else(|| "JWT missing kid header".to_string())?;
let jwks = fetch_or_get_jwks(state).await?; // First try against whatever's currently cached. If the kid isn't
// there or the signature doesn't verify, the cached JWKS is most
// likely stale (KC rotated keys) — refresh once and retry before
// giving up. Without this every key rotation produces a silent 401
// storm that only goes away when the agent restarts.
let jwks = fetch_or_get_jwks(state, false).await?;
match try_validate(token, &header, &kid, &jwks) {
Ok(ctx) => Ok(ctx),
Err(ValidationError::Permanent(e)) => Err(e),
Err(ValidationError::Stale(reason)) => {
tracing::info!(
kid = %kid,
reason = %reason,
"JWKS appears stale — forcing refresh and retrying"
);
let jwks = fetch_or_get_jwks(state, true).await?;
try_validate(token, &header, &kid, &jwks).map_err(|e| match e {
ValidationError::Stale(s) | ValidationError::Permanent(s) => s,
})
}
}
}
let jwk = jwks #[derive(Debug)]
enum ValidationError {
/// Refresh-eligible: cached JWKS may be stale.
Stale(String),
/// Refusing the token regardless of JWKS freshness.
Permanent(String),
}
fn try_validate(
token: &str,
header: &jsonwebtoken::Header,
kid: &str,
jwks: &JwkSet,
) -> Result<TenantContext, ValidationError> {
let jwk = match jwks
.keys .keys
.iter() .iter()
.find(|k| k.common.key_id.as_deref() == Some(&kid)) .find(|k| k.common.key_id.as_deref() == Some(kid))
.ok_or_else(|| "no matching key found in JWKS".to_string())?; {
Some(j) => j,
None => {
return Err(ValidationError::Stale(
"no matching key found in JWKS".to_string(),
))
}
};
let decoding_key = let decoding_key = DecodingKey::from_jwk(jwk)
DecodingKey::from_jwk(jwk).map_err(|e| format!("failed to create decoding key: {e}"))?; .map_err(|e| ValidationError::Permanent(format!("failed to create decoding key: {e}")))?;
let mut validation = Validation::new(header.alg); let mut validation = Validation::new(header.alg);
validation.validate_exp = true; validation.validate_exp = true;
validation.validate_aud = false; validation.validate_aud = false;
let data = decode::<Claims>(token, &decoding_key, &validation) let data = match decode::<Claims>(token, &decoding_key, &validation) {
.map_err(|e| format!("token validation failed: {e}"))?; Ok(d) => d,
Err(e) => {
// Signature mismatch is the other refresh-eligible failure:
// the matching kid is present but the key bytes don't match.
// Everything else (expired, malformed, etc.) is permanent.
return Err(
if matches!(e.kind(), jsonwebtoken::errors::ErrorKind::InvalidSignature) {
ValidationError::Stale(format!("token validation failed: {e}"))
} else {
ValidationError::Permanent(format!("token validation failed: {e}"))
},
);
}
};
claims_to_context(data.claims) claims_to_context(data.claims).map_err(ValidationError::Permanent)
} }
/// Map the decoded JWT payload into the platform-wide `TenantContext`. /// Map the decoded JWT payload into the platform-wide `TenantContext`.
@@ -198,14 +254,25 @@ fn claims_to_context(c: Claims) -> Result<TenantContext, String> {
}) })
} }
async fn fetch_or_get_jwks(state: &JwksState) -> Result<JwkSet, String> { async fn fetch_or_get_jwks(state: &JwksState, force: bool) -> Result<JwkSet, String> {
{ if !force {
let cached = state.jwks.read().await; let cached = state.jwks.read().await;
if let Some(ref jwks) = *cached { if let Some(ref jwks) = *cached {
return Ok(jwks.clone()); return Ok(jwks.clone());
} }
} }
// Hold the write lock across the fetch so concurrent refreshers
// don't all hammer Keycloak when keys rotate. If another writer
// already populated a fresh JWKS while we were waiting (and we
// weren't asked to force), use theirs.
let mut cached = state.jwks.write().await;
if !force {
if let Some(ref jwks) = *cached {
return Ok(jwks.clone());
}
}
let resp = reqwest::get(&state.jwks_url) let resp = reqwest::get(&state.jwks_url)
.await .await
.map_err(|e| format!("failed to fetch JWKS: {e}"))?; .map_err(|e| format!("failed to fetch JWKS: {e}"))?;
@@ -215,7 +282,6 @@ async fn fetch_or_get_jwks(state: &JwksState) -> Result<JwkSet, String> {
.await .await
.map_err(|e| format!("failed to parse JWKS: {e}"))?; .map_err(|e| format!("failed to parse JWKS: {e}"))?;
let mut cached = state.jwks.write().await;
*cached = Some(jwks.clone()); *cached = Some(jwks.clone());
Ok(jwks) Ok(jwks)
@@ -293,6 +359,24 @@ mod tests {
); );
} }
#[test]
fn try_validate_returns_stale_when_kid_missing_from_jwks() {
// Empty JWKS — the kid we ask for can't possibly match. The error
// must classify as Stale so the caller refreshes JWKS and retries.
let jwks = JwkSet { keys: vec![] };
let header = jsonwebtoken::Header {
alg: jsonwebtoken::Algorithm::RS256,
kid: Some("kid-rotated-out".to_string()),
..Default::default()
};
let err = try_validate("ignored.token.value", &header, "kid-rotated-out", &jwks)
.expect_err("should fail");
match err {
ValidationError::Stale(s) => assert!(s.contains("no matching key")),
ValidationError::Permanent(s) => panic!("must be Stale, got Permanent: {s}"),
}
}
#[test] #[test]
fn is_write_detects_methods() { fn is_write_detects_methods() {
assert!(!is_write(&Method::GET)); assert!(!is_write(&Method::GET));