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

feat(security): harden gateway boundaries and usage policies

Browse Source 
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/src/auth.rs
+597 -105
View File
@@ -4,9 +4,9 @@ use sqlx::{postgres::PgRow, PgPool, Row};
use aether_data_contracts::repository::auth::{
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
AuthApiKeyWriteRepository, CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord,
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
AuthApiKeyWriteRepository, CompareAndSwapAuthApiKeyCiphertext, CreateStandaloneApiKeyRecord,
CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord,
StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
};
use aether_data_contracts::DataLayerError;
@@ -197,6 +197,37 @@ WHERE api_keys.id = ANY($1::TEXT[])
ORDER BY api_keys.id ASC
"#;
const RESTORE_API_KEY_SELECT_SQL: &str = r#"
SELECT
api_keys.user_id,
api_keys.id AS api_key_id,
api_keys.key_hash,
api_keys.key_encrypted,
api_keys.name,
api_keys.allowed_providers,
api_keys.allowed_api_formats,
api_keys.allowed_models,
api_keys.ip_rules,
api_keys.rate_limit,
api_keys.concurrent_limit,
api_keys.force_capabilities,
api_keys.feature_settings,
api_keys.is_active,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS expires_at_unix_secs,
api_keys.auto_delete_on_expiry,
api_keys.total_requests,
COALESCE(api_keys.total_tokens, 0)::BIGINT AS total_tokens,
COALESCE(CAST(api_keys.total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
CAST(EXTRACT(EPOCH FROM api_keys.last_used_at) AS BIGINT) AS last_used_at_unix_secs,
CAST(EXTRACT(EPOCH FROM api_keys.created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM api_keys.updated_at) AS BIGINT) AS updated_at_unix_secs,
api_keys.is_standalone
FROM api_keys
WHERE api_keys.id = $1
LIMIT 1
FOR UPDATE
"#;
const LIST_EXPORT_BY_NAME_SEARCH_SQL: &str = r#"
SELECT
api_keys.user_id,
@@ -407,15 +438,15 @@ VALUES (
$10,
$11,
$12,
NULL,
$13,
$14,
$15,
FALSE,
FALSE,
$16,
FALSE,
FALSE,
$17,
$18,
$19,
NOW(),
NOW()
)
@@ -525,14 +556,17 @@ RETURNING
const UPDATE_USER_API_KEY_BASIC_SQL: &str = r#"
UPDATE api_keys
SET
name = COALESCE($3, name),
rate_limit = COALESCE($4, rate_limit),
concurrent_limit = COALESCE($5, concurrent_limit),
ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END,
key_encrypted = CASE WHEN $3 THEN $4 ELSE key_encrypted END,
name = CASE WHEN $5 THEN $6 ELSE name END,
rate_limit = CASE WHEN $7 THEN $8 ELSE rate_limit END,
concurrent_limit = CASE WHEN $9 THEN $10 ELSE concurrent_limit END,
ip_rules = CASE WHEN $11 THEN $12::jsonb ELSE ip_rules END,
feature_settings = CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END,
updated_at = NOW()
WHERE user_id = $1
AND id = $2
AND is_standalone = FALSE
AND ($15 = FALSE OR is_locked = FALSE)
RETURNING
user_id,
id AS api_key_id,
@@ -562,15 +596,17 @@ RETURNING
const UPDATE_STANDALONE_API_KEY_BASIC_SQL: &str = r#"
UPDATE api_keys
SET
name = COALESCE($2, name),
rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END,
concurrent_limit = CASE WHEN $5 THEN $6 ELSE concurrent_limit END,
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,
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,
key_encrypted = CASE WHEN $2 THEN $3 ELSE key_encrypted END,
name = CASE WHEN $4 THEN $5 ELSE name END,
force_capabilities = CASE WHEN $6 THEN $7::json ELSE force_capabilities END,
rate_limit = CASE WHEN $8 THEN $9 ELSE rate_limit END,
concurrent_limit = CASE WHEN $10 THEN $11 ELSE concurrent_limit END,
allowed_providers = CASE WHEN $12 THEN $13::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $14 THEN $15::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $16 THEN $17::json ELSE allowed_models END,
ip_rules = CASE WHEN $18 THEN $19::jsonb ELSE ip_rules END,
expires_at = CASE WHEN $20 THEN $21::timestamptz ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $22 THEN $23 ELSE auto_delete_on_expiry END,
updated_at = NOW()
WHERE id = $1
AND is_standalone = TRUE
@@ -608,6 +644,7 @@ SET
WHERE user_id = $1
AND id = $2
AND is_standalone = FALSE
AND ($4 = FALSE OR is_locked = FALSE)
RETURNING
user_id,
id AS api_key_id,
@@ -719,6 +756,7 @@ SET
WHERE user_id = $1
AND id = $2
AND is_standalone = FALSE
AND ($4 = FALSE OR is_locked = FALSE)
RETURNING
user_id,
id AS api_key_id,
@@ -753,6 +791,7 @@ SET
WHERE user_id = $1
AND id = $2
AND is_standalone = FALSE
AND ($4 = FALSE OR is_locked = FALSE)
RETURNING
user_id,
id AS api_key_id,
@@ -762,6 +801,7 @@ RETURNING
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
@@ -786,6 +826,7 @@ SET
WHERE user_id = $1
AND id = $2
AND is_standalone = FALSE
AND ($4 = FALSE OR is_locked = FALSE)
"#;
const SET_STANDALONE_API_KEY_FEATURE_SETTINGS_SQL: &str = r#"
@@ -797,26 +838,20 @@ WHERE id = $1
AND is_standalone = TRUE
"#;
const DISABLE_WALLET_BY_API_KEY_ID_SQL: &str = r#"
UPDATE wallets
SET status = 'disabled',
updated_at = NOW()
WHERE api_key_id = $1
AND status <> 'disabled'
"#;
const POSTGRES_ANONYMIZE_API_KEY_HISTORY_SQL: &[&str] = &[
"UPDATE request_candidates SET api_key_name = NULL WHERE api_key_id = $1",
"UPDATE video_tasks SET api_key_name = NULL WHERE api_key_id = $1",
"UPDATE usage SET api_key_name = NULL WHERE api_key_id = $1",
"UPDATE stats_daily_api_key SET api_key_name = NULL WHERE api_key_id = $1",
"UPDATE audit_logs SET description = 'deleted API key event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE api_key_id = $1",
"UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = $1)",
"UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.api_key_id = $1 AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))",
"UPDATE payment_orders SET gateway_response = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = $1)",
"UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = $1)",
];
const DELETE_USER_API_KEY_SQL: &str = r#"
DELETE FROM api_keys
WHERE user_id = $1
AND id = $2
AND is_standalone = FALSE
"#;
const DELETE_STANDALONE_API_KEY_SQL: &str = r#"
DELETE FROM api_keys
WHERE id = $1
AND is_standalone = TRUE
"#;
const POSTGRES_DELETE_API_KEY_DEPENDENTS_SQL: &[&str] =
&["DELETE FROM api_key_provider_mappings WHERE api_key_id = $1"];
#[derive(Debug, Clone)]
pub struct SqlxAuthApiKeySnapshotReadRepository {
@@ -952,10 +987,12 @@ impl SqlxAuthApiKeySnapshotReadRepository {
if user_ids.is_empty() {
return Ok(AuthApiKeyExportSummary::default());
}
let now_unix_secs =
datetime_from_unix_secs(now_unix_secs, "api_keys.summary_now")?.timestamp() as f64;
let row = sqlx::query(SUMMARIZE_EXPORT_BY_USER_IDS_SQL)
.bind(user_ids)
.bind(now_unix_secs as f64)
.bind(now_unix_secs)
.fetch_one(&self.pool)
.await
.map_postgres_err()?;
@@ -969,8 +1006,10 @@ impl SqlxAuthApiKeySnapshotReadRepository {
&self,
now_unix_secs: u64,
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
let now_unix_secs =
datetime_from_unix_secs(now_unix_secs, "api_keys.summary_now")?.timestamp() as f64;
let row = sqlx::query(SUMMARIZE_EXPORT_NON_STANDALONE_SQL)
.bind(now_unix_secs as f64)
.bind(now_unix_secs)
.fetch_one(&self.pool)
.await
.map_postgres_err()?;
@@ -1025,8 +1064,10 @@ impl SqlxAuthApiKeySnapshotReadRepository {
&self,
now_unix_secs: u64,
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
let now_unix_secs =
datetime_from_unix_secs(now_unix_secs, "api_keys.summary_now")?.timestamp() as f64;
let row = sqlx::query(SUMMARIZE_EXPORT_STANDALONE_SQL)
.bind(now_unix_secs as f64)
.bind(now_unix_secs)
.fetch_one(&self.pool)
.await
.map_postgres_err()?;
@@ -1173,12 +1214,20 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let expires_at = record
.expires_at_unix_secs
.map(|value| {
chrono::DateTime::<chrono::Utc>::from_timestamp(value as i64, 0).ok_or_else(|| {
DataLayerError::UnexpectedValue(format!("invalid api_keys.expires_at: {value}"))
})
})
.map(|value| datetime_from_unix_secs(value, "api_keys.expires_at"))
.transpose()?;
let mut tx = self.pool.begin().await.map_postgres_err()?;
let owner_exists: Option<String> = sqlx::query_scalar(
"SELECT id FROM users WHERE id = $1 AND is_deleted IS FALSE FOR UPDATE",
)
.bind(&record.user_id)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?;
if owner_exists.is_none() {
tx.rollback().await.map_postgres_err()?;
return Ok(None);
}
let row = sqlx::query(CREATE_USER_API_KEY_SQL)
.bind(record.api_key_id)
.bind(record.user_id)
@@ -1192,16 +1241,22 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(record.rate_limit)
.bind(record.concurrent_limit)
.bind(record.force_capabilities)
.bind(record.feature_settings)
.bind(record.is_active)
.bind(expires_at)
.bind(record.auto_delete_on_expiry)
.bind(record.total_requests as i64)
.bind(record.total_tokens as i64)
.bind(i64_from_u64(
record.total_requests,
"api_keys.total_requests",
)?)
.bind(i64_from_u64(record.total_tokens, "api_keys.total_tokens")?)
.bind(record.total_cost_usd)
.fetch_optional(&self.pool)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
let record = row.as_ref().map(map_auth_api_key_export_row).transpose()?;
tx.commit().await.map_err(postgres_error)?;
Ok(record)
}
async fn create_standalone_api_key(
@@ -1230,12 +1285,20 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let expires_at = record
.expires_at_unix_secs
.map(|value| {
chrono::DateTime::<chrono::Utc>::from_timestamp(value as i64, 0).ok_or_else(|| {
DataLayerError::UnexpectedValue(format!("invalid api_keys.expires_at: {value}"))
})
})
.map(|value| datetime_from_unix_secs(value, "api_keys.expires_at"))
.transpose()?;
let mut tx = self.pool.begin().await.map_postgres_err()?;
let owner_exists: Option<String> = sqlx::query_scalar(
"SELECT id FROM users WHERE id = $1 AND is_deleted IS FALSE FOR UPDATE",
)
.bind(&record.user_id)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?;
if owner_exists.is_none() {
tx.rollback().await.map_postgres_err()?;
return Ok(None);
}
let row = sqlx::query(CREATE_STANDALONE_API_KEY_SQL)
.bind(record.api_key_id)
.bind(record.user_id)
@@ -1252,13 +1315,18 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(record.is_active)
.bind(expires_at)
.bind(record.auto_delete_on_expiry)
.bind(record.total_requests as i64)
.bind(record.total_tokens as i64)
.bind(i64_from_u64(
record.total_requests,
"api_keys.total_requests",
)?)
.bind(i64_from_u64(record.total_tokens, "api_keys.total_tokens")?)
.bind(record.total_cost_usd)
.fetch_optional(&self.pool)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
let record = row.as_ref().map(map_auth_api_key_export_row).transpose()?;
tx.commit().await.map_err(postgres_error)?;
Ok(record)
}
async fn update_user_api_key_basic(
@@ -1272,14 +1340,84 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let feature_settings = record.feature_settings.clone().flatten();
let row = sqlx::query(UPDATE_USER_API_KEY_BASIC_SQL)
.bind(record.user_id)
.bind(record.api_key_id)
.bind(record.key_encrypted_present)
.bind(record.key_encrypted)
.bind(record.name_present)
.bind(record.name)
.bind(record.rate_limit_present)
.bind(record.rate_limit)
.bind(record.concurrent_limit_present)
.bind(record.concurrent_limit)
.bind(record.ip_rules.is_some())
.bind(ip_rules)
.bind(record.feature_settings.is_some())
.bind(feature_settings)
.bind(false)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn compare_and_swap_api_key_ciphertext(
&self,
mutation: &CompareAndSwapAuthApiKeyCiphertext,
) -> Result<bool, DataLayerError> {
let result = sqlx::query(
r#"
UPDATE api_keys
SET key_encrypted = $1
WHERE id = $2
AND user_id = $3
AND key_hash = $4
AND is_standalone = $5
AND key_encrypted = $6
"#,
)
.bind(&mutation.key_encrypted)
.bind(&mutation.api_key_id)
.bind(&mutation.user_id)
.bind(&mutation.key_hash)
.bind(mutation.is_standalone)
.bind(&mutation.expected_key_encrypted)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(result.rows_affected() == 1)
}
async fn update_user_api_key_basic_if_unlocked(
&self,
record: UpdateUserApiKeyBasicRecord,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let ip_rules = record
.ip_rules
.clone()
.flatten()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let feature_settings = record.feature_settings.clone().flatten();
let row = sqlx::query(UPDATE_USER_API_KEY_BASIC_SQL)
.bind(record.user_id)
.bind(record.api_key_id)
.bind(record.key_encrypted_present)
.bind(record.key_encrypted)
.bind(record.name_present)
.bind(record.name)
.bind(record.rate_limit_present)
.bind(record.rate_limit)
.bind(record.concurrent_limit_present)
.bind(record.concurrent_limit)
.bind(record.ip_rules.is_some())
.bind(ip_rules)
.bind(record.feature_settings.is_some())
.bind(feature_settings)
.bind(true)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1320,15 +1458,16 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let expires_at = record
.expires_at_unix_secs
.map(|value| {
chrono::DateTime::<chrono::Utc>::from_timestamp(value as i64, 0).ok_or_else(|| {
DataLayerError::UnexpectedValue(format!("invalid api_keys.expires_at: {value}"))
})
})
.map(|value| datetime_from_unix_secs(value, "api_keys.expires_at"))
.transpose()?;
let row = sqlx::query(UPDATE_STANDALONE_API_KEY_BASIC_SQL)
.bind(record.api_key_id)
.bind(record.key_encrypted_present)
.bind(record.key_encrypted)
.bind(record.name_present)
.bind(record.name)
.bind(record.force_capabilities.is_some())
.bind(record.force_capabilities.clone().flatten())
.bind(record.rate_limit_present)
.bind(record.rate_limit)
.bind(record.concurrent_limit_present)
@@ -1351,6 +1490,133 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn restore_api_key_if_matches(
&self,
expected: &StoredAuthApiKeyExportRecord,
restored: &StoredAuthApiKeyExportRecord,
) -> Result<bool, DataLayerError> {
if restored.api_key_id != expected.api_key_id
|| restored.user_id != expected.user_id
|| restored.key_hash != expected.key_hash
|| restored.is_standalone != expected.is_standalone
{
return Ok(false);
}
let mut tx = self.pool.begin().await.map_postgres_err()?;
let row = sqlx::query(RESTORE_API_KEY_SELECT_SQL)
.bind(&expected.api_key_id)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?;
let Some(row) = row else {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
};
let current = map_auth_api_key_export_row(&row)?;
if current != *expected {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
let expires_at = restored
.expires_at_unix_secs
.map(|value| datetime_from_unix_secs(value, "api_keys.expires_at"))
.transpose()?;
let last_used_at = restored
.last_used_at_unix_secs
.map(|value| datetime_from_unix_secs(value, "api_keys.last_used_at"))
.transpose()?;
let allowed_providers = restored
.allowed_providers
.as_ref()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let allowed_api_formats = restored
.allowed_api_formats
.as_ref()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let allowed_models = restored
.allowed_models
.as_ref()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let ip_rules = restored
.ip_rules
.as_ref()
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let result = sqlx::query(
r#"
UPDATE api_keys
SET key_encrypted = $1,
name = $2,
allowed_providers = $3::json,
allowed_api_formats = $4::json,
allowed_models = $5::json,
ip_rules = $6::jsonb,
rate_limit = $7,
concurrent_limit = $8,
force_capabilities = $9::json,
feature_settings = $10::jsonb,
is_active = $11,
expires_at = $12,
auto_delete_on_expiry = $13,
total_requests = $14,
total_tokens = $15,
total_cost_usd = $16,
last_used_at = $17,
updated_at = NOW()
WHERE id = $18
AND user_id = $19
AND key_hash = $20
AND is_standalone = $21
"#,
)
.bind(&restored.key_encrypted)
.bind(&restored.name)
.bind(allowed_providers)
.bind(allowed_api_formats)
.bind(allowed_models)
.bind(ip_rules)
.bind(restored.rate_limit)
.bind(restored.concurrent_limit)
.bind(&restored.force_capabilities)
.bind(&restored.feature_settings)
.bind(restored.is_active)
.bind(expires_at)
.bind(restored.auto_delete_on_expiry)
.bind(i64_from_u64(
restored.total_requests,
"api_keys.total_requests",
)?)
.bind(i64_from_u64(
restored.total_tokens,
"api_keys.total_tokens",
)?)
.bind(restored.total_cost_usd)
.bind(last_used_at)
.bind(&restored.api_key_id)
.bind(&restored.user_id)
.bind(&restored.key_hash)
.bind(restored.is_standalone)
.execute(&mut *tx)
.await
.map_postgres_err()?;
if result.rows_affected() != 1 {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
tx.commit().await.map_postgres_err()?;
Ok(true)
}
async fn set_user_api_key_active(
&self,
user_id: &str,
@@ -1361,6 +1627,24 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(user_id)
.bind(api_key_id)
.bind(is_active)
.bind(false)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn set_user_api_key_active_if_unlocked(
&self,
user_id: &str,
api_key_id: &str,
is_active: bool,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let row = sqlx::query(SET_USER_API_KEY_ACTIVE_SQL)
.bind(user_id)
.bind(api_key_id)
.bind(is_active)
.bind(true)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1411,6 +1695,28 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(user_id)
.bind(api_key_id)
.bind(allowed_providers)
.bind(false)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn set_user_api_key_allowed_providers_if_unlocked(
&self,
user_id: &str,
api_key_id: &str,
allowed_providers: Option<Vec<String>>,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let allowed_providers = allowed_providers
.map(serde_json::to_value)
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
let row = sqlx::query(SET_USER_API_KEY_ALLOWED_PROVIDERS_SQL)
.bind(user_id)
.bind(api_key_id)
.bind(allowed_providers)
.bind(true)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1427,6 +1733,24 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(user_id)
.bind(api_key_id)
.bind(force_capabilities)
.bind(false)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn set_user_api_key_force_capabilities_if_unlocked(
&self,
user_id: &str,
api_key_id: &str,
force_capabilities: Option<serde_json::Value>,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let row = sqlx::query(SET_USER_API_KEY_FORCE_CAPABILITIES_SQL)
.bind(user_id)
.bind(api_key_id)
.bind(force_capabilities)
.bind(true)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1443,6 +1767,32 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(user_id)
.bind(api_key_id)
.bind(feature_settings)
.bind(false)
.execute(&self.pool)
.await
.map_postgres_err()?;
if result.rows_affected() == 0 {
return Ok(None);
}
let api_key_ids = [api_key_id.to_string()];
Ok(self
.list_export_api_keys_by_ids(&api_key_ids)
.await?
.into_iter()
.find(|record| record.user_id == user_id && !record.is_standalone))
}
async fn set_user_api_key_feature_settings_if_unlocked(
&self,
user_id: &str,
api_key_id: &str,
feature_settings: Option<serde_json::Value>,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let result = sqlx::query(SET_USER_API_KEY_FEATURE_SETTINGS_SQL)
.bind(user_id)
.bind(api_key_id)
.bind(feature_settings)
.bind(true)
.execute(&self.pool)
.await
.map_postgres_err()?;
@@ -1464,10 +1814,15 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
if !total_cost_usd.is_finite() {
return Err(DataLayerError::InvalidInput(
"api_keys.total_cost_usd is not finite".to_string(),
));
}
let row = sqlx::query(SET_API_KEY_USAGE_TOTALS_SQL)
.bind(api_key_id)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(i64_from_u64(total_requests, "api_keys.total_requests")?)
.bind(i64_from_u64(total_tokens, "api_keys.total_tokens")?)
.bind(total_cost_usd)
.fetch_optional(&self.pool)
.await
@@ -1480,20 +1835,17 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
user_id: &str,
api_key_id: &str,
) -> Result<bool, DataLayerError> {
let mut tx = self.pool.begin().await.map_postgres_err()?;
sqlx::query(DISABLE_WALLET_BY_API_KEY_ID_SQL)
.bind(api_key_id)
.execute(&mut *tx)
self.delete_api_key(api_key_id, Some(user_id), false, false)
.await
.map_postgres_err()?;
let result = sqlx::query(DELETE_USER_API_KEY_SQL)
.bind(user_id)
.bind(api_key_id)
.execute(&mut *tx)
}
async fn delete_user_api_key_if_unlocked(
&self,
user_id: &str,
api_key_id: &str,
) -> Result<bool, DataLayerError> {
self.delete_api_key(api_key_id, Some(user_id), false, true)
.await
.map_postgres_err()?;
tx.commit().await.map_err(postgres_error)?;
Ok(result.rows_affected() > 0)
}
async fn set_standalone_api_key_feature_settings(
@@ -1519,22 +1871,106 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
}
async fn delete_standalone_api_key(&self, api_key_id: &str) -> Result<bool, DataLayerError> {
let mut tx = self.pool.begin().await.map_postgres_err()?;
sqlx::query(DISABLE_WALLET_BY_API_KEY_ID_SQL)
.bind(api_key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
let result = sqlx::query(DELETE_STANDALONE_API_KEY_SQL)
.bind(api_key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
tx.commit().await.map_err(postgres_error)?;
Ok(result.rows_affected() > 0)
self.delete_api_key(api_key_id, None, true, false).await
}
}
impl SqlxAuthApiKeySnapshotReadRepository {
async fn delete_api_key(
&self,
api_key_id: &str,
user_id: Option<&str>,
is_standalone: bool,
require_unlocked: bool,
) -> Result<bool, DataLayerError> {
let mut tx = self.pool.begin().await.map_postgres_err()?;
let matching_api_key = if let Some(user_id) = user_id {
if require_unlocked {
sqlx::query_scalar::<_, String>(
"SELECT id FROM api_keys WHERE id = $1 AND user_id = $2 AND is_standalone IS FALSE AND is_locked IS FALSE FOR UPDATE",
)
.bind(api_key_id)
.bind(user_id)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?
} else {
sqlx::query_scalar::<_, String>(
"SELECT id FROM api_keys WHERE id = $1 AND user_id = $2 AND is_standalone IS FALSE FOR UPDATE",
)
.bind(api_key_id)
.bind(user_id)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?
}
} else {
sqlx::query_scalar::<_, String>(
"SELECT id FROM api_keys WHERE id = $1 AND is_standalone IS TRUE FOR UPDATE",
)
.bind(api_key_id)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?
};
if matching_api_key.is_none() {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
sqlx::query(
"UPDATE wallets SET status = 'disabled', updated_at = NOW() WHERE api_key_id = $1 AND status <> 'disabled'",
)
.bind(api_key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
for sql in POSTGRES_ANONYMIZE_API_KEY_HISTORY_SQL {
sqlx::query(sql)
.bind(api_key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
for sql in POSTGRES_DELETE_API_KEY_DEPENDENTS_SQL {
sqlx::query(sql)
.bind(api_key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
let result = sqlx::query("DELETE FROM api_keys WHERE id = $1 AND is_standalone = $2")
.bind(api_key_id)
.bind(is_standalone)
.execute(&mut *tx)
.await
.map_postgres_err()?;
if result.rows_affected() != 1 {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
tx.commit().await.map_err(postgres_error)?;
Ok(true)
}
}
fn i64_from_u64(value: u64, field_name: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}")))
}
fn datetime_from_unix_secs(
value: u64,
field_name: &str,
) -> Result<chrono::DateTime<chrono::Utc>, DataLayerError> {
let unix_secs = i64_from_u64(value, field_name)?;
chrono::DateTime::<chrono::Utc>::from_timestamp(unix_secs, 0).ok_or_else(|| {
DataLayerError::InvalidInput(format!(
"{field_name} is outside the supported timestamp range: {value}"
))
})
}
fn row_get<T>(row: &sqlx::postgres::PgRow, column: &str) -> Result<T, DataLayerError>
where
for<'r> T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
@@ -1610,51 +2046,107 @@ fn map_auth_api_key_export_row(
#[cfg(test)]
mod tests {
use super::{
datetime_from_unix_secs, i64_from_u64, DataLayerError,
SqlxAuthApiKeySnapshotReadRepository, CREATE_STANDALONE_API_KEY_SQL,
CREATE_USER_API_KEY_SQL, UPDATE_STANDALONE_API_KEY_BASIC_SQL,
CREATE_USER_API_KEY_SQL, POSTGRES_ANONYMIZE_API_KEY_HISTORY_SQL,
POSTGRES_DELETE_API_KEY_DEPENDENTS_SQL, UPDATE_STANDALONE_API_KEY_BASIC_SQL,
UPDATE_USER_API_KEY_BASIC_SQL,
};
use crate::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
fn checked_u64_conversions_reject_counter_and_timestamp_overflow() {
assert_eq!(
i64_from_u64(i64::MAX as u64, "api_keys.total_requests")
.expect("i64 maximum should fit"),
i64::MAX
);
assert!(matches!(
i64_from_u64(u64::MAX, "api_keys.total_requests"),
Err(DataLayerError::InvalidInput(_))
));
assert!(matches!(
datetime_from_unix_secs(i64::MAX as u64, "api_keys.expires_at"),
Err(DataLayerError::InvalidInput(_))
));
}
#[test]
fn create_api_key_sql_orders_expiry_before_standalone_flags() {
assert!(CREATE_USER_API_KEY_SQL
.contains("expires_at,\n auto_delete_on_expiry,\n is_locked,\n is_standalone,"));
assert!(
CREATE_USER_API_KEY_SQL.contains("$13,\n $14,\n $15,\n FALSE,\n FALSE,\n $16,")
);
assert!(CREATE_USER_API_KEY_SQL
.contains("$13,\n $14,\n $15,\n $16,\n FALSE,\n FALSE,\n $17,"));
assert!(CREATE_STANDALONE_API_KEY_SQL
.contains("expires_at,\n auto_delete_on_expiry,\n is_locked,\n is_standalone,"));
assert!(CREATE_STANDALONE_API_KEY_SQL
.contains("$13,\n $14,\n $15,\n FALSE,\n TRUE,\n $16,"));
}
#[test]
fn api_key_delete_sql_preserves_ids_and_removes_private_snapshots() {
for table in [
"request_candidates",
"video_tasks",
"usage",
"stats_daily_api_key",
] {
assert!(POSTGRES_ANONYMIZE_API_KEY_HISTORY_SQL.iter().any(|sql| {
sql.starts_with(&format!("UPDATE {table} "))
&& sql.contains("SET api_key_name = NULL")
&& sql.ends_with("WHERE api_key_id = $1")
}));
}
assert!(POSTGRES_ANONYMIZE_API_KEY_HISTORY_SQL
.iter()
.any(|sql| sql
.starts_with("UPDATE audit_logs SET description = 'deleted API key event'")));
assert!(POSTGRES_ANONYMIZE_API_KEY_HISTORY_SQL.iter().any(|sql| sql
.starts_with("UPDATE payment_callbacks SET payload = NULL, error_message = NULL")));
assert_eq!(
POSTGRES_DELETE_API_KEY_DEPENDENTS_SQL,
&["DELETE FROM api_key_provider_mappings WHERE api_key_id = $1"]
);
}
#[test]
fn update_standalone_api_key_basic_sql_casts_json_case_values() {
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("concurrent_limit = CASE WHEN $5 THEN $6 ELSE concurrent_limit END"));
.contains("key_encrypted = CASE WHEN $2 THEN $3 ELSE key_encrypted END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END"));
.contains("concurrent_limit = CASE WHEN $10 THEN $11 ELSE concurrent_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END"
"allowed_providers = CASE WHEN $12 THEN $13::json ELSE allowed_providers END"
));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"allowed_api_formats = CASE WHEN $14 THEN $15::json ELSE allowed_api_formats END"
));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END"));
.contains("allowed_models = CASE WHEN $16 THEN $17::json ELSE allowed_models END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("ip_rules = CASE WHEN $13 THEN $14::jsonb ELSE ip_rules END"));
.contains("ip_rules = CASE WHEN $18 THEN $19::jsonb ELSE ip_rules END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END"));
.contains("rate_limit = CASE WHEN $8 THEN $9 ELSE rate_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END"));
.contains("expires_at = CASE WHEN $20 THEN $21::timestamptz ELSE expires_at END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END"
"auto_delete_on_expiry = CASE WHEN $22 THEN $23 ELSE auto_delete_on_expiry END"
));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"force_capabilities = CASE WHEN $6 THEN $7::json ELSE force_capabilities END"
));
}
#[test]
fn update_user_api_key_basic_sql_casts_ip_rules_as_jsonb() {
fn update_user_api_key_basic_sql_casts_json_patches_and_fences_locked_keys() {
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("ip_rules = CASE WHEN $6 THEN $7::jsonb ELSE ip_rules END"));
.contains("key_encrypted = CASE WHEN $3 THEN $4 ELSE key_encrypted END"));
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("ip_rules = CASE WHEN $11 THEN $12::jsonb ELSE ip_rules END"));
assert!(UPDATE_USER_API_KEY_BASIC_SQL.contains(
"feature_settings = CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END"
));
assert!(UPDATE_USER_API_KEY_BASIC_SQL.contains("AND ($15 = FALSE OR is_locked = FALSE)"));
}
#[tokio::test]