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

feat(admin): 完善代理节点与 OAuth 授权管理

Browse Source
This commit is contained in:
fawney19
2026-04-14 14:09:24 +08:00
parent 593640ac19
commit 861ae81ff0
44 changed files with 2757 additions and 365 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -508,7 +508,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
api_key_is_active: record.is_active,
api_key_is_locked: false,
api_key_is_standalone: true,
api_key_rate_limit: Some(record.rate_limit),
api_key_rate_limit: record.rate_limit,
api_key_concurrent_limit: Some(record.concurrent_limit),
api_key_expires_at_unix_secs: record.expires_at_unix_secs,
api_key_allowed_providers: record.allowed_providers.clone(),
@@ -536,7 +536,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
record.is_active,
false,
true,
Some(record.rate_limit),
record.rate_limit,
Some(record.concurrent_limit),
record.expires_at_unix_secs.map(|value| value as i64),
record
@@ -572,7 +572,7 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
.allowed_models
.as_ref()
.map(|value| serde_json::json!(value)),
Some(record.rate_limit),
record.rate_limit,
Some(record.concurrent_limit),
record.force_capabilities,
record.is_active,
@@ -650,12 +650,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.name = Some(name);
}
}
if let Some(rate_limit) = record.rate_limit {
if record.rate_limit_present {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_rate_limit = Some(rate_limit);
snapshot.api_key_rate_limit = record.rate_limit;
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.rate_limit = Some(rate_limit);
export.rate_limit = record.rate_limit;
}
}
if let Some(allowed_providers) = record.allowed_providers {
@@ -682,6 +682,19 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.allowed_models = allowed_models;
}
}
if record.expires_at_present {
if let Some(snapshot) = index.by_api_key_id.get_mut(&record.api_key_id) {
snapshot.api_key_expires_at_unix_secs = record.expires_at_unix_secs;
}
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.expires_at_unix_secs = record.expires_at_unix_secs;
}
}
if record.auto_delete_on_expiry_present {
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
export.auto_delete_on_expiry = record.auto_delete_on_expiry;
}
}
Ok(index.export_by_api_key_id.get(&record.api_key_id).cloned())
}

36
crates/aether-data/src/repository/auth/sql.rs
View File

@@ -468,10 +468,12 @@ const UPDATE_STANDALONE_API_KEY_BASIC_SQL: &str = r#"
UPDATE api_keys
SET
name = COALESCE($2, name),
rate_limit = COALESCE($3, rate_limit),
allowed_providers = CASE WHEN $4 THEN $5::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $6 THEN $7::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $8 THEN $9::json ELSE allowed_models END,
rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END,
allowed_providers = CASE WHEN $5 THEN $6::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $7 THEN $8::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $9 THEN $10::json ELSE allowed_models END,
expires_at = CASE WHEN $11 THEN $12 ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $13 THEN $14 ELSE auto_delete_on_expiry END,
updated_at = NOW()
WHERE id = $1
AND is_standalone = TRUE
@@ -1085,9 +1087,18 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.map(serde_json::to_value)
.transpose()
.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}"))
})
})
.transpose()?;
let row = sqlx::query(UPDATE_STANDALONE_API_KEY_BASIC_SQL)
.bind(record.api_key_id)
.bind(record.name)
.bind(record.rate_limit_present)
.bind(record.rate_limit)
.bind(record.allowed_providers.is_some())
.bind(allowed_providers)
@@ -1095,6 +1106,10 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(allowed_api_formats)
.bind(record.allowed_models.is_some())
.bind(allowed_models)
.bind(record.expires_at_present)
.bind(expires_at)
.bind(record.auto_delete_on_expiry_present)
.bind(record.auto_delete_on_expiry)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1299,12 +1314,19 @@ mod tests {
#[test]
fn update_standalone_api_key_basic_sql_casts_json_case_values() {
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_providers = CASE WHEN $4 THEN $5::json ELSE allowed_providers END"));
.contains("allowed_providers = CASE WHEN $5 THEN $6::json ELSE allowed_providers END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"allowed_api_formats = CASE WHEN $6 THEN $7::json ELSE allowed_api_formats END"
"allowed_api_formats = CASE WHEN $7 THEN $8::json ELSE allowed_api_formats END"
));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_models = CASE WHEN $8 THEN $9::json ELSE allowed_models END"));
.contains("allowed_models = CASE WHEN $9 THEN $10::json ELSE allowed_models END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("expires_at = CASE WHEN $11 THEN $12 ELSE expires_at END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL.contains(
"auto_delete_on_expiry = CASE WHEN $13 THEN $14 ELSE auto_delete_on_expiry END"
));
}
#[tokio::test]

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

@@ -390,7 +390,7 @@ pub struct CreateStandaloneApiKeyRecord {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub rate_limit: i32,
pub rate_limit: Option<i32>,
pub concurrent_limit: i32,
pub force_capabilities: Option<serde_json::Value>,
pub is_active: bool,
@@ -404,10 +404,15 @@ pub struct CreateStandaloneApiKeyRecord {
pub struct UpdateStandaloneApiKeyBasicRecord {
pub api_key_id: String,
pub name: Option<String>,
pub rate_limit_present: bool,
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>>>,
pub expires_at_present: bool,
pub expires_at_unix_secs: Option<u64>,
pub auto_delete_on_expiry_present: bool,
pub auto_delete_on_expiry: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]

130
crates/aether-data/src/repository/proxy_nodes/memory.rs
View File

@@ -8,8 +8,9 @@ use uuid::Uuid;
use super::types::{
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
};
use crate::DataLayerError;
@@ -106,6 +107,13 @@ impl InMemoryProxyNodeRepository {
(!config.is_empty()).then_some(Value::Object(config))
}
fn duplicate_proxy_node_error(node: &StoredProxyNode) -> DataLayerError {
DataLayerError::InvalidInput(format!(
"已存在相同地址的代理节点: {} ({}:{})",
node.name, node.ip, node.port
))
}
}
#[async_trait]
@@ -170,6 +178,109 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
Ok(updated)
}
async fn create_manual_node(
&self,
mutation: &ProxyNodeManualCreateMutation,
) -> Result<StoredProxyNode, DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
if let Some(existing) = nodes
.values()
.find(|node| node.ip == mutation.ip && node.port == mutation.port)
{
return Err(Self::duplicate_proxy_node_error(existing));
}
let now = Self::now_unix_secs();
let node = StoredProxyNode::new(
Uuid::new_v4().to_string(),
mutation.name.clone(),
mutation.ip.clone(),
mutation.port,
true,
"online".to_string(),
0,
0,
0,
0,
0,
0,
false,
false,
0,
)?
.with_manual_proxy_fields(
Some(mutation.proxy_url.clone()),
mutation.proxy_username.clone(),
mutation.proxy_password.clone(),
)
.with_runtime_fields(
mutation.region.clone(),
mutation.registered_by.clone(),
None,
None,
None,
None,
None,
None,
None,
now,
now,
);
nodes.insert(node.id.clone(), node.clone());
Ok(node)
}
async fn update_manual_node(
&self,
mutation: &ProxyNodeManualUpdateMutation,
) -> Result<Option<StoredProxyNode>, DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
let Some(existing) = nodes.get(&mutation.node_id).cloned() else {
return Ok(None);
};
if !existing.is_manual {
return Err(DataLayerError::InvalidInput(
"只能编辑手动添加的代理节点".to_string(),
));
}
let next_ip = mutation.ip.as_deref().unwrap_or(existing.ip.as_str());
let next_port = mutation.port.unwrap_or(existing.port);
if let Some(duplicate) = nodes.values().find(|node| {
node.id != mutation.node_id && node.ip == next_ip && node.port == next_port
}) {
return Err(Self::duplicate_proxy_node_error(duplicate));
}
let node = nodes
.get_mut(&mutation.node_id)
.expect("manual proxy node should be present");
if let Some(name) = mutation.name.as_ref() {
node.name = name.clone();
}
if let Some(ip) = mutation.ip.as_ref() {
node.ip = ip.clone();
}
if let Some(port) = mutation.port {
node.port = port;
}
if let Some(region) = mutation.region.as_ref() {
node.region = Some(region.clone());
}
if let Some(proxy_url) = mutation.proxy_url.as_ref() {
node.proxy_url = Some(proxy_url.clone());
}
if let Some(proxy_username) = mutation.proxy_username.as_ref() {
node.proxy_username = Some(proxy_username.clone());
}
if let Some(proxy_password) = mutation.proxy_password.as_ref() {
node.proxy_password = Some(proxy_password.clone());
}
node.updated_at_unix_secs = Self::now_unix_secs();
Ok(Some(node.clone()))
}
async fn register_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
@@ -404,6 +515,21 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
Ok(Some(node.clone()))
}
async fn delete_node(&self, node_id: &str) -> Result<Option<StoredProxyNode>, DataLayerError> {
let removed = self
.nodes
.write()
.expect("proxy node repository lock")
.remove(node_id);
if removed.is_some() {
self.events
.write()
.expect("proxy node repository lock")
.retain(|event| event.node_id != node_id);
}
Ok(removed)
}
async fn update_remote_config(
&self,
mutation: &ProxyNodeRemoteConfigMutation,

9
crates/aether-data/src/repository/proxy_nodes/mod.rs
View File

@@ -7,8 +7,9 @@ pub use sql::SqlxProxyNodeRepository;
pub use types::{
normalize_proxy_node_scheduling_state, proxy_node_accepts_new_tunnels, proxy_reported_version,
reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
remote_config_upgrade_target, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
PROXY_NODE_SCHEDULING_STATE_CORDONED, PROXY_NODE_SCHEDULING_STATE_DRAINING,
remote_config_upgrade_target, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyNode, StoredProxyNodeEvent, PROXY_NODE_SCHEDULING_STATE_CORDONED,
PROXY_NODE_SCHEDULING_STATE_DRAINING,
};

243
crates/aether-data/src/repository/proxy_nodes/sql.rs
View File

@@ -5,8 +5,9 @@ use sqlx::{postgres::PgRow, PgPool, Row};
use super::types::{
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
};
use crate::{
error::{postgres_error, SqlxResultExt},
@@ -213,6 +214,76 @@ VALUES (
)
"#;
const FIND_DUPLICATE_PROXY_NODE_SQL: &str = r#"
SELECT
id,
name,
ip,
port
FROM proxy_nodes
WHERE ip = $1
AND port = $2
LIMIT 1
FOR UPDATE
"#;
const FIND_DUPLICATE_PROXY_NODE_EXCLUDING_ID_SQL: &str = r#"
SELECT
id,
name,
ip,
port
FROM proxy_nodes
WHERE ip = $1
AND port = $2
AND id <> $3
LIMIT 1
FOR UPDATE
"#;
const INSERT_MANUAL_PROXY_NODE_SQL: &str = r#"
INSERT INTO proxy_nodes (
id,
name,
ip,
port,
region,
is_manual,
proxy_url,
proxy_username,
proxy_password,
status,
registered_by,
last_heartbeat_at,
heartbeat_interval,
active_connections,
total_requests,
tunnel_mode,
tunnel_connected,
config_version
)
VALUES (
$1,
$2,
$3,
$4,
$5,
TRUE,
$6,
$7,
$8,
'online'::proxynodestatus,
$9,
NULL,
0,
0,
0,
FALSE,
FALSE,
0
)
"#;
const UPDATE_PROXY_NODE_REGISTRATION_SQL: &str = r#"
UPDATE proxy_nodes
SET
@@ -234,6 +305,21 @@ SET
WHERE id = $1
"#;
const UPDATE_MANUAL_PROXY_NODE_SQL: &str = r#"
UPDATE proxy_nodes
SET
name = COALESCE($2, name),
ip = COALESCE($3, ip),
port = COALESCE($4, port),
region = COALESCE($5, region),
proxy_url = COALESCE($6, proxy_url),
proxy_username = COALESCE($7, proxy_username),
proxy_password = COALESCE($8, proxy_password),
updated_at = NOW()
WHERE id = $1
AND is_manual = TRUE
"#;
const UNREGISTER_PROXY_NODE_SQL: &str = r#"
UPDATE proxy_nodes
SET
@@ -244,6 +330,11 @@ SET
WHERE id = $1
"#;
const DELETE_PROXY_NODE_SQL: &str = r#"
DELETE FROM proxy_nodes
WHERE id = $1
"#;
const UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL: &str = r#"
UPDATE proxy_nodes
SET
@@ -410,6 +501,43 @@ impl SqlxProxyNodeRepository {
(!config.is_empty()).then_some(serde_json::Value::Object(config))
}
fn duplicate_proxy_node_detail(name: &str, ip: &str, port: i32) -> String {
format!("已存在相同地址的代理节点: {name} ({ip}:{port})")
}
async fn find_duplicate_proxy_node_locked(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
ip: &str,
port: i32,
exclude_node_id: Option<&str>,
) -> Result<Option<(String, String, i32)>, DataLayerError> {
let row = if let Some(exclude_node_id) = exclude_node_id {
sqlx::query(FIND_DUPLICATE_PROXY_NODE_EXCLUDING_ID_SQL)
.bind(ip)
.bind(port)
.bind(exclude_node_id)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
} else {
sqlx::query(FIND_DUPLICATE_PROXY_NODE_SQL)
.bind(ip)
.bind(port)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
};
row.map(|row| {
Ok((
row.try_get("name").map_postgres_err()?,
row.try_get("ip").map_postgres_err()?,
row.try_get("port").map_postgres_err()?,
))
})
.transpose()
}
}
#[async_trait]
@@ -462,6 +590,102 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
Ok(result.rows_affected() as usize)
}
async fn create_manual_node(
&self,
mutation: &ProxyNodeManualCreateMutation,
) -> Result<StoredProxyNode, DataLayerError> {
let lock_key = Self::registration_lock_key(&mutation.ip, mutation.port);
let mut tx = self.pool.begin().await.map_postgres_err()?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(lock_key)
.execute(&mut *tx)
.await
.map_postgres_err()?;
if let Some((name, ip, port)) =
Self::find_duplicate_proxy_node_locked(&mut tx, &mutation.ip, mutation.port, None)
.await?
{
return Err(DataLayerError::InvalidInput(
Self::duplicate_proxy_node_detail(&name, &ip, port),
));
}
let node_id = uuid::Uuid::new_v4().to_string();
sqlx::query(INSERT_MANUAL_PROXY_NODE_SQL)
.bind(&node_id)
.bind(&mutation.name)
.bind(&mutation.ip)
.bind(mutation.port)
.bind(mutation.region.as_deref())
.bind(&mutation.proxy_url)
.bind(mutation.proxy_username.as_deref())
.bind(mutation.proxy_password.as_deref())
.bind(mutation.registered_by.as_deref())
.execute(&mut *tx)
.await
.map_postgres_err()?;
tx.commit().await.map_err(postgres_error)?;
self.find_proxy_node(&node_id).await?.ok_or_else(|| {
DataLayerError::UnexpectedValue("created manual proxy node missing".to_string())
})
}
async fn update_manual_node(
&self,
mutation: &ProxyNodeManualUpdateMutation,
) -> Result<Option<StoredProxyNode>, DataLayerError> {
let existing = self.find_proxy_node(&mutation.node_id).await?;
let Some(existing) = existing else {
return Ok(None);
};
if !existing.is_manual {
return Err(DataLayerError::InvalidInput(
"只能编辑手动添加的代理节点".to_string(),
));
}
let next_ip = mutation.ip.as_deref().unwrap_or(existing.ip.as_str());
let next_port = mutation.port.unwrap_or(existing.port);
let lock_key = Self::registration_lock_key(next_ip, next_port);
let mut tx = self.pool.begin().await.map_postgres_err()?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(lock_key)
.execute(&mut *tx)
.await
.map_postgres_err()?;
if let Some((name, ip, port)) = Self::find_duplicate_proxy_node_locked(
&mut tx,
next_ip,
next_port,
Some(&mutation.node_id),
)
.await?
{
return Err(DataLayerError::InvalidInput(
Self::duplicate_proxy_node_detail(&name, &ip, port),
));
}
sqlx::query(UPDATE_MANUAL_PROXY_NODE_SQL)
.bind(&mutation.node_id)
.bind(mutation.name.as_deref())
.bind(mutation.ip.as_deref())
.bind(mutation.port)
.bind(mutation.region.as_deref())
.bind(mutation.proxy_url.as_deref())
.bind(mutation.proxy_username.as_deref())
.bind(mutation.proxy_password.as_deref())
.execute(&mut *tx)
.await
.map_postgres_err()?;
tx.commit().await.map_err(postgres_error)?;
self.find_proxy_node(&mutation.node_id).await
}
async fn register_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
@@ -721,6 +945,21 @@ VALUES (
self.find_proxy_node(&existing.id).await
}
async fn delete_node(&self, node_id: &str) -> Result<Option<StoredProxyNode>, DataLayerError> {
let existing = self.find_proxy_node(node_id).await?;
let Some(existing) = existing else {
return Ok(None);
};
sqlx::query(DELETE_PROXY_NODE_SQL)
.bind(node_id)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(Some(existing))
}
async fn update_remote_config(
&self,
mutation: &ProxyNodeRemoteConfigMutation,

39
crates/aether-data/src/repository/proxy_nodes/types.rs
View File

@@ -180,6 +180,30 @@ pub struct ProxyNodeRegistrationMutation {
pub tunnel_mode: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProxyNodeManualCreateMutation {
pub name: String,
pub ip: String,
pub port: i32,
pub region: Option<String>,
pub proxy_url: String,
pub proxy_username: Option<String>,
pub proxy_password: Option<String>,
pub registered_by: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProxyNodeManualUpdateMutation {
pub node_id: String,
pub name: Option<String>,
pub ip: Option<String>,
pub port: Option<i32>,
pub region: Option<String>,
pub proxy_url: Option<String>,
pub proxy_username: Option<String>,
pub proxy_password: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProxyNodeTunnelStatusMutation {
pub node_id: String,
@@ -348,6 +372,16 @@ pub trait ProxyNodeReadRepository: Send + Sync {
pub trait ProxyNodeWriteRepository: Send + Sync {
async fn reset_stale_tunnel_statuses(&self) -> Result<usize, crate::DataLayerError>;
async fn create_manual_node(
&self,
mutation: &ProxyNodeManualCreateMutation,
) -> Result<StoredProxyNode, crate::DataLayerError>;
async fn update_manual_node(
&self,
mutation: &ProxyNodeManualUpdateMutation,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
async fn register_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
@@ -368,6 +402,11 @@ pub trait ProxyNodeWriteRepository: Send + Sync {
node_id: &str,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
async fn delete_node(
&self,
node_id: &str,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
async fn update_remote_config(
&self,
mutation: &ProxyNodeRemoteConfigMutation,

88
crates/aether-provider-transport/src/auth.rs
View File

@@ -6,6 +6,7 @@ use super::headers::{
use super::snapshot::GatewayProviderTransportSnapshot;
const DEFAULT_ANTHROPIC_VERSION: &str = "2023-06-01";
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
fn collect_passthrough_headers(
headers: &http::HeaderMap,
@@ -233,10 +234,7 @@ pub fn resolve_local_openai_chat_auth(
if !matches!(auth_type.as_str(), "api_key" | "bearer") {
return None;
}
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
let secret = resolved_local_secret(transport)?;
Some(("authorization".to_string(), format!("Bearer {secret}")))
}
@@ -245,10 +243,7 @@ pub fn resolve_local_standard_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
let secret = resolved_local_secret(transport)?;
match auth_type.as_str() {
"api_key" => Some(("x-api-key".to_string(), secret.to_string())),
@@ -261,10 +256,7 @@ pub fn resolve_local_gemini_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
let secret = resolved_local_secret(transport)?;
match auth_type.as_str() {
"api_key" => Some(("x-goog-api-key".to_string(), secret.to_string())),
@@ -273,11 +265,76 @@ pub fn resolve_local_gemini_auth(
}
}
fn resolved_local_secret(transport: &GatewayProviderTransportSnapshot) -> Option<&str> {
let secret = transport.key.decrypted_api_key.trim();
(!secret.is_empty() && secret != PLACEHOLDER_API_KEY).then_some(secret)
}
#[cfg(test)]
mod tests {
use super::{build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth};
use super::{
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
resolve_local_standard_auth,
};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use std::collections::BTreeMap;
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "custom".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "claude:chat".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn claude_passthrough_headers_restore_stripped_anthropic_headers() {
let mut headers = http::HeaderMap::new();
@@ -380,4 +437,9 @@ mod tests {
Some("sk-upstream")
);
}
#[test]
fn local_standard_auth_rejects_placeholder_secret() {
assert!(resolve_local_standard_auth(&sample_transport()).is_none());
}
}

98
crates/aether-provider-transport/src/generic_oauth.rs
View File

@@ -3,10 +3,11 @@ use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use serde_json::{json, Value};
use url::form_urlencoded;
use super::oauth_refresh::{
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshError,
LocalResolvedOAuthRequestAuth,
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthRefreshAdapter,
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
};
use super::snapshot::GatewayProviderTransportSnapshot;
@@ -247,7 +248,7 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
async fn refresh(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
@@ -265,7 +266,6 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
let token_url = self.token_url_for_template(template);
let scope = (!template.scopes.is_empty()).then(|| template.scopes.join(" "));
let request = client.post(token_url);
let response = if template.uses_json_payload {
let mut body = serde_json::Map::from_iter([
(
@@ -284,44 +284,60 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
if let Some(scope) = scope.as_ref() {
body.insert("scope".to_string(), Value::String(scope.clone()));
}
request
.header("Content-Type", "application/json")
.header("Accept", "application/json")
.json(&Value::Object(body))
.send()
.await
executor
.execute(
template.provider_type,
transport,
&LocalOAuthHttpRequest {
request_id: "provider-oauth:local-refresh-token",
method: reqwest::Method::POST,
url: token_url,
headers: BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("accept".to_string(), "application/json".to_string()),
]),
json_body: Some(Value::Object(body)),
body_bytes: None,
},
)
.await?
} else {
let mut form = vec![
("grant_type", "refresh_token".to_string()),
("client_id", template.client_id.to_string()),
("refresh_token", refresh_token.clone()),
];
if let Some(scope) = scope.as_ref() {
form.push(("scope", scope.clone()));
}
if !template.client_secret.trim().is_empty() {
form.push(("client_secret", template.client_secret.to_string()));
}
request
.header("Content-Type", "application/x-www-form-urlencoded")
.header("Accept", "application/json")
.form(&form)
.send()
.await
}
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: template.provider_type,
source,
})?;
let status = response.status();
let body = response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: template.provider_type,
source,
})?;
let form_body = {
let mut form = form_urlencoded::Serializer::new(String::new());
form.append_pair("grant_type", "refresh_token");
form.append_pair("client_id", template.client_id);
form.append_pair("refresh_token", refresh_token.as_str());
if let Some(scope) = scope.as_ref() {
form.append_pair("scope", scope);
}
if !template.client_secret.trim().is_empty() {
form.append_pair("client_secret", template.client_secret);
}
form.finish().into_bytes()
};
executor
.execute(
template.provider_type,
transport,
&LocalOAuthHttpRequest {
request_id: "provider-oauth:local-refresh-token",
method: reqwest::Method::POST,
url: token_url,
headers: BTreeMap::from([
(
"content-type".to_string(),
"application/x-www-form-urlencoded".to_string(),
),
("accept".to_string(), "application/json".to_string()),
]),
json_body: None,
body_bytes: Some(form_body),
},
)
.await?
};
let status = reqwest::StatusCode::from_u16(response.status_code).unwrap_or_default();
let body = response.body_text;
if !status.is_success() {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: template.provider_type,

180
crates/aether-provider-transport/src/kiro/refresh.rs
View File

@@ -4,8 +4,8 @@ use async_trait::async_trait;
use serde_json::{json, Value};
use super::super::oauth_refresh::{
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshError,
LocalResolvedOAuthRequestAuth,
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthRefreshAdapter,
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
};
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::auth::{
@@ -34,13 +34,16 @@ impl KiroOAuthRefreshAdapter {
pub async fn refresh_auth_config(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
auth_config: &KiroAuthConfig,
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
if auth_config.is_idc_auth() {
self.refresh_idc_token(client, auth_config).await
self.refresh_idc_token(executor, transport, auth_config)
.await
} else {
self.refresh_social_token(client, auth_config).await
self.refresh_social_token(executor, transport, auth_config)
.await
}
}
@@ -101,7 +104,8 @@ impl KiroOAuthRefreshAdapter {
async fn refresh_social_token(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
auth_config: &KiroAuthConfig,
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
let url = self.social_refresh_url(auth_config);
@@ -122,36 +126,42 @@ impl KiroOAuthRefreshAdapter {
})?;
let kiro_version = auth_config.effective_kiro_version();
let user_agent = build_kiro_ide_tag(kiro_version, &machine_id);
let response = client
.post(url)
.header("User-Agent", user_agent)
.header("Host", host)
.header("Accept", "application/json, text/plain, */*")
.header("Content-Type", "application/json")
.header("Connection", "close")
.header("Accept-Encoding", "gzip, compress, deflate, br")
.json(&json!({
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default()
}))
.send()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
let response = executor
.execute(
PROVIDER_TYPE,
transport,
&LocalOAuthHttpRequest {
request_id: "provider-oauth:kiro-social-refresh",
method: reqwest::Method::POST,
url,
headers: std::collections::BTreeMap::from([
("user-agent".to_string(), user_agent),
("host".to_string(), host),
(
"accept".to_string(),
"application/json, text/plain, */*".to_string(),
),
("content-type".to_string(), "application/json".to_string()),
("connection".to_string(), "close".to_string()),
(
"accept-encoding".to_string(),
"gzip, compress, deflate, br".to_string(),
),
]),
json_body: Some(json!({
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default()
})),
body_bytes: None,
},
)
.await?;
let status = response.status();
let body = response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
let status = reqwest::StatusCode::from_u16(response.status_code).unwrap_or_default();
let body = response.body_text;
if !status.is_success() {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: PROVIDER_TYPE,
@@ -208,7 +218,8 @@ impl KiroOAuthRefreshAdapter {
async fn refresh_idc_token(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
auth_config: &KiroAuthConfig,
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
let url = self.idc_refresh_url(auth_config);
@@ -218,46 +229,49 @@ impl KiroOAuthRefreshAdapter {
.unwrap_or_else(|| {
format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region())
});
let response = client
.post(url)
.header("Content-Type", "application/json")
.header("Host", host)
.header("x-amz-user-agent", IDC_AMZ_USER_AGENT)
.header("User-Agent", "node")
.header("Accept", "*/*")
.json(&json!({
"clientId": auth_config
.client_id
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"clientSecret": auth_config
.client_secret
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"grantType": "refresh_token"
}))
.send()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
let response = executor
.execute(
PROVIDER_TYPE,
transport,
&LocalOAuthHttpRequest {
request_id: "provider-oauth:kiro-idc-refresh",
method: reqwest::Method::POST,
url,
headers: std::collections::BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("host".to_string(), host),
(
"x-amz-user-agent".to_string(),
IDC_AMZ_USER_AGENT.to_string(),
),
("user-agent".to_string(), "node".to_string()),
("accept".to_string(), "*/*".to_string()),
]),
json_body: Some(json!({
"clientId": auth_config
.client_id
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"clientSecret": auth_config
.client_secret
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"grantType": "refresh_token"
})),
body_bytes: None,
},
)
.await?;
let status = response.status();
let body = response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
let status = reqwest::StatusCode::from_u16(response.status_code).unwrap_or_default();
let body = response.body_text;
if !status.is_success() {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: PROVIDER_TYPE,
@@ -353,7 +367,7 @@ impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter {
async fn refresh(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
@@ -361,9 +375,11 @@ impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter {
return Ok(None);
};
let refreshed = if auth_config.is_idc_auth() {
self.refresh_idc_token(client, &auth_config).await?
self.refresh_idc_token(executor, transport, &auth_config)
.await?
} else {
self.refresh_social_token(client, &auth_config).await?
self.refresh_social_token(executor, transport, &auth_config)
.await?
};
Ok(Self::build_cached_entry(&refreshed))
}
@@ -410,7 +426,7 @@ mod tests {
use std::sync::{Arc, Mutex};
use super::super::super::oauth_refresh::{
LocalOAuthRefreshAdapter, LocalResolvedOAuthRequestAuth,
LocalOAuthRefreshAdapter, LocalResolvedOAuthRequestAuth, ReqwestLocalOAuthHttpExecutor,
};
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -565,9 +581,10 @@ mod tests {
"kiro_version":"1.2.3"
}"#,
);
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let entry = adapter
.refresh(&reqwest::Client::new(), &transport, None)
.refresh(&executor, &transport, None)
.await
.expect("refresh should succeed")
.expect("cached entry should exist");
@@ -665,9 +682,10 @@ mod tests {
"profile_arn":"arn:aws:bedrock:demo"
}"#,
);
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let entry = adapter
.refresh(&reqwest::Client::new(), &transport, None)
.refresh(&executor, &transport, None)
.await
.expect("refresh should succeed")
.expect("cached entry should exist");

5
crates/aether-provider-transport/src/lib.rs
View File

@@ -29,8 +29,9 @@ pub use network::{
TransportTunnelAttachmentOwner,
};
pub use oauth_refresh::{
supports_local_oauth_request_auth_resolution, CachedOAuthEntry, LocalOAuthRefreshCoordinator,
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
supports_local_oauth_request_auth_resolution, CachedOAuthEntry, LocalOAuthHttpExecutor,
LocalOAuthHttpRequest, LocalOAuthHttpResponse, LocalOAuthRefreshCoordinator,
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth, ReqwestLocalOAuthHttpExecutor,
};
pub use policy::{
local_gemini_transport_unsupported_reason,

100
crates/aether-provider-transport/src/oauth_refresh.rs
View File

@@ -42,6 +42,22 @@ pub struct CachedOAuthEntry {
pub metadata: Option<Value>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct LocalOAuthHttpRequest {
pub request_id: &'static str,
pub method: reqwest::Method,
pub url: String,
pub headers: BTreeMap<String, String>,
pub json_body: Option<Value>,
pub body_bytes: Option<Vec<u8>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LocalOAuthHttpResponse {
pub status_code: u16,
pub body_text: String,
}
#[derive(Debug, Error)]
pub enum LocalOAuthRefreshError {
#[error("{provider_type} oauth refresh request failed: {source}")]
@@ -63,6 +79,71 @@ pub enum LocalOAuthRefreshError {
},
}
#[async_trait]
pub trait LocalOAuthHttpExecutor: Send + Sync {
async fn execute(
&self,
provider_type: &'static str,
transport: &GatewayProviderTransportSnapshot,
request: &LocalOAuthHttpRequest,
) -> Result<LocalOAuthHttpResponse, LocalOAuthRefreshError>;
}
#[derive(Debug, Clone)]
pub struct ReqwestLocalOAuthHttpExecutor {
client: reqwest::Client,
}
impl ReqwestLocalOAuthHttpExecutor {
pub fn new(client: reqwest::Client) -> Self {
Self { client }
}
}
#[async_trait]
impl LocalOAuthHttpExecutor for ReqwestLocalOAuthHttpExecutor {
async fn execute(
&self,
provider_type: &'static str,
_transport: &GatewayProviderTransportSnapshot,
request: &LocalOAuthHttpRequest,
) -> Result<LocalOAuthHttpResponse, LocalOAuthRefreshError> {
let mut builder = self
.client
.request(request.method.clone(), request.url.as_str());
for (name, value) in &request.headers {
builder = builder.header(name, value);
}
if let Some(json_body) = request.json_body.as_ref() {
builder = builder.json(json_body);
} else if let Some(body_bytes) = request.body_bytes.as_ref() {
builder = builder.body(body_bytes.clone());
}
let response =
builder
.send()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type,
source,
})?;
let status_code = response.status().as_u16();
let body_text =
response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type,
source,
})?;
Ok(LocalOAuthHttpResponse {
status_code,
body_text,
})
}
}
#[async_trait]
pub trait LocalOAuthRefreshAdapter: Send + Sync {
fn provider_type(&self) -> &'static str;
@@ -94,7 +175,7 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync {
async fn refresh(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError>;
@@ -152,7 +233,7 @@ impl LocalOAuthRefreshCoordinator {
pub async fn resolve_with_result(
&self,
client: &reqwest::Client,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RedisLockRunner>,
distributed_owner: Option<&str>,
@@ -232,7 +313,7 @@ impl LocalOAuthRefreshCoordinator {
};
let refresh_result = adapter
.refresh(client, transport, cached_entry.as_ref())
.refresh(executor, transport, cached_entry.as_ref())
.await;
if let (Some(lock), Some(lease)) = (distributed_lock, distributed_lease.as_ref()) {
if let Err(err) = lock.release(lease).await {
@@ -300,8 +381,9 @@ mod tests {
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshCoordinator,
LocalOAuthRefreshError, LocalOAuthResolution, LocalResolvedOAuthRequestAuth,
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthRefreshAdapter,
LocalOAuthRefreshCoordinator, LocalOAuthRefreshError, LocalOAuthResolution,
LocalResolvedOAuthRequestAuth, ReqwestLocalOAuthHttpExecutor,
};
use async_trait::async_trait;
use std::sync::Arc;
@@ -351,7 +433,7 @@ mod tests {
async fn refresh(
&self,
_client: &reqwest::Client,
_executor: &dyn LocalOAuthHttpExecutor,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
@@ -427,14 +509,14 @@ mod tests {
refresh_hits: Arc::clone(&refresh_hits),
})]);
let transport = sample_transport();
let client = reqwest::Client::new();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let first = coordinator
.resolve_with_result(&client, &transport, None, None)
.resolve_with_result(&executor, &transport, None, None)
.await
.expect("first resolve should succeed");
let second = coordinator
.resolve_with_result(&client, &transport, None, None)
.resolve_with_result(&executor, &transport, None, None)
.await
.expect("second resolve should succeed");