i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support api key ip restriction rules

Browse Source
This commit is contained in:
fawney19
2026-05-20 16:11:49 +08:00
parent b6bdc08267
commit f76bbaab52
55 changed files with 698 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/aether-admin/src/observability/stats.rs
View File

@@ -2121,7 +2121,7 @@ mod tests {
api_key_allowed_providers: None,
api_key_allowed_api_formats: None,
api_key_allowed_models: None,
api_key_allowed_ips: None,
api_key_ip_rules: None,
}
}

2
crates/aether-data/migrations/mysql/20260518100000_add_api_key_allowed_ips.sql
View File

@@ -1,2 +0,0 @@
ALTER TABLE api_keys
ADD COLUMN allowed_ips TEXT NULL AFTER allowed_models;

2
crates/aether-data/migrations/mysql/20260520000000_add_api_key_ip_rules.sql Normal file
View File

@@ -0,0 +1,2 @@
ALTER TABLE api_keys
ADD COLUMN ip_rules TEXT NULL AFTER allowed_models;

5
crates/aether-data/migrations/postgres/20260518100000_add_api_key_allowed_ips.sql
View File

@@ -1,5 +0,0 @@
ALTER TABLE api_keys
ADD COLUMN IF NOT EXISTS allowed_ips jsonb NULL;
ALTER TABLE api_keys
ALTER COLUMN allowed_ips TYPE jsonb USING allowed_ips::jsonb;

5
crates/aether-data/migrations/postgres/20260520000000_add_api_key_ip_rules.sql Normal file
View File

@@ -0,0 +1,5 @@
ALTER TABLE api_keys
ADD COLUMN IF NOT EXISTS ip_rules jsonb NULL;
ALTER TABLE api_keys
ALTER COLUMN ip_rules TYPE jsonb USING ip_rules::jsonb;

1
crates/aether-data/migrations/sqlite/20260518100000_add_api_key_allowed_ips.sql
View File

@@ -1 +0,0 @@
ALTER TABLE api_keys ADD COLUMN allowed_ips TEXT;

1
crates/aether-data/migrations/sqlite/20260520000000_add_api_key_ip_rules.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE api_keys ADD COLUMN ip_rules TEXT;

2
crates/aether-data/schema/bootstrap/postgres/001_types_and_tables.sql
View File

@@ -176,7 +176,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys (
allowed_providers json,
allowed_api_formats json,
allowed_models json,
allowed_ips jsonb,
ip_rules jsonb,
rate_limit integer DEFAULT 100,
concurrent_limit integer,
force_capabilities json,

2
crates/aether-data/schema/generated/mysql/baseline/001_identity.sql
View File

@@ -75,7 +75,7 @@ CREATE TABLE IF NOT EXISTS api_keys (
`allowed_models` JSON,
`allowed_providers` JSON,
`allowed_api_formats` JSON,
`allowed_ips` JSON,
`ip_rules` JSON,
`rate_limit` INT DEFAULT 100,
`concurrent_limit` INT,
`force_capabilities` JSON,

2
crates/aether-data/schema/generated/postgres/baseline/001_identity.sql
View File

@@ -78,7 +78,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys (
allowed_models jsonb,
allowed_providers jsonb,
allowed_api_formats jsonb,
allowed_ips jsonb,
ip_rules jsonb,
rate_limit integer DEFAULT 100,
concurrent_limit integer,
force_capabilities jsonb,

2
crates/aether-data/schema/generated/sqlite/baseline/001_identity.sql
View File

@@ -73,7 +73,7 @@ CREATE TABLE IF NOT EXISTS api_keys (
allowed_models TEXT,
allowed_providers TEXT,
allowed_api_formats TEXT,
allowed_ips TEXT,
ip_rules TEXT,
rate_limit INTEGER DEFAULT 100,
concurrent_limit INTEGER,
force_capabilities TEXT,

2
crates/aether-data/schema/logical/001_identity.toml
View File

@@ -334,7 +334,7 @@ type = "json"
nullable = true
[[table.api_keys.columns]]
name = "allowed_ips"
name = "ip_rules"
type = "json"
nullable = true

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
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 = 20260519130000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

34
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -307,10 +307,10 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260515000000,
20260516000000,
20260518000000,
20260518100000,
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
}
@@ -471,31 +471,27 @@ fn management_tokens_json_columns_are_normalized_to_jsonb_in_postgres_schema_pat
}
#[test]
fn api_key_allowed_ips_is_jsonb_in_postgres_schema_paths() {
let api_key_allowed_ips_migration = POSTGRES_MIGRATOR
fn api_key_ip_rules_is_jsonb_in_postgres_schema_paths() {
let api_key_ip_rules_migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == 20260518100000)
.expect("api key allowed IPs migration should be embedded");
assert!(api_key_allowed_ips_migration
.find(|migration| migration.version == 20260520000000)
.expect("api key IP rules migration should be embedded");
assert!(api_key_ip_rules_migration
.sql
.contains("ADD COLUMN IF NOT EXISTS allowed_ips jsonb NULL"));
assert!(api_key_allowed_ips_migration
.contains("ADD COLUMN IF NOT EXISTS ip_rules jsonb NULL"));
assert!(api_key_ip_rules_migration
.sql
.contains("ALTER COLUMN allowed_ips TYPE jsonb USING allowed_ips::jsonb"));
.contains("ALTER COLUMN ip_rules TYPE jsonb USING ip_rules::jsonb"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("allowed_ips jsonb,"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ip_rules jsonb,"));
let bootstrap_schema =
include_str!("../../../schema/bootstrap/postgres/001_types_and_tables.sql");
assert!(bootstrap_schema.contains("allowed_ips jsonb,"));
let driver_schema =
include_str!("../../../schema/drivers/postgres/baseline/001_types_and_tables.sql");
assert!(driver_schema.contains("allowed_ips jsonb,"));
assert!(bootstrap_schema.contains("ip_rules jsonb,"));
let generated_identity =
include_str!("../../../schema/generated/postgres/baseline/001_identity.sql");
assert!(generated_identity.contains("allowed_ips jsonb,"));
assert!(generated_identity.contains("ip_rules jsonb,"));
}
#[test]
@@ -631,10 +627,10 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260512110000,
20260516000000,
20260518000000,
20260518100000,
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
assert_eq!(
@@ -653,10 +649,10 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260512110000,
20260516000000,
20260518000000,
20260518100000,
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
}
@@ -1175,10 +1171,10 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260515000000,
20260516000000,
20260518000000,
20260518100000,
20260519000000,
20260519120000,
20260519130000,
20260520000000,
]
);
}

40
crates/aether-data/src/repository/auth/memory.rs
View File

@@ -79,9 +79,9 @@ impl InMemoryAuthApiKeySnapshotRepository {
snapshot.api_key_is_standalone,
)
.and_then(|record| {
record.with_allowed_ips(
record.with_ip_rules(
snapshot
.api_key_allowed_ips
.api_key_ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)
@@ -504,7 +504,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_allowed_providers: record.allowed_providers.clone(),
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
api_key_allowed_models: record.allowed_models.clone(),
api_key_allowed_ips: record.allowed_ips.clone(),
api_key_ip_rules: record.ip_rules.clone(),
..template
}
} else {
@@ -543,9 +543,9 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.as_ref()
.map(|value| serde_json::json!(value)),
)?
.with_api_key_allowed_ips(
.with_api_key_ip_rules(
record
.allowed_ips
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
@@ -581,9 +581,9 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
record.total_cost_usd,
false,
)?
.with_allowed_ips(
.with_ip_rules(
record
.allowed_ips
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
@@ -640,7 +640,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_allowed_providers: record.allowed_providers.clone(),
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
api_key_allowed_models: record.allowed_models.clone(),
api_key_allowed_ips: record.allowed_ips.clone(),
api_key_ip_rules: record.ip_rules.clone(),
..template
}
} else {
@@ -679,9 +679,9 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.as_ref()
.map(|value| serde_json::json!(value)),
)?
.with_api_key_allowed_ips(
.with_api_key_ip_rules(
record
.allowed_ips
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
@@ -717,9 +717,9 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
record.total_cost_usd,
true,
)?
.with_allowed_ips(
.with_ip_rules(
record
.allowed_ips
.ip_rules
.as_ref()
.map(|value| serde_json::json!(value)),
)?
@@ -775,12 +775,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.concurrent_limit = Some(concurrent_limit);
}
}
if let Some(allowed_ips) = record.allowed_ips {
if let Some(ip_rules) = record.ip_rules {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_allowed_ips = allowed_ips.clone();
snapshot.api_key_ip_rules = ip_rules.clone();
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.allowed_ips = allowed_ips;
export.ip_rules = ip_rules;
}
}
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
@@ -848,12 +848,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.allowed_models = allowed_models;
}
}
if let Some(allowed_ips) = record.allowed_ips {
if let Some(ip_rules) = record.ip_rules {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_allowed_ips = allowed_ips.clone();
snapshot.api_key_ip_rules = ip_rules.clone();
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.allowed_ips = allowed_ips;
export.ip_rules = ip_rules;
}
}
if record.expires_at_present {
@@ -1289,7 +1289,7 @@ mod tests {
name: None,
rate_limit: None,
concurrent_limit: Some(11),
allowed_ips: None,
ip_rules: None,
})
.await
.expect("update should succeed")
@@ -1324,7 +1324,7 @@ mod tests {
allowed_providers: None,
allowed_api_formats: None,
allowed_models: None,
allowed_ips: None,
ip_rules: None,
expires_at_present: false,
expires_at_unix_secs: None,
auto_delete_on_expiry_present: false,

44
crates/aether-data/src/repository/auth/mysql.rs
View File

@@ -35,7 +35,7 @@ SELECT
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models,
api_keys.allowed_ips AS api_key_allowed_ips
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
"#;
@@ -50,7 +50,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -114,7 +114,7 @@ impl MysqlAuthApiKeyReadRepository {
r#"
INSERT INTO api_keys (
id, user_id, key_hash, key_encrypted, name, allowed_providers,
allowed_api_formats, allowed_models, allowed_ips, rate_limit, concurrent_limit,
allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit,
force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry,
total_requests, total_tokens, total_cost_usd, is_standalone,
created_at, updated_at
@@ -140,8 +140,8 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"api_keys.allowed_models",
)?)
.bind(json_string_from_string_list(
record.allowed_ips.as_ref(),
"api_keys.allowed_ips",
record.ip_rules.as_ref(),
"api_keys.ip_rules",
)?)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
@@ -181,7 +181,7 @@ struct CreateApiKeyInsertRecord {
allowed_providers: Option<Vec<String>>,
allowed_api_formats: Option<Vec<String>>,
allowed_models: Option<Vec<String>>,
allowed_ips: Option<Vec<String>>,
ip_rules: Option<Vec<String>>,
rate_limit: Option<i32>,
concurrent_limit: Option<i32>,
force_capabilities: Option<serde_json::Value>,
@@ -424,7 +424,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
allowed_ips: record.allowed_ips,
ip_rules: record.ip_rules,
rate_limit: Some(record.rate_limit),
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -452,7 +452,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
allowed_ips: record.allowed_ips,
ip_rules: record.ip_rules,
rate_limit: record.rate_limit,
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -478,7 +478,7 @@ UPDATE api_keys
SET name = COALESCE(?, name),
rate_limit = COALESCE(?, rate_limit),
concurrent_limit = COALESCE(?, concurrent_limit),
allowed_ips = CASE WHEN ? THEN ? ELSE allowed_ips END,
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
updated_at = ?
WHERE id = ?
AND user_id = ?
@@ -488,10 +488,10 @@ WHERE id = ?
.bind(record.name.as_deref())
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.allowed_ips.is_some())
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.allowed_ips,
"api_keys.allowed_ips",
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(now)
.bind(&record.api_key_id)
@@ -516,7 +516,7 @@ SET name = COALESCE(?, name),
allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END,
allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END,
allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END,
allowed_ips = CASE WHEN ? THEN ? ELSE allowed_ips END,
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
expires_at = CASE WHEN ? THEN ? ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN ? THEN ? ELSE auto_delete_on_expiry END,
updated_at = ?
@@ -544,10 +544,10 @@ WHERE id = ?
&record.allowed_models,
"api_keys.allowed_models",
)?)
.bind(record.allowed_ips.is_some())
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.allowed_ips,
"api_keys.allowed_ips",
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(record.expires_at_present)
.bind(optional_i64_from_u64(
@@ -944,9 +944,9 @@ fn map_auth_api_key_snapshot_row(
"api_keys.allowed_models",
)?,
)?
.with_api_key_allowed_ips(optional_json_from_string(
row.try_get("api_key_allowed_ips").map_sql_err()?,
"api_keys.allowed_ips",
.with_api_key_ip_rules(optional_json_from_string(
row.try_get("api_key_ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)?;
Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?))
}
@@ -991,9 +991,9 @@ fn map_auth_api_key_export_row(
row.try_get("is_standalone").map_sql_err()?,
)
.and_then(|record| {
record.with_allowed_ips(optional_json_from_string(
row.try_get("allowed_ips").map_sql_err()?,
"api_keys.allowed_ips",
record.with_ip_rules(optional_json_from_string(
row.try_get("ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)
})
.map(|record| record.with_feature_settings(feature_settings))

80
crates/aether-data/src/repository/auth/postgres.rs
View File

@@ -37,7 +37,7 @@ SELECT
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models,
api_keys.allowed_ips AS api_key_allowed_ips
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.key_hash = $1
@@ -68,7 +68,7 @@ SELECT
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models,
api_keys.allowed_ips AS api_key_allowed_ips
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1
@@ -99,7 +99,7 @@ SELECT
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models,
api_keys.allowed_ips AS api_key_allowed_ips
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1 AND users.id = $2
@@ -130,7 +130,7 @@ SELECT
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models,
api_keys.allowed_ips AS api_key_allowed_ips
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = ANY($1::TEXT[])
@@ -147,7 +147,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -178,7 +178,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -208,7 +208,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -238,7 +238,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -268,7 +268,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -342,7 +342,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -379,7 +379,7 @@ INSERT INTO api_keys (
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -429,7 +429,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -456,7 +456,7 @@ INSERT INTO api_keys (
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -506,7 +506,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -529,7 +529,7 @@ SET
name = COALESCE($3, name),
rate_limit = COALESCE($4, rate_limit),
concurrent_limit = COALESCE($5, concurrent_limit),
allowed_ips = CASE WHEN $6 THEN $7::jsonb ELSE allowed_ips END,
ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END,
updated_at = NOW()
WHERE user_id = $1
AND id = $2
@@ -543,7 +543,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -569,7 +569,7 @@ SET
allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END,
allowed_ips = CASE WHEN $13 THEN $14::jsonb ELSE allowed_ips END,
ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END,
expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END,
updated_at = NOW()
@@ -584,7 +584,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -618,7 +618,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -651,7 +651,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -695,7 +695,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
allowed_ips,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -1145,8 +1145,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let allowed_ips = record
.allowed_ips
let ip_rules = record
.ip_rules
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
@@ -1167,7 +1167,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_providers)
.bind(allowed_api_formats)
.bind(allowed_models)
.bind(allowed_ips)
.bind(ip_rules)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.force_capabilities)
@@ -1202,8 +1202,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let allowed_ips = record
.allowed_ips
let ip_rules = record
.ip_rules
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
@@ -1224,7 +1224,7 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_providers)
.bind(allowed_api_formats)
.bind(allowed_models)
.bind(allowed_ips)
.bind(ip_rules)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.force_capabilities)
@@ -1244,8 +1244,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
&self,
record: UpdateUserApiKeyBasicRecord,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let allowed_ips = record
.allowed_ips
let ip_rules = record
.ip_rules
.clone()
.flatten()
.map(serde_json::to_value)
@@ -1257,8 +1257,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(record.name)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.allowed_ips.is_some())
.bind(allowed_ips)
.bind(record.ip_rules.is_some())
.bind(ip_rules)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1290,8 +1290,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let allowed_ips = record
.allowed_ips
let ip_rules = record
.ip_rules
.clone()
.flatten()
.map(serde_json::to_value)
@@ -1318,8 +1318,8 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_api_formats)
.bind(record.allowed_models.is_some())
.bind(allowed_models)
.bind(record.allowed_ips.is_some())
.bind(allowed_ips)
.bind(record.ip_rules.is_some())
.bind(ip_rules)
.bind(record.expires_at_present)
.bind(expires_at)
.bind(record.auto_delete_on_expiry_present)
@@ -1549,7 +1549,7 @@ fn map_auth_api_key_snapshot_row(
row_get(row, "api_key_allowed_api_formats")?,
row_get(row, "api_key_allowed_models")?,
)?
.with_api_key_allowed_ips(row_get(row, "api_key_allowed_ips")?)?;
.with_api_key_ip_rules(row_get(row, "api_key_ip_rules")?)?;
Ok(snapshot.with_user_rate_limit(row_get(row, "user_rate_limit")?))
}
@@ -1577,7 +1577,7 @@ fn map_auth_api_key_export_row(
row_get(row, "total_cost_usd")?,
row_get(row, "is_standalone")?,
)
.and_then(|record| record.with_allowed_ips(row_get(row, "allowed_ips")?))
.and_then(|record| record.with_ip_rules(row_get(row, "ip_rules")?))
.map(|record| record.with_feature_settings(feature_settings))
.and_then(|record| {
record.with_activity_timestamps(
@@ -1622,7 +1622,7 @@ mod tests {
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_ips = CASE WHEN $13 THEN $14::jsonb ELSE allowed_ips END"));
.contains("ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
@@ -1633,9 +1633,9 @@ mod tests {
}
#[test]
fn update_user_api_key_basic_sql_casts_allowed_ips_as_jsonb() {
fn update_user_api_key_basic_sql_casts_ip_rules_as_jsonb() {
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("allowed_ips = CASE WHEN $6 THEN $7::jsonb ELSE allowed_ips END"));
.contains("ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END"));
}
#[tokio::test]

52
crates/aether-data/src/repository/auth/sqlite.rs
View File

@@ -35,7 +35,7 @@ SELECT
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models,
api_keys.allowed_ips AS api_key_allowed_ips
api_keys.ip_rules AS api_key_ip_rules
FROM api_keys
JOIN users ON users.id = api_keys.user_id
"#;
@@ -50,7 +50,7 @@ SELECT
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.allowed_ips,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
@@ -114,7 +114,7 @@ impl SqliteAuthApiKeyReadRepository {
r#"
INSERT INTO api_keys (
id, user_id, key_hash, key_encrypted, name, allowed_providers,
allowed_api_formats, allowed_models, allowed_ips, rate_limit, concurrent_limit,
allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit,
force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry,
total_requests, total_tokens, total_cost_usd, is_standalone,
created_at, updated_at
@@ -140,8 +140,8 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"api_keys.allowed_models",
)?)
.bind(json_string_from_string_list(
record.allowed_ips.as_ref(),
"api_keys.allowed_ips",
record.ip_rules.as_ref(),
"api_keys.ip_rules",
)?)
.bind(record.rate_limit)
.bind(record.concurrent_limit)
@@ -181,7 +181,7 @@ struct CreateApiKeyInsertRecord {
allowed_providers: Option<Vec<String>>,
allowed_api_formats: Option<Vec<String>>,
allowed_models: Option<Vec<String>>,
allowed_ips: Option<Vec<String>>,
ip_rules: Option<Vec<String>>,
rate_limit: Option<i32>,
concurrent_limit: Option<i32>,
force_capabilities: Option<serde_json::Value>,
@@ -424,7 +424,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
allowed_ips: record.allowed_ips,
ip_rules: record.ip_rules,
rate_limit: Some(record.rate_limit),
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -452,7 +452,7 @@ WHERE id = ?
allowed_providers: record.allowed_providers,
allowed_api_formats: record.allowed_api_formats,
allowed_models: record.allowed_models,
allowed_ips: record.allowed_ips,
ip_rules: record.ip_rules,
rate_limit: record.rate_limit,
concurrent_limit: record.concurrent_limit,
force_capabilities: record.force_capabilities,
@@ -478,7 +478,7 @@ UPDATE api_keys
SET name = COALESCE(?, name),
rate_limit = COALESCE(?, rate_limit),
concurrent_limit = COALESCE(?, concurrent_limit),
allowed_ips = CASE WHEN ? THEN ? ELSE allowed_ips END,
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
updated_at = ?
WHERE id = ?
AND user_id = ?
@@ -488,10 +488,10 @@ WHERE id = ?
.bind(record.name.as_deref())
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.allowed_ips.is_some())
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.allowed_ips,
"api_keys.allowed_ips",
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(now)
.bind(&record.api_key_id)
@@ -516,7 +516,7 @@ SET name = COALESCE(?, name),
allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END,
allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END,
allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END,
allowed_ips = CASE WHEN ? THEN ? ELSE allowed_ips END,
ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END,
expires_at = CASE WHEN ? THEN ? ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN ? THEN ? ELSE auto_delete_on_expiry END,
updated_at = ?
@@ -544,10 +544,10 @@ WHERE id = ?
&record.allowed_models,
"api_keys.allowed_models",
)?)
.bind(record.allowed_ips.is_some())
.bind(record.ip_rules.is_some())
.bind(json_string_from_nested_string_list(
&record.allowed_ips,
"api_keys.allowed_ips",
&record.ip_rules,
"api_keys.ip_rules",
)?)
.bind(record.expires_at_present)
.bind(optional_i64_from_u64(
@@ -944,9 +944,9 @@ fn map_auth_api_key_snapshot_row(
"api_keys.allowed_models",
)?,
)?
.with_api_key_allowed_ips(optional_json_from_string(
row.try_get("api_key_allowed_ips").map_sql_err()?,
"api_keys.allowed_ips",
.with_api_key_ip_rules(optional_json_from_string(
row.try_get("api_key_ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)?;
Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?))
}
@@ -991,9 +991,9 @@ fn map_auth_api_key_export_row(
row.try_get("is_standalone").map_sql_err()?,
)
.and_then(|record| {
record.with_allowed_ips(optional_json_from_string(
row.try_get("allowed_ips").map_sql_err()?,
"api_keys.allowed_ips",
record.with_ip_rules(optional_json_from_string(
row.try_get("ip_rules").map_sql_err()?,
"api_keys.ip_rules",
)?)
})
.map(|record| record.with_feature_settings(feature_settings))
@@ -1112,7 +1112,7 @@ mod tests {
allowed_providers: Some(vec!["openai".to_string()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec!["gpt-4.1".to_string()]),
allowed_ips: Some(vec!["203.0.113.10".to_string()]),
ip_rules: Some(vec!["203.0.113.10".to_string()]),
rate_limit: 100,
concurrent_limit: Some(5),
force_capabilities: Some(json!({"cache": true})),
@@ -1136,7 +1136,7 @@ mod tests {
name: Some("Updated User".to_string()),
rate_limit: Some(150),
concurrent_limit: Some(6),
allowed_ips: Some(Some(vec!["10.0.0.0/24".to_string()])),
ip_rules: Some(Some(vec!["10.0.0.0/24".to_string()])),
})
.await
.expect("user key should update")
@@ -1212,7 +1212,7 @@ mod tests {
allowed_providers: Some(vec!["openai".to_string()]),
allowed_api_formats: None,
allowed_models: None,
allowed_ips: None,
ip_rules: None,
rate_limit: None,
concurrent_limit: Some(2),
force_capabilities: None,
@@ -1239,7 +1239,7 @@ mod tests {
allowed_providers: Some(None),
allowed_api_formats: Some(Some(vec!["openai:responses".to_string()])),
allowed_models: Some(Some(vec!["gpt-4.1-mini".to_string()])),
allowed_ips: None,
ip_rules: None,
expires_at_present: true,
expires_at_unix_secs: Some(2_100_000_000),
auto_delete_on_expiry_present: true,

32
crates/aether-data/src/repository/auth/types.rs
View File

@@ -24,7 +24,7 @@ pub struct StoredAuthApiKeySnapshot {
pub api_key_allowed_providers: Option<Vec<String>>,
pub api_key_allowed_api_formats: Option<Vec<String>>,
pub api_key_allowed_models: Option<Vec<String>>,
pub api_key_allowed_ips: Option<Vec<String>>,
pub api_key_ip_rules: Option<Vec<String>>,
}
impl StoredAuthApiKeySnapshot {
@@ -98,15 +98,15 @@ impl StoredAuthApiKeySnapshot {
api_key_allowed_models,
"api_keys.allowed_models",
)?,
api_key_allowed_ips: None,
api_key_ip_rules: None,
})
}
pub fn with_api_key_allowed_ips(
pub fn with_api_key_ip_rules(
mut self,
api_key_allowed_ips: Option<serde_json::Value>,
api_key_ip_rules: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
self.api_key_allowed_ips = parse_string_list(api_key_allowed_ips, "api_keys.allowed_ips")?;
self.api_key_ip_rules = parse_string_list(api_key_ip_rules, "api_keys.ip_rules")?;
Ok(self)
}
@@ -158,7 +158,7 @@ pub struct ResolvedAuthApiKeySnapshot {
pub api_key_allowed_providers: Option<Vec<String>>,
pub api_key_allowed_api_formats: Option<Vec<String>>,
pub api_key_allowed_models: Option<Vec<String>>,
pub api_key_allowed_ips: Option<Vec<String>>,
pub api_key_ip_rules: Option<Vec<String>>,
pub currently_usable: bool,
}
@@ -188,7 +188,7 @@ impl ResolvedAuthApiKeySnapshot {
api_key_allowed_providers: snapshot.api_key_allowed_providers,
api_key_allowed_api_formats: snapshot.api_key_allowed_api_formats,
api_key_allowed_models: snapshot.api_key_allowed_models,
api_key_allowed_ips: snapshot.api_key_allowed_ips,
api_key_ip_rules: snapshot.api_key_ip_rules,
currently_usable,
};
resolved.constrain_non_standalone_api_key_policy_to_user_policy();
@@ -344,7 +344,7 @@ pub struct StoredAuthApiKeyExportRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub allowed_ips: Option<Vec<String>>,
pub ip_rules: Option<Vec<String>>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
@@ -416,7 +416,7 @@ impl StoredAuthApiKeyExportRecord {
"api_keys.allowed_api_formats",
)?,
allowed_models: parse_string_list(allowed_models, "api_keys.allowed_models")?,
allowed_ips: None,
ip_rules: None,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -459,11 +459,11 @@ impl StoredAuthApiKeyExportRecord {
self
}
pub fn with_allowed_ips(
pub fn with_ip_rules(
mut self,
allowed_ips: Option<serde_json::Value>,
ip_rules: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
self.allowed_ips = parse_string_list(allowed_ips, "api_keys.allowed_ips")?;
self.ip_rules = parse_string_list(ip_rules, "api_keys.ip_rules")?;
Ok(self)
}
}
@@ -491,7 +491,7 @@ pub struct CreateUserApiKeyRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub allowed_ips: Option<Vec<String>>,
pub ip_rules: Option<Vec<String>>,
pub rate_limit: i32,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
@@ -510,7 +510,7 @@ pub struct UpdateUserApiKeyBasicRecord {
pub name: Option<String>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
pub allowed_ips: Option<Option<Vec<String>>>,
pub ip_rules: Option<Option<Vec<String>>>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -523,7 +523,7 @@ pub struct CreateStandaloneApiKeyRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub allowed_ips: Option<Vec<String>>,
pub ip_rules: Option<Vec<String>>,
pub rate_limit: Option<i32>,
pub concurrent_limit: Option<i32>,
pub force_capabilities: Option<serde_json::Value>,
@@ -546,7 +546,7 @@ pub struct UpdateStandaloneApiKeyBasicRecord {
pub allowed_providers: Option<Option<Vec<String>>>,
pub allowed_api_formats: Option<Option<Vec<String>>>,
pub allowed_models: Option<Option<Vec<String>>>,
pub allowed_ips: Option<Option<Vec<String>>>,
pub ip_rules: Option<Option<Vec<String>>>,
pub expires_at_present: bool,
pub expires_at_unix_secs: Option<u64>,
pub auto_delete_on_expiry_present: bool,

2
crates/aether-data/src/repository/management_tokens/postgres.rs
View File

@@ -561,7 +561,7 @@ fn map_management_token_write_error(
.unwrap_or_else(|| "Management Token 名称已存在".to_string()),
),
(Some("23514"), Some("check_allowed_ips_not_empty")) => {
Some("IP 白名单不能为空，如需取消限制请不提供此字段".to_string())
Some("IP 限制规则不能为空，如需取消限制请不提供此字段".to_string())
}
_ => None,
}

12
crates/aether-data/src/repository/management_tokens/types.rs
View File

@@ -208,17 +208,17 @@ impl CreateManagementTokenRecord {
if let Some(allowed_ips) = &self.allowed_ips {
let Some(items) = allowed_ips.as_array() else {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must be an array".to_string(),
"IP 限制规则必须是数组".to_string(),
));
};
if items.is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must not be empty".to_string(),
"IP 限制规则不能为空".to_string(),
));
}
if items.iter().any(|value| value.as_str().is_none()) {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must contain only strings".to_string(),
"IP 限制规则只能包含字符串".to_string(),
));
}
}
@@ -258,17 +258,17 @@ impl UpdateManagementTokenRecord {
if let Some(allowed_ips) = &self.allowed_ips {
let Some(items) = allowed_ips.as_array() else {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must be an array".to_string(),
"IP 限制规则必须是数组".to_string(),
));
};
if items.is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must not be empty".to_string(),
"IP 限制规则不能为空".to_string(),
));
}
if items.iter().any(|value| value.as_str().is_none()) {
return Err(crate::DataLayerError::InvalidInput(
"allowed_ips must contain only strings".to_string(),
"IP 限制规则只能包含字符串".to_string(),
));
}
}