mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat(admin): 完善代理节点与 OAuth 授权管理
This commit is contained in:
@@ -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())
|
||||
}
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user