Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| cdfbb62f9d | |||
| 003835764e | |||
| e3aabe7d18 | |||
| 183234f9af | |||
| dbadff0aac | |||
| 116293519d |
Generated
+19
@@ -687,6 +687,7 @@ dependencies = [
|
||||
"tokio-cron-scheduler",
|
||||
"tokio-stream",
|
||||
"tokio-tungstenite 0.26.2",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
@@ -700,19 +701,23 @@ dependencies = [
|
||||
name = "compliance-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"bson",
|
||||
"chrono",
|
||||
"hex",
|
||||
"jsonwebtoken",
|
||||
"mongodb",
|
||||
"opentelemetry",
|
||||
"opentelemetry-appender-tracing",
|
||||
"opentelemetry-otlp",
|
||||
"opentelemetry_sdk",
|
||||
"reqwest",
|
||||
"secrecy",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-opentelemetry",
|
||||
"tracing-subscriber",
|
||||
@@ -826,6 +831,20 @@ dependencies = [
|
||||
"tracing-subscriber",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "compliance-smoke"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"compliance-core",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "console_error_panic_hook"
|
||||
version = "0.1.7"
|
||||
|
||||
@@ -6,6 +6,7 @@ members = [
|
||||
"compliance-graph",
|
||||
"compliance-dast",
|
||||
"compliance-mcp",
|
||||
"compliance-smoke",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -6,7 +6,7 @@ use tokio::sync::{broadcast, watch, Semaphore};
|
||||
use compliance_core::models::pentest::PentestEvent;
|
||||
use compliance_core::AgentConfig;
|
||||
|
||||
use crate::database::Database;
|
||||
use crate::database::{Database, DatabasePool};
|
||||
use crate::llm::LlmClient;
|
||||
use crate::pipeline::orchestrator::PipelineOrchestrator;
|
||||
|
||||
@@ -16,7 +16,13 @@ const DEFAULT_MAX_CONCURRENT_SESSIONS: usize = 5;
|
||||
#[derive(Clone)]
|
||||
pub struct ComplianceAgent {
|
||||
pub config: AgentConfig,
|
||||
/// Transitional single-database handle. Used by handlers that have
|
||||
/// not yet been migrated to `db_pool.for_tenant(&ctx)` (M7.2-B/C).
|
||||
/// Will be removed once every call site is tenant-scoped (M7.2-D).
|
||||
pub db: Database,
|
||||
/// Per-tenant Mongo broker introduced in M7.2-A. Handlers should
|
||||
/// prefer this and obtain a tenant-scoped [`Database`] from it.
|
||||
pub db_pool: DatabasePool,
|
||||
pub llm: Arc<LlmClient>,
|
||||
pub http: reqwest::Client,
|
||||
/// Per-session broadcast senders for SSE streaming.
|
||||
@@ -28,7 +34,7 @@ pub struct ComplianceAgent {
|
||||
}
|
||||
|
||||
impl ComplianceAgent {
|
||||
pub fn new(config: AgentConfig, db: Database) -> Self {
|
||||
pub fn new(config: AgentConfig, db: Database, db_pool: DatabasePool) -> Self {
|
||||
let llm = Arc::new(LlmClient::new(
|
||||
config.litellm_url.clone(),
|
||||
config.litellm_api_key.clone(),
|
||||
@@ -43,6 +49,7 @@ impl ComplianceAgent {
|
||||
Self {
|
||||
config,
|
||||
db,
|
||||
db_pool,
|
||||
llm,
|
||||
http,
|
||||
session_streams: Arc::new(DashMap::new()),
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -7,11 +7,13 @@ use mongodb::bson::doc;
|
||||
|
||||
use compliance_core::models::chat::{ChatRequest, ChatResponse, SourceReference};
|
||||
use compliance_core::models::embedding::EmbeddingBuildRun;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
use compliance_graph::graph::embedding_store::EmbeddingStore;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
use crate::rag::pipeline::RagPipeline;
|
||||
|
||||
use super::dto::tenant_db;
|
||||
use super::ApiResponse;
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
@@ -20,10 +22,12 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn chat(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
Json(req): Json<ChatRequest>,
|
||||
) -> Result<Json<ApiResponse<ChatResponse>>, StatusCode> {
|
||||
let pipeline = RagPipeline::new(agent.llm.clone(), agent.db.inner());
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let pipeline = RagPipeline::new(agent.llm.clone(), db.inner());
|
||||
|
||||
// Step 1: Embed the user's message
|
||||
let query_vectors = agent
|
||||
@@ -133,12 +137,15 @@ pub async fn chat(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn build_embeddings(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
// Resolve the tenant DB up front so we can move it into the spawn;
|
||||
// the JWT/dev context isn't available inside detached tasks.
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let agent_clone = (*agent).clone();
|
||||
tokio::spawn(async move {
|
||||
let repo = match agent_clone
|
||||
.db
|
||||
let repo = match db
|
||||
.repositories()
|
||||
.find_one(doc! { "_id": mongodb::bson::oid::ObjectId::parse_str(&repo_id).ok() })
|
||||
.await
|
||||
@@ -151,8 +158,7 @@ pub async fn build_embeddings(
|
||||
};
|
||||
|
||||
// Get latest graph build
|
||||
let build = match agent_clone
|
||||
.db
|
||||
let build = match db
|
||||
.graph_builds()
|
||||
.find_one(doc! { "repo_id": &repo_id })
|
||||
.sort(doc! { "started_at": -1 })
|
||||
@@ -171,26 +177,22 @@ pub async fn build_embeddings(
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
|
||||
// Get nodes
|
||||
let nodes: Vec<compliance_core::models::graph::CodeNode> = match agent_clone
|
||||
.db
|
||||
.graph_nodes()
|
||||
.find(doc! { "repo_id": &repo_id })
|
||||
.await
|
||||
{
|
||||
Ok(cursor) => {
|
||||
use futures_util::StreamExt;
|
||||
let mut items = Vec::new();
|
||||
let mut cursor = cursor;
|
||||
while let Some(Ok(item)) = cursor.next().await {
|
||||
items.push(item);
|
||||
let nodes: Vec<compliance_core::models::graph::CodeNode> =
|
||||
match db.graph_nodes().find(doc! { "repo_id": &repo_id }).await {
|
||||
Ok(cursor) => {
|
||||
use futures_util::StreamExt;
|
||||
let mut items = Vec::new();
|
||||
let mut cursor = cursor;
|
||||
while let Some(Ok(item)) = cursor.next().await {
|
||||
items.push(item);
|
||||
}
|
||||
items
|
||||
}
|
||||
items
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("[{repo_id}] Failed to fetch nodes: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
Err(e) => {
|
||||
tracing::error!("[{repo_id}] Failed to fetch nodes: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let creds = crate::pipeline::git::RepoCredentials {
|
||||
ssh_key_path: Some(agent_clone.config.ssh_key_path.clone()),
|
||||
@@ -207,7 +209,7 @@ pub async fn build_embeddings(
|
||||
}
|
||||
};
|
||||
|
||||
let pipeline = RagPipeline::new(agent_clone.llm.clone(), agent_clone.db.inner());
|
||||
let pipeline = RagPipeline::new(agent_clone.llm.clone(), db.inner());
|
||||
match pipeline
|
||||
.build_embeddings(&repo_id, &repo_path, &graph_build_id, &nodes)
|
||||
.await
|
||||
@@ -234,9 +236,11 @@ pub async fn build_embeddings(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn embedding_status(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Option<EmbeddingBuildRun>>>, StatusCode> {
|
||||
let store = EmbeddingStore::new(agent.db.inner());
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let store = EmbeddingStore::new(db.inner());
|
||||
let build = store.get_latest_build(&repo_id).await.map_err(|e| {
|
||||
tracing::error!("Failed to get embedding status: {e}");
|
||||
StatusCode::INTERNAL_SERVER_ERROR
|
||||
|
||||
@@ -7,9 +7,11 @@ use mongodb::bson::doc;
|
||||
use serde::Deserialize;
|
||||
|
||||
use compliance_core::models::dast::{DastFinding, DastScanRun, DastTarget, DastTargetType};
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
|
||||
use super::dto::tenant_db;
|
||||
use super::{collect_cursor_async, ApiResponse, PaginationParams};
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
@@ -45,9 +47,11 @@ fn default_rate_limit() -> u32 {
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_targets(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<DastTarget>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db
|
||||
.dast_targets()
|
||||
@@ -80,6 +84,7 @@ pub async fn list_targets(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn add_target(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Json(req): Json<AddTargetRequest>,
|
||||
) -> Result<Json<ApiResponse<DastTarget>>, StatusCode> {
|
||||
let mut target = DastTarget::new(req.name, req.base_url, req.target_type);
|
||||
@@ -89,9 +94,8 @@ pub async fn add_target(
|
||||
target.rate_limit = req.rate_limit;
|
||||
target.allow_destructive = req.allow_destructive;
|
||||
|
||||
agent
|
||||
.db
|
||||
.dast_targets()
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
db.dast_targets()
|
||||
.insert_one(&target)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
@@ -107,19 +111,19 @@ pub async fn add_target(
|
||||
#[tracing::instrument(skip_all, fields(target_id = %id))]
|
||||
pub async fn trigger_scan(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let target = agent
|
||||
.db
|
||||
let target = db
|
||||
.dast_targets()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
.ok_or(StatusCode::NOT_FOUND)?;
|
||||
|
||||
let db = agent.db.clone();
|
||||
tokio::spawn(async move {
|
||||
let orchestrator = compliance_dast::DastOrchestrator::new(100);
|
||||
match orchestrator.run_scan(&target, Vec::new()).await {
|
||||
@@ -147,9 +151,11 @@ pub async fn trigger_scan(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_scan_runs(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<DastScanRun>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db
|
||||
.dast_scan_runs()
|
||||
@@ -183,9 +189,11 @@ pub async fn list_scan_runs(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_findings(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<DastFinding>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db
|
||||
.dast_findings()
|
||||
@@ -219,12 +227,13 @@ pub async fn list_findings(
|
||||
#[tracing::instrument(skip_all, fields(finding_id = %id))]
|
||||
pub async fn get_finding(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<DastFinding>>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let finding = agent
|
||||
.db
|
||||
let finding = db
|
||||
.dast_findings()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
|
||||
@@ -180,6 +180,27 @@ pub struct SbomVersionDiff {
|
||||
pub(crate) type AgentExt = axum::extract::Extension<std::sync::Arc<crate::agent::ComplianceAgent>>;
|
||||
pub(crate) type ApiResult<T> = Result<axum::Json<ApiResponse<T>>, axum::http::StatusCode>;
|
||||
|
||||
/// Resolve a tenant-scoped [`Database`] from the request's
|
||||
/// [`TenantContext`] (inserted by the M7.1 JWT middleware, or by the
|
||||
/// dev fallback in unsecured environments). The pool ensures the
|
||||
/// tenant's indexes idempotently.
|
||||
///
|
||||
/// Returns 500 on the rare path where Mongo refuses the database
|
||||
/// handle — the M7.1 auth/status middleware already rejects every
|
||||
/// other failure mode with 4xx before we get here.
|
||||
pub(crate) async fn tenant_db(
|
||||
agent: &crate::agent::ComplianceAgent,
|
||||
tenant: &compliance_core::tenant_ctx::TenantCtx,
|
||||
) -> Result<crate::database::Database, axum::http::StatusCode> {
|
||||
agent.db_pool.for_tenant(&tenant.0).await.map_err(|e| {
|
||||
tracing::error!(
|
||||
tenant_id = %tenant.0.tenant_id,
|
||||
"Failed to acquire tenant database: {e}"
|
||||
);
|
||||
axum::http::StatusCode::INTERNAL_SERVER_ERROR
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn collect_cursor_async<T: serde::de::DeserializeOwned + Unpin + Send>(
|
||||
mut cursor: mongodb::Cursor<T>,
|
||||
) -> Vec<T> {
|
||||
|
||||
@@ -5,13 +5,16 @@ use mongodb::bson::doc;
|
||||
|
||||
use super::dto::*;
|
||||
use compliance_core::models::Finding;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
#[tracing::instrument(skip_all, fields(repo_id = ?filter.repo_id, severity = ?filter.severity, scan_type = ?filter.scan_type))]
|
||||
pub async fn list_findings(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(filter): Query<FindingsFilter>,
|
||||
) -> ApiResult<Vec<Finding>> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let mut query = doc! {};
|
||||
if let Some(repo_id) = &filter.repo_id {
|
||||
query.insert("repo_id", repo_id);
|
||||
@@ -81,11 +84,12 @@ pub async fn list_findings(
|
||||
#[tracing::instrument(skip_all, fields(finding_id = %id))]
|
||||
pub async fn get_finding(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Finding>>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let finding = agent
|
||||
.db
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let finding = db
|
||||
.findings()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -102,14 +106,14 @@ pub async fn get_finding(
|
||||
#[tracing::instrument(skip_all, fields(finding_id = %id))]
|
||||
pub async fn update_finding_status(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Json(req): Json<UpdateStatusRequest>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
agent
|
||||
.db
|
||||
.findings()
|
||||
db.findings()
|
||||
.update_one(
|
||||
doc! { "_id": oid },
|
||||
doc! { "$set": { "status": &req.status, "updated_at": mongodb::bson::DateTime::now() } },
|
||||
@@ -123,6 +127,7 @@ pub async fn update_finding_status(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn bulk_update_finding_status(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Json(req): Json<BulkUpdateStatusRequest>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oids: Vec<mongodb::bson::oid::ObjectId> = req
|
||||
@@ -135,8 +140,8 @@ pub async fn bulk_update_finding_status(
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
let result = agent
|
||||
.db
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let result = db
|
||||
.findings()
|
||||
.update_many(
|
||||
doc! { "_id": { "$in": oids } },
|
||||
@@ -153,14 +158,14 @@ pub async fn bulk_update_finding_status(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn update_finding_feedback(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Json(req): Json<UpdateFeedbackRequest>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
agent
|
||||
.db
|
||||
.findings()
|
||||
db.findings()
|
||||
.update_one(
|
||||
doc! { "_id": oid },
|
||||
doc! { "$set": { "developer_feedback": &req.feedback, "updated_at": mongodb::bson::DateTime::now() } },
|
||||
|
||||
@@ -7,9 +7,11 @@ use mongodb::bson::doc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use compliance_core::models::graph::{CodeEdge, CodeNode, GraphBuildRun, ImpactAnalysis};
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
|
||||
use super::dto::tenant_db;
|
||||
use super::{collect_cursor_async, ApiResponse};
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
@@ -36,9 +38,11 @@ fn default_search_limit() -> usize {
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn get_graph(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<GraphData>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
// Get latest build
|
||||
let build: Option<GraphBuildRun> = db
|
||||
@@ -98,9 +102,11 @@ pub async fn get_graph(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn get_nodes(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Vec<CodeNode>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let filter = doc! { "repo_id": &repo_id };
|
||||
|
||||
let nodes: Vec<CodeNode> = match db.graph_nodes().find(filter).await {
|
||||
@@ -123,9 +129,11 @@ pub async fn get_nodes(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn get_communities(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Vec<CommunityInfo>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let filter = doc! { "repo_id": &repo_id };
|
||||
|
||||
let nodes: Vec<CodeNode> = match db.graph_nodes().find(filter).await {
|
||||
@@ -176,9 +184,11 @@ pub struct CommunityInfo {
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id, finding_id = %finding_id))]
|
||||
pub async fn get_impact(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path((repo_id, finding_id)): Path<(String, String)>,
|
||||
) -> Result<Json<ApiResponse<Option<ImpactAnalysis>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let filter = doc! { "repo_id": &repo_id, "finding_id": &finding_id };
|
||||
|
||||
let impact = db
|
||||
@@ -198,10 +208,12 @@ pub async fn get_impact(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id, query = %params.q))]
|
||||
pub async fn search_symbols(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
Query(params): Query<SearchParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<CodeNode>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
// Simple text search on qualified_name and name fields
|
||||
let filter = doc! {
|
||||
@@ -234,10 +246,12 @@ pub async fn search_symbols(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn get_file_content(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
Query(params): Query<FileContentParams>,
|
||||
) -> Result<Json<ApiResponse<FileContent>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
// Look up the repository to get repo name
|
||||
let repo = db
|
||||
@@ -296,12 +310,13 @@ pub struct FileContent {
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
|
||||
pub async fn trigger_build(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(repo_id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let agent_clone = (*agent).clone();
|
||||
tokio::spawn(async move {
|
||||
let repo = match agent_clone
|
||||
.db
|
||||
let repo = match db
|
||||
.repositories()
|
||||
.find_one(doc! { "_id": mongodb::bson::oid::ObjectId::parse_str(&repo_id).ok() })
|
||||
.await
|
||||
@@ -333,8 +348,7 @@ pub async fn trigger_build(
|
||||
|
||||
match engine.build_graph(&repo_path, &repo_id, &graph_build_id) {
|
||||
Ok((code_graph, build_run)) => {
|
||||
let store =
|
||||
compliance_graph::graph::persistence::GraphStore::new(agent_clone.db.inner());
|
||||
let store = compliance_graph::graph::persistence::GraphStore::new(db.inner());
|
||||
let _ = store.delete_repo_graph(&repo_id).await;
|
||||
let _ = store
|
||||
.store_graph(&build_run, &code_graph.nodes, &code_graph.edges)
|
||||
|
||||
@@ -3,6 +3,7 @@ use mongodb::bson::doc;
|
||||
|
||||
use super::dto::*;
|
||||
use compliance_core::models::ScanRun;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn health() -> Json<serde_json::Value> {
|
||||
@@ -10,8 +11,12 @@ pub async fn health() -> Json<serde_json::Value> {
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn stats_overview(axum::extract::Extension(agent): AgentExt) -> ApiResult<OverviewStats> {
|
||||
let db = &agent.db;
|
||||
pub async fn stats_overview(
|
||||
axum::extract::Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
) -> ApiResult<OverviewStats> {
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
let total_repositories = db
|
||||
.repositories()
|
||||
|
||||
@@ -4,13 +4,16 @@ use mongodb::bson::doc;
|
||||
|
||||
use super::dto::*;
|
||||
use compliance_core::models::TrackerIssue;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_issues(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> ApiResult<Vec<TrackerIssue>> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db
|
||||
.tracker_issues()
|
||||
|
||||
@@ -5,15 +5,18 @@ use mongodb::bson::doc;
|
||||
use serde::Deserialize;
|
||||
|
||||
use compliance_core::models::notification::CveNotification;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use super::dto::{AgentExt, ApiResponse};
|
||||
use super::dto::{tenant_db, AgentExt, ApiResponse};
|
||||
|
||||
/// GET /api/v1/notifications — List CVE notifications (newest first)
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_notifications(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
axum::extract::Query(params): axum::extract::Query<NotificationFilter>,
|
||||
) -> Result<Json<ApiResponse<Vec<CveNotification>>>, StatusCode> {
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let mut filter = doc! {};
|
||||
|
||||
// Filter by status (default: show new + read, exclude dismissed)
|
||||
@@ -41,15 +44,13 @@ pub async fn list_notifications(
|
||||
let limit = params.limit.unwrap_or(50).min(200);
|
||||
let skip = (page - 1) * limit as u64;
|
||||
|
||||
let total = agent
|
||||
.db
|
||||
let total = db
|
||||
.cve_notifications()
|
||||
.count_documents(filter.clone())
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
let notifications: Vec<CveNotification> = match agent
|
||||
.db
|
||||
let notifications: Vec<CveNotification> = match db
|
||||
.cve_notifications()
|
||||
.find(filter)
|
||||
.sort(doc! { "created_at": -1 })
|
||||
@@ -83,9 +84,10 @@ pub async fn list_notifications(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn notification_count(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let count = agent
|
||||
.db
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let count = db
|
||||
.cve_notifications()
|
||||
.count_documents(doc! { "status": "new" })
|
||||
.await
|
||||
@@ -98,12 +100,13 @@ pub async fn notification_count(
|
||||
#[tracing::instrument(skip_all, fields(id = %id))]
|
||||
pub async fn mark_read(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
axum::extract::Path(id): axum::extract::Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let result = agent
|
||||
.db
|
||||
let result = db
|
||||
.cve_notifications()
|
||||
.update_one(
|
||||
doc! { "_id": oid },
|
||||
@@ -125,12 +128,13 @@ pub async fn mark_read(
|
||||
#[tracing::instrument(skip_all, fields(id = %id))]
|
||||
pub async fn dismiss_notification(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
axum::extract::Path(id): axum::extract::Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let result = agent
|
||||
.db
|
||||
let result = db
|
||||
.cve_notifications()
|
||||
.update_one(
|
||||
doc! { "_id": oid },
|
||||
@@ -149,9 +153,10 @@ pub async fn dismiss_notification(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn mark_all_read(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let result = agent
|
||||
.db
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let result = db
|
||||
.cve_notifications()
|
||||
.update_many(
|
||||
doc! { "status": "new" },
|
||||
|
||||
@@ -13,10 +13,11 @@ use compliance_core::models::dast::DastFinding;
|
||||
use compliance_core::models::finding::Finding;
|
||||
use compliance_core::models::pentest::*;
|
||||
use compliance_core::models::sbom::SbomEntry;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
|
||||
use super::super::dto::collect_cursor_async;
|
||||
use super::super::dto::{collect_cursor_async, tenant_db};
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
|
||||
@@ -35,11 +36,15 @@ pub struct ExportBody {
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn export_session_report(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Json(body): Json<ExportBody>,
|
||||
) -> Result<axum::response::Response, (StatusCode, String)> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
|
||||
if body.password.len() < 8 {
|
||||
return Err((
|
||||
@@ -49,8 +54,7 @@ pub async fn export_session_report(
|
||||
}
|
||||
|
||||
// Fetch session
|
||||
let session = agent
|
||||
.db
|
||||
let session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -64,9 +68,7 @@ pub async fn export_session_report(
|
||||
|
||||
// Resolve target name
|
||||
let target = if let Ok(tid) = mongodb::bson::oid::ObjectId::parse_str(&session.target_id) {
|
||||
agent
|
||||
.db
|
||||
.dast_targets()
|
||||
db.dast_targets()
|
||||
.find_one(doc! { "_id": tid })
|
||||
.await
|
||||
.ok()
|
||||
@@ -84,8 +86,7 @@ pub async fn export_session_report(
|
||||
.unwrap_or_default();
|
||||
|
||||
// Fetch attack chain nodes
|
||||
let nodes: Vec<AttackChainNode> = match agent
|
||||
.db
|
||||
let nodes: Vec<AttackChainNode> = match db
|
||||
.attack_chain_nodes()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "started_at": 1 })
|
||||
@@ -96,8 +97,7 @@ pub async fn export_session_report(
|
||||
};
|
||||
|
||||
// Fetch DAST findings for this session, then deduplicate
|
||||
let raw_findings: Vec<DastFinding> = match agent
|
||||
.db
|
||||
let raw_findings: Vec<DastFinding> = match db
|
||||
.dast_findings()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "severity": -1, "created_at": -1 })
|
||||
@@ -122,8 +122,7 @@ pub async fn export_session_report(
|
||||
.or_else(|| target.as_ref().and_then(|t| t.repo_id.clone()));
|
||||
|
||||
let (sast_findings, sbom_entries, code_context) = if let Some(ref rid) = repo_id {
|
||||
let sast: Vec<Finding> = match agent
|
||||
.db
|
||||
let sast: Vec<Finding> = match db
|
||||
.findings()
|
||||
.find(doc! {
|
||||
"repo_id": rid,
|
||||
@@ -143,8 +142,7 @@ pub async fn export_session_report(
|
||||
Err(_) => Vec::new(),
|
||||
};
|
||||
|
||||
let sbom: Vec<SbomEntry> = match agent
|
||||
.db
|
||||
let sbom: Vec<SbomEntry> = match db
|
||||
.sbom_entries()
|
||||
.find(doc! {
|
||||
"repo_id": rid,
|
||||
@@ -164,8 +162,7 @@ pub async fn export_session_report(
|
||||
};
|
||||
|
||||
// Build code context from graph nodes
|
||||
let code_ctx: Vec<CodeContextHint> = match agent
|
||||
.db
|
||||
let code_ctx: Vec<CodeContextHint> = match db
|
||||
.graph_nodes()
|
||||
.find(doc! { "repo_id": rid, "is_entry_point": true })
|
||||
.limit(50)
|
||||
|
||||
@@ -7,11 +7,12 @@ use mongodb::bson::doc;
|
||||
use serde::Deserialize;
|
||||
|
||||
use compliance_core::models::pentest::*;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
use crate::pentest::PentestOrchestrator;
|
||||
|
||||
use super::super::dto::{collect_cursor_async, ApiResponse, PaginationParams};
|
||||
use super::super::dto::{collect_cursor_async, tenant_db, ApiResponse, PaginationParams};
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
|
||||
@@ -43,6 +44,7 @@ pub struct LookupRepoQuery {
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn create_session(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Json(req): Json<CreateSessionRequest>,
|
||||
) -> Result<Json<ApiResponse<PentestSession>>, (StatusCode, String)> {
|
||||
// Try to acquire a concurrency permit
|
||||
@@ -57,6 +59,10 @@ pub async fn create_session(
|
||||
)
|
||||
})?;
|
||||
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
|
||||
if let Some(ref config) = req.config {
|
||||
// ── Wizard path ──────────────────────────────────────────────
|
||||
if !config.disclaimer_accepted {
|
||||
@@ -67,8 +73,7 @@ pub async fn create_session(
|
||||
}
|
||||
|
||||
// Look up or auto-create DastTarget by app_url
|
||||
let target = match agent
|
||||
.db
|
||||
let target = match db
|
||||
.dast_targets()
|
||||
.find_one(doc! { "base_url": &config.app_url })
|
||||
.await
|
||||
@@ -87,7 +92,7 @@ pub async fn create_session(
|
||||
}
|
||||
t.allow_destructive = config.allow_destructive;
|
||||
t.excluded_paths = config.scope_exclusions.clone();
|
||||
let res = agent.db.dast_targets().insert_one(&t).await.map_err(|e| {
|
||||
let res = db.dast_targets().insert_one(&t).await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Failed to create target: {e}"),
|
||||
@@ -110,8 +115,7 @@ pub async fn create_session(
|
||||
|
||||
// Resolve repo_id from git_repo_url if provided
|
||||
if let Some(ref git_url) = config.git_repo_url {
|
||||
if let Ok(Some(repo)) = agent
|
||||
.db
|
||||
if let Ok(Some(repo)) = db
|
||||
.repositories()
|
||||
.find_one(doc! { "git_url": git_url })
|
||||
.await
|
||||
@@ -120,8 +124,7 @@ pub async fn create_session(
|
||||
}
|
||||
}
|
||||
|
||||
let insert_result = agent
|
||||
.db
|
||||
let insert_result = db
|
||||
.pentest_sessions()
|
||||
.insert_one(&session)
|
||||
.await
|
||||
@@ -212,8 +215,7 @@ pub async fn create_session(
|
||||
// Persist encrypted credentials to DB
|
||||
if session_for_task.config.is_some() {
|
||||
if let Some(sid) = session.id {
|
||||
let _ = agent
|
||||
.db
|
||||
let _ = db
|
||||
.pentest_sessions()
|
||||
.update_one(
|
||||
doc! { "_id": sid },
|
||||
@@ -245,12 +247,13 @@ pub async fn create_session(
|
||||
});
|
||||
|
||||
let llm = agent.llm.clone();
|
||||
let db = agent.db.clone();
|
||||
let db_for_orchestrator = db.clone();
|
||||
let session_clone = session.clone();
|
||||
let target_clone = target.clone();
|
||||
let agent_ref = agent.clone();
|
||||
tokio::spawn(async move {
|
||||
let orchestrator = PentestOrchestrator::new(llm, db, event_tx, Some(pause_rx));
|
||||
let orchestrator =
|
||||
PentestOrchestrator::new(llm, db_for_orchestrator, event_tx, Some(pause_rx));
|
||||
orchestrator
|
||||
.run_session_guarded(&session_clone, &target_clone, &initial_message)
|
||||
.await;
|
||||
@@ -292,8 +295,7 @@ pub async fn create_session(
|
||||
)
|
||||
})?;
|
||||
|
||||
let target = agent
|
||||
.db
|
||||
let target = db
|
||||
.dast_targets()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -310,8 +312,7 @@ pub async fn create_session(
|
||||
let mut session = PentestSession::new(target_id, strategy);
|
||||
session.repo_id = target.repo_id.clone();
|
||||
|
||||
let insert_result = agent
|
||||
.db
|
||||
let insert_result = db
|
||||
.pentest_sessions()
|
||||
.insert_one(&session)
|
||||
.await
|
||||
@@ -338,12 +339,13 @@ pub async fn create_session(
|
||||
});
|
||||
|
||||
let llm = agent.llm.clone();
|
||||
let db = agent.db.clone();
|
||||
let db_for_orchestrator = db.clone();
|
||||
let session_clone = session.clone();
|
||||
let target_clone = target.clone();
|
||||
let agent_ref = agent.clone();
|
||||
tokio::spawn(async move {
|
||||
let orchestrator = PentestOrchestrator::new(llm, db, event_tx, Some(pause_rx));
|
||||
let orchestrator =
|
||||
PentestOrchestrator::new(llm, db_for_orchestrator, event_tx, Some(pause_rx));
|
||||
orchestrator
|
||||
.run_session_guarded(&session_clone, &target_clone, &initial_message)
|
||||
.await;
|
||||
@@ -373,10 +375,11 @@ fn parse_strategy(s: &str) -> PentestStrategy {
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn lookup_repo(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<LookupRepoQuery>,
|
||||
) -> Result<Json<ApiResponse<serde_json::Value>>, StatusCode> {
|
||||
let repo = agent
|
||||
.db
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let repo = db
|
||||
.repositories()
|
||||
.find_one(doc! { "git_url": ¶ms.url })
|
||||
.await
|
||||
@@ -402,9 +405,11 @@ pub async fn lookup_repo(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_sessions(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<PentestSession>>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db
|
||||
.pentest_sessions()
|
||||
@@ -438,12 +443,13 @@ pub async fn list_sessions(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn get_session(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<PentestSession>>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let mut session = agent
|
||||
.db
|
||||
let mut session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -471,15 +477,18 @@ pub async fn get_session(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn send_message(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Json(req): Json<SendMessageRequest>,
|
||||
) -> Result<Json<ApiResponse<PentestMessage>>, (StatusCode, String)> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
|
||||
// Verify session exists and is running
|
||||
let session = agent
|
||||
.db
|
||||
let session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -506,8 +515,7 @@ pub async fn send_message(
|
||||
)
|
||||
})?;
|
||||
|
||||
let target = agent
|
||||
.db
|
||||
let target = db
|
||||
.dast_targets()
|
||||
.find_one(doc! { "_id": target_oid })
|
||||
.await
|
||||
@@ -527,13 +535,13 @@ pub async fn send_message(
|
||||
// Store user message
|
||||
let session_id = id.clone();
|
||||
let user_msg = PentestMessage::user(session_id.clone(), req.message.clone());
|
||||
let _ = agent.db.pentest_messages().insert_one(&user_msg).await;
|
||||
let _ = db.pentest_messages().insert_one(&user_msg).await;
|
||||
|
||||
let response_msg = user_msg.clone();
|
||||
|
||||
// Spawn orchestrator to continue the session
|
||||
let llm = agent.llm.clone();
|
||||
let db = agent.db.clone();
|
||||
let db_for_orchestrator = db.clone();
|
||||
let message = req.message.clone();
|
||||
|
||||
// Use existing broadcast sender if available, otherwise create a new one
|
||||
@@ -548,7 +556,7 @@ pub async fn send_message(
|
||||
.unwrap_or_else(|| agent.register_session_stream(&session_id));
|
||||
|
||||
tokio::spawn(async move {
|
||||
let orchestrator = PentestOrchestrator::new(llm, db, event_tx, None);
|
||||
let orchestrator = PentestOrchestrator::new(llm, db_for_orchestrator, event_tx, None);
|
||||
orchestrator
|
||||
.run_session_guarded(&session, &target, &message)
|
||||
.await;
|
||||
@@ -565,13 +573,16 @@ pub async fn send_message(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn stop_session(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<PentestSession>>, (StatusCode, String)> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
|
||||
let session = agent
|
||||
.db
|
||||
let session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -590,9 +601,7 @@ pub async fn stop_session(
|
||||
));
|
||||
}
|
||||
|
||||
agent
|
||||
.db
|
||||
.pentest_sessions()
|
||||
db.pentest_sessions()
|
||||
.update_one(
|
||||
doc! { "_id": oid },
|
||||
doc! { "$set": {
|
||||
@@ -612,8 +621,7 @@ pub async fn stop_session(
|
||||
// Clean up session resources
|
||||
agent.cleanup_session(&id);
|
||||
|
||||
let updated = agent
|
||||
.db
|
||||
let updated = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -641,13 +649,16 @@ pub async fn stop_session(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn pause_session(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, String)> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
|
||||
let session = agent
|
||||
.db
|
||||
let session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -684,13 +695,16 @@ pub async fn pause_session(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn resume_session(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, String)> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
|
||||
let session = agent
|
||||
.db
|
||||
let session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -727,12 +741,13 @@ pub async fn resume_session(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn get_attack_chain(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<ApiResponse<Vec<AttackChainNode>>>, StatusCode> {
|
||||
let _oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let nodes = match agent
|
||||
.db
|
||||
let nodes = match db
|
||||
.attack_chain_nodes()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "started_at": 1 })
|
||||
@@ -757,21 +772,21 @@ pub async fn get_attack_chain(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn get_messages(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<PentestMessage>>>, StatusCode> {
|
||||
let _oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = agent
|
||||
.db
|
||||
let total = db
|
||||
.pentest_messages()
|
||||
.count_documents(doc! { "session_id": &id })
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
let messages = match agent
|
||||
.db
|
||||
let messages = match db
|
||||
.pentest_messages()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "created_at": 1 })
|
||||
@@ -797,21 +812,21 @@ pub async fn get_messages(
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn get_session_findings(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> Result<Json<ApiResponse<Vec<compliance_core::models::dast::DastFinding>>>, StatusCode> {
|
||||
let _oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = agent
|
||||
.db
|
||||
let total = db
|
||||
.dast_findings()
|
||||
.count_documents(doc! { "session_id": &id })
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
let findings = match agent
|
||||
.db
|
||||
let findings = match db
|
||||
.dast_findings()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "created_at": -1 })
|
||||
|
||||
@@ -6,10 +6,11 @@ use axum::Json;
|
||||
use mongodb::bson::doc;
|
||||
|
||||
use compliance_core::models::pentest::*;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
|
||||
use super::super::dto::{collect_cursor_async, ApiResponse};
|
||||
use super::super::dto::{collect_cursor_async, tenant_db, ApiResponse};
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
|
||||
@@ -17,8 +18,10 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn pentest_stats(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
) -> Result<Json<ApiResponse<PentestStats>>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
let running_sessions = db
|
||||
.pentest_sessions()
|
||||
|
||||
@@ -11,10 +11,11 @@ use tokio_stream::wrappers::BroadcastStream;
|
||||
use tokio_stream::StreamExt;
|
||||
|
||||
use compliance_core::models::pentest::*;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
|
||||
use super::super::dto::collect_cursor_async;
|
||||
use super::super::dto::{collect_cursor_async, tenant_db};
|
||||
|
||||
type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
|
||||
@@ -25,13 +26,14 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn session_stream(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
// Verify session exists
|
||||
let _session = agent
|
||||
.db
|
||||
let _session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -43,8 +45,7 @@ pub async fn session_stream(
|
||||
let mut initial_events: Vec<Result<Event, Infallible>> = Vec::new();
|
||||
|
||||
// Fetch recent messages for this session
|
||||
let messages: Vec<PentestMessage> = match agent
|
||||
.db
|
||||
let messages: Vec<PentestMessage> = match db
|
||||
.pentest_messages()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "created_at": 1 })
|
||||
@@ -56,8 +57,7 @@ pub async fn session_stream(
|
||||
};
|
||||
|
||||
// Fetch recent attack chain nodes
|
||||
let nodes: Vec<AttackChainNode> = match agent
|
||||
.db
|
||||
let nodes: Vec<AttackChainNode> = match db
|
||||
.attack_chain_nodes()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "started_at": 1 })
|
||||
@@ -94,8 +94,7 @@ pub async fn session_stream(
|
||||
}
|
||||
|
||||
// Add current session status event
|
||||
let session = agent
|
||||
.db
|
||||
let session = db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
|
||||
@@ -5,13 +5,16 @@ use mongodb::bson::doc;
|
||||
|
||||
use super::dto::*;
|
||||
use compliance_core::models::*;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_repositories(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> ApiResult<Vec<TrackedRepository>> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db
|
||||
.repositories()
|
||||
@@ -43,6 +46,7 @@ pub async fn list_repositories(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn add_repository(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Json(req): Json<AddRepositoryRequest>,
|
||||
) -> Result<Json<ApiResponse<TrackedRepository>>, (StatusCode, String)> {
|
||||
// Validate repository access before saving
|
||||
@@ -69,17 +73,15 @@ pub async fn add_repository(
|
||||
repo.tracker_token = req.tracker_token;
|
||||
repo.scan_schedule = req.scan_schedule;
|
||||
|
||||
agent
|
||||
.db
|
||||
.repositories()
|
||||
.insert_one(&repo)
|
||||
let db = tenant_db(&agent, &tenant)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
(
|
||||
StatusCode::CONFLICT,
|
||||
"Repository already exists".to_string(),
|
||||
)
|
||||
})?;
|
||||
.map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
|
||||
db.repositories().insert_one(&repo).await.map_err(|_| {
|
||||
(
|
||||
StatusCode::CONFLICT,
|
||||
"Repository already exists".to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(Json(ApiResponse {
|
||||
data: repo,
|
||||
@@ -91,10 +93,12 @@ pub async fn add_repository(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %id))]
|
||||
pub async fn update_repository(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
Json(req): Json<UpdateRepositoryRequest>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
|
||||
let mut set_doc = doc! { "updated_at": mongodb::bson::DateTime::now() };
|
||||
|
||||
@@ -126,8 +130,7 @@ pub async fn update_repository(
|
||||
set_doc.insert("scan_schedule", schedule);
|
||||
}
|
||||
|
||||
let result = agent
|
||||
.db
|
||||
let result = db
|
||||
.repositories()
|
||||
.update_one(doc! { "_id": oid }, doc! { "$set": set_doc })
|
||||
.await
|
||||
@@ -170,11 +173,12 @@ pub async fn trigger_scan(
|
||||
/// Return the webhook secret for a repository (used by dashboard to display it)
|
||||
pub async fn get_webhook_config(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let repo = agent
|
||||
.db
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let repo = db
|
||||
.repositories()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
@@ -196,10 +200,12 @@ pub async fn get_webhook_config(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = %id))]
|
||||
pub async fn delete_repository(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Path(id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
// Delete the repository
|
||||
let result = db
|
||||
|
||||
@@ -6,6 +6,7 @@ use mongodb::bson::doc;
|
||||
|
||||
use super::dto::*;
|
||||
use compliance_core::models::SbomEntry;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
const COPYLEFT_LICENSES: &[&str] = &[
|
||||
"GPL-2.0",
|
||||
@@ -29,8 +30,10 @@ const COPYLEFT_LICENSES: &[&str] = &[
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn sbom_filters(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
let managers: Vec<String> = db
|
||||
.sbom_entries()
|
||||
@@ -61,9 +64,11 @@ pub async fn sbom_filters(
|
||||
#[tracing::instrument(skip_all, fields(repo_id = ?filter.repo_id, package_manager = ?filter.package_manager))]
|
||||
pub async fn list_sbom(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(filter): Query<SbomFilter>,
|
||||
) -> ApiResult<Vec<SbomEntry>> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let mut query = doc! {};
|
||||
|
||||
if let Some(repo_id) = &filter.repo_id {
|
||||
@@ -120,9 +125,11 @@ pub async fn list_sbom(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn export_sbom(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<SbomExportParams>,
|
||||
) -> Result<impl IntoResponse, StatusCode> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let entries: Vec<SbomEntry> = match db
|
||||
.sbom_entries()
|
||||
.find(doc! { "repo_id": ¶ms.repo_id })
|
||||
@@ -236,9 +243,11 @@ pub async fn export_sbom(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn license_summary(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<SbomFilter>,
|
||||
) -> ApiResult<Vec<LicenseSummary>> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let mut query = doc! {};
|
||||
if let Some(repo_id) = ¶ms.repo_id {
|
||||
query.insert("repo_id", repo_id);
|
||||
@@ -285,9 +294,11 @@ pub async fn license_summary(
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn sbom_diff(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<SbomDiffParams>,
|
||||
) -> ApiResult<SbomDiffResult> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
|
||||
let entries_a: Vec<SbomEntry> = match db
|
||||
.sbom_entries()
|
||||
|
||||
@@ -4,13 +4,16 @@ use mongodb::bson::doc;
|
||||
|
||||
use super::dto::*;
|
||||
use compliance_core::models::ScanRun;
|
||||
use compliance_core::tenant_ctx::TenantCtx;
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn list_scan_runs(
|
||||
Extension(agent): AgentExt,
|
||||
tenant: TenantCtx,
|
||||
Query(params): Query<PaginationParams>,
|
||||
) -> ApiResult<Vec<ScanRun>> {
|
||||
let db = &agent.db;
|
||||
let db = tenant_db(&agent, &tenant).await?;
|
||||
let db = &db;
|
||||
let skip = (params.page.saturating_sub(1)) * params.limit as u64;
|
||||
let total = db.scan_runs().count_documents(doc! {}).await.unwrap_or(0);
|
||||
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
pub mod auth_middleware;
|
||||
pub mod handlers;
|
||||
pub mod routes;
|
||||
pub mod server;
|
||||
|
||||
@@ -1,17 +1,54 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::Request;
|
||||
use axum::http::HeaderValue;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use axum::{middleware, Extension};
|
||||
use tokio::sync::RwLock;
|
||||
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 compliance_core::{TenantContext, TenantStatus};
|
||||
|
||||
use crate::agent::ComplianceAgent;
|
||||
use crate::api::auth_middleware::{require_jwt_auth, JwksState};
|
||||
use crate::api::routes;
|
||||
use crate::error::AgentError;
|
||||
|
||||
/// Synthetic tenant id used when Keycloak isn't configured (local dev,
|
||||
/// `cargo run` against a bare Mongo). Lets the handler stack stay
|
||||
/// uniformly tenant-scoped without the operator having to spin up KC
|
||||
/// just to poke at the API. Override via `DEV_TENANT_ID`.
|
||||
const DEFAULT_DEV_TENANT_ID: &str = "dev";
|
||||
|
||||
/// Inject a synthetic [`TenantContext`] for any request that lacks one.
|
||||
/// Only mounted when Keycloak is NOT configured; with KC, the real
|
||||
/// `require_jwt_auth` middleware owns this and we never reach here
|
||||
/// without a context.
|
||||
///
|
||||
/// Public so the integration-test harness can mount it without
|
||||
/// duplicating the synthetic-context shape.
|
||||
pub async fn inject_dev_tenant(mut request: Request, next: Next) -> Response {
|
||||
if request.extensions().get::<TenantContext>().is_none() {
|
||||
let tenant_id =
|
||||
std::env::var("DEV_TENANT_ID").unwrap_or_else(|_| DEFAULT_DEV_TENANT_ID.to_string());
|
||||
let ctx = TenantContext {
|
||||
tenant_slug: tenant_id.clone(),
|
||||
tenant_id,
|
||||
org_roles: vec![],
|
||||
products: vec![],
|
||||
plan: "dev".to_string(),
|
||||
status: TenantStatus::Active,
|
||||
user_id: "dev-user".to_string(),
|
||||
user_name: None,
|
||||
};
|
||||
request.extensions_mut().insert(ctx);
|
||||
}
|
||||
next.run(request).await
|
||||
}
|
||||
|
||||
pub async fn start_api_server(agent: ComplianceAgent, port: u16) -> Result<(), AgentError> {
|
||||
let mut app = routes::build_router()
|
||||
.layer(Extension(Arc::new(agent.clone())))
|
||||
@@ -44,11 +81,22 @@ 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_jwt_auth));
|
||||
.layer(middleware::from_fn(require_tenant_status))
|
||||
.layer(middleware::from_fn(require_jwt_auth))
|
||||
.layer(Extension(jwks_state));
|
||||
} else {
|
||||
tracing::warn!("Keycloak not configured - API endpoints are unprotected");
|
||||
let tenant_id =
|
||||
std::env::var("DEV_TENANT_ID").unwrap_or_else(|_| DEFAULT_DEV_TENANT_ID.to_string());
|
||||
tracing::warn!(
|
||||
tenant_id = %tenant_id,
|
||||
"Keycloak not configured — running unauthenticated against the dev tenant. \
|
||||
DO NOT use in any environment with real customer data."
|
||||
);
|
||||
app = app.layer(middleware::from_fn(inject_dev_tenant));
|
||||
}
|
||||
|
||||
let addr = format!("0.0.0.0:{port}");
|
||||
|
||||
@@ -1,11 +1,151 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use mongodb::bson::doc;
|
||||
use mongodb::options::IndexOptions;
|
||||
use mongodb::{Client, Collection, IndexModel};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use compliance_core::models::*;
|
||||
use compliance_core::TenantContext;
|
||||
|
||||
use crate::error::AgentError;
|
||||
|
||||
/// Mongo enforces a 63-byte cap on database names (older clusters: 64
|
||||
/// on Linux, 63 on Windows; we target the conservative limit).
|
||||
const MAX_DB_NAME_LEN: usize = 63;
|
||||
|
||||
/// Hex length of the SHA-256 truncation used for the hash fallback
|
||||
/// tenant DB name (16 bytes → 32 hex chars). 16 bytes gives ~2^64
|
||||
/// birthday-collision resistance — at our 10s-100s tenant scale this
|
||||
/// is effectively impossible to hit.
|
||||
const HASH_HEX_LEN: usize = 32;
|
||||
|
||||
/// Largest `db_prefix` that still guarantees the hash-fallback name
|
||||
/// fits in the 63-byte cap: `prefix + "_" + 32 hex chars`.
|
||||
const MAX_PREFIX_LEN: usize = MAX_DB_NAME_LEN - 1 - HASH_HEX_LEN;
|
||||
|
||||
/// Per-tenant Mongo connection broker (M7.2 isolation model).
|
||||
///
|
||||
/// Holds one [`Client`] and hands out [`Database`] handles physically
|
||||
/// scoped to `<db_prefix>_<tenant_id>`. The driver is the isolation
|
||||
/// boundary — a handle for tenant A cannot see tenant B's documents
|
||||
/// because it is connected to a different database, not because of an
|
||||
/// application-level filter.
|
||||
///
|
||||
/// Index creation runs idempotently the first time each tenant is seen
|
||||
/// in the process's lifetime. Mongo's `createIndex` is itself idempotent
|
||||
/// by index name; the in-memory `ensured` set just skips the round-trip.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct DatabasePool {
|
||||
client: Client,
|
||||
db_prefix: String,
|
||||
ensured: Arc<DashMap<String, ()>>,
|
||||
}
|
||||
|
||||
impl DatabasePool {
|
||||
/// Connect to the cluster and prepare to hand out tenant databases
|
||||
/// named `<db_prefix>_<tenant_id>`.
|
||||
///
|
||||
/// Validates `db_prefix.len() <= MAX_PREFIX_LEN` so the
|
||||
/// hash-fallback path is provably within Mongo's 63-byte db-name
|
||||
/// cap. Refuses to construct a pool that could ever produce an
|
||||
/// over-long name.
|
||||
pub async fn connect(uri: &str, db_prefix: &str) -> Result<Self, AgentError> {
|
||||
if db_prefix.len() > MAX_PREFIX_LEN {
|
||||
return Err(AgentError::Other(format!(
|
||||
"db_prefix '{db_prefix}' is {} chars; max is {MAX_PREFIX_LEN} so the \
|
||||
hash-fallback tenant DB name fits Mongo's {MAX_DB_NAME_LEN}-byte cap",
|
||||
db_prefix.len()
|
||||
)));
|
||||
}
|
||||
let client = Client::with_uri_str(uri).await?;
|
||||
client
|
||||
.database("admin")
|
||||
.run_command(doc! { "ping": 1 })
|
||||
.await?;
|
||||
tracing::info!(
|
||||
"MongoDB cluster reachable; per-tenant pool ready (db prefix '{db_prefix}')"
|
||||
);
|
||||
Ok(Self {
|
||||
client,
|
||||
db_prefix: db_prefix.to_string(),
|
||||
ensured: Arc::new(DashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Return a [`Database`] scoped to this tenant. Ensures indexes on
|
||||
/// first call per tenant (per process). Cheap on the hot path —
|
||||
/// subsequent calls skip the round-trip.
|
||||
pub async fn for_tenant(&self, ctx: &TenantContext) -> Result<Database, AgentError> {
|
||||
let db_name = self.tenant_db_name(&ctx.tenant_id);
|
||||
let db = Database::from_database(self.client.database(&db_name));
|
||||
// `DashMap::insert` returns the previous value; `None` means we
|
||||
// were the first writer for this tenant_id and own the
|
||||
// index-ensure work.
|
||||
if self.ensured.insert(ctx.tenant_id.clone(), ()).is_none() {
|
||||
if let Err(e) = db.ensure_indexes().await {
|
||||
// Roll the marker back so the next request retries.
|
||||
self.ensured.remove(&ctx.tenant_id);
|
||||
return Err(e);
|
||||
}
|
||||
tracing::debug!(
|
||||
tenant_id = %ctx.tenant_id,
|
||||
db_name = %db_name,
|
||||
"Indexes ensured for tenant database"
|
||||
);
|
||||
}
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
/// Compute the Mongo database name for a tenant. Public for tests
|
||||
/// and tenant offboarding (`pool.client().database(name).drop()`).
|
||||
///
|
||||
/// Format: `<prefix>_<sanitized_tenant_id>` if it fits the 63-byte
|
||||
/// cap, else `<prefix>_<sha256-16-byte-hex-of-tenant_id>`. The
|
||||
/// `db_prefix` length invariant established at [`Self::connect`]
|
||||
/// guarantees the hash-fallback name always fits — no runtime
|
||||
/// assertion needed.
|
||||
///
|
||||
/// Collision resistance: the hash fallback is a 16-byte SHA-256
|
||||
/// truncation, which gives ~2^64 birthday-collision resistance. At
|
||||
/// our 10s–100s tenant scale the probability of two tenant_ids
|
||||
/// colliding is effectively zero. (8-byte truncation would have
|
||||
/// been ~2^32 — too close for comfort on a regulated product.)
|
||||
pub fn tenant_db_name(&self, tenant_id: &str) -> String {
|
||||
let sanitized = sanitize_tenant_id(tenant_id);
|
||||
let natural = format!("{}_{}", self.db_prefix, sanitized);
|
||||
if natural.len() <= MAX_DB_NAME_LEN {
|
||||
natural
|
||||
} else {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(tenant_id.as_bytes());
|
||||
let digest = hasher.finalize();
|
||||
let suffix = hex::encode(&digest[..HASH_HEX_LEN / 2]);
|
||||
format!("{}_{}", self.db_prefix, suffix)
|
||||
}
|
||||
}
|
||||
|
||||
/// Raw client handle. Reserved for cross-tenant admin flows that
|
||||
/// must opt in explicitly (tenant listing, drop-on-offboard).
|
||||
pub fn client(&self) -> &Client {
|
||||
&self.client
|
||||
}
|
||||
}
|
||||
|
||||
/// Mongo database names disallow `/`, `\`, `.`, `"`, `$`, ` `, and NUL.
|
||||
/// breakpilot-dev tenant_ids are UUIDs so this is belt-and-braces, but
|
||||
/// it lets the pool tolerate any future tenant_id shape without surprise.
|
||||
fn sanitize_tenant_id(tenant_id: &str) -> String {
|
||||
tenant_id
|
||||
.chars()
|
||||
.map(|c| match c {
|
||||
'/' | '\\' | '.' | '"' | '$' | ' ' | '\0' => '_',
|
||||
c => c,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Database {
|
||||
inner: mongodb::Database,
|
||||
@@ -20,6 +160,12 @@ impl Database {
|
||||
Ok(Self { inner: db })
|
||||
}
|
||||
|
||||
/// Wrap an already-resolved Mongo database. Used by [`DatabasePool`]
|
||||
/// to hand out tenant-scoped handles without a fresh client per tenant.
|
||||
pub(crate) fn from_database(inner: mongodb::Database) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
pub async fn ensure_indexes(&self) -> Result<(), AgentError> {
|
||||
// repositories: unique git_url
|
||||
self.repositories()
|
||||
|
||||
@@ -28,7 +28,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let db = database::Database::connect(&config.mongodb_uri, &config.mongodb_database).await?;
|
||||
db.ensure_indexes().await?;
|
||||
|
||||
let agent = agent::ComplianceAgent::new(config.clone(), db.clone());
|
||||
// M7.2-A: per-tenant pool. Uses `mongodb_database` as the db-name
|
||||
// prefix so tenant databases land as `<prefix>_<tenant_id>` next to
|
||||
// the legacy single-tenant database.
|
||||
let db_pool =
|
||||
database::DatabasePool::connect(&config.mongodb_uri, &config.mongodb_database).await?;
|
||||
|
||||
let agent = agent::ComplianceAgent::new(config.clone(), db.clone(), db_pool);
|
||||
|
||||
tracing::info!("Starting scheduler...");
|
||||
let scheduler_agent = agent.clone();
|
||||
|
||||
@@ -7,7 +7,7 @@ use std::sync::Arc;
|
||||
|
||||
use compliance_agent::agent::ComplianceAgent;
|
||||
use compliance_agent::api;
|
||||
use compliance_agent::database::Database;
|
||||
use compliance_agent::database::{Database, DatabasePool};
|
||||
use compliance_core::AgentConfig;
|
||||
use secrecy::SecretString;
|
||||
|
||||
@@ -33,6 +33,10 @@ impl TestServer {
|
||||
.expect("Failed to connect to MongoDB — is it running?");
|
||||
db.ensure_indexes().await.expect("Failed to create indexes");
|
||||
|
||||
let db_pool = DatabasePool::connect(&mongodb_uri, &db_name)
|
||||
.await
|
||||
.expect("Failed to build DatabasePool");
|
||||
|
||||
let config = AgentConfig {
|
||||
mongodb_uri: mongodb_uri.clone(),
|
||||
mongodb_database: db_name.clone(),
|
||||
@@ -69,11 +73,15 @@ impl TestServer {
|
||||
pentest_imap_password: None,
|
||||
};
|
||||
|
||||
let agent = ComplianceAgent::new(config, db);
|
||||
let agent = ComplianceAgent::new(config, db, db_pool);
|
||||
|
||||
// Build the router with the agent extension
|
||||
// Build the router with the agent extension. After M7.2-B every
|
||||
// handler takes a TenantCtx extractor; without KC in the test
|
||||
// harness, the dev-tenant injector mounts a synthetic context so
|
||||
// tests run end-to-end against `<db_name>_dev`.
|
||||
let app = api::routes::build_router()
|
||||
.layer(axum::extract::Extension(Arc::new(agent)))
|
||||
.layer(axum::middleware::from_fn(api::server::inject_dev_tenant))
|
||||
.layer(tower_http::cors::CorsLayer::permissive());
|
||||
|
||||
// Bind to port 0 to get a random available port
|
||||
@@ -156,10 +164,20 @@ impl TestServer {
|
||||
&self.db_name
|
||||
}
|
||||
|
||||
/// Drop the test database on cleanup
|
||||
/// Drop the test database on cleanup. Post-M7.2-B the actual data
|
||||
/// lives in `<db_name>_<tenant>` per-tenant databases; list those
|
||||
/// off the cluster and drop them too.
|
||||
pub async fn cleanup(&self) {
|
||||
if let Ok(client) = mongodb::Client::with_uri_str(&self.mongodb_uri).await {
|
||||
client.database(&self.db_name).drop().await.ok();
|
||||
if let Ok(names) = client.list_database_names().await {
|
||||
let prefix = format!("{}_", self.db_name);
|
||||
for name in names {
|
||||
if name.starts_with(&prefix) {
|
||||
client.database(&name).drop().await.ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
//! M7.2-A — `DatabasePool` isolation proof.
|
||||
//!
|
||||
//! Two `TenantContext`s, two databases, one client. Insert on A, query
|
||||
//! on B → empty. Insert on B, query on A → only A's docs. Proves that
|
||||
//! the per-tenant database split actually isolates at the driver level
|
||||
//! and not at "we hope we filter."
|
||||
//!
|
||||
//! Requires MongoDB. Set `TEST_MONGODB_URI` to override the default
|
||||
//! `mongodb://root:example@localhost:27017/?authSource=admin`.
|
||||
|
||||
#![allow(clippy::expect_used, clippy::unwrap_used)]
|
||||
|
||||
use compliance_agent::database::DatabasePool;
|
||||
use compliance_core::models::TrackedRepository;
|
||||
use compliance_core::{OrgRole, TenantContext, TenantStatus};
|
||||
use mongodb::bson::doc;
|
||||
|
||||
fn ctx(tenant_id: &str, slug: &str) -> TenantContext {
|
||||
TenantContext {
|
||||
tenant_id: tenant_id.to_string(),
|
||||
tenant_slug: slug.to_string(),
|
||||
org_roles: vec![OrgRole::ItAdmin],
|
||||
products: vec!["compliance-scanner".to_string()],
|
||||
plan: "starter".to_string(),
|
||||
status: TenantStatus::Active,
|
||||
user_id: "u-1".to_string(),
|
||||
user_name: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn fixture_repo(name: &str, git_url: &str) -> TrackedRepository {
|
||||
TrackedRepository {
|
||||
id: None,
|
||||
name: name.to_string(),
|
||||
git_url: git_url.to_string(),
|
||||
default_branch: "main".to_string(),
|
||||
local_path: None,
|
||||
scan_schedule: None,
|
||||
webhook_enabled: false,
|
||||
webhook_secret: None,
|
||||
tracker_type: None,
|
||||
tracker_owner: None,
|
||||
tracker_repo: None,
|
||||
tracker_token: None,
|
||||
auth_token: None,
|
||||
auth_username: None,
|
||||
last_scanned_commit: None,
|
||||
findings_count: 0,
|
||||
created_at: chrono::Utc::now(),
|
||||
updated_at: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_isolates_tenants_at_driver_level() {
|
||||
let uri = std::env::var("TEST_MONGODB_URI")
|
||||
.unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into());
|
||||
// Unique per run so parallel test invocations don't collide. Kept
|
||||
// short because Mongo caps db names at 63 bytes (prefix + tenant_id).
|
||||
let prefix = format!("m72a_{}", short_id());
|
||||
|
||||
let pool = DatabasePool::connect(&uri, &prefix)
|
||||
.await
|
||||
.expect("Failed to connect to MongoDB — is it running?");
|
||||
|
||||
let acme = ctx("00000000-0000-0000-0000-00000000acme", "acme");
|
||||
let globex = ctx("00000000-0000-0000-0000-0000globex000", "globex");
|
||||
|
||||
let acme_db = pool.for_tenant(&acme).await.expect("acme db");
|
||||
let globex_db = pool.for_tenant(&globex).await.expect("globex db");
|
||||
|
||||
// Write distinct repos into each tenant's database.
|
||||
acme_db
|
||||
.repositories()
|
||||
.insert_one(fixture_repo("acme-app", "git@example.com:acme/app.git"))
|
||||
.await
|
||||
.expect("insert acme");
|
||||
globex_db
|
||||
.repositories()
|
||||
.insert_one(fixture_repo(
|
||||
"globex-platform",
|
||||
"git@example.com:globex/platform.git",
|
||||
))
|
||||
.await
|
||||
.expect("insert globex");
|
||||
|
||||
// The point of the whole exercise: acme can ONLY see acme's repo
|
||||
// and globex can ONLY see globex's, with no filter doc anywhere
|
||||
// because the isolation is at the database handle, not in the query.
|
||||
let acme_seen = collect(&acme_db).await;
|
||||
let globex_seen = collect(&globex_db).await;
|
||||
|
||||
assert_eq!(acme_seen.len(), 1, "acme should see exactly its own repo");
|
||||
assert_eq!(acme_seen[0].name, "acme-app");
|
||||
assert_eq!(
|
||||
globex_seen.len(),
|
||||
1,
|
||||
"globex should see exactly its own repo"
|
||||
);
|
||||
assert_eq!(globex_seen[0].name, "globex-platform");
|
||||
|
||||
// Sanity: the two databases really are different by name.
|
||||
let acme_db_name = pool.tenant_db_name(&acme.tenant_id);
|
||||
let globex_db_name = pool.tenant_db_name(&globex.tenant_id);
|
||||
assert_ne!(acme_db_name, globex_db_name);
|
||||
assert!(acme_db_name.starts_with(&prefix));
|
||||
|
||||
// Cleanup — drop both per-tenant databases.
|
||||
pool.client()
|
||||
.database(&acme_db_name)
|
||||
.drop()
|
||||
.await
|
||||
.expect("drop acme");
|
||||
pool.client()
|
||||
.database(&globex_db_name)
|
||||
.drop()
|
||||
.await
|
||||
.expect("drop globex");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn for_tenant_is_idempotent_index_creation() {
|
||||
let uri = std::env::var("TEST_MONGODB_URI")
|
||||
.unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into());
|
||||
let prefix = format!("m72a_{}", short_id());
|
||||
let pool = DatabasePool::connect(&uri, &prefix).await.expect("connect");
|
||||
|
||||
let acme = ctx("00000000-0000-0000-0000-00000000acme", "acme");
|
||||
|
||||
// Second call must not fail (ensure_indexes already ran, in-memory
|
||||
// marker is set, Mongo's createIndex is idempotent by name anyway).
|
||||
let _ = pool.for_tenant(&acme).await.expect("first call");
|
||||
let _ = pool.for_tenant(&acme).await.expect("second call");
|
||||
let _ = pool.for_tenant(&acme).await.expect("third call");
|
||||
|
||||
// Cleanup
|
||||
let db_name = pool.tenant_db_name(&acme.tenant_id);
|
||||
pool.client().database(&db_name).drop().await.expect("drop");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tenant_db_name_sanitizes_unsafe_characters() {
|
||||
let uri = std::env::var("TEST_MONGODB_URI")
|
||||
.unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into());
|
||||
let pool = DatabasePool::connect(&uri, "m72a_sanitize")
|
||||
.await
|
||||
.expect("connect");
|
||||
|
||||
// Mongo db names cannot contain `/ \ . " $ <space> NUL`. The pool
|
||||
// must rewrite these without exploding on connect.
|
||||
let funky = "te/n.a\\nt$id\" with spaces";
|
||||
let name = pool.tenant_db_name(funky);
|
||||
for c in ['/', '\\', '.', '"', '$', ' '] {
|
||||
assert!(
|
||||
!name.contains(c),
|
||||
"sanitized db name still contains {c:?}: {name}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tenant_db_name_falls_back_to_hash_when_too_long() {
|
||||
let uri = std::env::var("TEST_MONGODB_URI")
|
||||
.unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into());
|
||||
let pool = DatabasePool::connect(&uri, "m72a_long")
|
||||
.await
|
||||
.expect("connect");
|
||||
|
||||
// 100-byte tenant_id would overflow the 63-byte db-name cap with
|
||||
// any reasonable prefix. The pool must hash it down.
|
||||
let huge = "x".repeat(100);
|
||||
let name = pool.tenant_db_name(&huge);
|
||||
assert!(name.len() <= 63, "hashed name should fit: {name}");
|
||||
assert!(name.starts_with("m72a_long_"));
|
||||
// The hash suffix is 32 hex chars (16-byte SHA-256 truncation).
|
||||
let suffix = name.trim_start_matches("m72a_long_");
|
||||
assert_eq!(
|
||||
suffix.len(),
|
||||
32,
|
||||
"expected 32-hex suffix (16-byte hash), got {suffix:?}"
|
||||
);
|
||||
assert!(suffix.chars().all(|c| c.is_ascii_hexdigit()));
|
||||
|
||||
// Stable: same input → same output.
|
||||
assert_eq!(name, pool.tenant_db_name(&huge));
|
||||
|
||||
// Different inputs → different outputs (collision check on a tiny
|
||||
// sample — full birthday-resistance is a proof not a test).
|
||||
let huge2 = "y".repeat(100);
|
||||
assert_ne!(pool.tenant_db_name(&huge), pool.tenant_db_name(&huge2));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connect_rejects_overlong_db_prefix() {
|
||||
let uri = std::env::var("TEST_MONGODB_URI")
|
||||
.unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into());
|
||||
|
||||
// MAX_PREFIX_LEN is 30 (= 63 - 1 - 32). A 31-char prefix MUST be
|
||||
// rejected at construction so the hash-fallback path can never
|
||||
// produce an over-long db name at runtime.
|
||||
let too_long = "a".repeat(31);
|
||||
let err = DatabasePool::connect(&uri, &too_long).await.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(
|
||||
msg.contains("max is 30") || msg.contains(&too_long),
|
||||
"error should explain the cap: {msg}"
|
||||
);
|
||||
|
||||
// Exactly 30 chars is the inclusive bound — must succeed.
|
||||
let just_right = "a".repeat(30);
|
||||
let _ = DatabasePool::connect(&uri, &just_right)
|
||||
.await
|
||||
.expect("30-char prefix should be accepted");
|
||||
}
|
||||
|
||||
/// Short UUID slug for keeping test prefixes well under Mongo's 63-byte
|
||||
/// db-name cap.
|
||||
fn short_id() -> String {
|
||||
uuid::Uuid::new_v4().simple().to_string()[..8].to_string()
|
||||
}
|
||||
|
||||
/// Drain a `repositories` find cursor on the given tenant database.
|
||||
async fn collect(db: &compliance_agent::database::Database) -> Vec<TrackedRepository> {
|
||||
let mut cursor = db
|
||||
.repositories()
|
||||
.find(doc! {})
|
||||
.await
|
||||
.expect("find repositories");
|
||||
let mut out = Vec::new();
|
||||
while cursor.advance().await.expect("advance") {
|
||||
out.push(cursor.deserialize_current().expect("deserialize"));
|
||||
}
|
||||
out
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -18,6 +18,15 @@ telemetry = [
|
||||
"dep:tracing-subscriber",
|
||||
"dep:tracing",
|
||||
]
|
||||
# Pulls in the M7.1 Axum middleware + extractor. Consumers that don't
|
||||
# embed an HTTP server (e.g. the wasm dashboard frontend) leave it off.
|
||||
axum = [
|
||||
"dep:axum",
|
||||
"dep:jsonwebtoken",
|
||||
"dep:reqwest",
|
||||
"dep:tokio",
|
||||
"dep:tracing",
|
||||
]
|
||||
|
||||
[dependencies]
|
||||
serde = { workspace = true }
|
||||
@@ -37,3 +46,7 @@ opentelemetry-appender-tracing = { version = "0.29", optional = true }
|
||||
tracing-opentelemetry = { version = "0.30", optional = true }
|
||||
tracing-subscriber = { workspace = true, optional = true }
|
||||
tracing = { workspace = true, optional = true }
|
||||
axum = { version = "0.8", optional = true }
|
||||
jsonwebtoken = { version = "9", optional = true }
|
||||
reqwest = { workspace = true, optional = true }
|
||||
tokio = { workspace = true, optional = true }
|
||||
|
||||
@@ -0,0 +1,390 @@
|
||||
//! M7.1 — JWT validation + tenant context propagation.
|
||||
//!
|
||||
//! `require_jwt_auth` validates a Bearer JWT against Keycloak's JWKS and
|
||||
//! attaches a [`TenantContext`] to the request extensions. Downstream
|
||||
//! middleware ([`require_tenant_status`]) and Axum extractors
|
||||
//! ([`crate::tenant_ctx::TenantCtx`]) read it from there.
|
||||
//!
|
||||
//! Skipped paths:
|
||||
//! * `/api/v1/health` — Kubernetes liveness; never authenticated.
|
||||
//!
|
||||
//! Failure modes:
|
||||
//! * No `JwksState` extension → pass-through (single-tenant dev mode).
|
||||
//! * Missing / malformed Bearer header → 401.
|
||||
//! * Signature / expiry invalid → 401.
|
||||
//! * Claims present but tenant_id missing → 401 (treated as a malformed
|
||||
//! token; the realm must always issue tenant_id).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::Request,
|
||||
http::Method,
|
||||
middleware::Next,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use jsonwebtoken::{decode, decode_header, jwk::JwkSet, DecodingKey, Validation};
|
||||
use reqwest::StatusCode;
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::{OrgRole, TenantContext, TenantStatus};
|
||||
|
||||
/// Cached JWKS from Keycloak for token validation.
|
||||
#[derive(Clone)]
|
||||
pub struct JwksState {
|
||||
pub jwks: Arc<RwLock<Option<JwkSet>>>,
|
||||
pub jwks_url: String,
|
||||
}
|
||||
|
||||
/// Raw shape of the JWT payload — matches the breakpilot-dev realm's
|
||||
/// protocol-mapper output. Missing fields default to "" / empty so a
|
||||
/// realm that hasn't been fully wired yet still validates.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct Claims {
|
||||
sub: String,
|
||||
#[serde(default)]
|
||||
name: Option<String>,
|
||||
#[serde(default)]
|
||||
preferred_username: Option<String>,
|
||||
#[serde(default)]
|
||||
tenant_id: String,
|
||||
#[serde(default)]
|
||||
tenant_slug: String,
|
||||
#[serde(default)]
|
||||
org_roles: Vec<String>,
|
||||
#[serde(default)]
|
||||
products: Vec<String>,
|
||||
#[serde(default)]
|
||||
plan: String,
|
||||
#[serde(default)]
|
||||
tenant_status: Option<TenantStatus>,
|
||||
}
|
||||
|
||||
const PUBLIC_ENDPOINTS: &[&str] = &["/api/v1/health"];
|
||||
|
||||
/// Middleware that validates Bearer JWT tokens against Keycloak's JWKS
|
||||
/// and attaches a `TenantContext` extension on success.
|
||||
///
|
||||
/// Skips validation for the health endpoint.
|
||||
/// If `JwksState` is not present (Keycloak not configured), requests
|
||||
/// pass through and downstream code must handle the missing context.
|
||||
pub async fn require_jwt_auth(mut 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(ctx) => {
|
||||
request.extensions_mut().insert(ctx);
|
||||
next.run(request).await
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("JWT validation failed: {e}");
|
||||
(StatusCode::UNAUTHORIZED, "Invalid token").into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Middleware that enforces the M7.1 `tenant_status` contract.
|
||||
///
|
||||
/// * `Active` / `Trial` / `Demo` — pass through.
|
||||
/// * `Frozen` — read-only after cancel / non-payment. Writes return 402.
|
||||
/// * `Archived` — data-retention window closed. Every request returns 410.
|
||||
///
|
||||
/// Pass-through when no `TenantContext` is present (single-tenant dev or
|
||||
/// the upstream JWT middleware ran without `JwksState`).
|
||||
pub async fn require_tenant_status(request: Request, next: Next) -> Response {
|
||||
let ctx = match request.extensions().get::<TenantContext>() {
|
||||
Some(c) => c.clone(),
|
||||
None => return next.run(request).await,
|
||||
};
|
||||
|
||||
if ctx.status.is_archived() {
|
||||
return (
|
||||
StatusCode::GONE,
|
||||
"Tenant archived — data retention window closed",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
if ctx.status.is_frozen() && is_write(request.method()) {
|
||||
return (
|
||||
StatusCode::PAYMENT_REQUIRED,
|
||||
"Tenant frozen — read-only. Re-activate to resume writes.",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
next.run(request).await
|
||||
}
|
||||
|
||||
/// Treat anything other than GET/HEAD/OPTIONS as a write. Good enough for
|
||||
/// REST. The few exceptions (e.g. read-side POSTs) can opt out at the
|
||||
/// handler level once we have them.
|
||||
fn is_write(m: &Method) -> bool {
|
||||
!matches!(m, &Method::GET | &Method::HEAD | &Method::OPTIONS)
|
||||
}
|
||||
|
||||
async fn validate_token(token: &str, state: &JwksState) -> Result<TenantContext, String> {
|
||||
let header = decode_header(token).map_err(|e| format!("failed to decode JWT header: {e}"))?;
|
||||
|
||||
let kid = header
|
||||
.kid
|
||||
.clone()
|
||||
.ok_or_else(|| "JWT missing kid header".to_string())?;
|
||||
|
||||
// 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,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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
|
||||
.iter()
|
||||
.find(|k| k.common.key_id.as_deref() == Some(kid))
|
||||
{
|
||||
Some(j) => j,
|
||||
None => {
|
||||
return Err(ValidationError::Stale(
|
||||
"no matching key found in JWKS".to_string(),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
let decoding_key = DecodingKey::from_jwk(jwk)
|
||||
.map_err(|e| ValidationError::Permanent(format!("failed to create decoding key: {e}")))?;
|
||||
|
||||
let mut validation = Validation::new(header.alg);
|
||||
validation.validate_exp = true;
|
||||
validation.validate_aud = false;
|
||||
|
||||
let data = match decode::<Claims>(token, &decoding_key, &validation) {
|
||||
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).map_err(ValidationError::Permanent)
|
||||
}
|
||||
|
||||
/// Map the decoded JWT payload into the platform-wide `TenantContext`.
|
||||
/// Pulled out for unit testing — no I/O.
|
||||
fn claims_to_context(c: Claims) -> Result<TenantContext, String> {
|
||||
if c.tenant_id.is_empty() {
|
||||
return Err("JWT is missing tenant_id claim".to_string());
|
||||
}
|
||||
|
||||
let status = c.tenant_status.unwrap_or_else(|| {
|
||||
tracing::warn!(
|
||||
"JWT missing tenant_status claim for tenant {} — defaulting to Trial",
|
||||
c.tenant_id
|
||||
);
|
||||
TenantStatus::Trial
|
||||
});
|
||||
|
||||
Ok(TenantContext {
|
||||
tenant_id: c.tenant_id,
|
||||
tenant_slug: c.tenant_slug,
|
||||
org_roles: c.org_roles.iter().map(|r| OrgRole::parse(r)).collect(),
|
||||
products: c.products,
|
||||
plan: c.plan,
|
||||
status,
|
||||
user_id: c.sub,
|
||||
user_name: c.name.or(c.preferred_username),
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch_or_get_jwks(state: &JwksState, force: bool) -> Result<JwkSet, String> {
|
||||
if !force {
|
||||
let cached = state.jwks.read().await;
|
||||
if let Some(ref jwks) = *cached {
|
||||
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)
|
||||
.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}"))?;
|
||||
|
||||
*cached = Some(jwks.clone());
|
||||
|
||||
Ok(jwks)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::expect_used, clippy::unwrap_used)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn base_claims() -> Claims {
|
||||
Claims {
|
||||
sub: "user-123".to_string(),
|
||||
name: Some("Alice Acme".to_string()),
|
||||
preferred_username: None,
|
||||
tenant_id: "00000000-0000-0000-0000-000000000001".to_string(),
|
||||
tenant_slug: "acme".to_string(),
|
||||
org_roles: vec!["IT_ADMIN".to_string()],
|
||||
products: vec!["compliance".to_string()],
|
||||
plan: "professional".to_string(),
|
||||
tenant_status: Some(TenantStatus::Active),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_to_context_happy_path() {
|
||||
let ctx = claims_to_context(base_claims()).expect("should map");
|
||||
assert_eq!(ctx.tenant_id, "00000000-0000-0000-0000-000000000001");
|
||||
assert_eq!(ctx.tenant_slug, "acme");
|
||||
assert_eq!(ctx.org_roles, vec![OrgRole::ItAdmin]);
|
||||
assert_eq!(ctx.products, vec!["compliance"]);
|
||||
assert_eq!(ctx.plan, "professional");
|
||||
assert_eq!(ctx.status, TenantStatus::Active);
|
||||
assert_eq!(ctx.user_id, "user-123");
|
||||
assert_eq!(ctx.user_name.as_deref(), Some("Alice Acme"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_to_context_rejects_missing_tenant_id() {
|
||||
let mut c = base_claims();
|
||||
c.tenant_id = "".to_string();
|
||||
let err = claims_to_context(c).expect_err("should reject");
|
||||
assert!(err.contains("tenant_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_to_context_defaults_status_when_missing() {
|
||||
let mut c = base_claims();
|
||||
c.tenant_status = None;
|
||||
let ctx = claims_to_context(c).expect("should map");
|
||||
assert_eq!(ctx.status, TenantStatus::Trial);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_to_context_falls_back_to_preferred_username() {
|
||||
let mut c = base_claims();
|
||||
c.name = None;
|
||||
c.preferred_username = Some("alice@acme.dev".to_string());
|
||||
let ctx = claims_to_context(c).expect("should map");
|
||||
assert_eq!(ctx.user_name.as_deref(), Some("alice@acme.dev"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claims_to_context_parses_multiple_roles() {
|
||||
let mut c = base_claims();
|
||||
c.org_roles = vec![
|
||||
"IT_ADMIN".to_string(),
|
||||
"CXO".to_string(),
|
||||
"GARBAGE".to_string(),
|
||||
];
|
||||
let ctx = claims_to_context(c).expect("should map");
|
||||
assert_eq!(
|
||||
ctx.org_roles,
|
||||
vec![OrgRole::ItAdmin, OrgRole::Cxo, OrgRole::Unknown]
|
||||
);
|
||||
}
|
||||
|
||||
#[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]
|
||||
fn is_write_detects_methods() {
|
||||
assert!(!is_write(&Method::GET));
|
||||
assert!(!is_write(&Method::HEAD));
|
||||
assert!(!is_write(&Method::OPTIONS));
|
||||
assert!(is_write(&Method::POST));
|
||||
assert!(is_write(&Method::PUT));
|
||||
assert!(is_write(&Method::PATCH));
|
||||
assert!(is_write(&Method::DELETE));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
//! Database helpers shared across the workspace.
|
||||
//!
|
||||
//! `tenant_filter` returns the BSON filter that every query and update
|
||||
//! against a tenant-scoped collection MUST include. Centralising it here
|
||||
//! makes the rule grep-able and keeps query call-sites from accidentally
|
||||
//! omitting it.
|
||||
//!
|
||||
//! Future work (M7.2+): each collection model grows a `tenant_id` field
|
||||
//! and every `find` / `update_*` / `delete_*` call gets this filter
|
||||
//! merged in. The migration to per-collection scoping is tracked
|
||||
//! separately — this helper is the building block.
|
||||
|
||||
use bson::{doc, Document};
|
||||
|
||||
use crate::TenantContext;
|
||||
|
||||
/// Returns `{ "tenant_id": <ctx.tenant_id> }`. Merge this into every
|
||||
/// query filter against a tenant-scoped collection.
|
||||
///
|
||||
/// Use [`tenant_filter_merge`] when you need to combine it with other
|
||||
/// query conditions — it preserves both halves without overwriting.
|
||||
pub fn tenant_filter(ctx: &TenantContext) -> Document {
|
||||
doc! { "tenant_id": &ctx.tenant_id }
|
||||
}
|
||||
|
||||
/// Returns the tenant filter merged with caller-supplied conditions.
|
||||
/// The tenant_id always wins on key conflict — callers cannot
|
||||
/// accidentally override the scoping.
|
||||
pub fn tenant_filter_merge(ctx: &TenantContext, mut extra: Document) -> Document {
|
||||
extra.insert("tenant_id", &ctx.tenant_id);
|
||||
extra
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::TenantStatus;
|
||||
|
||||
fn ctx() -> TenantContext {
|
||||
TenantContext {
|
||||
tenant_id: "t-abc".to_string(),
|
||||
tenant_slug: "acme".to_string(),
|
||||
org_roles: vec![],
|
||||
products: vec![],
|
||||
plan: "starter".to_string(),
|
||||
status: TenantStatus::Active,
|
||||
user_id: "u-1".to_string(),
|
||||
user_name: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn produces_tenant_id_filter() {
|
||||
let f = tenant_filter(&ctx());
|
||||
assert_eq!(f.get_str("tenant_id"), Ok("t-abc"));
|
||||
assert_eq!(f.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_preserves_extra_conditions() {
|
||||
let extra = doc! { "status": "open", "severity": "high" };
|
||||
let f = tenant_filter_merge(&ctx(), extra);
|
||||
assert_eq!(f.get_str("tenant_id"), Ok("t-abc"));
|
||||
assert_eq!(f.get_str("status"), Ok("open"));
|
||||
assert_eq!(f.get_str("severity"), Ok("high"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_overrides_caller_tenant_id() {
|
||||
let extra = doc! { "tenant_id": "evil-other", "status": "open" };
|
||||
let f = tenant_filter_merge(&ctx(), extra);
|
||||
assert_eq!(f.get_str("tenant_id"), Ok("t-abc"));
|
||||
assert_eq!(f.get_str("status"), Ok("open"));
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,17 @@
|
||||
pub mod config;
|
||||
pub mod db;
|
||||
pub mod error;
|
||||
pub mod models;
|
||||
#[cfg(feature = "telemetry")]
|
||||
pub mod telemetry;
|
||||
pub mod tenant;
|
||||
pub mod traits;
|
||||
|
||||
#[cfg(feature = "axum")]
|
||||
pub mod auth;
|
||||
#[cfg(feature = "axum")]
|
||||
pub mod tenant_ctx;
|
||||
|
||||
pub use config::{AgentConfig, DashboardConfig};
|
||||
pub use error::CoreError;
|
||||
pub use tenant::{OrgRole, TenantContext, TenantStatus};
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
//! Tenant context propagated through every authenticated request.
|
||||
//!
|
||||
//! M7.1 single source of truth for "who is this request for". Claims come
|
||||
//! from a Keycloak-issued JWT and land here via [`crate::auth::require_jwt_auth`]
|
||||
//! (enabled with the `axum` feature). Handlers reach into the request
|
||||
//! extensions with the [`crate::tenant_ctx::TenantCtx`] extractor.
|
||||
//!
|
||||
//! The shape mirrors the JWT claim names the breakpilot-platform realm
|
||||
//! emits (see `platform/orca-platform/dev/keycloak/realm-export.json`).
|
||||
//! Stable contract — adding fields is fine; renaming is a breaking
|
||||
//! change for every downstream product.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Tenant lifecycle status from `PLATFORM_ARCHITECTURE.md §5c`.
|
||||
///
|
||||
/// Drives the `tenant_status` middleware:
|
||||
/// * `Demo` / `Trial` / `Active` — full access.
|
||||
/// * `Frozen` — read-only after cancel / non-payment. Mutating endpoints
|
||||
/// return 402.
|
||||
/// * `Archived` — data-retention window closed. Every endpoint returns 410.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum TenantStatus {
|
||||
Demo,
|
||||
Trial,
|
||||
Active,
|
||||
Frozen,
|
||||
Archived,
|
||||
}
|
||||
|
||||
impl TenantStatus {
|
||||
/// True for statuses that block write paths.
|
||||
pub fn is_frozen(&self) -> bool {
|
||||
matches!(self, TenantStatus::Frozen)
|
||||
}
|
||||
/// True for statuses that block every request.
|
||||
pub fn is_archived(&self) -> bool {
|
||||
matches!(self, TenantStatus::Archived)
|
||||
}
|
||||
/// True for the shared demo tenant — metering, billing, and audit
|
||||
/// export are skipped.
|
||||
pub fn is_demo(&self) -> bool {
|
||||
matches!(self, TenantStatus::Demo)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TenantStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Demo => write!(f, "demo"),
|
||||
Self::Trial => write!(f, "trial"),
|
||||
Self::Active => write!(f, "active"),
|
||||
Self::Frozen => write!(f, "frozen"),
|
||||
Self::Archived => write!(f, "archived"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Org-level role baked into the JWT by the realm's protocol mapper.
|
||||
/// `PLATFORM_ARCHITECTURE.md §6` is the canonical list.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "UPPERCASE")]
|
||||
pub enum OrgRole {
|
||||
ItAdmin,
|
||||
Cxo,
|
||||
Finance,
|
||||
Legal,
|
||||
User,
|
||||
/// Anything we haven't enumerated yet — forwards-compatible.
|
||||
#[serde(other)]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl OrgRole {
|
||||
/// Parses a single role string (Keycloak emits these as `IT_ADMIN`,
|
||||
/// `CXO`, etc.). Round-trips with the JSON layer.
|
||||
pub fn parse(s: &str) -> Self {
|
||||
match s {
|
||||
"IT_ADMIN" => OrgRole::ItAdmin,
|
||||
"CXO" => OrgRole::Cxo,
|
||||
"FINANCE" => OrgRole::Finance,
|
||||
"LEGAL" => OrgRole::Legal,
|
||||
"USER" => OrgRole::User,
|
||||
_ => OrgRole::Unknown,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Everything we know about the requesting tenant at the moment a request
|
||||
/// lands. Cheap to clone (every field is owned + small).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TenantContext {
|
||||
/// `tenants.id` from the platform's tenant-registry (UUID).
|
||||
pub tenant_id: String,
|
||||
/// Lowercase URL-safe slug. Useful for log lines + audit emit.
|
||||
pub tenant_slug: String,
|
||||
/// Org-level roles the authenticated user holds inside this tenant.
|
||||
/// Drives the per-handler RBAC in `M7.1-followup` PRs.
|
||||
pub org_roles: Vec<OrgRole>,
|
||||
/// Products this tenant is currently entitled to. Used to short-circuit
|
||||
/// MCP / API calls for unsubscribed products.
|
||||
pub products: Vec<String>,
|
||||
/// Customer plan (`starter` / `professional` / `enterprise`) — gates
|
||||
/// per-plan feature flags (e.g., MCP server is enterprise-only).
|
||||
pub plan: String,
|
||||
/// Lifecycle status — read by `require_tenant_status` middleware.
|
||||
pub status: TenantStatus,
|
||||
/// Keycloak user id of the requester (`sub` claim). Required for audit
|
||||
/// emit so we know WHO did the thing, not just WHICH tenant.
|
||||
pub user_id: String,
|
||||
/// Optional user-facing name from the `name` / `preferred_username`
|
||||
/// claim. Only used in audit + log lines.
|
||||
pub user_name: Option<String>,
|
||||
}
|
||||
|
||||
impl TenantContext {
|
||||
/// True if the caller holds at least one of the listed roles.
|
||||
pub fn has_any_role(&self, roles: &[OrgRole]) -> bool {
|
||||
self.org_roles.iter().any(|r| roles.contains(r))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn org_role_parses_known_values() {
|
||||
assert_eq!(OrgRole::parse("IT_ADMIN"), OrgRole::ItAdmin);
|
||||
assert_eq!(OrgRole::parse("CXO"), OrgRole::Cxo);
|
||||
assert_eq!(OrgRole::parse("USER"), OrgRole::User);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn org_role_unknown_is_forward_compat() {
|
||||
assert_eq!(OrgRole::parse("FUTURE_ROLE"), OrgRole::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tenant_status_predicates() {
|
||||
assert!(TenantStatus::Frozen.is_frozen());
|
||||
assert!(!TenantStatus::Active.is_frozen());
|
||||
assert!(TenantStatus::Archived.is_archived());
|
||||
assert!(TenantStatus::Demo.is_demo());
|
||||
assert!(!TenantStatus::Active.is_demo());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn has_any_role_matches() {
|
||||
let ctx = TenantContext {
|
||||
tenant_id: "t1".into(),
|
||||
tenant_slug: "acme".into(),
|
||||
org_roles: vec![OrgRole::ItAdmin],
|
||||
products: vec![],
|
||||
plan: "professional".into(),
|
||||
status: TenantStatus::Active,
|
||||
user_id: "u".into(),
|
||||
user_name: None,
|
||||
};
|
||||
assert!(ctx.has_any_role(&[OrgRole::ItAdmin]));
|
||||
assert!(ctx.has_any_role(&[OrgRole::Cxo, OrgRole::ItAdmin]));
|
||||
assert!(!ctx.has_any_role(&[OrgRole::User, OrgRole::Cxo]));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
//! Axum extractor for the per-request `TenantContext`.
|
||||
//!
|
||||
//! Handlers consume it as a normal extractor argument:
|
||||
//!
|
||||
//! ```ignore
|
||||
//! async fn list_findings(TenantCtx(ctx): TenantCtx) -> Json<...> {
|
||||
//! let filter = compliance_core::db::tenant_filter(&ctx);
|
||||
//! ...
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! The middleware ([`crate::auth::require_jwt_auth`]) is responsible for
|
||||
//! inserting the context into the request extensions. If it's missing on
|
||||
//! a route that uses this extractor, that's a bug in the wiring — we
|
||||
//! return 401 so the caller sees an auth failure rather than a 500.
|
||||
|
||||
use axum::{
|
||||
extract::FromRequestParts,
|
||||
http::{request::Parts, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
|
||||
use crate::TenantContext;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TenantCtx(pub TenantContext);
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TenantCtxRejection;
|
||||
|
||||
impl IntoResponse for TenantCtxRejection {
|
||||
fn into_response(self) -> Response {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Missing tenant context — request was not authenticated",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> FromRequestParts<S> for TenantCtx
|
||||
where
|
||||
S: Send + Sync,
|
||||
{
|
||||
type Rejection = TenantCtxRejection;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
|
||||
parts
|
||||
.extensions
|
||||
.get::<TenantContext>()
|
||||
.cloned()
|
||||
.map(TenantCtx)
|
||||
.ok_or(TenantCtxRejection)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::expect_used, clippy::unwrap_used)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::TenantStatus;
|
||||
use axum::http::Request;
|
||||
|
||||
fn ctx() -> TenantContext {
|
||||
TenantContext {
|
||||
tenant_id: "t-1".to_string(),
|
||||
tenant_slug: "acme".to_string(),
|
||||
org_roles: vec![],
|
||||
products: vec![],
|
||||
plan: "starter".to_string(),
|
||||
status: TenantStatus::Active,
|
||||
user_id: "u-1".to_string(),
|
||||
user_name: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn extracts_context_when_present() {
|
||||
let mut req = Request::new(());
|
||||
req.extensions_mut().insert(ctx());
|
||||
let (mut parts, _) = req.into_parts();
|
||||
let TenantCtx(found) = TenantCtx::from_request_parts(&mut parts, &())
|
||||
.await
|
||||
.expect("extractor should succeed");
|
||||
assert_eq!(found.tenant_id, "t-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_when_missing() {
|
||||
let req: Request<()> = Request::new(());
|
||||
let (mut parts, _) = req.into_parts();
|
||||
let err = TenantCtx::from_request_parts(&mut parts, &()).await;
|
||||
assert!(err.is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
[package]
|
||||
name = "compliance-smoke"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "Tiny Axum service exercising compliance-core M7.1 tenant gating. Run smoke.sh against it before merging anything that touches the auth/tenant path."
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "compliance-smoke"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
compliance-core = { workspace = true, features = ["axum"] }
|
||||
axum = "0.8"
|
||||
tokio = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
@@ -0,0 +1,111 @@
|
||||
//! M7.1 smoke service.
|
||||
//!
|
||||
//! A standalone Axum binary whose only job is to host the
|
||||
//! [`compliance_core::auth`] middleware + [`compliance_core::tenant_ctx`]
|
||||
//! extractor on three endpoints, so `scripts/smoke.sh` can prove the
|
||||
//! tenant-gating contract end-to-end before any auth-path PR merges.
|
||||
//!
|
||||
//! Endpoints:
|
||||
//! * `GET /api/v1/health` — public, never authenticated.
|
||||
//! * `GET /api/v1/echo` — protected read; returns the [`TenantContext`].
|
||||
//! * `POST /api/v1/echo` — protected write; exercises the `Frozen → 402`
|
||||
//! gate on the same handler.
|
||||
//!
|
||||
//! Configuration (env):
|
||||
//! * `KEYCLOAK_URL` — e.g. `http://localhost:8080`. Required.
|
||||
//! * `KEYCLOAK_REALM` — e.g. `certifai`. Required.
|
||||
//! * `SMOKE_PORT` — defaults to `3010`.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{middleware, routing::get, Extension, Json, Router};
|
||||
use compliance_core::{
|
||||
auth::{require_jwt_auth, require_tenant_status, JwksState},
|
||||
tenant_ctx::TenantCtx,
|
||||
};
|
||||
use serde::Serialize;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct EchoResponse {
|
||||
method: &'static str,
|
||||
tenant_id: String,
|
||||
tenant_slug: String,
|
||||
plan: String,
|
||||
status: String,
|
||||
products: Vec<String>,
|
||||
org_roles: Vec<String>,
|
||||
user_id: String,
|
||||
user_name: Option<String>,
|
||||
}
|
||||
|
||||
async fn health() -> Json<serde_json::Value> {
|
||||
Json(serde_json::json!({ "ok": true }))
|
||||
}
|
||||
|
||||
async fn echo_read(TenantCtx(ctx): TenantCtx) -> Json<EchoResponse> {
|
||||
Json(echo(ctx, "GET"))
|
||||
}
|
||||
|
||||
async fn echo_write(TenantCtx(ctx): TenantCtx) -> Json<EchoResponse> {
|
||||
Json(echo(ctx, "POST"))
|
||||
}
|
||||
|
||||
fn echo(ctx: compliance_core::TenantContext, method: &'static str) -> EchoResponse {
|
||||
EchoResponse {
|
||||
method,
|
||||
tenant_id: ctx.tenant_id,
|
||||
tenant_slug: ctx.tenant_slug,
|
||||
plan: ctx.plan,
|
||||
status: ctx.status.to_string(),
|
||||
products: ctx.products,
|
||||
org_roles: ctx.org_roles.iter().map(|r| format!("{r:?}")).collect(),
|
||||
user_id: ctx.user_id,
|
||||
user_name: ctx.user_name,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||||
)
|
||||
.init();
|
||||
|
||||
let kc_url = std::env::var("KEYCLOAK_URL")
|
||||
.map_err(|_| "KEYCLOAK_URL is required (e.g. http://localhost:8080)")?;
|
||||
let kc_realm = std::env::var("KEYCLOAK_REALM")
|
||||
.map_err(|_| "KEYCLOAK_REALM is required (e.g. certifai)")?;
|
||||
let port: u16 = std::env::var("SMOKE_PORT")
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(3010);
|
||||
|
||||
let jwks_url = format!("{kc_url}/realms/{kc_realm}/protocol/openid-connect/certs");
|
||||
let jwks_state = JwksState {
|
||||
jwks: Arc::new(RwLock::new(None)),
|
||||
jwks_url: jwks_url.clone(),
|
||||
};
|
||||
|
||||
// Layers execute outermost-first. The Extension must be registered
|
||||
// before `require_jwt_auth` so the middleware can read JwksState; the
|
||||
// status gate must run after JWT so `TenantContext` is in extensions.
|
||||
let app = Router::new()
|
||||
.route("/api/v1/health", get(health))
|
||||
.route("/api/v1/echo", get(echo_read).post(echo_write))
|
||||
.layer(middleware::from_fn(require_tenant_status))
|
||||
.layer(middleware::from_fn(require_jwt_auth))
|
||||
.layer(Extension(jwks_state));
|
||||
|
||||
let addr = format!("0.0.0.0:{port}");
|
||||
let listener = tokio::net::TcpListener::bind(&addr).await?;
|
||||
tracing::info!(
|
||||
port,
|
||||
jwks = %jwks_url,
|
||||
"compliance-smoke listening — try `scripts/smoke.sh`"
|
||||
);
|
||||
axum::serve(listener, app).await?;
|
||||
Ok(())
|
||||
}
|
||||
Executable
+136
@@ -0,0 +1,136 @@
|
||||
#!/usr/bin/env bash
|
||||
# M7.1 tenant-gating smoke test.
|
||||
#
|
||||
# Drives compliance-smoke against a live Keycloak realm with five test
|
||||
# users (one per tenant_status), asserts the response code on each
|
||||
# endpoint, and exits non-zero on any mismatch.
|
||||
#
|
||||
# Pre-reqs (one-time):
|
||||
# * KC up at $KC_URL with realm $KC_REALM
|
||||
# * Client $KC_CLIENT has direct-access-grants enabled
|
||||
# * Users + tenant_status mappers per certifai/keycloak/realm-export.json
|
||||
# * compliance-smoke binary running and reachable at $SMOKE_URL
|
||||
#
|
||||
# Usage:
|
||||
# scripts/smoke.sh # uses defaults below
|
||||
# SMOKE_URL=... scripts/smoke.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
KC_URL="${KC_URL:-http://localhost:8080}"
|
||||
KC_REALM="${KC_REALM:-certifai}"
|
||||
KC_CLIENT="${KC_CLIENT:-certifai-dashboard}"
|
||||
SMOKE_URL="${SMOKE_URL:-http://localhost:3010}"
|
||||
|
||||
readonly TOKEN_ENDPOINT="${KC_URL}/realms/${KC_REALM}/protocol/openid-connect/token"
|
||||
|
||||
PASS=0
|
||||
FAIL=0
|
||||
|
||||
red() { printf '\033[31m%s\033[0m' "$*"; }
|
||||
green() { printf '\033[32m%s\033[0m' "$*"; }
|
||||
yellow() { printf '\033[33m%s\033[0m' "$*"; }
|
||||
|
||||
# Fetches an access token via direct access grant. Echoes the raw token.
|
||||
get_token() {
|
||||
local user="$1" pass="$2"
|
||||
curl -sS -X POST "$TOKEN_ENDPOINT" \
|
||||
-H 'Content-Type: application/x-www-form-urlencoded' \
|
||||
-d "grant_type=password" \
|
||||
-d "client_id=${KC_CLIENT}" \
|
||||
-d "username=${user}" \
|
||||
-d "password=${pass}" \
|
||||
-d "scope=openid" \
|
||||
| sed -n 's/.*"access_token":"\([^"]*\)".*/\1/p'
|
||||
}
|
||||
|
||||
# Hits SMOKE_URL$path with the given method and (optional) bearer token,
|
||||
# asserts the response status code matches $want.
|
||||
assert_status() {
|
||||
local label="$1" method="$2" path="$3" want="$4" token="${5:-}"
|
||||
local args=(-sS -o /dev/null -w '%{http_code}' -X "$method" "${SMOKE_URL}${path}")
|
||||
if [[ -n "$token" ]]; then
|
||||
args+=(-H "Authorization: Bearer ${token}")
|
||||
fi
|
||||
local got
|
||||
got=$(curl "${args[@]}")
|
||||
if [[ "$got" == "$want" ]]; then
|
||||
printf ' %s %s %-4s %-15s → %s\n' "$(green PASS)" "$label" "$method" "$path" "$got"
|
||||
PASS=$((PASS + 1))
|
||||
else
|
||||
printf ' %s %s %-4s %-15s → got %s, want %s\n' "$(red FAIL)" "$label" "$method" "$path" "$got" "$want"
|
||||
FAIL=$((FAIL + 1))
|
||||
fi
|
||||
}
|
||||
|
||||
header() {
|
||||
printf '\n%s %s\n' "$(yellow '##')" "$1"
|
||||
}
|
||||
|
||||
# ---- Pre-flight ----------------------------------------------------------
|
||||
header "Pre-flight"
|
||||
if ! curl -sS -o /dev/null -w '%{http_code}\n' "${SMOKE_URL}/api/v1/health" | grep -q '^200$'; then
|
||||
printf ' %s smoke service not reachable at %s\n' "$(red ERR)" "$SMOKE_URL"
|
||||
exit 2
|
||||
fi
|
||||
if ! curl -sS -o /dev/null -w '%{http_code}\n' "${KC_URL}/realms/${KC_REALM}/.well-known/openid-configuration" | grep -q '^200$'; then
|
||||
printf ' %s Keycloak realm %s not reachable at %s\n' "$(red ERR)" "$KC_REALM" "$KC_URL"
|
||||
exit 2
|
||||
fi
|
||||
printf ' %s smoke service + Keycloak both up\n' "$(green OK)"
|
||||
|
||||
# ---- Public endpoint --------------------------------------------------
|
||||
header "Public endpoint (no auth required)"
|
||||
assert_status anon GET /api/v1/health 200
|
||||
|
||||
# ---- Anonymous access to protected endpoints ----------------------------
|
||||
header "Anonymous → 401 on protected endpoints"
|
||||
assert_status anon GET /api/v1/echo 401
|
||||
assert_status anon POST /api/v1/echo 401
|
||||
|
||||
# ---- Bad token ----------------------------------------------------------
|
||||
header "Bad token → 401"
|
||||
assert_status bogus GET /api/v1/echo 401 "not-a-real-jwt"
|
||||
assert_status bogus POST /api/v1/echo 401 "not-a-real-jwt"
|
||||
|
||||
# ---- Active tenant (admin user) -----------------------------------------
|
||||
header "admin@certifai.local (active) → full access"
|
||||
TOKEN=$(get_token admin@certifai.local admin)
|
||||
if [[ -z "$TOKEN" ]]; then
|
||||
printf ' %s failed to fetch token for admin\n' "$(red ERR)"
|
||||
exit 2
|
||||
fi
|
||||
assert_status active GET /api/v1/echo 200 "$TOKEN"
|
||||
assert_status active POST /api/v1/echo 200 "$TOKEN"
|
||||
|
||||
# ---- Active tenant (USER role) ------------------------------------------
|
||||
header "user@certifai.local (active) → full access"
|
||||
TOKEN=$(get_token user@certifai.local user)
|
||||
assert_status active GET /api/v1/echo 200 "$TOKEN"
|
||||
assert_status active POST /api/v1/echo 200 "$TOKEN"
|
||||
|
||||
# ---- Trial tenant -------------------------------------------------------
|
||||
header "trial@acme.local (trial) → full access"
|
||||
TOKEN=$(get_token trial@acme.local trial)
|
||||
assert_status trial GET /api/v1/echo 200 "$TOKEN"
|
||||
assert_status trial POST /api/v1/echo 200 "$TOKEN"
|
||||
|
||||
# ---- Frozen tenant ------------------------------------------------------
|
||||
header "frozen@acme.local (frozen) → read-only, writes 402"
|
||||
TOKEN=$(get_token frozen@acme.local frozen)
|
||||
assert_status frozen GET /api/v1/echo 200 "$TOKEN"
|
||||
assert_status frozen POST /api/v1/echo 402 "$TOKEN"
|
||||
|
||||
# ---- Archived tenant ----------------------------------------------------
|
||||
header "archived@acme.local (archived) → 410 everywhere"
|
||||
TOKEN=$(get_token archived@acme.local archived)
|
||||
assert_status archived GET /api/v1/echo 410 "$TOKEN"
|
||||
assert_status archived POST /api/v1/echo 410 "$TOKEN"
|
||||
|
||||
# ---- Summary ------------------------------------------------------------
|
||||
printf '\n'
|
||||
if [[ "$FAIL" -gt 0 ]]; then
|
||||
printf '%s %d passed, %d failed\n' "$(red FAIL)" "$PASS" "$FAIL"
|
||||
exit 1
|
||||
fi
|
||||
printf '%s %d/%d assertions passed\n' "$(green PASS)" "$PASS" "$PASS"
|
||||
Reference in New Issue
Block a user