fix: exempt admins from default user group limits

This commit is contained in:
fawney19
2026-05-11 18:20:01 +08:00
parent 9bed5e9f83
commit fed676f54f
19 changed files with 264 additions and 20 deletions

View File

@@ -290,6 +290,7 @@ mod tests {
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: Some(allowed_models),
});

View File

@@ -43,6 +43,8 @@ pub(crate) struct GatewayControlAuthContext {
#[serde(skip)]
pub(crate) api_key_is_standalone: bool,
#[serde(skip)]
pub(crate) admin_bypass_limits: bool,
#[serde(skip)]
pub(crate) local_rejection: Option<GatewayLocalAuthRejection>,
#[serde(skip)]
pub(crate) allowed_models: Option<Vec<String>>,
@@ -512,6 +514,7 @@ pub(super) async fn resolve_data_backed_auth_context(
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey),
allowed_models: None,
}));
@@ -568,6 +571,7 @@ async fn resolve_trusted_auth_context(
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey),
allowed_models: None,
}));
@@ -656,6 +660,8 @@ async fn build_data_backed_auth_context(
user_rate_limit: snapshot.user_rate_limit,
api_key_rate_limit: snapshot.api_key_rate_limit,
api_key_is_standalone: snapshot.api_key_is_standalone,
admin_bypass_limits: snapshot.user_role.eq_ignore_ascii_case("admin")
&& !snapshot.api_key_is_standalone,
local_rejection,
allowed_models,
}

View File

@@ -1652,12 +1652,21 @@ impl GatewayDataState {
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));
}
let Some(repository) = self.user_reader.as_ref() else {
return Ok(Some(snapshot));
};
let Some(user) = repository.find_user_auth_by_id(&snapshot.user_id).await? else {
return Ok(Some(snapshot));
};
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));
}
let export_row = repository.find_export_user_by_id(&snapshot.user_id).await?;
let mut groups = repository
.list_user_groups_for_user(&snapshot.user_id)
@@ -1739,6 +1748,18 @@ impl GatewayDataState {
}
}
fn apply_admin_unrestricted_auth_snapshot(snapshot: &mut GatewayAuthApiKeySnapshot) {
snapshot.user_allowed_providers = None;
snapshot.user_allowed_api_formats = None;
snapshot.user_allowed_models = None;
snapshot.user_rate_limit = None;
snapshot.api_key_allowed_providers = None;
snapshot.api_key_allowed_api_formats = None;
snapshot.api_key_allowed_models = None;
snapshot.api_key_rate_limit = None;
snapshot.api_key_concurrent_limit = None;
}
fn resolve_effective_list_policy(
user_values: Option<Vec<String>>,
user_mode: &str,
@@ -1875,16 +1896,27 @@ mod tests {
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord,
StoredAuthApiKeySnapshot,
};
use aether_data::repository::users::StoredUserGroup;
use aether_data::repository::users::{
InMemoryUserReadRepository, StoredUserAuthRecord, StoredUserGroup, UpsertUserGroupRecord,
UserReadRepository,
};
use crate::data::GatewayDataState;
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
sample_snapshot_with_role(api_key_id, user_id, "user")
}
fn sample_snapshot_with_role(
api_key_id: &str,
user_id: &str,
role: &str,
) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
role.to_string(),
"local".to_string(),
true,
false,
@@ -1906,6 +1938,26 @@ mod tests {
.expect("snapshot should build")
}
fn sample_auth_user(user_id: &str, role: &str) -> StoredUserAuthRecord {
StoredUserAuthRecord::new(
user_id.to_string(),
Some("alice@example.com".to_string()),
true,
"alice".to_string(),
Some("hash".to_string()),
role.to_string(),
"local".to_string(),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
true,
false,
None,
None,
)
.expect("auth user should build")
}
fn sample_group(
id: &str,
priority: i32,
@@ -2053,6 +2105,61 @@ mod tests {
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")
.with_user_rate_limit(Some(120));
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()]);
snapshot.api_key_rate_limit = Some(5);
snapshot.api_key_concurrent_limit = Some(1);
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 group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "Restricted".to_string(),
description: None,
priority: 10,
allowed_providers: Some(vec!["openai".to_string()]),
allowed_providers_mode: "specific".to_string(),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_api_formats_mode: "specific".to_string(),
allowed_models: Some(vec!["gpt-4.1".to_string()]),
allowed_models_mode: "specific".to_string(),
rate_limit: Some(1),
rate_limit_mode: "custom".to_string(),
})
.await
.expect("group should create")
.expect("group should exist");
user_repository
.add_user_to_group(&group.id, "admin-1")
.await
.expect("group membership should create");
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.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 data_state_lists_auth_api_key_export_records() {
let repository = Arc::new(

View File

@@ -152,6 +152,16 @@ impl<'a> AdminAppState<'a> {
self.app.include_default_user_group_ids(group_ids).await
}
pub(crate) async fn include_default_user_group_ids_for_role(
&self,
group_ids: &[String],
role: &str,
) -> Result<Vec<String>, GatewayError> {
self.app
.include_default_user_group_ids_for_role(group_ids, role)
.await
}
pub(crate) async fn effective_default_user_group_id(
&self,
) -> Result<Option<String>, GatewayError> {

View File

@@ -192,7 +192,7 @@ pub(in super::super) async fn build_admin_create_user_response(
};
let requested_group_ids = normalize_admin_user_group_ids(payload.group_ids);
let group_ids = state
.include_default_user_group_ids(&requested_group_ids)
.include_default_user_group_ids_for_role(&requested_group_ids, &role)
.await?;
let groups = if group_ids.is_empty() {
Vec::new()

View File

@@ -29,7 +29,7 @@ pub(in super::super) async fn build_admin_update_user_response(
let Some(user_id) = admin_user_id_from_detail_path(request_context.path()) else {
return Ok(build_admin_users_bad_request_response("缺少 user_id"));
};
let Some(_existing_user) = state.find_user_auth_by_id(&user_id).await? else {
let Some(existing_user) = state.find_user_auth_by_id(&user_id).await? else {
return Ok((
http::StatusCode::NOT_FOUND,
Json(json!({ "detail": "用户不存在" })),
@@ -130,6 +130,7 @@ pub(in super::super) async fn build_admin_update_user_response(
},
None => None,
};
let effective_role = role.as_deref().unwrap_or(existing_user.role.as_str());
if payload.rate_limit.is_some_and(|value| value < 0) {
return Ok((
http::StatusCode::BAD_REQUEST,
@@ -251,7 +252,19 @@ pub(in super::super) async fn build_admin_update_user_response(
let requested_group_ids = normalize_admin_user_group_ids(payload.group_ids);
Some(
state
.include_default_user_group_ids(&requested_group_ids)
.include_default_user_group_ids_for_role(&requested_group_ids, effective_role)
.await?,
)
} else if role.is_some() {
let requested_group_ids = state
.list_user_groups_for_user(&user_id)
.await?
.into_iter()
.map(|group| group.id)
.collect::<Vec<_>>();
Some(
state
.include_default_user_group_ids_for_role(&requested_group_ids, effective_role)
.await?,
)
} else {

View File

@@ -364,6 +364,9 @@ impl RpmPlan {
if auth.local_rejection.is_some() || auth.user_id.is_empty() || auth.api_key_id.is_empty() {
return None;
}
if auth.admin_bypass_limits {
return None;
}
let now_ts = current_unix_secs();
let bucket = config.current_bucket(now_ts);
@@ -488,6 +491,7 @@ mod tests {
user_rate_limit: Some(1),
api_key_rate_limit: Some(10),
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: None,
});
@@ -526,6 +530,7 @@ mod tests {
user_rate_limit: None,
api_key_rate_limit: Some(10),
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: None,
});
@@ -550,6 +555,39 @@ mod tests {
}
}
#[tokio::test]
async fn limiter_skips_admin_bypass_context() {
let limiter = FrontdoorUserRpmLimiter::new(FrontdoorUserRpmConfig::new(60, 120, false))
.with_system_default_limit_for_tests(1);
let decision = sample_decision(GatewayControlAuthContext {
user_id: "admin-1".to_string(),
api_key_id: "key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: Some(10.0),
access_allowed: true,
user_rate_limit: Some(1),
api_key_rate_limit: Some(1),
api_key_is_standalone: false,
admin_bypass_limits: true,
local_rejection: None,
allowed_models: None,
});
let state = AppState::new().expect("state should build for tests");
let first = limiter
.check_and_consume(&state, Some(&decision))
.await
.expect("check should succeed");
let second = limiter
.check_and_consume(&state, Some(&decision))
.await
.expect("check should succeed");
assert_eq!(first, FrontdoorUserRpmOutcome::NotApplicable);
assert_eq!(second, FrontdoorUserRpmOutcome::NotApplicable);
}
#[test]
fn config_normalizes_non_positive_values() {
let config = FrontdoorUserRpmConfig::new(0, 0, true);
@@ -574,6 +612,7 @@ mod tests {
user_rate_limit: Some(1),
api_key_rate_limit: Some(10),
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: None,
});

View File

@@ -56,12 +56,23 @@ impl AppState {
&self,
group_ids: &[String],
) -> Result<Vec<String>, GatewayError> {
let mut group_ids = group_ids
.iter()
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<BTreeSet<_>>();
self.include_default_user_group_ids_for_role(group_ids, "user")
.await
}
pub(crate) async fn include_default_user_group_ids_for_role(
&self,
group_ids: &[String],
role: &str,
) -> Result<Vec<String>, GatewayError> {
let mut group_ids = normalized_user_group_ids(group_ids);
if role.trim().eq_ignore_ascii_case("admin") {
if let Some(default_group_id) = self.configured_default_user_group_id().await? {
group_ids.remove(&default_group_id);
}
group_ids.remove(BUILTIN_DEFAULT_USER_GROUP_ID);
return Ok(group_ids.into_iter().collect());
}
if let Some(default_group_id) = self.effective_default_user_group_id().await? {
group_ids.insert(default_group_id);
}
@@ -69,7 +80,7 @@ impl AppState {
}
pub(crate) async fn add_all_users_to_group(&self, group_id: &str) -> Result<(), GatewayError> {
for user in self.list_export_users().await? {
for user in self.list_non_admin_export_users().await? {
self.add_user_to_group(group_id, &user.id).await?;
}
Ok(())
@@ -801,3 +812,12 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
}
fn normalized_user_group_ids(group_ids: &[String]) -> BTreeSet<String> {
group_ids
.iter()
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect()
}

View File

@@ -444,11 +444,11 @@ async fn gateway_allows_default_user_group_access_policy_updates() {
let user_repository = Arc::new(
InMemoryUserReadRepository::seed_auth_users(vec![
sample_admin_user("user-1"),
sample_admin_user_with_role("admin-1", "admin", "admin@example.com", "admin"),
sample_admin_user_with_role("user-2", "user", "bob@example.com", "bob"),
])
.with_export_users(vec![
sample_admin_export_user("user-1"),
sample_admin_export_user_with("admin", true, "admin-1", "admin@example.com", "admin"),
sample_admin_export_user_with("user", true, "user-2", "bob@example.com", "bob"),
]),
);
@@ -504,7 +504,8 @@ async fn gateway_allows_default_user_group_access_policy_updates() {
.list_user_group_members(&group_id)
.await
.expect("default members should list");
assert_eq!(members.len(), 2);
assert_eq!(members.len(), 1);
assert_eq!(members[0].user_id, "user-2");
let update_response = client
.put(format!("{gateway_url}/api/admin/user-groups/{group_id}"))

View File

@@ -51,6 +51,7 @@ pub(super) fn sample_auth_context() -> GatewayControlAuthContext {
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: None,
}

View File

@@ -99,4 +99,5 @@ VALUES (
INSERT IGNORE INTO user_group_members (group_id, user_id, created_at)
SELECT '00000000-0000-0000-0000-000000000001', id, UNIX_TIMESTAMP()
FROM users
WHERE is_deleted = 0;
WHERE is_deleted = 0
AND LOWER(role) <> 'admin';

View File

@@ -0,0 +1,12 @@
DELETE member
FROM user_group_members AS member
JOIN users ON users.id = member.user_id
WHERE LOWER(users.role) = 'admin'
AND (
member.group_id = '00000000-0000-0000-0000-000000000001'
OR member.group_id IN (
SELECT TRIM(BOTH '"' FROM value)
FROM system_configs
WHERE `key` = 'default_user_group_id'
)
);

View File

@@ -98,4 +98,5 @@ INSERT INTO public.user_group_members (group_id, user_id)
SELECT '00000000-0000-0000-0000-000000000001', id
FROM public.users
WHERE is_deleted IS FALSE
AND LOWER(role) <> 'admin'
ON CONFLICT (group_id, user_id) DO NOTHING;

View File

@@ -0,0 +1,12 @@
DELETE FROM public.user_group_members AS member
USING public.users AS users
WHERE member.user_id = users.id
AND LOWER(users.role) = 'admin'
AND (
member.group_id = '00000000-0000-0000-0000-000000000001'
OR member.group_id IN (
SELECT value #>> '{}'
FROM public.system_configs
WHERE key = 'default_user_group_id'
)
);

View File

@@ -97,4 +97,5 @@ VALUES (
INSERT OR IGNORE INTO user_group_members (group_id, user_id, created_at)
SELECT '00000000-0000-0000-0000-000000000001', id, CAST(strftime('%s', 'now') AS INTEGER)
FROM users
WHERE is_deleted = 0;
WHERE is_deleted = 0
AND LOWER(role) <> 'admin';

View File

@@ -0,0 +1,14 @@
DELETE FROM user_group_members
WHERE user_id IN (
SELECT id
FROM users
WHERE LOWER(role) = 'admin'
)
AND (
group_id = '00000000-0000-0000-0000-000000000001'
OR group_id IN (
SELECT TRIM(value, '"')
FROM system_configs
WHERE key = 'default_user_group_id'
)
);

View File

@@ -37,6 +37,7 @@ INSERT INTO public.user_group_members (group_id, user_id)
SELECT '00000000-0000-0000-0000-000000000001', id
FROM public.users
WHERE is_deleted IS FALSE
AND LOWER(role) <> 'admin'
ON CONFLICT (group_id, user_id) DO NOTHING;
SELECT pg_catalog.set_config('search_path', 'public', true);

View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260511000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260511120000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

View File

@@ -299,6 +299,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260510000000,
20260510120000,
20260511000000,
20260511120000,
]
);
}
@@ -562,7 +563,8 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260508000000,
20260509000000,
20260509120000,
20260510120000
20260510120000,
20260511120000
]
);
assert_eq!(
@@ -573,7 +575,8 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260508000000,
20260509000000,
20260509120000,
20260510120000
20260510120000,
20260511120000
]
);
}
@@ -1084,6 +1087,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260510000000,
20260510120000,
20260511000000,
20260511120000,
]
);
}