feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate,完善并发门控与审计系统

新增 crate:
- aether-runtime: 服务运行时基础设施(并发门控、分布式并发、指标、队列、优雅关闭、tracing)
- aether-cache: 通用 TTL 缓存与命名空间抽象
- aether-data: 数据访问层(PostgreSQL/Redis 后端、repository 模式)
- aether-http: HTTP 客户端封装(重试、配置)
- aether-testkit: 集成测试工具集(gateway/executor/hub/proxy fixture、等待、负载测试)

gateway 扩展:
- 引入 audit 模块(shadow 执行审计、决策链路追踪、请求审计 bundle)
- 引入 cache 模块(AuthContext 缓存、direct-plan bypass 缓存)
- 引入 data 模块(auth/candidates/config/usage/video_tasks 数据访问)
- 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控
- 新增本地 auth 拒绝、过载响应构建器
- 补充 control/auth_cache/video/concurrency 集成测试

aether-proxy 扩展:
- AppState 集成 stream_gate / distributed_stream_gate 并发门控
- 新增 ProxyAdmissionError 及准入拒绝流程
- stream_handler 补充门控饱和/不可用场景测试
- 配置与注册客户端逻辑完善

aether-hub 扩展:
- main.rs 引入运行时初始化、指标端点、健康检查
- local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
fawney19
2026-03-24 15:12:56 +08:00
parent eaf8475f9e
commit b5a0070023
157 changed files with 22097 additions and 448 deletions

View File

@@ -0,0 +1,133 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
use crate::DataLayerError;
#[derive(Debug, Default)]
struct MemoryAuthApiKeyIndex {
by_api_key_id: BTreeMap<String, StoredAuthApiKeySnapshot>,
by_key_hash: BTreeMap<String, String>,
}
#[derive(Debug, Default)]
pub struct InMemoryAuthApiKeySnapshotRepository {
index: RwLock<MemoryAuthApiKeyIndex>,
}
impl InMemoryAuthApiKeySnapshotRepository {
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = (Option<String>, StoredAuthApiKeySnapshot)>,
{
let mut by_api_key_id = BTreeMap::new();
let mut by_key_hash = BTreeMap::new();
for (key_hash, snapshot) in items {
if let Some(key_hash) = key_hash {
by_key_hash.insert(key_hash, snapshot.api_key_id.clone());
}
by_api_key_id.insert(snapshot.api_key_id.clone(), snapshot);
}
Self {
index: RwLock::new(MemoryAuthApiKeyIndex {
by_api_key_id,
by_key_hash,
}),
}
}
}
#[async_trait]
impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
let index = self
.index
.read()
.expect("auth api key snapshot repository lock");
Ok(match key {
AuthApiKeyLookupKey::KeyHash(key_hash) => index
.by_key_hash
.get(key_hash)
.and_then(|api_key_id| index.by_api_key_id.get(api_key_id))
.cloned(),
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
index.by_api_key_id.get(api_key_id).cloned()
}
AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
} => index
.by_api_key_id
.get(api_key_id)
.filter(|snapshot| snapshot.user_id == user_id)
.cloned(),
})
}
}
#[cfg(test)]
mod tests {
use super::InMemoryAuthApiKeySnapshotRepository;
use crate::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
};
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("snapshot should build")
}
#[tokio::test]
async fn reads_auth_snapshot_by_all_supported_keys() {
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_snapshot("key-1", "user-1"),
)]);
assert!(repository
.find_api_key_snapshot(AuthApiKeyLookupKey::KeyHash("hash-1"))
.await
.expect("find by hash should succeed")
.is_some());
assert!(repository
.find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-1"))
.await
.expect("find by api key id should succeed")
.is_some());
assert!(repository
.find_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
user_id: "user-1",
api_key_id: "key-1",
})
.await
.expect("find by user/api key ids should succeed")
.is_some());
}
}

View File

@@ -0,0 +1,9 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryAuthApiKeySnapshotRepository;
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
pub use types::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthRepository, StoredAuthApiKeySnapshot,
};

View File

@@ -0,0 +1,202 @@
use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
use crate::DataLayerError;
const FIND_BY_KEY_HASH_SQL: &str = r#"
SELECT
users.id AS user_id,
users.username,
users.email,
users.role::text AS user_role,
users.auth_source::text AS user_auth_source,
users.is_active AS user_is_active,
users.is_deleted AS user_is_deleted,
users.allowed_providers AS user_allowed_providers,
users.allowed_api_formats AS user_allowed_api_formats,
users.allowed_models AS user_allowed_models,
api_keys.id AS api_key_id,
api_keys.name AS api_key_name,
api_keys.is_active AS api_key_is_active,
api_keys.is_locked AS api_key_is_locked,
api_keys.is_standalone AS api_key_is_standalone,
api_keys.rate_limit AS api_key_rate_limit,
api_keys.concurrent_limit AS api_key_concurrent_limit,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
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
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.key_hash = $1
LIMIT 1
"#;
const FIND_BY_API_KEY_ID_SQL: &str = r#"
SELECT
users.id AS user_id,
users.username,
users.email,
users.role::text AS user_role,
users.auth_source::text AS user_auth_source,
users.is_active AS user_is_active,
users.is_deleted AS user_is_deleted,
users.allowed_providers AS user_allowed_providers,
users.allowed_api_formats AS user_allowed_api_formats,
users.allowed_models AS user_allowed_models,
api_keys.id AS api_key_id,
api_keys.name AS api_key_name,
api_keys.is_active AS api_key_is_active,
api_keys.is_locked AS api_key_is_locked,
api_keys.is_standalone AS api_key_is_standalone,
api_keys.rate_limit AS api_key_rate_limit,
api_keys.concurrent_limit AS api_key_concurrent_limit,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
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
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1
LIMIT 1
"#;
const FIND_BY_USER_API_KEY_IDS_SQL: &str = r#"
SELECT
users.id AS user_id,
users.username,
users.email,
users.role::text AS user_role,
users.auth_source::text AS user_auth_source,
users.is_active AS user_is_active,
users.is_deleted AS user_is_deleted,
users.allowed_providers AS user_allowed_providers,
users.allowed_api_formats AS user_allowed_api_formats,
users.allowed_models AS user_allowed_models,
api_keys.id AS api_key_id,
api_keys.name AS api_key_name,
api_keys.is_active AS api_key_is_active,
api_keys.is_locked AS api_key_is_locked,
api_keys.is_standalone AS api_key_is_standalone,
api_keys.rate_limit AS api_key_rate_limit,
api_keys.concurrent_limit AS api_key_concurrent_limit,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
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
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1 AND users.id = $2
LIMIT 1
"#;
#[derive(Debug, Clone)]
pub struct SqlxAuthApiKeySnapshotReadRepository {
pool: PgPool,
}
impl SqlxAuthApiKeySnapshotReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
let row = match key {
AuthApiKeyLookupKey::KeyHash(key_hash) => {
sqlx::query(FIND_BY_KEY_HASH_SQL)
.bind(key_hash)
.fetch_optional(&self.pool)
.await?
}
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
sqlx::query(FIND_BY_API_KEY_ID_SQL)
.bind(api_key_id)
.fetch_optional(&self.pool)
.await?
}
AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
} => {
sqlx::query(FIND_BY_USER_API_KEY_IDS_SQL)
.bind(api_key_id)
.bind(user_id)
.fetch_optional(&self.pool)
.await?
}
};
row.as_ref().map(map_auth_api_key_snapshot_row).transpose()
}
}
#[async_trait]
impl AuthApiKeyReadRepository for SqlxAuthApiKeySnapshotReadRepository {
async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
Self::find_api_key_snapshot(self, key).await
}
}
fn map_auth_api_key_snapshot_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredAuthApiKeySnapshot, DataLayerError> {
StoredAuthApiKeySnapshot::new(
row.try_get("user_id")?,
row.try_get("username")?,
row.try_get("email")?,
row.try_get("user_role")?,
row.try_get("user_auth_source")?,
row.try_get("user_is_active")?,
row.try_get("user_is_deleted")?,
row.try_get("user_allowed_providers")?,
row.try_get("user_allowed_api_formats")?,
row.try_get("user_allowed_models")?,
row.try_get("api_key_id")?,
row.try_get("api_key_name")?,
row.try_get("api_key_is_active")?,
row.try_get("api_key_is_locked")?,
row.try_get("api_key_is_standalone")?,
row.try_get("api_key_rate_limit")?,
row.try_get("api_key_concurrent_limit")?,
row.try_get("api_key_expires_at_unix_secs")?,
row.try_get("api_key_allowed_providers")?,
row.try_get("api_key_allowed_api_formats")?,
row.try_get("api_key_allowed_models")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxAuthApiKeySnapshotReadRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxAuthApiKeySnapshotReadRepository::new(pool);
let _ = repository.pool();
}
}

View File

@@ -0,0 +1,225 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredAuthApiKeySnapshot {
pub user_id: String,
pub username: String,
pub email: Option<String>,
pub user_role: String,
pub user_auth_source: String,
pub user_is_active: bool,
pub user_is_deleted: bool,
pub user_allowed_providers: Option<Vec<String>>,
pub user_allowed_api_formats: Option<Vec<String>>,
pub user_allowed_models: Option<Vec<String>>,
pub api_key_id: String,
pub api_key_name: Option<String>,
pub api_key_is_active: bool,
pub api_key_is_locked: bool,
pub api_key_is_standalone: bool,
pub api_key_rate_limit: Option<i32>,
pub api_key_concurrent_limit: Option<i32>,
pub api_key_expires_at_unix_secs: Option<u64>,
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>>,
}
impl StoredAuthApiKeySnapshot {
#[allow(clippy::too_many_arguments)]
pub fn new(
user_id: String,
username: String,
email: Option<String>,
user_role: String,
user_auth_source: String,
user_is_active: bool,
user_is_deleted: bool,
user_allowed_providers: Option<serde_json::Value>,
user_allowed_api_formats: Option<serde_json::Value>,
user_allowed_models: Option<serde_json::Value>,
api_key_id: String,
api_key_name: Option<String>,
api_key_is_active: bool,
api_key_is_locked: bool,
api_key_is_standalone: bool,
api_key_rate_limit: Option<i32>,
api_key_concurrent_limit: Option<i32>,
api_key_expires_at_unix_secs: Option<i64>,
api_key_allowed_providers: Option<serde_json::Value>,
api_key_allowed_api_formats: Option<serde_json::Value>,
api_key_allowed_models: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
Ok(Self {
user_id,
username,
email,
user_role,
user_auth_source,
user_is_active,
user_is_deleted,
user_allowed_providers: parse_string_list(
user_allowed_providers,
"users.allowed_providers",
)?,
user_allowed_api_formats: parse_string_list(
user_allowed_api_formats,
"users.allowed_api_formats",
)?,
user_allowed_models: parse_string_list(user_allowed_models, "users.allowed_models")?,
api_key_id,
api_key_name,
api_key_is_active,
api_key_is_locked,
api_key_is_standalone,
api_key_rate_limit,
api_key_concurrent_limit,
api_key_expires_at_unix_secs: api_key_expires_at_unix_secs
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid api_keys.expires_at_unix_secs: {value}"
))
})
})
.transpose()?,
api_key_allowed_providers: parse_string_list(
api_key_allowed_providers,
"api_keys.allowed_providers",
)?,
api_key_allowed_api_formats: parse_string_list(
api_key_allowed_api_formats,
"api_keys.allowed_api_formats",
)?,
api_key_allowed_models: parse_string_list(
api_key_allowed_models,
"api_keys.allowed_models",
)?,
})
}
pub fn is_currently_usable(&self, now_unix_secs: u64) -> bool {
if !self.user_is_active || self.user_is_deleted {
return false;
}
if !self.api_key_is_active {
return false;
}
if self.api_key_is_locked && !self.api_key_is_standalone {
return false;
}
if let Some(expires_at_unix_secs) = self.api_key_expires_at_unix_secs {
if expires_at_unix_secs < now_unix_secs {
return false;
}
}
true
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuthApiKeyLookupKey<'a> {
KeyHash(&'a str),
ApiKeyId(&'a str),
UserApiKeyIds {
user_id: &'a str,
api_key_id: &'a str,
},
}
#[async_trait]
pub trait AuthApiKeyReadRepository: Send + Sync {
async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
}
pub trait AuthRepository: AuthApiKeyReadRepository + Send + Sync {}
impl<T> AuthRepository for T where T: AuthApiKeyReadRepository + Send + Sync {}
fn parse_string_list(
value: Option<serde_json::Value>,
field_name: &str,
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
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"))
})?;
let mut items = Vec::with_capacity(array.len());
for item in array {
let Some(item) = item.as_str() else {
return Err(crate::DataLayerError::UnexpectedValue(format!(
"{field_name} contains a non-string item"
)));
};
items.push(item.to_string());
}
Ok(Some(items))
}
#[cfg(test)]
mod tests {
use super::StoredAuthApiKeySnapshot;
#[test]
fn rejects_non_array_allowed_providers() {
assert!(StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
None,
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!({"bad": true})),
None,
None,
"key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
None,
None,
None,
None,
)
.is_err());
}
#[test]
fn expired_non_standalone_key_is_not_usable() {
let snapshot = StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
None,
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(100),
None,
None,
None,
)
.expect("snapshot should build");
assert!(!snapshot.is_currently_usable(101));
}
}