mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: 扩展 Rust gateway 全功能模块,新增 billing/crypto/wallet crate 及完整数据层
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate - aether-data 扩展 repository 层:announcements、auth_modules、billing、 candidate_selection、gemini_file_mappings、global_models、management_tokens、 oauth_providers、proxy_nodes、quota、users、wallet 等模块 - aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/ video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块 - 重构 executor decision 和 gateway state 为模块目录结构 - 新增 gateway router、frontdoor 路由层及对应测试 - Python 侧 API 路由重构,新增 compat/support 模块 - 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
@@ -3,13 +3,20 @@ use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
|
||||
use super::types::{
|
||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||
AuthApiKeyWriteRepository, CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord,
|
||||
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryAuthApiKeyIndex {
|
||||
by_api_key_id: BTreeMap<String, StoredAuthApiKeySnapshot>,
|
||||
export_by_api_key_id: BTreeMap<String, StoredAuthApiKeyExportRecord>,
|
||||
by_key_hash: BTreeMap<String, String>,
|
||||
touch_counts: BTreeMap<String, usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -23,8 +30,46 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
I: IntoIterator<Item = (Option<String>, StoredAuthApiKeySnapshot)>,
|
||||
{
|
||||
let mut by_api_key_id = BTreeMap::new();
|
||||
let mut export_by_api_key_id = BTreeMap::new();
|
||||
let mut by_key_hash = BTreeMap::new();
|
||||
for (key_hash, snapshot) in items {
|
||||
let derived_key_hash = key_hash
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("memory-{}", snapshot.api_key_id));
|
||||
export_by_api_key_id.insert(
|
||||
snapshot.api_key_id.clone(),
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
snapshot.user_id.clone(),
|
||||
snapshot.api_key_id.clone(),
|
||||
derived_key_hash.clone(),
|
||||
None,
|
||||
snapshot.api_key_name.clone(),
|
||||
snapshot
|
||||
.api_key_allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
snapshot
|
||||
.api_key_allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
snapshot
|
||||
.api_key_allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
snapshot.api_key_rate_limit,
|
||||
snapshot.api_key_concurrent_limit,
|
||||
None,
|
||||
snapshot.api_key_is_active,
|
||||
snapshot
|
||||
.api_key_expires_at_unix_secs
|
||||
.map(|value| value as i64),
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
snapshot.api_key_is_standalone,
|
||||
)
|
||||
.expect("derived auth api key export record should build"),
|
||||
);
|
||||
if let Some(key_hash) = key_hash {
|
||||
by_key_hash.insert(key_hash, snapshot.api_key_id.clone());
|
||||
}
|
||||
@@ -33,10 +78,38 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
Self {
|
||||
index: RwLock::new(MemoryAuthApiKeyIndex {
|
||||
by_api_key_id,
|
||||
export_by_api_key_id,
|
||||
by_key_hash,
|
||||
touch_counts: BTreeMap::new(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_export_records<I>(mut self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAuthApiKeyExportRecord>,
|
||||
{
|
||||
let index = self
|
||||
.index
|
||||
.get_mut()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
for item in items {
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.insert(item.api_key_id.clone(), item);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn touch_count(&self, api_key_id: &str) -> usize {
|
||||
self.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock")
|
||||
.touch_counts
|
||||
.get(api_key_id)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -68,13 +141,698 @@ impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(api_key_ids
|
||||
.iter()
|
||||
.filter_map(|api_key_id| index.by_api_key_id.get(api_key_id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| {
|
||||
!record.is_standalone && user_ids.iter().any(|id| id == &record.user_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_api_keys_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(api_key_ids
|
||||
.iter()
|
||||
.filter_map(|api_key_id| index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_export_standalone_api_keys_page(
|
||||
&self,
|
||||
query: &StandaloneApiKeyExportListQuery,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| {
|
||||
record.is_standalone
|
||||
&& query
|
||||
.is_active
|
||||
.is_none_or(|is_active| record.is_active == is_active)
|
||||
})
|
||||
.skip(query.skip)
|
||||
.take(query.limit)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn count_export_standalone_api_keys(
|
||||
&self,
|
||||
is_active: Option<bool>,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| {
|
||||
record.is_standalone
|
||||
&& is_active.is_none_or(|expected| record.is_active == expected)
|
||||
})
|
||||
.count() as u64)
|
||||
}
|
||||
|
||||
async fn summarize_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let mut summary = AuthApiKeyExportSummary::default();
|
||||
for record in index.export_by_api_key_id.values().filter(|record| {
|
||||
!record.is_standalone && user_ids.iter().any(|id| id == &record.user_id)
|
||||
}) {
|
||||
summary.total = summary.total.saturating_add(1);
|
||||
if record.is_active
|
||||
&& record
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|expires_at_unix_secs| expires_at_unix_secs >= now_unix_secs)
|
||||
{
|
||||
summary.active = summary.active.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn summarize_export_non_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let mut summary = AuthApiKeyExportSummary::default();
|
||||
for record in index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| !record.is_standalone)
|
||||
{
|
||||
summary.total = summary.total.saturating_add(1);
|
||||
if record.is_active
|
||||
&& record
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|expires_at_unix_secs| expires_at_unix_secs >= now_unix_secs)
|
||||
{
|
||||
summary.active = summary.active.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn find_export_standalone_api_key_by_id(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.get(api_key_id)
|
||||
.filter(|record| record.is_standalone)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn summarize_export_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let mut summary = AuthApiKeyExportSummary::default();
|
||||
for record in index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| record.is_standalone)
|
||||
{
|
||||
summary.total = summary.total.saturating_add(1);
|
||||
if record.is_active
|
||||
&& record
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|expires_at_unix_secs| expires_at_unix_secs >= now_unix_secs)
|
||||
{
|
||||
summary.active = summary.active.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn list_export_standalone_api_keys(
|
||||
&self,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(index
|
||||
.export_by_api_key_id
|
||||
.values()
|
||||
.filter(|record| record.is_standalone)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
async fn touch_last_used_at(&self, api_key_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
if !index.by_api_key_id.contains_key(api_key_id) {
|
||||
return Ok(false);
|
||||
}
|
||||
let counter = index
|
||||
.touch_counts
|
||||
.entry(api_key_id.to_string())
|
||||
.or_insert(0);
|
||||
*counter += 1;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn create_user_api_key(
|
||||
&self,
|
||||
record: CreateUserApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
if index.by_api_key_id.contains_key(&record.api_key_id) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.id: {}",
|
||||
record.api_key_id
|
||||
)));
|
||||
}
|
||||
if index.by_key_hash.contains_key(&record.key_hash) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.key_hash: {}",
|
||||
record.key_hash
|
||||
)));
|
||||
}
|
||||
|
||||
let template = index
|
||||
.by_api_key_id
|
||||
.values()
|
||||
.find(|snapshot| snapshot.user_id == record.user_id)
|
||||
.cloned();
|
||||
let snapshot = if let Some(template) = template {
|
||||
StoredAuthApiKeySnapshot {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.clone(),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: Some(record.rate_limit),
|
||||
api_key_concurrent_limit: Some(record.concurrent_limit),
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
..template
|
||||
}
|
||||
} else {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
record.user_id.clone(),
|
||||
format!(
|
||||
"user-{}",
|
||||
&record.user_id.chars().take(8).collect::<String>()
|
||||
),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
record.api_key_id.clone(),
|
||||
record.name.clone(),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)?
|
||||
};
|
||||
|
||||
let export = StoredAuthApiKeyExportRecord::new(
|
||||
record.user_id.clone(),
|
||||
record.api_key_id.clone(),
|
||||
record.key_hash.clone(),
|
||||
record.key_encrypted,
|
||||
record.name,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
false,
|
||||
)?;
|
||||
|
||||
index
|
||||
.by_key_hash
|
||||
.insert(record.key_hash, record.api_key_id.clone());
|
||||
index
|
||||
.by_api_key_id
|
||||
.insert(record.api_key_id.clone(), snapshot);
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.insert(record.api_key_id, export.clone());
|
||||
Ok(Some(export))
|
||||
}
|
||||
|
||||
async fn create_standalone_api_key(
|
||||
&self,
|
||||
record: CreateStandaloneApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
if index.by_api_key_id.contains_key(&record.api_key_id) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.id: {}",
|
||||
record.api_key_id
|
||||
)));
|
||||
}
|
||||
if index.by_key_hash.contains_key(&record.key_hash) {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"duplicate api_keys.key_hash: {}",
|
||||
record.key_hash
|
||||
)));
|
||||
}
|
||||
|
||||
let template = index
|
||||
.by_api_key_id
|
||||
.values()
|
||||
.find(|snapshot| snapshot.user_id == record.user_id)
|
||||
.cloned();
|
||||
let snapshot = if let Some(template) = template {
|
||||
StoredAuthApiKeySnapshot {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.clone(),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: true,
|
||||
api_key_rate_limit: Some(record.rate_limit),
|
||||
api_key_concurrent_limit: Some(record.concurrent_limit),
|
||||
api_key_expires_at_unix_secs: None,
|
||||
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(),
|
||||
..template
|
||||
}
|
||||
} else {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
record.user_id.clone(),
|
||||
format!(
|
||||
"admin-{}",
|
||||
&record.user_id.chars().take(8).collect::<String>()
|
||||
),
|
||||
None,
|
||||
"admin".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
record.api_key_id.clone(),
|
||||
record.name.clone(),
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
)?
|
||||
};
|
||||
|
||||
let export = StoredAuthApiKeyExportRecord::new(
|
||||
record.user_id.clone(),
|
||||
record.api_key_id.clone(),
|
||||
record.key_hash.clone(),
|
||||
record.key_encrypted,
|
||||
record.name,
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
true,
|
||||
)?;
|
||||
|
||||
index
|
||||
.by_key_hash
|
||||
.insert(record.key_hash, record.api_key_id.clone());
|
||||
index
|
||||
.by_api_key_id
|
||||
.insert(record.api_key_id.clone(), snapshot);
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.insert(record.api_key_id, export.clone());
|
||||
Ok(Some(export))
|
||||
}
|
||||
|
||||
async fn update_user_api_key_basic(
|
||||
&self,
|
||||
record: UpdateUserApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(&record.api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != record.user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(name) = record.name {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_name = Some(name.clone());
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.name = Some(name);
|
||||
}
|
||||
}
|
||||
if let Some(rate_limit) = record.rate_limit {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_rate_limit = Some(rate_limit);
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.rate_limit = Some(rate_limit);
|
||||
}
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn update_standalone_api_key_basic(
|
||||
&self,
|
||||
record: UpdateStandaloneApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(&record.api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(name) = record.name {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_name = Some(name.clone());
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.name = Some(name);
|
||||
}
|
||||
}
|
||||
if let Some(rate_limit) = record.rate_limit {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_rate_limit = Some(rate_limit);
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.rate_limit = Some(rate_limit);
|
||||
}
|
||||
}
|
||||
if let Some(allowed_providers) = record.allowed_providers {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_allowed_providers = allowed_providers.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.allowed_providers = allowed_providers;
|
||||
}
|
||||
}
|
||||
if let Some(allowed_api_formats) = record.allowed_api_formats {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_allowed_api_formats = allowed_api_formats.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.allowed_api_formats = allowed_api_formats;
|
||||
}
|
||||
}
|
||||
if let Some(allowed_models) = record.allowed_models {
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
|
||||
snapshot.api_key_allowed_models = allowed_models.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
|
||||
export.allowed_models = allowed_models;
|
||||
}
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_user_api_key_active(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_is_active = is_active;
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) {
|
||||
export.is_active = is_active;
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_standalone_api_key_active(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_is_active = is_active;
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) {
|
||||
export.is_active = is_active;
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_user_api_key_locked(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_locked: bool,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(false);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_is_locked = is_locked;
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn set_user_api_key_allowed_providers(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
allowed_providers: Option<Vec<String>>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Some(snapshot) = index.by_api_key_id.get_mut(api_key_id) {
|
||||
snapshot.api_key_allowed_providers = allowed_providers.clone();
|
||||
}
|
||||
if let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) {
|
||||
export.allowed_providers = allowed_providers;
|
||||
}
|
||||
Ok(index.export_by_api_key_id.get(api_key_id).cloned())
|
||||
}
|
||||
|
||||
async fn set_user_api_key_force_capabilities(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
export.force_capabilities = force_capabilities;
|
||||
Ok(Some(export.clone()))
|
||||
}
|
||||
|
||||
async fn delete_user_api_key(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if snapshot.user_id != user_id || snapshot.api_key_is_standalone {
|
||||
return Ok(false);
|
||||
}
|
||||
index.by_api_key_id.remove(api_key_id);
|
||||
index.export_by_api_key_id.remove(api_key_id);
|
||||
index.by_key_hash.retain(|_, value| value != api_key_id);
|
||||
index.touch_counts.remove(api_key_id);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn delete_standalone_api_key(&self, api_key_id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
let Some(snapshot) = index.by_api_key_id.get(api_key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if !snapshot.api_key_is_standalone {
|
||||
return Ok(false);
|
||||
}
|
||||
index.by_api_key_id.remove(api_key_id);
|
||||
index.export_by_api_key_id.remove(api_key_id);
|
||||
index.by_key_hash.retain(|_, value| value != api_key_id);
|
||||
index.touch_counts.remove(api_key_id);
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryAuthApiKeySnapshotRepository;
|
||||
use crate::repository::auth::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository,
|
||||
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
@@ -129,5 +887,132 @@ mod tests {
|
||||
.await
|
||||
.expect("find by user/api key ids should succeed")
|
||||
.is_some());
|
||||
let snapshots = repository
|
||||
.list_api_key_snapshots_by_ids(&["key-1".to_string(), "missing".to_string()])
|
||||
.await
|
||||
.expect("batch lookup should succeed");
|
||||
assert_eq!(snapshots.len(), 1);
|
||||
assert_eq!(snapshots[0].api_key_id, "key-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn touches_last_used_for_existing_key() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_snapshot("key-1", "user-1"),
|
||||
)]);
|
||||
|
||||
assert!(repository
|
||||
.touch_last_used_at("key-1")
|
||||
.await
|
||||
.expect("touch should succeed"));
|
||||
assert_eq!(repository.touch_count("key-1"), 1);
|
||||
assert!(!repository
|
||||
.touch_last_used_at("missing")
|
||||
.await
|
||||
.expect("missing touch should succeed"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_export_records_for_user_bound_and_standalone_keys() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![
|
||||
(
|
||||
Some("hash-user".to_string()),
|
||||
sample_snapshot("key-user", "user-1"),
|
||||
),
|
||||
(
|
||||
Some("hash-standalone".to_string()),
|
||||
sample_snapshot("key-standalone", "admin-1"),
|
||||
),
|
||||
])
|
||||
.with_export_records(vec![
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
"user-1".to_string(),
|
||||
"key-user".to_string(),
|
||||
"hash-user".to_string(),
|
||||
Some("enc-user".to_string()),
|
||||
Some("default".to_string()),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
Some(120),
|
||||
Some(7),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
Some(200),
|
||||
false,
|
||||
14,
|
||||
1.5,
|
||||
false,
|
||||
)
|
||||
.expect("user export record should build"),
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
"admin-1".to_string(),
|
||||
"key-standalone".to_string(),
|
||||
"hash-standalone".to_string(),
|
||||
Some("enc-standalone".to_string()),
|
||||
Some("standalone".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(1),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
true,
|
||||
2,
|
||||
0.25,
|
||||
true,
|
||||
)
|
||||
.expect("standalone export record should build"),
|
||||
]);
|
||||
|
||||
let user_records = repository
|
||||
.list_export_api_keys_by_user_ids(&["user-1".to_string()])
|
||||
.await
|
||||
.expect("user export lookup should succeed");
|
||||
assert_eq!(user_records.len(), 1);
|
||||
assert_eq!(user_records[0].api_key_id, "key-user");
|
||||
assert_eq!(user_records[0].key_encrypted.as_deref(), Some("enc-user"));
|
||||
assert_eq!(user_records[0].total_requests, 14);
|
||||
|
||||
let standalone_records = repository
|
||||
.list_export_standalone_api_keys()
|
||||
.await
|
||||
.expect("standalone export lookup should succeed");
|
||||
assert_eq!(standalone_records.len(), 1);
|
||||
assert_eq!(standalone_records[0].api_key_id, "key-standalone");
|
||||
assert!(standalone_records[0].is_standalone);
|
||||
|
||||
let selected_records = repository
|
||||
.list_export_api_keys_by_ids(&[
|
||||
"key-standalone".to_string(),
|
||||
"missing".to_string(),
|
||||
"key-user".to_string(),
|
||||
])
|
||||
.await
|
||||
.expect("api key id export lookup should succeed");
|
||||
assert_eq!(selected_records.len(), 2);
|
||||
assert_eq!(selected_records[0].api_key_id, "key-standalone");
|
||||
assert_eq!(selected_records[1].api_key_id, "key-user");
|
||||
|
||||
let paged_records = repository
|
||||
.list_export_standalone_api_keys_page(&StandaloneApiKeyExportListQuery {
|
||||
skip: 0,
|
||||
limit: 10,
|
||||
is_active: Some(true),
|
||||
})
|
||||
.await
|
||||
.expect("standalone export page should succeed");
|
||||
assert_eq!(paged_records.len(), 1);
|
||||
assert_eq!(paged_records[0].api_key_id, "key-standalone");
|
||||
assert_eq!(
|
||||
repository
|
||||
.count_export_standalone_api_keys(Some(true))
|
||||
.await
|
||||
.expect("standalone export count should succeed"),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,5 +5,8 @@ mod types;
|
||||
pub use memory::InMemoryAuthApiKeySnapshotRepository;
|
||||
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
|
||||
pub use types::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthRepository, StoredAuthApiKeySnapshot,
|
||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||
AuthApiKeyWriteRepository, AuthRepository, CreateStandaloneApiKeyRecord,
|
||||
CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord,
|
||||
StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
|
||||
};
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -9,6 +9,7 @@ pub struct StoredAuthApiKeySnapshot {
|
||||
pub user_auth_source: String,
|
||||
pub user_is_active: bool,
|
||||
pub user_is_deleted: bool,
|
||||
pub user_rate_limit: Option<i32>,
|
||||
pub user_allowed_providers: Option<Vec<String>>,
|
||||
pub user_allowed_api_formats: Option<Vec<String>>,
|
||||
pub user_allowed_models: Option<Vec<String>>,
|
||||
@@ -58,6 +59,7 @@ impl StoredAuthApiKeySnapshot {
|
||||
user_auth_source,
|
||||
user_is_active,
|
||||
user_is_deleted,
|
||||
user_rate_limit: None,
|
||||
user_allowed_providers: parse_string_list(
|
||||
user_allowed_providers,
|
||||
"users.allowed_providers",
|
||||
@@ -115,6 +117,157 @@ impl StoredAuthApiKeySnapshot {
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
pub fn with_user_rate_limit(mut self, user_rate_limit: Option<i32>) -> Self {
|
||||
self.user_rate_limit = user_rate_limit;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAuthApiKeyExportRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: Option<i32>,
|
||||
pub concurrent_limit: Option<i32>,
|
||||
pub force_capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub auto_delete_on_expiry: bool,
|
||||
pub total_requests: u64,
|
||||
pub total_cost_usd: f64,
|
||||
pub is_standalone: bool,
|
||||
}
|
||||
|
||||
impl StoredAuthApiKeyExportRecord {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
user_id: String,
|
||||
api_key_id: String,
|
||||
key_hash: String,
|
||||
key_encrypted: Option<String>,
|
||||
name: Option<String>,
|
||||
allowed_providers: Option<serde_json::Value>,
|
||||
allowed_api_formats: Option<serde_json::Value>,
|
||||
allowed_models: Option<serde_json::Value>,
|
||||
rate_limit: Option<i32>,
|
||||
concurrent_limit: Option<i32>,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
is_active: bool,
|
||||
expires_at_unix_secs: Option<i64>,
|
||||
auto_delete_on_expiry: bool,
|
||||
total_requests: i64,
|
||||
total_cost_usd: f64,
|
||||
is_standalone: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if user_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.user_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if api_key_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if key_hash.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.key_hash is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"api_keys.total_cost_usd is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
user_id,
|
||||
api_key_id,
|
||||
key_hash,
|
||||
key_encrypted,
|
||||
name,
|
||||
allowed_providers: parse_string_list(allowed_providers, "api_keys.allowed_providers")?,
|
||||
allowed_api_formats: parse_string_list(
|
||||
allowed_api_formats,
|
||||
"api_keys.allowed_api_formats",
|
||||
)?,
|
||||
allowed_models: parse_string_list(allowed_models, "api_keys.allowed_models")?,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities,
|
||||
is_active,
|
||||
expires_at_unix_secs: expires_at_unix_secs
|
||||
.map(|value| parse_u64_i64(value, "api_keys.expires_at_unix_secs"))
|
||||
.transpose()?,
|
||||
auto_delete_on_expiry,
|
||||
total_requests: parse_u64_i64(total_requests, "api_keys.total_requests")?,
|
||||
total_cost_usd,
|
||||
is_standalone,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct AuthApiKeyExportSummary {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct StandaloneApiKeyExportListQuery {
|
||||
pub skip: usize,
|
||||
pub limit: usize,
|
||||
pub is_active: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateUserApiKeyRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub rate_limit: i32,
|
||||
pub concurrent_limit: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateUserApiKeyBasicRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub name: Option<String>,
|
||||
pub rate_limit: Option<i32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateStandaloneApiKeyRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: i32,
|
||||
pub concurrent_limit: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateStandaloneApiKeyBasicRecord {
|
||||
pub api_key_id: String,
|
||||
pub name: Option<String>,
|
||||
pub rate_limit: Option<i32>,
|
||||
pub allowed_providers: Option<Option<Vec<String>>>,
|
||||
pub allowed_api_formats: Option<Option<Vec<String>>>,
|
||||
pub allowed_models: Option<Option<Vec<String>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -133,11 +286,137 @@ pub trait AuthApiKeyReadRepository: Send + Sync {
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
|
||||
|
||||
async fn list_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_api_keys_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_standalone_api_keys_page(
|
||||
&self,
|
||||
query: &StandaloneApiKeyExportListQuery,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn count_export_standalone_api_keys(
|
||||
&self,
|
||||
is_active: Option<bool>,
|
||||
) -> Result<u64, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_non_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_export_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, crate::DataLayerError>;
|
||||
|
||||
async fn find_export_standalone_api_key_by_id(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn list_export_standalone_api_keys(
|
||||
&self,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait AuthRepository: AuthApiKeyReadRepository + Send + Sync {}
|
||||
#[async_trait]
|
||||
pub trait AuthApiKeyWriteRepository: Send + Sync {
|
||||
async fn touch_last_used_at(&self, api_key_id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
impl<T> AuthRepository for T where T: AuthApiKeyReadRepository + Send + Sync {}
|
||||
async fn create_user_api_key(
|
||||
&self,
|
||||
record: CreateUserApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn create_standalone_api_key(
|
||||
&self,
|
||||
record: CreateStandaloneApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn update_user_api_key_basic(
|
||||
&self,
|
||||
record: UpdateUserApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn update_standalone_api_key_basic(
|
||||
&self,
|
||||
record: UpdateStandaloneApiKeyBasicRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_active(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_standalone_api_key_active(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_locked(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_locked: bool,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_allowed_providers(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
allowed_providers: Option<Vec<String>>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn set_user_api_key_force_capabilities(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_user_api_key(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn delete_standalone_api_key(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait AuthRepository:
|
||||
AuthApiKeyReadRepository + AuthApiKeyWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> AuthRepository for T where
|
||||
T: AuthApiKeyReadRepository + AuthApiKeyWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
fn parse_string_list(
|
||||
value: Option<serde_json::Value>,
|
||||
@@ -146,9 +425,43 @@ fn parse_string_list(
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let array = value.as_array().ok_or_else(|| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("{field_name} is not a JSON array"))
|
||||
})?;
|
||||
parse_string_list_value(&value, field_name)
|
||||
}
|
||||
|
||||
fn parse_string_list_value(
|
||||
value: &serde_json::Value,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
serde_json::Value::Array(array) => parse_string_list_array(array, field_name).map(Some),
|
||||
serde_json::Value::String(raw) => parse_embedded_string_list(raw, field_name),
|
||||
_ => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} is not a JSON array"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_embedded_string_list(
|
||||
raw: &str,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
|
||||
return parse_string_list_value(&decoded, field_name);
|
||||
}
|
||||
|
||||
Ok(Some(vec![raw.to_string()]))
|
||||
}
|
||||
|
||||
fn parse_string_list_array(
|
||||
array: &[serde_json::Value],
|
||||
field_name: &str,
|
||||
) -> Result<Vec<String>, crate::DataLayerError> {
|
||||
let mut items = Vec::with_capacity(array.len());
|
||||
for item in array {
|
||||
let Some(item) = item.as_str() else {
|
||||
@@ -156,14 +469,23 @@ fn parse_string_list(
|
||||
"{field_name} contains a non-string item"
|
||||
)));
|
||||
};
|
||||
items.push(item.to_string());
|
||||
let item = item.trim();
|
||||
if !item.is_empty() {
|
||||
items.push(item.to_string());
|
||||
}
|
||||
}
|
||||
Ok(Some(items))
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
fn parse_u64_i64(value: i64, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredAuthApiKeySnapshot;
|
||||
use super::{StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot};
|
||||
|
||||
#[test]
|
||||
fn rejects_non_array_allowed_providers() {
|
||||
@@ -193,6 +515,72 @@ mod tests {
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_stringified_allowed_provider_array() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!("[\"openai\", \" gemini \"]")),
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("snapshot should build");
|
||||
|
||||
assert_eq!(
|
||||
snapshot.user_allowed_providers,
|
||||
Some(vec!["openai".to_string(), "gemini".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_single_string_allowed_provider() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!("openai")),
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("snapshot should build");
|
||||
|
||||
assert_eq!(
|
||||
snapshot.user_allowed_providers,
|
||||
Some(vec!["openai".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_non_standalone_key_is_not_usable() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
@@ -222,4 +610,59 @@ mod tests {
|
||||
|
||||
assert!(!snapshot.is_currently_usable(101));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_record_rejects_negative_totals() {
|
||||
assert!(StoredAuthApiKeyExportRecord::new(
|
||||
"user-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"hash-1".to_string(),
|
||||
Some("enc".to_string()),
|
||||
Some("default".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
-1,
|
||||
0.0,
|
||||
false,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn export_record_accepts_stringified_allowed_models() {
|
||||
let record = StoredAuthApiKeyExportRecord::new(
|
||||
"user-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"hash-1".to_string(),
|
||||
Some("enc".to_string()),
|
||||
Some("default".to_string()),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!("[\"gpt-5\", \" gpt-4.1 \"]")),
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
Some(200),
|
||||
false,
|
||||
12,
|
||||
1.25,
|
||||
false,
|
||||
)
|
||||
.expect("export record should build");
|
||||
|
||||
assert_eq!(
|
||||
record.allowed_models,
|
||||
Some(vec!["gpt-5".to_string(), "gpt-4.1".to_string()])
|
||||
);
|
||||
assert_eq!(record.total_requests, 12);
|
||||
assert_eq!(record.total_cost_usd, 1.25);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user