mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 13:10:21 +08:00
fix(auth): resolve group policy before key intersection
This commit is contained in:
@@ -12,10 +12,7 @@ use super::{
|
||||
StoredWalletSnapshot, UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord,
|
||||
};
|
||||
use crate::LocalMutationOutcome;
|
||||
use aether_data::repository::auth::{
|
||||
read_resolved_auth_api_key_snapshot_by_key_hash,
|
||||
read_resolved_auth_api_key_snapshot_by_user_api_key_ids,
|
||||
};
|
||||
use aether_data::repository::auth::ResolvedAuthApiKeySnapshotReader;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct GatewayUserEffectiveListPolicies {
|
||||
@@ -1751,15 +1748,14 @@ impl GatewayDataState {
|
||||
let snapshot = crate::request_diagnostics::observe_db_operation(
|
||||
"auth_api_key_snapshot",
|
||||
self.database_pool_summary(),
|
||||
read_resolved_auth_api_key_snapshot_by_user_api_key_ids(
|
||||
self,
|
||||
self.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
now_unix_secs,
|
||||
),
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
self.apply_user_group_effective_policies(snapshot).await
|
||||
self.apply_user_group_effective_policies(snapshot, now_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_auth_api_key_snapshot_by_key_hash(
|
||||
@@ -1770,25 +1766,33 @@ impl GatewayDataState {
|
||||
let snapshot = crate::request_diagnostics::observe_db_operation(
|
||||
"auth_api_key_snapshot_by_hash",
|
||||
self.database_pool_summary(),
|
||||
read_resolved_auth_api_key_snapshot_by_key_hash(self, key_hash, now_unix_secs),
|
||||
self.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::KeyHash(key_hash)),
|
||||
)
|
||||
.await?;
|
||||
self.apply_user_group_effective_policies(snapshot).await
|
||||
self.apply_user_group_effective_policies(snapshot, now_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn apply_user_group_effective_policies(
|
||||
&self,
|
||||
snapshot: Option<GatewayAuthApiKeySnapshot>,
|
||||
snapshot: Option<StoredAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<GatewayAuthApiKeySnapshot>, DataLayerError> {
|
||||
let Some(mut snapshot) = snapshot else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone {
|
||||
apply_admin_unrestricted_auth_snapshot(&mut snapshot);
|
||||
return Ok(Some(snapshot));
|
||||
return Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
|
||||
snapshot,
|
||||
now_unix_secs,
|
||||
)));
|
||||
}
|
||||
let Some(repository) = self.user_reader.as_ref() else {
|
||||
return Ok(Some(snapshot));
|
||||
return Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
|
||||
snapshot,
|
||||
now_unix_secs,
|
||||
)));
|
||||
};
|
||||
let Some(user) = crate::request_diagnostics::observe_db_operation(
|
||||
"auth_user_policy",
|
||||
@@ -1797,57 +1801,50 @@ impl GatewayDataState {
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(snapshot));
|
||||
return Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
|
||||
snapshot,
|
||||
now_unix_secs,
|
||||
)));
|
||||
};
|
||||
if user.role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone {
|
||||
snapshot.user_role = user.role;
|
||||
apply_admin_unrestricted_auth_snapshot(&mut snapshot);
|
||||
return Ok(Some(snapshot));
|
||||
return Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
|
||||
snapshot,
|
||||
now_unix_secs,
|
||||
)));
|
||||
}
|
||||
let groups = self
|
||||
.effective_user_groups_for_user(&snapshot.user_id)
|
||||
.await?;
|
||||
|
||||
let mut allowed_providers =
|
||||
let allowed_providers =
|
||||
resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
|
||||
(
|
||||
&group.allowed_providers_mode,
|
||||
group.allowed_providers.clone(),
|
||||
)
|
||||
});
|
||||
let mut allowed_api_formats =
|
||||
let allowed_api_formats =
|
||||
resolve_effective_api_format_policy(None, "unrestricted", &groups, |group| {
|
||||
(
|
||||
&group.allowed_api_formats_mode,
|
||||
group.allowed_api_formats.clone(),
|
||||
)
|
||||
});
|
||||
let mut allowed_models =
|
||||
let allowed_models =
|
||||
resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
|
||||
(&group.allowed_models_mode, group.allowed_models.clone())
|
||||
});
|
||||
let user_rate_limit = resolve_effective_rate_limit_policy(None, "system", &groups);
|
||||
if !snapshot.api_key_is_standalone {
|
||||
constrain_api_key_list_policy_to_user_policy(
|
||||
&mut allowed_providers,
|
||||
&mut snapshot.api_key_allowed_providers,
|
||||
);
|
||||
constrain_api_key_api_format_policy_to_user_policy(
|
||||
&mut allowed_api_formats,
|
||||
&mut snapshot.api_key_allowed_api_formats,
|
||||
);
|
||||
constrain_api_key_list_policy_to_user_policy(
|
||||
&mut allowed_models,
|
||||
&mut snapshot.api_key_allowed_models,
|
||||
);
|
||||
}
|
||||
snapshot.apply_user_policy(
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
user_rate_limit,
|
||||
);
|
||||
Ok(Some(snapshot))
|
||||
snapshot.user_allowed_providers = allowed_providers;
|
||||
snapshot.user_allowed_api_formats = allowed_api_formats;
|
||||
snapshot.user_allowed_models = allowed_models;
|
||||
snapshot.user_rate_limit = user_rate_limit;
|
||||
Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
|
||||
snapshot,
|
||||
now_unix_secs,
|
||||
)))
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_user_effective_list_policies(
|
||||
@@ -1970,7 +1967,7 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_admin_unrestricted_auth_snapshot(snapshot: &mut GatewayAuthApiKeySnapshot) {
|
||||
fn apply_admin_unrestricted_auth_snapshot(snapshot: &mut StoredAuthApiKeySnapshot) {
|
||||
snapshot.user_allowed_providers = None;
|
||||
snapshot.user_allowed_api_formats = None;
|
||||
snapshot.user_allowed_models = None;
|
||||
@@ -2159,38 +2156,6 @@ fn rate_limit_policy_value(policy: Option<RateLimitRestriction>) -> Option<i32>
|
||||
}
|
||||
}
|
||||
|
||||
fn constrain_api_key_list_policy_to_user_policy(
|
||||
user_policy: &mut Option<Vec<String>>,
|
||||
api_key_policy: &mut Option<Vec<String>>,
|
||||
) {
|
||||
let Some(api_key_values) = api_key_policy.as_ref().filter(|values| !values.is_empty()) else {
|
||||
return;
|
||||
};
|
||||
let Some(user_values) = user_policy.clone() else {
|
||||
return;
|
||||
};
|
||||
let effective = intersect_list_policies(Some(api_key_values.to_vec()), Some(user_values))
|
||||
.unwrap_or_default();
|
||||
*user_policy = Some(effective.clone());
|
||||
*api_key_policy = Some(effective);
|
||||
}
|
||||
|
||||
fn constrain_api_key_api_format_policy_to_user_policy(
|
||||
user_policy: &mut Option<Vec<String>>,
|
||||
api_key_policy: &mut Option<Vec<String>>,
|
||||
) {
|
||||
let Some(api_key_values) = api_key_policy.as_ref().filter(|values| !values.is_empty()) else {
|
||||
return;
|
||||
};
|
||||
let Some(user_values) = user_policy.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let effective =
|
||||
aether_ai_formats::intersect_api_format_allowed_lists(api_key_values, user_values);
|
||||
*user_policy = Some(effective.clone());
|
||||
*api_key_policy = Some(effective);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
@@ -2529,17 +2494,6 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_specific_policy_cannot_expand_user_policy() {
|
||||
let mut user_policy = Some(vec!["gpt-5".to_string()]);
|
||||
let mut api_key_policy = Some(vec!["gpt-4.1".to_string()]);
|
||||
|
||||
constrain_api_key_list_policy_to_user_policy(&mut user_policy, &mut api_key_policy);
|
||||
|
||||
assert_eq!(user_policy, Some(Vec::<String>::new()));
|
||||
assert_eq!(api_key_policy, Some(Vec::<String>::new()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_non_standalone_snapshot_bypasses_group_and_key_policies() {
|
||||
let mut snapshot = sample_snapshot_with_role("key-admin", "admin-1", "admin")
|
||||
@@ -2595,6 +2549,38 @@ mod tests {
|
||||
assert_eq!(resolved.api_key_concurrent_limit, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn current_admin_role_bypasses_stored_user_and_key_policies() {
|
||||
let mut snapshot = sample_snapshot("key-admin", "admin-1");
|
||||
snapshot.api_key_allowed_providers = Some(vec!["anthropic".to_string()]);
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["anthropic:messages".to_string()]);
|
||||
snapshot.api_key_allowed_models = Some(vec!["claude-sonnet-4-5".to_string()]);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-admin".to_string()),
|
||||
snapshot,
|
||||
)]));
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
|
||||
sample_auth_user("admin-1", "admin"),
|
||||
]));
|
||||
let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository)
|
||||
.with_user_reader(user_repository);
|
||||
|
||||
let resolved = state
|
||||
.read_auth_api_key_snapshot_by_key_hash("hash-admin", 100)
|
||||
.await
|
||||
.expect("snapshot should resolve")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(resolved.user_role, "admin");
|
||||
assert_eq!(resolved.effective_allowed_providers(), None);
|
||||
assert_eq!(resolved.effective_allowed_api_formats(), None);
|
||||
assert_eq!(resolved.effective_allowed_models(), None);
|
||||
assert_eq!(resolved.user_rate_limit, None);
|
||||
assert_eq!(resolved.api_key_rate_limit, None);
|
||||
assert_eq!(resolved.api_key_concurrent_limit, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn user_personal_policy_fields_are_ignored_when_groups_are_applied() {
|
||||
let mut snapshot = sample_snapshot("key-user", "user-1").with_user_rate_limit(Some(200));
|
||||
@@ -2701,6 +2687,46 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn snapshot_without_user_reader_uses_stored_policy_intersection() {
|
||||
let mut snapshot = sample_snapshot("key-search", "user-search");
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["openai:search".to_string()]);
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-search".to_string()),
|
||||
snapshot,
|
||||
)]));
|
||||
let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository);
|
||||
|
||||
let resolved = state
|
||||
.read_auth_api_key_snapshot_by_key_hash("hash-search", 100)
|
||||
.await
|
||||
.expect("snapshot should resolve")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(resolved.effective_allowed_api_formats(), Some(&[][..]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_current_user_uses_stored_policy_intersection() {
|
||||
let mut snapshot = sample_snapshot("key-search", "missing-user");
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["openai:search".to_string()]);
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-search".to_string()),
|
||||
snapshot,
|
||||
)]));
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::default());
|
||||
let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository)
|
||||
.with_user_reader(user_repository);
|
||||
|
||||
let resolved = state
|
||||
.read_auth_api_key_snapshot_by_key_hash("hash-search", 100)
|
||||
.await
|
||||
.expect("snapshot should resolve")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(resolved.effective_allowed_api_formats(), Some(&[][..]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_lists_auth_api_key_export_records() {
|
||||
let repository = Arc::new(
|
||||
|
||||
Reference in New Issue
Block a user