Cargo.lock

@@ -687,6 +687,7 @@ dependencies = [
  "tokio-cron-scheduler",
  "tokio-stream",
  "tokio-tungstenite 0.26.2",
+ "tower",
  "tower-http",
  "tracing",
  "tracing-subscriber",

compliance-agent/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 workspace = true
 
 [dependencies]
-compliance-core = { workspace = true, features = ["mongodb", "telemetry"] }
+compliance-core = { workspace = true, features = ["mongodb", "telemetry", "axum"] }
 compliance-graph = { path = "../compliance-graph" }
 compliance-dast = { path = "../compliance-dast" }
 serde = { workspace = true }
@@ -44,7 +44,8 @@ dashmap = { workspace = true }
 tokio-stream = { workspace = true }
 
 [dev-dependencies]
-compliance-core = { workspace = true, features = ["mongodb"] }
+compliance-core = { workspace = true, features = ["mongodb", "axum"] }
+tower = { version = "0.5", features = ["util"] }
 reqwest = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true }

compliance-agent/src/api/auth_middleware.rs

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

compliance-agent/src/api/mod.rs

@@ -1,4 +1,3 @@
-pub mod auth_middleware;
 pub mod handlers;
 pub mod routes;
 pub mod server;

compliance-agent/src/api/server.rs

@@ -7,8 +7,9 @@ use tower_http::cors::CorsLayer;
 use tower_http::set_header::SetResponseHeaderLayer;
 use tower_http::trace::TraceLayer;
 
+use compliance_core::auth::{require_jwt_auth, require_tenant_status, JwksState};
+
 use crate::agent::ComplianceAgent;
-use crate::api::auth_middleware::{require_jwt_auth, JwksState};
 use crate::api::routes;
 use crate::error::AgentError;
 
@@ -44,9 +45,13 @@ pub async fn start_api_server(agent: ComplianceAgent, port: u16) -> Result<(), A
             jwks_url,
         };
         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
-            .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 {
         tracing::warn!("Keycloak not configured - API endpoints are unprotected");
     }

compliance-agent/tests/tenant_status_middleware.rs

@@ -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);
+}