fix(auth): resolve group policy before key intersection

This commit is contained in:
MMEXA
2026-07-12 03:35:42 +08:00
parent 2316df5c9a
commit 8d4d42a887
+110 -84
View File
@@ -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(