diff --git a/Cargo.lock b/Cargo.lock index f18788e1..b7224480 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1177,6 +1177,12 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "http" version = "1.4.1" @@ -2675,6 +2681,7 @@ dependencies = [ "chrono", "dashmap", "futures", + "hex", "parking_lot", "rand 0.8.6", "ring", @@ -3498,6 +3505,7 @@ dependencies = [ "chrono", "encryption-inmemory", "futures", + "hex", "reqwest 0.12.28", "ring", "rook-core", diff --git a/apps/rook/dashboard/e2e/api-keys.spec.ts b/apps/rook/dashboard/e2e/api-keys.spec.ts index 1f46f708..1b11879d 100644 --- a/apps/rook/dashboard/e2e/api-keys.spec.ts +++ b/apps/rook/dashboard/e2e/api-keys.spec.ts @@ -107,7 +107,7 @@ async function revokeKeysByLabelViaApi(page: Page, label: string): Promise async function createApiKeyViaApi( page: Page, label: string, - scopes: string[] = ['read'], + scopes: string[] = ['chat:read'], tier: string = 'free' ): Promise<{ id: string; plaintextKey: string }> { const csrf = await getCsrfToken(page) @@ -235,7 +235,7 @@ test.describe('API Keys - Create Flow', () => { // Fill in the form await page.getByLabel(/label/i).fill(createLabel) - await page.getByRole('dialog').getByText(/read/i).click() + await page.getByRole('dialog').getByText(/^chat read$/i).click() // Submit await page.getByRole('button', { name: /create key/i }).click() @@ -283,7 +283,7 @@ test.describe('API Keys - Edit Flow', () => { // Revoke any leftover key for this worker, then create a fresh one. await revokeKeysByLabelViaApi(page, editLabel) await revokeKeysByLabelViaApi(page, `updated-key-label-${testInfo.workerIndex}`) - await createApiKeyViaApi(page, editLabel, ['read'], 'free') + await createApiKeyViaApi(page, editLabel, ['chat:read'], 'free') }) test('opens edit modal when clicking edit button', async ({ page }) => { @@ -366,7 +366,7 @@ test.describe('API Keys - Revoke Flow', () => { // Revoke only OUR worker's label — avoids touching keys belonging to // concurrent browser workers, which caused a race condition on reload. await revokeKeysByLabelViaApi(page, revokeLabel) - await createApiKeyViaApi(page, revokeLabel, ['read'], 'free') + await createApiKeyViaApi(page, revokeLabel, ['chat:read'], 'free') }) test('shows confirmation dialog when revoking', async ({ page }) => { diff --git a/apps/rook/dashboard/src/views/ApiKeysView.vue b/apps/rook/dashboard/src/views/ApiKeysView.vue index 03fa05e2..351307d7 100644 --- a/apps/rook/dashboard/src/views/ApiKeysView.vue +++ b/apps/rook/dashboard/src/views/ApiKeysView.vue @@ -170,8 +170,8 @@ const hasNextPage = computed(() => offset.value + limit.value < total.value) const hasPrevPage = computed(() => offset.value > 0) const scopesOptions = [ - { value: 'read', label: 'Read' }, - { value: 'write', label: 'Write' }, + { value: 'chat:read', label: 'Chat Read' }, + { value: 'chat:write', label: 'Chat Write' }, ] const tierOptions = [ diff --git a/crates/application/rook-usecases/Cargo.toml b/crates/application/rook-usecases/Cargo.toml index afc724cd..1c8b1802 100644 --- a/crates/application/rook-usecases/Cargo.toml +++ b/crates/application/rook-usecases/Cargo.toml @@ -46,3 +46,6 @@ subtle = "2" # Base64 for token encoding base64 = "0.22" + +# Hex encoding for digest output +hex = "0.4" diff --git a/crates/application/rook-usecases/src/auth/bootstrap_status.rs b/crates/application/rook-usecases/src/auth/bootstrap_status.rs index a78934e9..a34e83b2 100644 --- a/crates/application/rook-usecases/src/auth/bootstrap_status.rs +++ b/crates/application/rook-usecases/src/auth/bootstrap_status.rs @@ -74,6 +74,8 @@ impl BootstrapStatus { scopes: vec![ApiKeyScope::parse("admin").expect("static admin scope")], tier: ApiKeyTier::Enterprise, expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], }) .await?; diff --git a/crates/application/rook-usecases/src/auth/login.rs b/crates/application/rook-usecases/src/auth/login.rs index 3a6ace83..450768b4 100644 --- a/crates/application/rook-usecases/src/auth/login.rs +++ b/crates/application/rook-usecases/src/auth/login.rs @@ -87,7 +87,7 @@ impl Login { // Step 5: SHA-256 hash the token → token_hash let mut hasher = Sha256::new(); hasher.update(token_bytes); - let token_hash = format!("{:x}", hasher.finalize()); + let token_hash = hex::encode(hasher.finalize()); // Step 6: Create session let new_session = rook_core::NewSession { diff --git a/crates/application/rook-usecases/src/auth/validate_session.rs b/crates/application/rook-usecases/src/auth/validate_session.rs index 6a0743a5..6b4a4b71 100644 --- a/crates/application/rook-usecases/src/auth/validate_session.rs +++ b/crates/application/rook-usecases/src/auth/validate_session.rs @@ -54,7 +54,7 @@ impl ValidateSession { // Step 2: SHA-256 hash the token bytes let mut hasher = Sha256::new(); hasher.update(&token_bytes); - let token_hash = format!("{:x}", hasher.finalize()); + let token_hash = hex::encode(hasher.finalize()); // Step 3: Look up session by token hash let session = self diff --git a/crates/application/rook-usecases/src/authenticate_client_api.rs b/crates/application/rook-usecases/src/authenticate_client_api.rs index 1da3f64d..174533c8 100644 --- a/crates/application/rook-usecases/src/authenticate_client_api.rs +++ b/crates/application/rook-usecases/src/authenticate_client_api.rs @@ -152,8 +152,10 @@ mod tests { ApiKeySubject { id: ApiKeyId::new("key_1"), label: "Production".to_string(), - scopes: vec![ApiKeyScope::parse("read").expect("scope")], + scopes: vec![ApiKeyScope::parse("chat:read").expect("scope")], tier: ApiKeyTier::Pro, + allowed_models: vec![], + allowed_providers: vec![], } } diff --git a/crates/application/rook-usecases/src/manage_api_keys.rs b/crates/application/rook-usecases/src/manage_api_keys.rs index 6c4b2b26..a347d5df 100644 --- a/crates/application/rook-usecases/src/manage_api_keys.rs +++ b/crates/application/rook-usecases/src/manage_api_keys.rs @@ -5,6 +5,7 @@ use chrono::{DateTime, Utc}; use rand::RngCore; use rook_core::{ ApiKeyId, ApiKeyRecord, ApiKeyRepositoryError, ApiKeyRepositoryPort, ApiKeyScope, ApiKeyTier, + ApiKeyValidationError, ModelId, ProviderId, }; #[derive(Debug, thiserror::Error)] @@ -68,6 +69,9 @@ impl ManageApiKeys { } } + // Validate all requested scopes are canonical. + validate_scopes(&request.scopes)?; + let raw_key = generate_api_key(); let key_hash = hash_api_key(&raw_key, &self.hash_secret); let key_prefix: String = raw_key.chars().take(8).collect(); @@ -86,6 +90,8 @@ impl ManageApiKeys { expires_at: request.expires_at, created_at: now, last_used_at: None, + allowed_models: request.allowed_models, + allowed_providers: request.allowed_providers, }; self.repo.create(&record).await?; @@ -104,7 +110,14 @@ impl ManageApiKeys { .ok_or_else(|| ManageApiKeysError::NotFound(id.clone()))?; let label = request.label.unwrap_or(existing.label); - let scopes = request.scopes.unwrap_or(existing.scopes); + let scopes = match request.scopes { + Some(new_scopes) => { + // Validate incoming scopes before applying the update. + validate_scopes(&new_scopes)?; + new_scopes + } + None => existing.scopes, + }; let tier = request.tier.unwrap_or(existing.tier); let is_active = request.is_active.unwrap_or(existing.is_active); @@ -133,6 +146,10 @@ impl ManageApiKeys { expires_at, created_at: existing.created_at, last_used_at: existing.last_used_at, + allowed_models: request.allowed_models.unwrap_or(existing.allowed_models), + allowed_providers: request + .allowed_providers + .unwrap_or(existing.allowed_providers), }; self.repo.update(&updated).await?; @@ -161,6 +178,8 @@ pub struct CreateApiKeyRequest { pub scopes: Vec, pub tier: ApiKeyTier, pub expires_at: Option>, + pub allowed_models: Vec, + pub allowed_providers: Vec, } #[derive(Debug, Clone, Default)] @@ -170,6 +189,27 @@ pub struct UpdateApiKeyRequest { pub tier: Option, pub is_active: Option, pub expires_at: Option>>, + pub allowed_models: Option>, + pub allowed_providers: Option>, +} + +/// Validates that every scope in the slice is a known canonical value. +/// Returns `ManageApiKeysError::Validation` on the first unknown scope found. +fn validate_scopes(scopes: &[ApiKeyScope]) -> ManageApiKeysResult<()> { + for scope in scopes { + ApiKeyScope::parse(scope.as_str()).map_err(|e| match e { + ApiKeyValidationError::UnknownScope(s) => { + ManageApiKeysError::Validation(format!("unknown scope: {s}")) + } + ApiKeyValidationError::EmptyScope => { + ManageApiKeysError::Validation("scope must not be empty".into()) + } + ApiKeyValidationError::InvalidTier(t) => { + ManageApiKeysError::Validation(format!("invalid tier: {t}")) + } + })?; + } + Ok(()) } fn generate_api_key() -> String { @@ -307,9 +347,11 @@ mod tests { // 1. Create a key let create_req = CreateApiKeyRequest { label: "Dev Key".to_string(), - scopes: vec![ApiKeyScope::parse("read").unwrap()], + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], }; let (record, raw_key) = usecase.create(create_req).await.unwrap(); assert_eq!(record.label, "Dev Key"); @@ -333,6 +375,8 @@ mod tests { tier: Some(ApiKeyTier::Enterprise), is_active: Some(false), expires_at: None, + allowed_models: None, + allowed_providers: None, }; let updated = usecase.update(&record.id, update_req).await.unwrap(); assert_eq!(updated.label, "Prod Key"); @@ -354,9 +398,11 @@ mod tests { let create_req = CreateApiKeyRequest { label: "Test Key".to_string(), - scopes: vec![ApiKeyScope::parse("read").unwrap()], + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], }; let (record, _) = usecase.create(create_req).await.unwrap(); assert!(record.is_active); @@ -377,9 +423,11 @@ mod tests { for i in 0..5 { let create_req = CreateApiKeyRequest { label: format!("Key {}", i), - scopes: vec![ApiKeyScope::parse("read").unwrap()], + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], }; usecase.create(create_req).await.unwrap(); } @@ -402,9 +450,11 @@ mod tests { let create_req = CreateApiKeyRequest { label: "Expired Key".to_string(), - scopes: vec![ApiKeyScope::parse("read").unwrap()], + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, expires_at: Some(Utc::now() - chrono::Duration::days(1)), + allowed_models: vec![], + allowed_providers: vec![], }; let result = usecase.create(create_req).await; @@ -424,9 +474,11 @@ mod tests { let create_req = CreateApiKeyRequest { label: "To Delete".to_string(), - scopes: vec![ApiKeyScope::parse("read").unwrap()], + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], }; let (record, _) = usecase.create(create_req).await.unwrap(); @@ -438,4 +490,146 @@ mod tests { assert!(!found.is_active); assert!(found.revoked_at.is_some()); } + + #[tokio::test] + async fn test_create_with_unknown_scope_is_rejected() { + let repo = Arc::new(FakeApiKeyRepository::default()); + let usecase = ManageApiKeys::new(repo.clone(), "test-secret"); + + // Use a pre-built ApiKeyScope via parse_lenient to bypass the strict check + // and simulate a caller passing an unknown scope in the request struct directly. + // We can only reach this code path by constructing the scope via parse_lenient. + let bad_scope = ApiKeyScope::parse_lenient("legacy:scope"); + let create_req = CreateApiKeyRequest { + label: "Bad Scope Key".to_string(), + scopes: vec![bad_scope], + tier: ApiKeyTier::Free, + expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + + let result = usecase.create(create_req).await; + assert!(result.is_err()); + match result { + Err(ManageApiKeysError::Validation(msg)) => { + assert!(msg.contains("unknown scope"), "message was: {msg}"); + } + other => panic!("expected Validation error, got {:?}", other), + } + } + + #[tokio::test] + async fn test_update_with_unknown_scope_is_rejected() { + let repo = Arc::new(FakeApiKeyRepository::default()); + let usecase = ManageApiKeys::new(repo.clone(), "test-secret"); + + // Create a valid key first. + let create_req = CreateApiKeyRequest { + label: "Valid Key".to_string(), + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], + tier: ApiKeyTier::Free, + expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + let (record, _) = usecase.create(create_req).await.unwrap(); + + // Attempt update with an unknown scope. + let bad_scope = ApiKeyScope::parse_lenient("legacy:scope"); + let update_req = UpdateApiKeyRequest { + scopes: Some(vec![bad_scope]), + ..Default::default() + }; + + let result = usecase.update(&record.id, update_req).await; + assert!(result.is_err()); + match result { + Err(ManageApiKeysError::Validation(msg)) => { + assert!(msg.contains("unknown scope"), "message was: {msg}"); + } + other => panic!("expected Validation error, got {:?}", other), + } + } + + #[tokio::test] + async fn test_create_with_allowed_models_and_providers() { + use rook_core::{ModelId, ProviderId}; + let repo = Arc::new(FakeApiKeyRepository::default()); + let usecase = ManageApiKeys::new(repo.clone(), "test-secret"); + + let create_req = CreateApiKeyRequest { + label: "Restricted Key".to_string(), + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], + tier: ApiKeyTier::Free, + expires_at: None, + allowed_models: vec![ModelId::new("gpt-4"), ModelId::new("claude-3")], + allowed_providers: vec![ProviderId::new("openai")], + }; + let (record, _) = usecase.create(create_req).await.unwrap(); + + assert_eq!(record.allowed_models.len(), 2); + assert_eq!(record.allowed_models[0].as_str(), "gpt-4"); + assert_eq!(record.allowed_providers.len(), 1); + assert_eq!(record.allowed_providers[0].as_str(), "openai"); + } + + #[tokio::test] + async fn test_update_allowed_models_and_providers() { + use rook_core::{ModelId, ProviderId}; + let repo = Arc::new(FakeApiKeyRepository::default()); + let usecase = ManageApiKeys::new(repo.clone(), "test-secret"); + + let create_req = CreateApiKeyRequest { + label: "Key".to_string(), + scopes: vec![ApiKeyScope::parse("chat:read").unwrap()], + tier: ApiKeyTier::Free, + expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + let (record, _) = usecase.create(create_req).await.unwrap(); + assert!(record.allowed_models.is_empty()); + + // Update to add restrictions + let update_req = UpdateApiKeyRequest { + allowed_models: Some(vec![ModelId::new("gpt-4o")]), + allowed_providers: Some(vec![ProviderId::new("anthropic")]), + ..Default::default() + }; + let updated = usecase.update(&record.id, update_req).await.unwrap(); + assert_eq!(updated.allowed_models.len(), 1); + assert_eq!(updated.allowed_models[0].as_str(), "gpt-4o"); + assert_eq!(updated.allowed_providers[0].as_str(), "anthropic"); + + // Update to clear restrictions (set to empty = all allowed) + let clear_req = UpdateApiKeyRequest { + allowed_models: Some(vec![]), + allowed_providers: Some(vec![]), + ..Default::default() + }; + let cleared = usecase.update(&record.id, clear_req).await.unwrap(); + assert!(cleared.allowed_models.is_empty()); + assert!(cleared.allowed_providers.is_empty()); + } + + #[tokio::test] + async fn test_empty_restrictions_means_all_allowed() { + let repo = Arc::new(FakeApiKeyRepository::default()); + let usecase = ManageApiKeys::new(repo.clone(), "test-secret"); + + let create_req = CreateApiKeyRequest { + label: "Unrestricted Key".to_string(), + scopes: vec![ApiKeyScope::parse("chat:write").unwrap()], + tier: ApiKeyTier::Pro, + expires_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + let (record, _) = usecase.create(create_req).await.unwrap(); + + // Empty = unrestricted: no allowed_models means all models OK + assert!(record.allowed_models.is_empty(), "empty = unrestricted"); + assert!(record.allowed_providers.is_empty(), "empty = unrestricted"); + } } diff --git a/crates/application/rook-usecases/src/route_request.rs b/crates/application/rook-usecases/src/route_request.rs index 1a09a894..f10ea31e 100644 --- a/crates/application/rook-usecases/src/route_request.rs +++ b/crates/application/rook-usecases/src/route_request.rs @@ -55,6 +55,16 @@ impl RouteRequest { let cache_key = req.cache_key(); let start = Instant::now(); + // 0. Model restriction check (before any provider interaction) + if !req.restrictions.allowed_models.is_empty() + && !req.restrictions.allowed_models.contains(&req.model) + { + return Err(CortexError::forbidden(format!( + "model '{}' is not permitted by this API key", + req.model.as_str() + ))); + } + // 1. Cache hit? if req.metadata.cacheable { if let Some(cached) = self.cache.get(&cache_key).await? { @@ -67,6 +77,16 @@ impl RouteRequest { let provider = self.router.select(&req).await?; let provider_id = provider.id().clone(); + // 2a. Provider restriction check (after selection, before execution) + if !req.restrictions.allowed_providers.is_empty() + && !req.restrictions.allowed_providers.contains(&provider_id) + { + return Err(CortexError::forbidden(format!( + "provider '{}' is not permitted by this API key", + provider_id.as_str() + ))); + } + let provider_format = provider.api_format(); let provider_req = self.format_translator.translate_request( client_format, @@ -130,8 +150,29 @@ impl RouteRequest { ) -> Result>, CortexError> { let start = Instant::now(); + + // 0. Model restriction check + if !req.restrictions.allowed_models.is_empty() + && !req.restrictions.allowed_models.contains(&req.model) + { + return Err(CortexError::forbidden(format!( + "model '{}' is not permitted by this API key", + req.model.as_str() + ))); + } + let provider = self.router.select(&req).await?; let provider_id = provider.id().clone(); + + // 0a. Provider restriction check + if !req.restrictions.allowed_providers.is_empty() + && !req.restrictions.allowed_providers.contains(&provider_id) + { + return Err(CortexError::forbidden(format!( + "provider '{}' is not permitted by this API key", + provider_id.as_str() + ))); + } let provider_format = provider.api_format(); let provider_req = self.format_translator.translate_request( client_format, @@ -403,6 +444,7 @@ mod tests { cacheable: true, priority: 1, }, + restrictions: rook_core::ApiKeyRestrictions::default(), } } @@ -452,4 +494,98 @@ mod tests { assert_eq!(entries[0].status, RequestStatus::Success); assert_eq!(entries[0].usage.as_ref().unwrap().total_tokens, 5); } + + fn make_usecase() -> RouteRequest { + let provider: Arc = Arc::new(TestProvider { + id: ProviderId::new("test-provider"), + }); + RouteRequest::new( + Arc::new(TestRouter { provider }), + Arc::new(TestCache { + get_calls: Mutex::new(0), + set_calls: Mutex::new(0), + }), + Arc::new(TestAudit { + entries: Mutex::new(Vec::new()), + }), + Arc::new(TestFormatTranslator), + ) + } + + #[tokio::test] + async fn execute_is_forbidden_when_model_not_in_allowed_list() { + let usecase = make_usecase(); + let mut req = request(); + req.restrictions.allowed_models = + vec![ModelId::new("gpt-4"), ModelId::new("claude-3-opus")]; + // TEST_MODEL is "gpt-test" — not in the allowed list + + let result = usecase.execute(req).await; + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!(err.is_forbidden(), "expected forbidden, got: {err}"); + } + + #[tokio::test] + async fn execute_succeeds_when_model_is_in_allowed_list() { + let usecase = make_usecase(); + let mut req = request(); + req.restrictions.allowed_models = vec![TEST_MODEL.clone()]; + + let result = usecase.execute(req).await; + assert!(result.is_ok(), "expected success, got: {:?}", result.err()); + } + + #[tokio::test] + async fn execute_succeeds_when_allowed_models_is_empty() { + let usecase = make_usecase(); + // default restrictions — empty = unrestricted + let result = usecase.execute(request()).await; + assert!(result.is_ok(), "expected success, got: {:?}", result.err()); + } + + #[tokio::test] + async fn execute_is_forbidden_when_provider_not_in_allowed_list() { + let usecase = make_usecase(); + let mut req = request(); + // test-provider is the provider selected; restrict to a different one + req.restrictions.allowed_providers = vec![ProviderId::new("anthropic")]; + + let result = usecase.execute(req).await; + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!(err.is_forbidden(), "expected forbidden, got: {err}"); + } + + #[tokio::test] + async fn execute_succeeds_when_provider_is_in_allowed_list() { + let usecase = make_usecase(); + let mut req = request(); + req.restrictions.allowed_providers = vec![ProviderId::new("test-provider")]; + + let result = usecase.execute(req).await; + assert!(result.is_ok(), "expected success, got: {:?}", result.err()); + } + + #[tokio::test] + async fn execute_stream_is_forbidden_when_model_not_allowed() { + let usecase = make_usecase(); + let mut req = request(); + req.restrictions.allowed_models = vec![ModelId::new("gpt-4")]; + + let result = usecase.execute_stream(req).await; + assert!(result.is_err()); + assert!(result.err().unwrap().is_forbidden()); + } + + #[tokio::test] + async fn execute_stream_is_forbidden_when_provider_not_allowed() { + let usecase = make_usecase(); + let mut req = request(); + req.restrictions.allowed_providers = vec![ProviderId::new("openai")]; + + let result = usecase.execute_stream(req).await; + assert!(result.is_err()); + assert!(result.err().unwrap().is_forbidden()); + } } diff --git a/crates/application/rook-usecases/src/router_impl.rs b/crates/application/rook-usecases/src/router_impl.rs index df9f6082..9ccf1fe8 100644 --- a/crates/application/rook-usecases/src/router_impl.rs +++ b/crates/application/rook-usecases/src/router_impl.rs @@ -322,6 +322,7 @@ mod tests { cacheable: true, priority: 1, }, + restrictions: rook_core::ApiKeyRestrictions::default(), } } diff --git a/crates/domain/rook-core/src/api_key.rs b/crates/domain/rook-core/src/api_key.rs index 63ebb62b..795ae775 100644 --- a/crates/domain/rook-core/src/api_key.rs +++ b/crates/domain/rook-core/src/api_key.rs @@ -2,6 +2,7 @@ use std::str::FromStr; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use shared_kernel::{ModelId, ProviderId}; use smol_str::SmolStr; #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -23,18 +24,73 @@ impl std::fmt::Display for ApiKeyId { } } +/// Canonical set of valid API key scopes. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum KnownScope { + ChatRead, + ChatWrite, + ProvidersRead, + ProvidersWrite, + Admin, +} + +impl KnownScope { + pub fn as_str(self) -> &'static str { + match self { + Self::ChatRead => "chat:read", + Self::ChatWrite => "chat:write", + Self::ProvidersRead => "providers:read", + Self::ProvidersWrite => "providers:write", + Self::Admin => "admin", + } + } +} + +impl FromStr for KnownScope { + type Err = (); + + fn from_str(value: &str) -> Result { + match value { + "chat:read" => Ok(Self::ChatRead), + "chat:write" => Ok(Self::ChatWrite), + "providers:read" => Ok(Self::ProvidersRead), + "providers:write" => Ok(Self::ProvidersWrite), + "admin" => Ok(Self::Admin), + _ => Err(()), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ApiKeyScope(SmolStr); impl ApiKeyScope { + /// Strict parse — only canonical scope values are accepted. + /// Returns `Err(UnknownScope)` for anything outside the known set. pub fn parse(value: &str) -> Result { let value = value.trim(); if value.is_empty() { return Err(ApiKeyValidationError::EmptyScope); } + if value.parse::().is_err() { + return Err(ApiKeyValidationError::UnknownScope(value.to_string())); + } Ok(Self(value.into())) } + /// Lenient parse for reading from the database. + /// Accepts any non-empty string and logs a warning for unknown scopes + /// so existing DB rows are never rejected. + pub fn parse_lenient(value: &str) -> Self { + let value = value.trim(); + if value.is_empty() { + tracing::warn!(scope = "", "parse_lenient received empty scope"); + } else if value.parse::().is_err() { + tracing::warn!(scope = value, "unknown API key scope loaded from database"); + } + Self(value.into()) + } + pub fn as_str(&self) -> &str { &self.0 } @@ -76,12 +132,18 @@ pub struct ApiKeySubject { pub label: String, pub scopes: Vec, pub tier: ApiKeyTier, + /// Empty vec means unrestricted (all models allowed). + pub allowed_models: Vec, + /// Empty vec means unrestricted (all providers allowed). + pub allowed_providers: Vec, } #[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)] pub enum ApiKeyValidationError { #[error("API key scope must not be empty")] EmptyScope, + #[error("unknown API key scope: {0}")] + UnknownScope(String), #[error("invalid API key tier: {0}")] InvalidTier(String), } @@ -109,4 +171,108 @@ pub struct ApiKeyRecord { pub expires_at: Option>, pub created_at: DateTime, pub last_used_at: Option>, + /// Empty vec means unrestricted (all models allowed). + pub allowed_models: Vec, + /// Empty vec means unrestricted (all providers allowed). + pub allowed_providers: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn all_five_known_scopes_parse_ok() { + assert!(ApiKeyScope::parse("chat:read").is_ok()); + assert!(ApiKeyScope::parse("chat:write").is_ok()); + assert!(ApiKeyScope::parse("providers:read").is_ok()); + assert!(ApiKeyScope::parse("providers:write").is_ok()); + assert!(ApiKeyScope::parse("admin").is_ok()); + } + + #[test] + fn unknown_scope_returns_unknown_scope_error() { + let err = ApiKeyScope::parse("read").unwrap_err(); + assert_eq!(err, ApiKeyValidationError::UnknownScope("read".to_string())); + } + + #[test] + fn empty_scope_returns_empty_scope_error() { + let err = ApiKeyScope::parse("").unwrap_err(); + assert_eq!(err, ApiKeyValidationError::EmptyScope); + + let err_ws = ApiKeyScope::parse(" ").unwrap_err(); + assert_eq!(err_ws, ApiKeyValidationError::EmptyScope); + } + + #[test] + fn uppercase_scope_is_rejected() { + // Scope matching is case-sensitive; uppercase variants are unknown. + let err = ApiKeyScope::parse("Chat:Read").unwrap_err(); + assert!( + matches!(err, ApiKeyValidationError::UnknownScope(_)), + "expected UnknownScope, got {:?}", + err + ); + + let err2 = ApiKeyScope::parse("ADMIN").unwrap_err(); + assert!(matches!(err2, ApiKeyValidationError::UnknownScope(_))); + } + + #[test] + fn parse_lenient_accepts_unknown_scope_without_error() { + // Should not panic or return an error; just log a warning. + let scope = ApiKeyScope::parse_lenient("legacy:custom"); + assert_eq!(scope.as_str(), "legacy:custom"); + } + + #[test] + fn parse_lenient_accepts_known_scope() { + let scope = ApiKeyScope::parse_lenient("chat:read"); + assert_eq!(scope.as_str(), "chat:read"); + } + + #[test] + fn api_key_record_has_allowed_models_and_providers_fields() { + use shared_kernel::{ModelId, ProviderId}; + let record = ApiKeyRecord { + id: ApiKeyId::new("test-key"), + label: "Test".to_string(), + key_hash: "hash".to_string(), + key_prefix: "prefix".to_string(), + scopes: vec![], + tier: ApiKeyTier::Free, + is_active: true, + revoked_at: None, + expires_at: None, + created_at: chrono::Utc::now(), + last_used_at: None, + allowed_models: vec![ModelId::new("gpt-4"), ModelId::new("claude-3")], + allowed_providers: vec![ProviderId::new("openai")], + }; + assert_eq!(record.allowed_models.len(), 2); + assert_eq!(record.allowed_providers.len(), 1); + } + + #[test] + fn api_key_record_empty_restrictions_means_unrestricted() { + let record = ApiKeyRecord { + id: ApiKeyId::new("test-key"), + label: "Test".to_string(), + key_hash: "hash".to_string(), + key_prefix: "prefix".to_string(), + scopes: vec![], + tier: ApiKeyTier::Free, + is_active: true, + revoked_at: None, + expires_at: None, + created_at: chrono::Utc::now(), + last_used_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + // Empty means unrestricted — both fields are present and empty + assert!(record.allowed_models.is_empty()); + assert!(record.allowed_providers.is_empty()); + } } diff --git a/crates/domain/rook-core/src/lib.rs b/crates/domain/rook-core/src/lib.rs index 8f4b3d7b..67e19184 100644 --- a/crates/domain/rook-core/src/lib.rs +++ b/crates/domain/rook-core/src/lib.rs @@ -41,8 +41,10 @@ mod api_key_tests { #[test] fn api_key_scope_trims_and_rejects_empty_values() { - let scope = ApiKeyScope::parse(" read ").expect("scope"); - assert_eq!(scope.as_str(), "read"); + // Only canonical scope values are accepted by parse(). + let scope = ApiKeyScope::parse(" chat:read ").expect("scope"); + assert_eq!(scope.as_str(), "chat:read"); assert!(ApiKeyScope::parse(" ").is_err()); + assert!(ApiKeyScope::parse("read").is_err()); // non-canonical → UnknownScope } } diff --git a/crates/domain/rook-core/src/model.rs b/crates/domain/rook-core/src/model.rs index a809dd95..df70dfad 100644 --- a/crates/domain/rook-core/src/model.rs +++ b/crates/domain/rook-core/src/model.rs @@ -125,6 +125,14 @@ pub struct NewSession { /// --------------------------------------------------------------------------- /// Request / Response /// --------------------------------------------------------------------------- +/// API key restriction filters. Empty vecs mean "unrestricted" (all allowed). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ApiKeyRestrictions { + /// If non-empty, only these model IDs are allowed. + pub allowed_models: Vec, + /// If non-empty, only these provider IDs are allowed. + pub allowed_providers: Vec, +} #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CompletionRequest { @@ -137,6 +145,9 @@ pub struct CompletionRequest { pub tools: Option, pub tool_choice: Option, pub metadata: RequestMetadata, + /// API key restrictions — empty vecs mean unrestricted. + #[serde(default)] + pub restrictions: ApiKeyRestrictions, } impl CompletionRequest { diff --git a/crates/domain/shared-kernel/src/error.rs b/crates/domain/shared-kernel/src/error.rs index 1fe193bb..8749e892 100644 --- a/crates/domain/shared-kernel/src/error.rs +++ b/crates/domain/shared-kernel/src/error.rs @@ -63,6 +63,32 @@ impl CortexError { self.source.is::() } + pub fn forbidden(msg: impl Into) -> Self { + Self { + source: Box::new(ForbiddenError(msg.into())), + } + } + + pub fn is_forbidden(&self) -> bool { + self.source.is::() + } + + /// Error code for `MODEL_NOT_ALLOWED` and `PROVIDER_NOT_ALLOWED` so the + /// HTTP layer can return a specific `code` to the client. + pub fn forbidden_code(&self) -> Option<&'static str> { + if let Some(ForbiddenError(msg)) = self.source.downcast_ref::() { + if msg.starts_with("model ") { + Some("model_not_allowed") + } else if msg.starts_with("provider ") { + Some("provider_not_allowed") + } else { + None + } + } else { + None + } + } + pub fn is_rate_limited(&self) -> bool { self.source.is::() } @@ -147,6 +173,17 @@ impl fmt::Display for InvalidRequestError { impl std::error::Error for InvalidRequestError {} +#[derive(Debug)] +pub struct ForbiddenError(pub String); + +impl fmt::Display for ForbiddenError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "forbidden: {}", self.0) + } +} + +impl std::error::Error for ForbiddenError {} + // --------------------------------------------------------------------------- // Result type alias // --------------------------------------------------------------------------- diff --git a/crates/infrastructure/auth-sqlite/src/lib.rs b/crates/infrastructure/auth-sqlite/src/lib.rs index adf568c4..b1564aab 100644 --- a/crates/infrastructure/auth-sqlite/src/lib.rs +++ b/crates/infrastructure/auth-sqlite/src/lib.rs @@ -6,8 +6,8 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; use rook_core::{ ApiKeyId, ApiKeyRecord, ApiKeyRepositoryError, ApiKeyRepositoryPort, ApiKeyScope, - ApiKeySubject, ApiKeyTier, NewSession, NewUser, PasswordHash, Session, SessionId, - SessionRepositoryError, SessionRepositoryPort, User, UserId, UserRepositoryError, + ApiKeySubject, ApiKeyTier, ModelId, NewSession, NewUser, PasswordHash, ProviderId, Session, + SessionId, SessionRepositoryError, SessionRepositoryPort, User, UserId, UserRepositoryError, UserRepositoryPort, }; use rusqlite::{params, Connection, ErrorCode, OptionalExtension}; @@ -50,8 +50,9 @@ impl SqliteApiKeyRepository { conn.execute( "INSERT INTO api_keys ( id, label, key_hash, key_prefix, scopes_json, tier, is_active, - revoked_at, expires_at, created_at, last_used_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + revoked_at, expires_at, created_at, last_used_at, + allowed_models_json, allowed_providers_json + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)", params![ record.id.to_string(), record.label, @@ -64,6 +65,8 @@ impl SqliteApiKeyRepository { optional_datetime(record.expires_at), record.created_at.to_rfc3339(), optional_datetime(record.last_used_at), + "[]", + "[]", ], ) .map_err(db_error)?; @@ -117,7 +120,7 @@ impl TestApiKeyRecord { label: "Production".to_string(), key_hash: key_hash.to_string(), key_prefix: key_hash.chars().take(8).collect(), - scopes: vec![ApiKeyScope::parse("read").expect("scope")], + scopes: vec![ApiKeyScope::parse("chat:read").expect("scope")], tier: ApiKeyTier::Free, is_active: true, revoked_at: None, @@ -137,7 +140,7 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { let conn = self.lock()?; let now = Utc::now().to_rfc3339(); conn.query_row( - "SELECT id, label, scopes_json, tier + "SELECT id, label, scopes_json, tier, allowed_models_json, allowed_providers_json FROM api_keys WHERE key_hash = ?1 AND is_active = 1 @@ -173,7 +176,8 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { let mut stmt = conn .prepare( "SELECT id, label, key_hash, key_prefix, scopes_json, tier, is_active, - revoked_at, expires_at, created_at, last_used_at + revoked_at, expires_at, created_at, last_used_at, + allowed_models_json, allowed_providers_json FROM api_keys ORDER BY created_at DESC", ) @@ -190,7 +194,8 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { let conn = self.lock()?; conn.query_row( "SELECT id, label, key_hash, key_prefix, scopes_json, tier, is_active, - revoked_at, expires_at, created_at, last_used_at + revoked_at, expires_at, created_at, last_used_at, + allowed_models_json, allowed_providers_json FROM api_keys WHERE id = ?1", params![id.to_string()], @@ -205,8 +210,9 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { conn.execute( "INSERT INTO api_keys ( id, label, key_hash, key_prefix, scopes_json, tier, is_active, - revoked_at, expires_at, created_at, last_used_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + revoked_at, expires_at, created_at, last_used_at, + allowed_models_json, allowed_providers_json + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)", params![ record.id.to_string(), record.label, @@ -219,6 +225,8 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { optional_datetime(record.expires_at), record.created_at.to_rfc3339(), optional_datetime(record.last_used_at), + models_to_json(&record.allowed_models)?, + providers_to_json(&record.allowed_providers)?, ], ) .map_err(db_error)?; @@ -236,8 +244,10 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { is_active = ?4, revoked_at = ?5, expires_at = ?6, - last_used_at = ?7 - WHERE id = ?8", + last_used_at = ?7, + allowed_models_json = ?8, + allowed_providers_json = ?9 + WHERE id = ?10", params![ record.label, scopes_to_json(&record.scopes)?, @@ -246,6 +256,8 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { optional_datetime(record.revoked_at), optional_datetime(record.expires_at), optional_datetime(record.last_used_at), + models_to_json(&record.allowed_models)?, + providers_to_json(&record.allowed_providers)?, record.id.to_string(), ], ) @@ -309,7 +321,8 @@ impl ApiKeyRepositoryPort for SqliteApiKeyRepository { let mut stmt = conn .prepare( "SELECT id, label, key_hash, key_prefix, scopes_json, tier, is_active, - revoked_at, expires_at, created_at, last_used_at + revoked_at, expires_at, created_at, last_used_at, + allowed_models_json, allowed_providers_json FROM api_keys WHERE is_active = 1 ORDER BY created_at DESC @@ -340,11 +353,15 @@ fn row_to_subject(row: &rusqlite::Row<'_>) -> rusqlite::Result { let label: String = row.get("label")?; let scopes_json: String = row.get("scopes_json")?; let tier: String = row.get("tier")?; + let allowed_models_json: String = row.get("allowed_models_json")?; + let allowed_providers_json: String = row.get("allowed_providers_json")?; Ok(ApiKeySubject { id: ApiKeyId::new(id), label, scopes: scopes_from_json(&scopes_json).map_err(invalid_data)?, tier: ApiKeyTier::from_str(&tier).map_err(|error| invalid_data(error.to_string()))?, + allowed_models: models_from_json(&allowed_models_json).map_err(invalid_data)?, + allowed_providers: providers_from_json(&allowed_providers_json).map_err(invalid_data)?, }) } @@ -360,6 +377,8 @@ fn row_to_record(row: &rusqlite::Row<'_>) -> rusqlite::Result { let expires_at_str: Option = row.get("expires_at")?; let created_at_str: String = row.get("created_at")?; let last_used_at_str: Option = row.get("last_used_at")?; + let allowed_models_json: String = row.get("allowed_models_json")?; + let allowed_providers_json: String = row.get("allowed_providers_json")?; let revoked_at = revoked_at_str.map(|s| parse_datetime(&s)).transpose()?; let expires_at = expires_at_str.map(|s| parse_datetime(&s)).transpose()?; @@ -378,15 +397,17 @@ fn row_to_record(row: &rusqlite::Row<'_>) -> rusqlite::Result { expires_at, created_at, last_used_at, + allowed_models: models_from_json(&allowed_models_json).map_err(invalid_data)?, + allowed_providers: providers_from_json(&allowed_providers_json).map_err(invalid_data)?, }) } fn scopes_from_json(value: &str) -> Result, String> { let values = serde_json::from_str::>(value).map_err(|error| error.to_string())?; - values + Ok(values .iter() - .map(|scope| ApiKeyScope::parse(scope).map_err(|error| error.to_string())) - .collect() + .map(|scope| ApiKeyScope::parse_lenient(scope)) + .collect()) } fn scopes_to_json(scopes: &[ApiKeyScope]) -> Result { @@ -395,6 +416,26 @@ fn scopes_to_json(scopes: &[ApiKeyScope]) -> Result Result, String> { + let values = serde_json::from_str::>(value).map_err(|e| e.to_string())?; + Ok(values.into_iter().map(ModelId::new).collect()) +} + +fn models_to_json(models: &[ModelId]) -> Result { + let values: Vec<&str> = models.iter().map(|m| m.as_str()).collect(); + serde_json::to_string(&values).map_err(|e| ApiKeyRepositoryError::Database(e.to_string())) +} + +fn providers_from_json(value: &str) -> Result, String> { + let values = serde_json::from_str::>(value).map_err(|e| e.to_string())?; + Ok(values.into_iter().map(ProviderId::new).collect()) +} + +fn providers_to_json(providers: &[ProviderId]) -> Result { + let values: Vec<&str> = providers.iter().map(|p| p.as_str()).collect(); + serde_json::to_string(&values).map_err(|e| ApiKeyRepositoryError::Database(e.to_string())) +} + fn bool_to_i64(value: bool) -> i64 { i64::from(value) } @@ -714,7 +755,7 @@ mod tests { use chrono::{Duration, Utc}; use rook_core::{ - ApiKeyId, ApiKeyRepositoryPort, ApiKeyTier, NewSession as CoreNewSession, + ApiKeyId, ApiKeyRepositoryPort, ApiKeyScope, ApiKeyTier, NewSession as CoreNewSession, NewUser as CoreNewUser, PasswordHash as CorePasswordHash, SessionRepositoryPort, UserRepositoryPort, }; @@ -771,7 +812,7 @@ mod tests { let mut active = TestApiKeyRecord::active("active", "hash-active"); active .scopes - .push(rook_core::ApiKeyScope::parse("write").expect("scope")); + .push(rook_core::ApiKeyScope::parse("chat:write").expect("scope")); active.tier = ApiKeyTier::Pro; active.expires_at = Some(Utc::now() + Duration::days(1)); repo.insert_test_key(active).await.expect("insert active"); @@ -787,7 +828,7 @@ mod tests { .expect("active subject"); assert_eq!(active.id, ApiKeyId::new("active")); assert_eq!(active.tier, ApiKeyTier::Pro); - assert_eq!(active.scopes[0].as_str(), "read"); + assert_eq!(active.scopes[0].as_str(), "chat:read"); assert!(repo .find_active_by_hash("hash-revoked") .await @@ -832,13 +873,15 @@ mod tests { label: "Development Key".to_string(), key_hash: "hash-123".to_string(), key_prefix: "rk-123".to_string(), - scopes: vec![rook_core::ApiKeyScope::parse("read").unwrap()], + scopes: vec![rook_core::ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, is_active: true, revoked_at: None, expires_at: None, created_at: Utc::now(), last_used_at: None, + allowed_models: vec![], + allowed_providers: vec![], }; repo.create(&record).await.expect("create"); @@ -933,13 +976,15 @@ mod tests { label: format!("Key {}", i), key_hash: format!("hash-{}", i), key_prefix: format!("rk-{}", i), - scopes: vec![rook_core::ApiKeyScope::parse("read").unwrap()], + scopes: vec![rook_core::ApiKeyScope::parse("chat:read").unwrap()], tier: ApiKeyTier::Free, is_active: true, revoked_at: None, expires_at: None, created_at: Utc::now(), last_used_at: None, + allowed_models: vec![], + allowed_providers: vec![], }; repo.create(&record).await.expect("create"); } @@ -1189,4 +1234,157 @@ mod tests { cleanup_test_db(db_path.as_ref()); } + + #[test] + fn create_api_key_with_allowed_models_and_providers_persists_correctly() { + use rook_core::{ModelId, ProviderId}; + runtime().block_on(async { + let repo = SqliteApiKeyRepository::new(":memory:").expect("repo"); + + let record = rook_core::ApiKeyRecord { + id: ApiKeyId::new("key-restricted"), + label: "Restricted Key".to_string(), + key_hash: "hash-restricted".to_string(), + key_prefix: "rk-rest".to_string(), + scopes: vec![rook_core::ApiKeyScope::parse("chat:read").unwrap()], + tier: ApiKeyTier::Free, + is_active: true, + revoked_at: None, + expires_at: None, + created_at: Utc::now(), + last_used_at: None, + allowed_models: vec![ModelId::new("gpt-4"), ModelId::new("claude-3")], + allowed_providers: vec![ProviderId::new("openai")], + }; + repo.create(&record).await.expect("create"); + + let found = repo + .find(&ApiKeyId::new("key-restricted")) + .await + .expect("find") + .expect("some"); + + assert_eq!(found.allowed_models.len(), 2); + assert_eq!(found.allowed_models[0].as_str(), "gpt-4"); + assert_eq!(found.allowed_models[1].as_str(), "claude-3"); + assert_eq!(found.allowed_providers.len(), 1); + assert_eq!(found.allowed_providers[0].as_str(), "openai"); + }); + } + + #[test] + fn update_api_key_restrictions_persists_correctly() { + use rook_core::{ModelId, ProviderId}; + runtime().block_on(async { + let repo = SqliteApiKeyRepository::new(":memory:").expect("repo"); + + // Start with empty restrictions + let record = rook_core::ApiKeyRecord { + id: ApiKeyId::new("key-update-test"), + label: "Update Test".to_string(), + key_hash: "hash-update".to_string(), + key_prefix: "rk-upd".to_string(), + scopes: vec![rook_core::ApiKeyScope::parse("chat:read").unwrap()], + tier: ApiKeyTier::Free, + is_active: true, + revoked_at: None, + expires_at: None, + created_at: Utc::now(), + last_used_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + repo.create(&record).await.expect("create"); + + // Update with restrictions + let mut updated = record.clone(); + updated.allowed_models = vec![ModelId::new("gpt-4o")]; + updated.allowed_providers = vec![ProviderId::new("anthropic")]; + repo.update(&updated).await.expect("update"); + + let found = repo + .find(&ApiKeyId::new("key-update-test")) + .await + .expect("find") + .expect("some"); + assert_eq!(found.allowed_models.len(), 1); + assert_eq!(found.allowed_models[0].as_str(), "gpt-4o"); + assert_eq!(found.allowed_providers[0].as_str(), "anthropic"); + + // Clear restrictions + let mut cleared = found; + cleared.allowed_models = vec![]; + cleared.allowed_providers = vec![]; + repo.update(&cleared).await.expect("update clear"); + + let refound = repo + .find(&ApiKeyId::new("key-update-test")) + .await + .expect("find") + .expect("some"); + assert!(refound.allowed_models.is_empty(), "should be cleared"); + assert!(refound.allowed_providers.is_empty(), "should be cleared"); + }); + } + + #[test] + fn empty_restrictions_means_unrestricted() { + runtime().block_on(async { + let repo = SqliteApiKeyRepository::new(":memory:").expect("repo"); + + let record = rook_core::ApiKeyRecord { + id: ApiKeyId::new("key-empty-restrictions"), + label: "Unrestricted".to_string(), + key_hash: "hash-unrestricted".to_string(), + key_prefix: "rk-unr".to_string(), + scopes: vec![rook_core::ApiKeyScope::parse("chat:write").unwrap()], + tier: ApiKeyTier::Pro, + is_active: true, + revoked_at: None, + expires_at: None, + created_at: Utc::now(), + last_used_at: None, + allowed_models: vec![], + allowed_providers: vec![], + }; + repo.create(&record).await.expect("create"); + + let found = repo + .find(&ApiKeyId::new("key-empty-restrictions")) + .await + .expect("find") + .expect("some"); + + assert!(found.allowed_models.is_empty(), "empty = unrestricted"); + assert!(found.allowed_providers.is_empty(), "empty = unrestricted"); + }); + } + + // Regression: scopes_from_json used to call ApiKeyScope::parse (strict) which + // rejected pre-#83 legacy scopes ('read', 'write') and broke reads of any + // existing key. parse_lenient accepts unknown strings and logs a warning. + #[test] + fn read_key_with_legacy_scope_string_is_preserved() { + runtime().block_on(async { + let repo = SqliteApiKeyRepository::new(":memory:").expect("repo"); + // Use the test helper and override scopes to a legacy 'read' value. + let mut record = TestApiKeyRecord::active("legacy_key_id", "hash:legacy"); + record.label = "Legacy Key".to_string(); + record.scopes = vec![ApiKeyScope::parse_lenient("read")]; + repo.insert_test_key(record) + .await + .expect("insert legacy key"); + + // Re-fetch — this exercises the scopes_from_json read path. + let result = repo + .find_active_by_hash("hash:legacy") + .await + .expect("query") + .expect("legacy key must be readable"); + + assert_eq!(result.label, "Legacy Key"); + assert_eq!(result.scopes.len(), 1); + assert_eq!(result.scopes[0].as_str(), "read"); + }); + } } diff --git a/crates/infrastructure/db-migration/src/migrations/V1__allowed_models_providers.sql b/crates/infrastructure/db-migration/src/migrations/V1__allowed_models_providers.sql new file mode 100644 index 00000000..5406e6ea --- /dev/null +++ b/crates/infrastructure/db-migration/src/migrations/V1__allowed_models_providers.sql @@ -0,0 +1,6 @@ +-- V1__allowed_models_providers.sql +-- Adds allowed_models_json and allowed_providers_json restriction columns to api_keys. +-- Empty JSON array ('[]') means unrestricted (all models/providers allowed). + +ALTER TABLE api_keys ADD COLUMN allowed_models_json TEXT NOT NULL DEFAULT '[]'; +ALTER TABLE api_keys ADD COLUMN allowed_providers_json TEXT NOT NULL DEFAULT '[]'; diff --git a/crates/infrastructure/providers-anthropic/tests/provider.rs b/crates/infrastructure/providers-anthropic/tests/provider.rs index f14a0abb..ce58a82f 100644 --- a/crates/infrastructure/providers-anthropic/tests/provider.rs +++ b/crates/infrastructure/providers-anthropic/tests/provider.rs @@ -78,6 +78,7 @@ async fn complete_returns_valid_response_from_mock_server() { cacheable: true, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; let result = provider.complete(&req).await; diff --git a/crates/infrastructure/providers-gemini/tests/provider.rs b/crates/infrastructure/providers-gemini/tests/provider.rs index 82d0ced9..f59f6fdf 100644 --- a/crates/infrastructure/providers-gemini/tests/provider.rs +++ b/crates/infrastructure/providers-gemini/tests/provider.rs @@ -67,6 +67,7 @@ async fn complete_returns_error_not_implemented() { cacheable: true, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; let result = provider.complete(&req).await; diff --git a/crates/infrastructure/providers-groq/tests/provider.rs b/crates/infrastructure/providers-groq/tests/provider.rs index 7239a3d2..8239fc39 100644 --- a/crates/infrastructure/providers-groq/tests/provider.rs +++ b/crates/infrastructure/providers-groq/tests/provider.rs @@ -48,6 +48,7 @@ async fn complete_returns_error_when_not_implemented() { cacheable: true, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; let result = provider.complete(&req).await; diff --git a/crates/infrastructure/providers-ollama/tests/provider.rs b/crates/infrastructure/providers-ollama/tests/provider.rs index 79646138..6fbf2440 100644 --- a/crates/infrastructure/providers-ollama/tests/provider.rs +++ b/crates/infrastructure/providers-ollama/tests/provider.rs @@ -88,6 +88,7 @@ async fn complete_returns_error_when_not_implemented() { cacheable: true, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; // Complete is not yet implemented — returns provider error diff --git a/crates/infrastructure/providers-openai/tests/provider.rs b/crates/infrastructure/providers-openai/tests/provider.rs index e253bb8a..0abe29c2 100644 --- a/crates/infrastructure/providers-openai/tests/provider.rs +++ b/crates/infrastructure/providers-openai/tests/provider.rs @@ -96,6 +96,7 @@ async fn complete_returns_response_on_success() { cacheable: true, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; let result = provider.complete(&req).await; @@ -145,6 +146,7 @@ async fn stream_returns_chunks_on_openai_sse_success() { cacheable: true, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; let chunks = provider diff --git a/crates/infrastructure/transport-axum/Cargo.toml b/crates/infrastructure/transport-axum/Cargo.toml index 3c6bce3c..b1abeda2 100644 --- a/crates/infrastructure/transport-axum/Cargo.toml +++ b/crates/infrastructure/transport-axum/Cargo.toml @@ -39,6 +39,9 @@ tracing = "0.1" # SHA-256 for session token hashing sha2 = "0.11" +# Hex encoding for digest output +hex = "0.4" + [dev-dependencies] axum-test = "20" encryption-inmemory = { path = "../encryption-inmemory" } diff --git a/crates/infrastructure/transport-axum/src/anthropic_adapter.rs b/crates/infrastructure/transport-axum/src/anthropic_adapter.rs index 6b7b8cc5..a5f0a9ee 100644 --- a/crates/infrastructure/transport-axum/src/anthropic_adapter.rs +++ b/crates/infrastructure/transport-axum/src/anthropic_adapter.rs @@ -1,6 +1,9 @@ // Anthropic adapter — translates between Anthropic wire format and domain model -use rook_core::{CompletionRequest, Message, MessageContent, RequestMetadata, Role, StreamChunk}; +use rook_core::{ + ApiKeyRestrictions, CompletionRequest, Message, MessageContent, RequestMetadata, Role, + StreamChunk, +}; use serde::{Deserialize, Serialize}; use shared_kernel::{CortexError, ModelId, RequestId}; @@ -124,6 +127,7 @@ impl From for CompletionRequest { cacheable: false, priority: 5, }, + restrictions: ApiKeyRestrictions::default(), } } } diff --git a/crates/infrastructure/transport-axum/src/authz.rs b/crates/infrastructure/transport-axum/src/authz.rs index 17527b5e..56a086b0 100644 --- a/crates/infrastructure/transport-axum/src/authz.rs +++ b/crates/infrastructure/transport-axum/src/authz.rs @@ -33,6 +33,8 @@ const TRUSTED_HEADERS: &[&str] = &[ "x-authz-auth-id", "x-authz-auth-label", "x-authz-auth-scopes", + "x-authz-allowed-models", + "x-authz-allowed-providers", ]; #[derive(Clone)] @@ -68,7 +70,7 @@ impl AuthzConfig { format!("key_{}", idx + 1), format!("Client API Key {}", idx + 1), key.to_string(), - ["read", "write"], + ["chat:read", "chat:write"], ) }) .collect(); @@ -263,6 +265,8 @@ pub struct Subject { pub id: String, pub label: String, pub scopes: Vec, + pub allowed_models: Vec, + pub allowed_providers: Vec, } impl Subject { @@ -272,6 +276,8 @@ impl Subject { id: "public".to_string(), label: "Public".to_string(), scopes: Vec::new(), + allowed_models: Vec::new(), + allowed_providers: Vec::new(), } } } @@ -423,7 +429,14 @@ pub async fn middleware( } remove_trusted_headers(request.headers_mut()); - let outcome = evaluate_policy(route_class, request.headers(), &config).await; + let outcome = evaluate_policy( + route_class, + request.method(), + request.uri().path(), + request.headers(), + &config, + ) + .await; if !outcome.allow { let mut resp = rejection_response(request.uri().path(), route_class, outcome).into_response(); @@ -467,12 +480,14 @@ pub fn classify_route(method: &Method, path: &str) -> AuthTier { pub async fn evaluate_policy( route_class: AuthTier, + method: &Method, + path: &str, headers: &HeaderMap, config: &AuthzConfig, ) -> AuthOutcome { match route_class { AuthTier::Public => AuthOutcome::allow(Subject::anonymous()), - AuthTier::ClientApi => client_api_policy(headers, config).await, + AuthTier::ClientApi => client_api_policy(method, path, headers, config).await, AuthTier::Management => management_policy(headers, config).await, } } @@ -525,6 +540,16 @@ pub fn stamp_trusted_headers( insert_header(headers, "x-authz-auth-id", &subject.id); insert_header(headers, "x-authz-auth-label", &subject.label); insert_header(headers, "x-authz-auth-scopes", &subject.scopes.join(",")); + insert_header( + headers, + "x-authz-allowed-models", + &subject.allowed_models.join(","), + ); + insert_header( + headers, + "x-authz-allowed-providers", + &subject.allowed_providers.join(","), + ); } pub struct PreflightResponse { @@ -592,7 +617,36 @@ pub fn body_size_rejection( }) } -async fn client_api_policy(headers: &HeaderMap, config: &AuthzConfig) -> AuthOutcome { +fn required_scope(method: &Method, path: &str) -> Option<&'static str> { + if !path.starts_with("/v1/") { + return None; + } + if path.starts_with("/v1/providers/") || path.starts_with("/v1/providers") { + return if *method == Method::GET { + Some("providers:read") + } else { + Some("providers:write") + }; + } + if path.starts_with("/v1/chat/") { + return match *method { + Method::GET => Some("chat:read"), + _ => Some("chat:write"), + }; + } + // Other /v1/* routes: GET defaults to chat:read, writes require chat:write + match *method { + Method::GET => Some("chat:read"), + _ => Some("chat:write"), + } +} + +async fn client_api_policy( + method: &Method, + path: &str, + headers: &HeaderMap, + config: &AuthzConfig, +) -> AuthOutcome { let Some(api_key) = extract_api_key(headers) else { return AuthOutcome::reject(StatusCode::UNAUTHORIZED, "MISSING_API_KEY"); }; @@ -609,7 +663,20 @@ async fn client_api_policy(headers: &HeaderMap, config: &AuthzConfig) -> AuthOut .iter() .map(|scope| scope.as_str().to_string()) .collect(), + allowed_models: api_key_subject + .allowed_models + .iter() + .map(|m| m.as_str().to_string()) + .collect(), + allowed_providers: api_key_subject + .allowed_providers + .iter() + .map(|p| p.as_str().to_string()) + .collect(), }; + if let Some(rejection) = check_scope(method, path, &subject) { + return rejection; + } return match config .rate_limiter .check(&subject.id, RateLimitTier::from(api_key_subject.tier)) @@ -651,7 +718,12 @@ async fn client_api_policy(headers: &HeaderMap, config: &AuthzConfig) -> AuthOut id: credential.id.clone(), label: credential.label.clone(), scopes: credential.scopes.clone(), + allowed_models: Vec::new(), + allowed_providers: Vec::new(), }; + if let Some(rejection) = check_scope(method, path, &subject) { + return rejection; + } match config.rate_limiter.check(&credential.id, credential.tier) { Ok(snapshot) => AuthOutcome::allow(subject).with_rate_limit(snapshot), Err(snapshot) => AuthOutcome::reject(StatusCode::TOO_MANY_REQUESTS, "RATE_LIMIT_EXCEEDED") @@ -659,6 +731,17 @@ async fn client_api_policy(headers: &HeaderMap, config: &AuthzConfig) -> AuthOut } } +fn check_scope(method: &Method, path: &str, subject: &Subject) -> Option { + let required = required_scope(method, path)?; + if subject.scopes.iter().any(|s| s == "admin" || s == required) { + return None; + } + Some(AuthOutcome::reject( + StatusCode::FORBIDDEN, + "INSUFFICIENT_SCOPE", + )) +} + async fn management_policy(headers: &HeaderMap, config: &AuthzConfig) -> AuthOutcome { // Extract auth_token cookie let Some(cookie_value) = extract_cookie(headers, "auth_token") else { @@ -679,6 +762,8 @@ async fn management_policy(headers: &HeaderMap, config: &AuthzConfig) -> AuthOut id: validated.session.user_id.to_string(), label: validated.username, scopes: vec!["admin".to_string()], + allowed_models: Vec::new(), + allowed_providers: Vec::new(), }; AuthOutcome::allow(subject) } @@ -741,6 +826,8 @@ fn verify_jwt(token: &str, secret: &str) -> Result { id: id.clone(), label: id, scopes: vec!["admin".to_string()], + allowed_models: Vec::new(), + allowed_providers: Vec::new(), }) } @@ -811,6 +898,7 @@ fn rejection_message(code: &str) -> &'static str { "AUTH_MISCONFIGURED" => "Authentication is not configured", "AUTH_BACKEND_ERROR" => "Authentication backend error", "RATE_LIMIT_EXCEEDED" => "Rate limit exceeded", + "INSUFFICIENT_SCOPE" => "Insufficient scope for this operation", _ => "Unauthorized", } } @@ -995,8 +1083,14 @@ mod tests { .expect("runtime") } - fn evaluate(route_class: AuthTier, headers: &HeaderMap, config: &AuthzConfig) -> AuthOutcome { - runtime().block_on(evaluate_policy(route_class, headers, config)) + fn evaluate( + route_class: AuthTier, + method: &Method, + path: &str, + headers: &HeaderMap, + config: &AuthzConfig, + ) -> AuthOutcome { + runtime().block_on(evaluate_policy(route_class, method, path, headers, config)) } #[test] @@ -1023,26 +1117,38 @@ mod tests { "key_1", "Production Key", "sk-live", - ["read", "write"], + ["chat:read", "chat:write"], )], "test-secret", ); let mut headers = HeaderMap::new(); headers.insert("authorization", HeaderValue::from_static("Bearer sk-live")); - let outcome = evaluate(AuthTier::ClientApi, &headers, &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); assert!(outcome.allow); let subject = outcome.subject.expect("subject"); assert_eq!(subject.kind, AuthKind::ApiKey); assert_eq!(subject.id, "key_1"); - assert_eq!(subject.scopes, vec!["read", "write"]); + assert_eq!(subject.scopes, vec!["chat:read", "chat:write"]); } #[test] fn client_api_rejects_missing_api_key() { let config = AuthzConfig::new(Vec::new(), "test-secret"); - let outcome = evaluate(AuthTier::ClientApi, &HeaderMap::new(), &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &HeaderMap::new(), + &config, + ); assert!(!outcome.allow); assert_eq!(outcome.status, Some(StatusCode::UNAUTHORIZED)); @@ -1163,14 +1269,20 @@ mod tests { "key_1", "Production Key", "sk-live", - ["read"], + ["chat:read"], )], "test-secret", ); let mut headers = HeaderMap::new(); headers.insert("x-api-key", HeaderValue::from_static("sk-live")); - let outcome = evaluate(AuthTier::ClientApi, &headers, &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); let snapshot = outcome.rate_limit.expect("rate limit snapshot"); assert_eq!(snapshot.limit, 100); @@ -1184,7 +1296,7 @@ mod tests { "key_1", "Production Key", "sk-live", - ["read"], + ["chat:read"], )], "test-secret", ); @@ -1192,10 +1304,22 @@ mod tests { headers.insert("x-api-key", HeaderValue::from_static("sk-live")); for _ in 0..100 { - let outcome = evaluate(AuthTier::ClientApi, &headers, &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); assert!(outcome.allow); } - let outcome = evaluate(AuthTier::ClientApi, &headers, &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); assert!(!outcome.allow); assert_eq!(outcome.status, Some(StatusCode::TOO_MANY_REQUESTS)); @@ -1209,21 +1333,29 @@ mod tests { *repo.subject.lock().expect("subject") = Some(ApiKeySubject { id: ApiKeyId::new("persisted-key"), label: "Persisted Key".to_string(), - scopes: vec![ApiKeyScope::parse("read").expect("scope")], + scopes: vec![ApiKeyScope::parse("chat:read").expect("scope")], tier: ApiKeyTier::Enterprise, + allowed_models: vec![], + allowed_providers: vec![], }); let auth = AuthenticateClientApi::new(repo, "hash-secret"); let config = AuthzConfig::with_client_auth(auth, false, "test-secret"); let mut headers = HeaderMap::new(); headers.insert("authorization", HeaderValue::from_static("Bearer sk-live")); - let outcome = evaluate(AuthTier::ClientApi, &headers, &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); assert!(outcome.allow); let subject = outcome.subject.expect("subject"); assert_eq!(subject.id, "persisted-key"); assert_eq!(subject.label, "Persisted Key"); - assert_eq!(subject.scopes, vec!["read"]); + assert_eq!(subject.scopes, vec!["chat:read"]); assert_eq!(outcome.rate_limit.expect("rate limit").limit, 10_000); } @@ -1234,10 +1366,157 @@ mod tests { let mut headers = HeaderMap::new(); headers.insert("x-api-key", HeaderValue::from_static("sk-invalid")); - let outcome = evaluate(AuthTier::ClientApi, &headers, &config); + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); assert!(!outcome.allow); assert_eq!(outcome.status, Some(StatusCode::UNAUTHORIZED)); assert_eq!(outcome.code, Some("INVALID_API_KEY")); } + + #[test] + fn client_api_with_chat_read_scope_allowed_on_get_route() { + let config = AuthzConfig::new( + vec![ApiKeyCredential::new( + "key_read", + "Read-only Key", + "sk-read", + ["chat:read"], + )], + "test-secret", + ); + let mut headers = HeaderMap::new(); + headers.insert("x-api-key", HeaderValue::from_static("sk-read")); + + let outcome = evaluate( + AuthTier::ClientApi, + &Method::GET, + "/v1/models", + &headers, + &config, + ); + + assert!( + outcome.allow, + "chat:read key must be allowed on GET /v1/models" + ); + } + + #[test] + fn client_api_with_chat_read_scope_rejected_on_write_route() { + let config = AuthzConfig::new( + vec![ApiKeyCredential::new( + "key_read", + "Read-only Key", + "sk-read", + ["chat:read"], + )], + "test-secret", + ); + let mut headers = HeaderMap::new(); + headers.insert("x-api-key", HeaderValue::from_static("sk-read")); + + let outcome = evaluate( + AuthTier::ClientApi, + &Method::POST, + "/v1/chat/completions", + &headers, + &config, + ); + + assert!(!outcome.allow); + assert_eq!(outcome.status, Some(StatusCode::FORBIDDEN)); + assert_eq!(outcome.code, Some("INSUFFICIENT_SCOPE")); + } + + #[test] + fn client_api_with_admin_scope_allowed_on_any_route() { + let config = AuthzConfig::new( + vec![ApiKeyCredential::new( + "key_admin", + "Admin Key", + "sk-admin", + ["admin"], + )], + "test-secret", + ); + let mut headers = HeaderMap::new(); + headers.insert("x-api-key", HeaderValue::from_static("sk-admin")); + + let outcome = evaluate( + AuthTier::ClientApi, + &Method::POST, + "/v1/chat/completions", + &headers, + &config, + ); + + assert!(outcome.allow, "admin scope must bypass scope enforcement"); + } + + #[test] + fn client_api_with_chat_write_scope_allowed_on_write_route() { + let config = AuthzConfig::new( + vec![ApiKeyCredential::new( + "key_write", + "Write Key", + "sk-write", + ["chat:write"], + )], + "test-secret", + ); + let mut headers = HeaderMap::new(); + headers.insert("x-api-key", HeaderValue::from_static("sk-write")); + + let outcome = evaluate( + AuthTier::ClientApi, + &Method::POST, + "/v1/chat/completions", + &headers, + &config, + ); + + assert!( + outcome.allow, + "chat:write key must be allowed on POST /v1/chat/completions" + ); + } + + // Regression: required_scope fallback used to return chat:read for ALL /v1/* + // methods, letting a read-only key hit POST /v1/messages. Now non-GET must + // require chat:write (or admin). + #[test] + fn client_api_with_chat_read_scope_rejected_on_post_to_messages() { + let config = AuthzConfig::new( + vec![ApiKeyCredential::new( + "key_read", + "Read-only Key", + "sk-read", + ["chat:read"], + )], + "test-secret", + ); + let mut headers = HeaderMap::new(); + headers.insert("x-api-key", HeaderValue::from_static("sk-read")); + + let outcome = evaluate( + AuthTier::ClientApi, + &Method::POST, + "/v1/messages", + &headers, + &config, + ); + + assert!( + !outcome.allow, + "chat:read must NOT be allowed on POST /v1/messages" + ); + assert_eq!(outcome.status, Some(StatusCode::FORBIDDEN)); + assert_eq!(outcome.code, Some("INSUFFICIENT_SCOPE")); + } } diff --git a/crates/infrastructure/transport-axum/src/format_registry.rs b/crates/infrastructure/transport-axum/src/format_registry.rs index 975d6ff3..4827e054 100644 --- a/crates/infrastructure/transport-axum/src/format_registry.rs +++ b/crates/infrastructure/transport-axum/src/format_registry.rs @@ -239,6 +239,7 @@ mod tests { cacheable: false, priority: 0, }, + restrictions: rook_core::ApiKeyRestrictions::default(), }; let error = registry diff --git a/crates/infrastructure/transport-axum/src/handlers/api_key.rs b/crates/infrastructure/transport-axum/src/handlers/api_key.rs index fa92f7d9..86a54cad 100644 --- a/crates/infrastructure/transport-axum/src/handlers/api_key.rs +++ b/crates/infrastructure/transport-axum/src/handlers/api_key.rs @@ -22,7 +22,7 @@ where Deserialize::deserialize(d).map(Some) } -use rook_core::{ApiKeyId, ApiKeyScope, ApiKeyTier}; +use rook_core::{ApiKeyId, ApiKeyScope, ApiKeyTier, ModelId, ProviderId}; use rook_usecases::{CreateApiKeyRequest, UpdateApiKeyRequest}; use crate::api_key_dto::ListApiKeysResponseDto; @@ -37,6 +37,10 @@ pub struct CreateApiKeyRequestDto { pub scopes: Vec, pub tier: String, pub expires_at: Option>, + #[serde(default)] + pub allowed_models: Vec, + #[serde(default)] + pub allowed_providers: Vec, } #[derive(Debug, Deserialize)] @@ -48,6 +52,8 @@ pub struct UpdateApiKeyRequestDto { pub is_active: Option, #[serde(default, deserialize_with = "double_option")] pub expires_at: Option>>, + pub allowed_models: Option>, + pub allowed_providers: Option>, } #[derive(Debug, Serialize)] @@ -63,6 +69,8 @@ pub struct ApiKeyRecordResponseDto { pub expires_at: Option>, pub created_at: DateTime, pub last_used_at: Option>, + pub allowed_models: Vec, + pub allowed_providers: Vec, } impl From<&rook_core::ApiKeyRecord> for ApiKeyRecordResponseDto { @@ -82,6 +90,16 @@ impl From<&rook_core::ApiKeyRecord> for ApiKeyRecordResponseDto { expires_at: record.expires_at, created_at: record.created_at, last_used_at: record.last_used_at, + allowed_models: record + .allowed_models + .iter() + .map(|m| m.as_str().to_string()) + .collect(), + allowed_providers: record + .allowed_providers + .iter() + .map(|p| p.as_str().to_string()) + .collect(), } } } @@ -152,6 +170,12 @@ pub async fn create_api_key( scopes, tier, expires_at: req.expires_at, + allowed_models: req.allowed_models.into_iter().map(ModelId::new).collect(), + allowed_providers: req + .allowed_providers + .into_iter() + .map(ProviderId::new) + .collect(), }; let (record, plaintext_key) = mak.create(domain_req).await.map_err(map_error)?; @@ -221,6 +245,12 @@ pub async fn update_api_key( tier, is_active: req.is_active, expires_at: req.expires_at, + allowed_models: req + .allowed_models + .map(|v| v.into_iter().map(ModelId::new).collect()), + allowed_providers: req + .allowed_providers + .map(|v| v.into_iter().map(ProviderId::new).collect()), }; let record = mak.update(&key_id, domain_req).await.map_err(map_error)?; diff --git a/crates/infrastructure/transport-axum/src/handlers/auth.rs b/crates/infrastructure/transport-axum/src/handlers/auth.rs index ca89508b..8b731242 100644 --- a/crates/infrastructure/transport-axum/src/handlers/auth.rs +++ b/crates/infrastructure/transport-axum/src/handlers/auth.rs @@ -115,7 +115,7 @@ pub async fn logout_handler( // SHA-256 hash to find session let mut hasher = Sha256::new(); hasher.update(&token_bytes); - let token_hash = format!("{:x}", hasher.finalize()); + let token_hash = hex::encode(hasher.finalize()); // Revoke the session. // "Session not found" and "already revoked" are both fine — the client's intent diff --git a/crates/infrastructure/transport-axum/src/openai_adapter.rs b/crates/infrastructure/transport-axum/src/openai_adapter.rs index acdb2d3d..351157b4 100644 --- a/crates/infrastructure/transport-axum/src/openai_adapter.rs +++ b/crates/infrastructure/transport-axum/src/openai_adapter.rs @@ -1,7 +1,8 @@ // OpenAI adapter — translates between OpenAI wire format and domain model use rook_core::{ - CompletionRequest, FinishReason, Message, MessageContent, RequestMetadata, Role, StreamChunk, + ApiKeyRestrictions, CompletionRequest, FinishReason, Message, MessageContent, RequestMetadata, + Role, StreamChunk, }; use serde::{Deserialize, Serialize}; use shared_kernel::{ModelId, RequestId}; @@ -128,6 +129,7 @@ impl From for CompletionRequest { cacheable: true, priority: 5, }, + restrictions: ApiKeyRestrictions::default(), } } } diff --git a/crates/infrastructure/transport-axum/src/routes.rs b/crates/infrastructure/transport-axum/src/routes.rs index 762e9418..1bfb083d 100644 --- a/crates/infrastructure/transport-axum/src/routes.rs +++ b/crates/infrastructure/transport-axum/src/routes.rs @@ -5,15 +5,15 @@ use std::{convert::Infallible, sync::Arc}; use axum::response::sse::Event; use axum::{ extract::{Json, State}, - http::{header, StatusCode}, + http::{header, HeaderMap, StatusCode}, middleware, response::{AppendHeaders, IntoResponse, Response, Sse}, routing::{delete, get, post, put}, Router, }; use futures::StreamExt; -use rook_core::{ApiFormat, CompletionRequest, HealthPort, HealthStatus}; -use shared_kernel::CortexError; +use rook_core::{ApiFormat, ApiKeyRestrictions, CompletionRequest, HealthPort, HealthStatus}; +use shared_kernel::{CortexError, ModelId, ProviderId}; use tower_http::limit::RequestBodyLimitLayer; use tracing::error; @@ -133,12 +133,70 @@ fn extract_client_ip(request: &axum::extract::Request) -> std::net::IpAddr { .unwrap_or_else(|| std::net::IpAddr::from([127, 0, 0, 1])) } +/// Extract `ApiKeyRestrictions` from the trusted authz headers stamped by the middleware. +/// +/// Headers are comma-separated lists — empty string means no restriction (unrestricted). +/// Fails closed if the headers are missing or non-UTF-8: the authz middleware must +/// always stamp these headers, so missing/invalid values indicate a routing bug or +/// an attempt to bypass the middleware. +fn restrictions_from_headers(headers: &HeaderMap) -> Result { + let allowed_models = parse_csv_header(headers, "x-authz-allowed-models")? + .into_iter() + .map(ModelId::new) + .collect(); + let allowed_providers = parse_csv_header(headers, "x-authz-allowed-providers")? + .into_iter() + .map(ProviderId::new) + .collect(); + Ok(ApiKeyRestrictions { + allowed_models, + allowed_providers, + }) +} + +/// Parse a comma-separated header value into a Vec. +/// +/// Returns `Err(HttpError)` if the header is missing or not valid UTF-8. +/// An empty value (present but empty) yields an empty Vec (unrestricted). +fn parse_csv_header(headers: &HeaderMap, name: &'static str) -> Result, HttpError> { + let value = headers + .get(name) + .ok_or_else(|| { + error!( + header = name, + "trusted authz header missing — middleware bypass?" + ); + HttpError { + status: StatusCode::INTERNAL_SERVER_ERROR, + code: "AUTHZ_HEADER_MISSING", + message: format!("authz header {name} missing"), + } + })? + .to_str() + .map_err(|error| { + error!(header = name, %error, "trusted authz header is not valid UTF-8"); + HttpError { + status: StatusCode::INTERNAL_SERVER_ERROR, + code: "AUTHZ_HEADER_INVALID", + message: format!("authz header {name} is not valid UTF-8: {error}"), + } + })?; + Ok(value + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .collect()) +} + /// POST /v1/chat/completions — OpenAI-compatible async fn chat_completions( State(usecases): State, + headers: HeaderMap, Json(body): Json, ) -> Result { - let req = CompletionRequest::from(body); + let mut req = CompletionRequest::from(body); + req.restrictions = restrictions_from_headers(&headers)?; if req.stream { return chat_completions_stream(usecases, req).await; @@ -153,6 +211,17 @@ async fn chat_completions( let openai_resp = OpenAIChatResponse::from(&resp); Ok(Json(openai_resp).into_response()) } + Err(e) if e.is_forbidden() => { + let body = OpenAIErrorResponse { + error: OpenAIErrorBody { + error_type: "invalid_request_error".to_string(), + code: Some("model_not_allowed".to_string()), + message: e.to_string(), + param: None, + }, + }; + Ok((StatusCode::FORBIDDEN, Json(body)).into_response()) + } Err(e) if e.is_all_providers_exhausted() => { let body = OpenAIErrorResponse { error: OpenAIErrorBody { @@ -206,6 +275,8 @@ async fn chat_completions_stream( .await { Ok(stream) => stream, + Err(error) if error.is_forbidden() => return Err(map_forbidden_openai(&error)), + Err(error) if error.is_rate_limited() => return Err(map_rate_limited(&error)), Err(error) => { let error_event = openai_error_event(error); let body = futures::stream::once(async move { Ok::(error_event) }); @@ -234,6 +305,29 @@ async fn chat_completions_stream( Ok(response) } +/// Map a `forbidden` `CortexError` to the OpenAI-shaped 403 response used by +/// `chat_completions` so streaming and non-streaming share identical behavior. +fn map_forbidden_openai(error: &CortexError) -> HttpError { + let code = error.forbidden_code().unwrap_or("model_not_allowed"); + HttpError { + status: StatusCode::FORBIDDEN, + code: match code { + "provider_not_allowed" => "PROVIDER_NOT_ALLOWED", + _ => "MODEL_NOT_ALLOWED", + }, + message: error.to_string(), + } +} + +/// Map a `rate_limited` `CortexError` to a 429 with the standard retry-after headers. +fn map_rate_limited(error: &CortexError) -> HttpError { + HttpError { + status: StatusCode::TOO_MANY_REQUESTS, + code: "RATE_LIMITED", + message: error.to_string(), + } +} + fn openai_error_event(error: shared_kernel::CortexError) -> Event { let body = OpenAIErrorResponse { error: OpenAIErrorBody { @@ -288,9 +382,11 @@ async fn list_models(State(_usecases): State) -> impl IntoResponse { /// POST /v1/messages — Anthropic-compatible async fn anthropic_messages( State(usecases): State, + headers: HeaderMap, Json(body): Json, ) -> Result { - let req = CompletionRequest::from(body); + let mut req = CompletionRequest::from(body); + req.restrictions = restrictions_from_headers(&headers)?; if req.stream { return anthropic_messages_stream(usecases, req).await; @@ -305,6 +401,7 @@ async fn anthropic_messages( let anthropic_resp = AnthropicMessagesResponse::from(&resp); Ok(Json(anthropic_resp).into_response()) } + Err(e) if e.is_forbidden() => Ok((StatusCode::FORBIDDEN, e.to_string()).into_response()), Err(e) if e.is_all_providers_exhausted() => { Ok((StatusCode::SERVICE_UNAVAILABLE, "All providers unavailable").into_response()) } @@ -325,6 +422,14 @@ async fn anthropic_messages_stream( .await { Ok(stream) => stream, + Err(error) if error.is_forbidden() => { + return Err(HttpError { + status: StatusCode::FORBIDDEN, + code: "MODEL_NOT_ALLOWED", + message: error.to_string(), + }); + } + Err(error) if error.is_rate_limited() => return Err(map_rate_limited(&error)), Err(error) => { let error_event = serde_json::to_string(&AnthropicSseEvent::from(error)) .map(|data| Event::default().data(data)) diff --git a/crates/infrastructure/transport-axum/tests/api_key_routes.rs b/crates/infrastructure/transport-axum/tests/api_key_routes.rs index 6db2c2b1..1c85e54f 100644 --- a/crates/infrastructure/transport-axum/tests/api_key_routes.rs +++ b/crates/infrastructure/transport-axum/tests/api_key_routes.rs @@ -14,8 +14,8 @@ fn test_record() -> ApiKeyRecord { key_hash: "hash_abc123".to_string(), key_prefix: "rook_test".to_string(), scopes: vec![ - ApiKeyScope::parse("read").unwrap(), - ApiKeyScope::parse("write").unwrap(), + ApiKeyScope::parse("chat:read").unwrap(), + ApiKeyScope::parse("chat:write").unwrap(), ], tier: ApiKeyTier::Pro, is_active: true, @@ -23,6 +23,8 @@ fn test_record() -> ApiKeyRecord { expires_at: Some(Utc::now() + chrono::Duration::days(30)), created_at: Utc::now(), last_used_at: None, + allowed_models: vec![], + allowed_providers: vec![], } } @@ -34,7 +36,10 @@ fn api_key_record_response_dto_converts_correctly() { assert_eq!(dto.id, "key_123"); assert_eq!(dto.label, "Production Key"); assert_eq!(dto.key_prefix, "rook_test"); - assert_eq!(dto.scopes, vec!["read".to_string(), "write".to_string()]); + assert_eq!( + dto.scopes, + vec!["chat:read".to_string(), "chat:write".to_string()] + ); assert_eq!(dto.tier, "pro"); assert!(dto.is_active); assert!(dto.expires_at.is_some()); diff --git a/crates/infrastructure/transport-axum/tests/format_translation_integration.rs b/crates/infrastructure/transport-axum/tests/format_translation_integration.rs index d53090ac..222c442a 100644 --- a/crates/infrastructure/transport-axum/tests/format_translation_integration.rs +++ b/crates/infrastructure/transport-axum/tests/format_translation_integration.rs @@ -387,6 +387,7 @@ fn registry_domain_request() -> CompletionRequest { cacheable: false, priority: 1, }, + restrictions: rook_core::ApiKeyRestrictions::default(), } }