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

feat(proxy): 实现代理节点批量升级回滚、隧道重定向跟随及远程配置管理

Browse Source 
核心功能：
- 新增代理节点批量升级回滚工作流，支持分批升级、健康探针、跳过/重试/取消等操作
- proxy 隧道流处理器支持 HTTP 重定向跟随（最多 10 跳），区分 307/308 可重播与不可重播请求体
- proxy 协议新增 follow_redirects / http1_only 字段，网关侧同步支持
- 新增代理节点远端配置变更接口（名称、允许端口、调度状态、升级目标等）
- 新增代理节点注册/反注册/心跳的 Admin API，及节点过期清理维护任务
- gateway 隧道 owner-relay 支持流式代理大请求体，新增 5 MiB 默认限制
- 新增 ProxyNodeRegistrationMutation / ProxyNodeRemoteConfigMutation 数据类型
- proxy 配置新增重定向重播预算、心跳间隔等参数，TUI 安装向导同步更新
- 前端 ProxyNodes 页面新增批量升级操作面板及滚动进度展示
This commit is contained in:
fawney19
2026-04-12 16:02:38 +08:00
parent 7c5bb7f383
commit 9703840a36
83 changed files with 11832 additions and 1520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
crates/aether-data/src/repository/management_tokens/memory.rs
View File

@@ -1,3 +1,4 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -13,6 +14,7 @@ use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryManagementTokenRepository {
items: RwLock<Vec<StoredManagementTokenWithUser>>,
hashes: RwLock<BTreeMap<String, String>>,
}
impl InMemoryManagementTokenRepository {
@@ -22,6 +24,18 @@ impl InMemoryManagementTokenRepository {
{
Self {
items: RwLock::new(items.into_iter().collect()),
hashes: RwLock::new(BTreeMap::new()),
}
}
pub fn seed_with_hashes<I, J>(items: I, hashes: J) -> Self
where
I: IntoIterator<Item = StoredManagementTokenWithUser>,
J: IntoIterator<Item = (String, String)>,
{
Self {
items: RwLock::new(items.into_iter().collect()),
hashes: RwLock::new(hashes.into_iter().collect()),
}
}
@@ -31,6 +45,10 @@ impl InMemoryManagementTokenRepository {
.ok()
.map(|duration| duration.as_secs())
}
fn remove_hash_for_token(hashes: &mut BTreeMap<String, String>, token_id: &str) {
hashes.retain(|_, existing_token_id| existing_token_id != token_id);
}
}
#[async_trait]
@@ -77,6 +95,24 @@ impl ManagementTokenReadRepository for InMemoryManagementTokenRepository {
let items = self.items.read().expect("management token repository lock");
Ok(items.iter().find(|item| item.token.id == token_id).cloned())
}
async fn get_management_token_with_user_by_hash(
&self,
token_hash: &str,
) -> Result<Option<StoredManagementTokenWithUser>, DataLayerError> {
let token_id = {
let hashes = self
.hashes
.read()
.expect("management token repository lock");
hashes.get(token_hash).cloned()
};
let Some(token_id) = token_id else {
return Ok(None);
};
let items = self.items.read().expect("management token repository lock");
Ok(items.iter().find(|item| item.token.id == token_id).cloned())
}
}
#[async_trait]
@@ -91,6 +127,10 @@ impl ManagementTokenWriteRepository for InMemoryManagementTokenRepository {
.items
.write()
.expect("management token repository lock");
let mut hashes = self
.hashes
.write()
.expect("management token repository lock");
if items
.iter()
.any(|item| item.token.user_id == record.user_id && item.token.name == record.name)
@@ -118,6 +158,7 @@ impl ManagementTokenWriteRepository for InMemoryManagementTokenRepository {
token.clone(),
record.user.clone(),
));
hashes.insert(record.token_hash.clone(), record.id.clone());
Ok(token)
}
@@ -183,8 +224,13 @@ impl ManagementTokenWriteRepository for InMemoryManagementTokenRepository {
.items
.write()
.expect("management token repository lock");
let mut hashes = self
.hashes
.write()
.expect("management token repository lock");
let original_len = items.len();
items.retain(|item| item.token.id != token_id);
Self::remove_hash_for_token(&mut hashes, token_id);
Ok(items.len() != original_len)
}
@@ -215,16 +261,41 @@ impl ManagementTokenWriteRepository for InMemoryManagementTokenRepository {
.items
.write()
.expect("management token repository lock");
let mut hashes = self
.hashes
.write()
.expect("management token repository lock");
let Some(item) = items
.iter_mut()
.find(|item| item.token.id == mutation.token_id)
else {
return Ok(None);
};
Self::remove_hash_for_token(&mut hashes, &mutation.token_id);
hashes.insert(mutation.token_hash.clone(), mutation.token_id.clone());
item.token.token_prefix = mutation.token_prefix.clone();
item.token.updated_at_unix_secs = Self::now_unix_secs();
Ok(Some(item.token.clone()))
}
async fn record_management_token_usage(
&self,
token_id: &str,
last_used_ip: Option<&str>,
) -> Result<Option<StoredManagementToken>, DataLayerError> {
let mut items = self
.items
.write()
.expect("management token repository lock");
let Some(item) = items.iter_mut().find(|item| item.token.id == token_id) else {
return Ok(None);
};
item.token.last_used_at_unix_secs = Self::now_unix_secs();
item.token.last_used_ip = last_used_ip.map(ToOwned::to_owned);
item.token.usage_count = item.token.usage_count.saturating_add(1);
item.token.updated_at_unix_secs = Self::now_unix_secs();
Ok(Some(item.token.clone()))
}
}
#[cfg(test)]
@@ -254,10 +325,16 @@ mod tests {
#[tokio::test]
async fn lists_filters_and_mutates_management_tokens() {
let repository = InMemoryManagementTokenRepository::seed(vec![
sample_token("token-1", "user-1", true),
sample_token("token-2", "user-2", false),
]);
let repository = InMemoryManagementTokenRepository::seed_with_hashes(
vec![
sample_token("token-1", "user-1", true),
sample_token("token-2", "user-2", false),
],
vec![
("hash-1".to_string(), "token-1".to_string()),
("hash-2".to_string(), "token-2".to_string()),
],
);
let page = repository
.list_management_tokens(&ManagementTokenListQuery {
@@ -333,10 +410,31 @@ mod tests {
.expect("token should exist");
assert_eq!(regenerated.token_prefix.as_deref(), Some("ae_5678"));
let by_hash = repository
.get_management_token_with_user_by_hash("hash-3b")
.await
.expect("lookup by hash should succeed")
.expect("token should exist");
assert_eq!(by_hash.token.id, "token-3");
let used = repository
.record_management_token_usage("token-3", Some("127.0.0.1"))
.await
.expect("usage update should succeed")
.expect("token should exist");
assert_eq!(used.last_used_ip.as_deref(), Some("127.0.0.1"));
assert_eq!(used.usage_count, 1);
let deleted = repository
.delete_management_token("token-1")
.await
.expect("delete should succeed");
assert!(deleted);
let deleted_by_hash = repository
.get_management_token_with_user_by_hash("hash-1")
.await
.expect("hash lookup should succeed");
assert!(deleted_by_hash.is_none());
}
}

74
crates/aether-data/src/repository/management_tokens/sql.rs
View File

@@ -69,6 +69,31 @@ WHERE mt.id = $1
LIMIT 1
"#;
const GET_MANAGEMENT_TOKEN_WITH_USER_BY_HASH_SQL: &str = r#"
SELECT
mt.id,
mt.user_id,
mt.name,
mt.description,
mt.token_prefix,
mt.allowed_ips,
EXTRACT(EPOCH FROM mt.expires_at)::bigint AS expires_at_unix_secs,
EXTRACT(EPOCH FROM mt.last_used_at)::bigint AS last_used_at_unix_secs,
mt.last_used_ip,
COALESCE(mt.usage_count, 0) AS usage_count,
mt.is_active,
EXTRACT(EPOCH FROM mt.created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM mt.updated_at)::bigint AS updated_at_unix_secs,
u.id AS user_row_id,
u.email AS user_email,
u.username AS user_username,
u.role::text AS user_role
FROM management_tokens mt
JOIN users u ON u.id = mt.user_id
WHERE mt.token_hash = $1
LIMIT 1
"#;
const DELETE_MANAGEMENT_TOKEN_SQL: &str = r#"
DELETE FROM management_tokens
WHERE id = $1
@@ -196,6 +221,29 @@ RETURNING
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
"#;
const RECORD_MANAGEMENT_TOKEN_USAGE_SQL: &str = r#"
UPDATE management_tokens
SET last_used_at = NOW(),
last_used_ip = $2,
usage_count = COALESCE(usage_count, 0) + 1,
updated_at = NOW()
WHERE id = $1
RETURNING
id,
user_id,
name,
description,
token_prefix,
allowed_ips,
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
last_used_ip,
COALESCE(usage_count, 0) AS usage_count,
is_active,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
"#;
#[derive(Debug, Clone)]
pub struct SqlxManagementTokenRepository {
pool: PgPool,
@@ -250,6 +298,18 @@ impl ManagementTokenReadRepository for SqlxManagementTokenRepository {
.map_postgres_err()?;
row.as_ref().map(map_token_with_user_row).transpose()
}
async fn get_management_token_with_user_by_hash(
&self,
token_hash: &str,
) -> Result<Option<StoredManagementTokenWithUser>, DataLayerError> {
let row = sqlx::query(GET_MANAGEMENT_TOKEN_WITH_USER_BY_HASH_SQL)
.bind(token_hash)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_token_with_user_row).transpose()
}
}
#[async_trait]
@@ -341,6 +401,20 @@ impl ManagementTokenWriteRepository for SqlxManagementTokenRepository {
.map_postgres_err()?;
row.as_ref().map(map_token_row).transpose()
}
async fn record_management_token_usage(
&self,
token_id: &str,
last_used_ip: Option<&str>,
) -> Result<Option<StoredManagementToken>, DataLayerError> {
let row = sqlx::query(RECORD_MANAGEMENT_TOKEN_USAGE_SQL)
.bind(token_id)
.bind(last_used_ip)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_token_row).transpose()
}
}
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {

11
crates/aether-data/src/repository/management_tokens/types.rs
View File

@@ -306,6 +306,11 @@ pub trait ManagementTokenReadRepository: Send + Sync {
&self,
token_id: &str,
) -> Result<Option<StoredManagementTokenWithUser>, crate::DataLayerError>;
async fn get_management_token_with_user_by_hash(
&self,
token_hash: &str,
) -> Result<Option<StoredManagementTokenWithUser>, crate::DataLayerError>;
}
#[async_trait]
@@ -332,4 +337,10 @@ pub trait ManagementTokenWriteRepository: Send + Sync {
&self,
mutation: &RegenerateManagementTokenSecret,
) -> Result<Option<StoredManagementToken>, crate::DataLayerError>;
async fn record_management_token_usage(
&self,
token_id: &str,
last_used_ip: Option<&str>,
) -> Result<Option<StoredManagementToken>, crate::DataLayerError>;
}

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

@@ -3,9 +3,12 @@ use std::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use serde_json::{json, Map, Value};
use uuid::Uuid;
use super::types::{
normalize_proxy_metadata, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
};
use crate::DataLayerError;
@@ -58,6 +61,51 @@ impl InMemoryProxyNodeRepository {
fn next_event_id(events: &[StoredProxyNodeEvent]) -> i64 {
events.iter().map(|event| event.id).max().unwrap_or(0) + 1
}
fn normalize_remote_config(
mutation: &ProxyNodeRemoteConfigMutation,
existing: Option<&Value>,
) -> Option<Value> {
let mut config = match existing {
Some(Value::Object(map)) => map.clone(),
_ => Map::new(),
};
if let Some(node_name) = mutation.node_name.as_ref() {
config.insert("node_name".to_string(), Value::String(node_name.clone()));
}
if let Some(allowed_ports) = mutation.allowed_ports.as_ref() {
config.insert("allowed_ports".to_string(), json!(allowed_ports));
}
if let Some(log_level) = mutation.log_level.as_ref() {
config.insert("log_level".to_string(), Value::String(log_level.clone()));
}
if let Some(heartbeat_interval) = mutation.heartbeat_interval {
config.insert("heartbeat_interval".to_string(), json!(heartbeat_interval));
}
if let Some(scheduling_state) = mutation.scheduling_state.as_ref() {
match scheduling_state {
Some(state) => {
config.insert("scheduling_state".to_string(), Value::String(state.clone()));
}
None => {
config.remove("scheduling_state");
}
}
}
if let Some(upgrade_to) = mutation.upgrade_to.as_ref() {
match upgrade_to {
Some(version) => {
config.insert("upgrade_to".to_string(), Value::String(version.clone()));
}
None => {
config.remove("upgrade_to");
}
}
}
(!config.is_empty()).then_some(Value::Object(config))
}
}
#[async_trait]
@@ -102,6 +150,117 @@ impl ProxyNodeReadRepository for InMemoryProxyNodeRepository {
#[async_trait]
impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
async fn reset_stale_tunnel_statuses(&self) -> Result<usize, DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
let now = Self::now_unix_secs();
let mut updated = 0usize;
for node in nodes.values_mut() {
if node.is_manual || !node.tunnel_connected {
continue;
}
node.tunnel_connected = false;
node.status = "offline".to_string();
node.active_connections = 0;
node.tunnel_connected_at_unix_secs = now;
node.updated_at_unix_secs = now;
updated = updated.saturating_add(1);
}
Ok(updated)
}
async fn register_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
) -> Result<StoredProxyNode, DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
let now = Self::now_unix_secs();
let normalized_proxy_metadata = normalize_proxy_metadata(
mutation.proxy_metadata.as_ref(),
mutation.proxy_version.as_deref(),
);
if let Some(existing_id) = nodes
.iter()
.find(|(_, node)| {
!node.is_manual && node.ip == mutation.ip && node.port == mutation.port
})
.map(|(node_id, _)| node_id.clone())
{
let node = nodes
.get_mut(&existing_id)
.expect("existing proxy node should be present");
node.name = mutation.name.clone();
node.ip = mutation.ip.clone();
node.port = mutation.port;
node.region = mutation.region.clone();
node.last_heartbeat_at_unix_secs = now;
node.heartbeat_interval = mutation.heartbeat_interval;
node.tunnel_mode = mutation.tunnel_mode;
node.registered_by = mutation.registered_by.clone();
if let Some(active_connections) = mutation.active_connections {
node.active_connections = active_connections;
}
if let Some(total_requests) = mutation.total_requests {
node.total_requests = total_requests;
}
if let Some(avg_latency_ms) = mutation.avg_latency_ms {
node.avg_latency_ms = Some(avg_latency_ms);
}
if let Some(hardware_info) = mutation.hardware_info.as_ref() {
node.hardware_info = Some(hardware_info.clone());
}
if let Some(estimated_max_concurrency) = mutation.estimated_max_concurrency {
node.estimated_max_concurrency = Some(estimated_max_concurrency);
}
if let Some(proxy_metadata) = normalized_proxy_metadata {
node.proxy_metadata = Some(proxy_metadata);
}
if node.created_at_unix_ms.is_none() {
node.created_at_unix_ms = now;
}
node.updated_at_unix_secs = now;
return Ok(node.clone());
}
let mut node = StoredProxyNode::new(
Uuid::new_v4().to_string(),
mutation.name.clone(),
mutation.ip.clone(),
mutation.port,
false,
"offline".to_string(),
mutation.heartbeat_interval,
mutation.active_connections.unwrap_or(0),
mutation.total_requests.unwrap_or(0),
0,
0,
0,
mutation.tunnel_mode,
false,
0,
)?
.with_runtime_fields(
mutation.region.clone(),
mutation.registered_by.clone(),
now,
mutation.avg_latency_ms,
normalized_proxy_metadata,
mutation.hardware_info.clone(),
mutation.estimated_max_concurrency,
None,
None,
now,
now,
);
node.avg_latency_ms = mutation.avg_latency_ms;
nodes.insert(node.id.clone(), node.clone());
Ok(node)
}
async fn apply_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,
@@ -154,6 +313,15 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
if let Some(value) = mutation.stream_errors_delta.filter(|value| *value > 0) {
node.stream_errors += value;
}
let reconciled_remote_config = reconcile_remote_config_after_heartbeat(
node.remote_config.as_ref(),
mutation.proxy_version.as_deref(),
);
if reconciled_remote_config != node.remote_config {
node.remote_config = reconciled_remote_config;
node.config_version = node.config_version.saturating_add(1);
node.updated_at_unix_secs = now;
}
Ok(Some(node.clone()))
}
@@ -204,6 +372,9 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
} else {
"offline".to_string()
};
if !mutation.connected {
node.active_connections = 0;
}
node.updated_at_unix_secs = Some(event_time);
let event_id = Self::next_event_id(&events);
events.push(StoredProxyNodeEvent {
@@ -215,14 +386,55 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
});
Ok(Some(node.clone()))
}
async fn unregister_node(
&self,
node_id: &str,
) -> Result<Option<StoredProxyNode>, DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
let Some(node) = nodes.get_mut(node_id) else {
return Ok(None);
};
let now = Self::now_unix_secs();
node.status = "offline".to_string();
node.tunnel_connected = false;
node.active_connections = 0;
node.tunnel_connected_at_unix_secs = now;
node.updated_at_unix_secs = now;
Ok(Some(node.clone()))
}
async fn update_remote_config(
&self,
mutation: &ProxyNodeRemoteConfigMutation,
) -> Result<Option<StoredProxyNode>, DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
let Some(node) = nodes.get_mut(&mutation.node_id) else {
return Ok(None);
};
if node.is_manual {
return Err(DataLayerError::InvalidInput(
"手动节点不支持远程配置下发".to_string(),
));
}
if let Some(node_name) = mutation.node_name.as_ref() {
node.name = node_name.clone();
}
node.remote_config = Self::normalize_remote_config(mutation, node.remote_config.as_ref());
node.config_version = node.config_version.saturating_add(1);
node.updated_at_unix_secs = Self::now_unix_secs();
Ok(Some(node.clone()))
}
}
#[cfg(test)]
mod tests {
use super::InMemoryProxyNodeRepository;
use crate::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyNode, StoredProxyNodeEvent,
};
use serde_json::json;
@@ -334,6 +546,7 @@ mod tests {
.expect("node should exist");
assert_eq!(updated.status, "offline");
assert!(!updated.tunnel_connected);
assert_eq!(updated.active_connections, 0);
let events = repository
.list_proxy_node_events("node-1", 10)
@@ -386,4 +599,133 @@ mod tests {
assert_eq!(events[0].id, 2);
assert_eq!(events[0].detail.as_deref(), Some("newer"));
}
#[tokio::test]
async fn resets_stale_tunnel_statuses_without_touching_manual_nodes() {
let mut stale_tunnel = sample_node();
stale_tunnel.tunnel_connected = true;
stale_tunnel.status = "online".to_string();
stale_tunnel.active_connections = 6;
let mut manual_node = sample_node();
manual_node.id = "manual-node".to_string();
manual_node.is_manual = true;
manual_node.tunnel_connected = true;
manual_node.status = "online".to_string();
manual_node.active_connections = 4;
let repository = InMemoryProxyNodeRepository::seed(vec![stale_tunnel, manual_node]);
let updated = repository
.reset_stale_tunnel_statuses()
.await
.expect("reset should succeed");
assert_eq!(updated, 1);
let stale = repository
.find_proxy_node("node-1")
.await
.expect("lookup should succeed")
.expect("stale node should exist");
assert_eq!(stale.status, "offline");
assert!(!stale.tunnel_connected);
assert_eq!(stale.active_connections, 0);
let manual = repository
.find_proxy_node("manual-node")
.await
.expect("lookup should succeed")
.expect("manual node should exist");
assert_eq!(manual.status, "online");
assert!(manual.tunnel_connected);
assert_eq!(manual.active_connections, 4);
}
#[tokio::test]
async fn registers_updates_config_and_unregisters_nodes() {
let repository = InMemoryProxyNodeRepository::default();
let registered = repository
.register_node(&ProxyNodeRegistrationMutation {
name: "proxy-01".to_string(),
ip: "127.0.0.1".to_string(),
port: 0,
region: Some("test".to_string()),
heartbeat_interval: 30,
active_connections: Some(1),
total_requests: Some(2),
avg_latency_ms: Some(3.5),
hardware_info: Some(json!({"cpu": "arm64"})),
estimated_max_concurrency: Some(128),
proxy_metadata: Some(json!({"arch": "arm64"})),
proxy_version: Some("2.1.0".to_string()),
registered_by: Some("admin-1".to_string()),
tunnel_mode: true,
})
.await
.expect("register should succeed");
assert_eq!(registered.status, "offline");
assert_eq!(registered.total_requests, 2);
assert_eq!(registered.config_version, 0);
let updated = repository
.update_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: registered.id.clone(),
node_name: Some("proxy-02".to_string()),
allowed_ports: Some(vec![443, 8443]),
log_level: Some("info".to_string()),
heartbeat_interval: Some(45),
scheduling_state: Some(Some("draining".to_string())),
upgrade_to: Some(Some("2.2.0".to_string())),
})
.await
.expect("config update should succeed")
.expect("node should exist");
assert_eq!(updated.name, "proxy-02");
assert_eq!(updated.config_version, 1);
assert_eq!(
updated
.remote_config
.as_ref()
.and_then(|value| value.get("scheduling_state")),
Some(&json!("draining"))
);
assert_eq!(
updated
.remote_config
.as_ref()
.and_then(|value| value.get("upgrade_to")),
Some(&json!("2.2.0"))
);
let after_upgrade = repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: registered.id.clone(),
heartbeat_interval: None,
active_connections: Some(2),
total_requests_delta: Some(1),
avg_latency_ms: Some(2.0),
failed_requests_delta: Some(0),
dns_failures_delta: Some(0),
stream_errors_delta: Some(0),
proxy_metadata: Some(json!({"arch": "arm64"})),
proxy_version: Some("proxy-v2.2.0".to_string()),
})
.await
.expect("heartbeat should succeed")
.expect("node should exist");
assert_eq!(after_upgrade.config_version, 2);
assert!(after_upgrade
.remote_config
.as_ref()
.and_then(|value| value.get("upgrade_to"))
.is_none());
let unregistered = repository
.unregister_node(&registered.id)
.await
.expect("unregister should succeed")
.expect("node should exist");
assert_eq!(unregistered.status, "offline");
assert!(!unregistered.tunnel_connected);
}
}

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

@@ -5,6 +5,10 @@ mod types;
pub use memory::InMemoryProxyNodeRepository;
pub use sql::SqlxProxyNodeRepository;
pub use types::{
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeTunnelStatusMutation,
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,
};

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

@@ -1,8 +1,10 @@
use async_trait::async_trait;
use sha2::{Digest, Sha256};
use sqlx::{postgres::PgRow, PgPool, Row};
use super::types::{
normalize_proxy_metadata, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
};
use crate::{
@@ -129,6 +131,140 @@ SET
WHERE id = $1
"#;
const FIND_EXISTING_TUNNEL_NODE_SQL: &str = r#"
SELECT
id,
name,
ip,
port,
region,
is_manual,
proxy_url,
proxy_username,
proxy_password,
CAST(status AS TEXT) AS status,
registered_by,
EXTRACT(EPOCH FROM last_heartbeat_at)::bigint AS last_heartbeat_at_unix_secs,
heartbeat_interval,
active_connections,
total_requests,
CAST(avg_latency_ms AS DOUBLE PRECISION) AS avg_latency_ms,
failed_requests,
dns_failures,
stream_errors,
proxy_metadata,
hardware_info,
estimated_max_concurrency,
tunnel_mode,
tunnel_connected,
EXTRACT(EPOCH FROM tunnel_connected_at)::bigint AS tunnel_connected_at_unix_secs,
remote_config,
config_version,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM proxy_nodes
WHERE ip = $1
AND port = $2
AND is_manual = FALSE
ORDER BY created_at ASC, id ASC
LIMIT 1
FOR UPDATE
"#;
const INSERT_PROXY_NODE_SQL: &str = r#"
INSERT INTO proxy_nodes (
id,
name,
ip,
port,
region,
status,
registered_by,
last_heartbeat_at,
heartbeat_interval,
active_connections,
total_requests,
avg_latency_ms,
hardware_info,
estimated_max_concurrency,
tunnel_mode,
tunnel_connected,
proxy_metadata
)
VALUES (
$1,
$2,
$3,
$4,
$5,
'offline'::proxynodestatus,
$6,
NOW(),
$7,
COALESCE($8, 0),
COALESCE($9, 0),
$10,
$11,
$12,
$13,
FALSE,
$14
)
"#;
const UPDATE_PROXY_NODE_REGISTRATION_SQL: &str = r#"
UPDATE proxy_nodes
SET
name = $2,
ip = $3,
port = $4,
region = $5,
registered_by = $6,
last_heartbeat_at = NOW(),
heartbeat_interval = $7,
active_connections = COALESCE($8, active_connections),
total_requests = COALESCE($9, total_requests),
avg_latency_ms = COALESCE($10, avg_latency_ms),
hardware_info = COALESCE($11, hardware_info),
estimated_max_concurrency = COALESCE($12, estimated_max_concurrency),
tunnel_mode = $13,
proxy_metadata = COALESCE($14, proxy_metadata),
updated_at = NOW()
WHERE id = $1
"#;
const UNREGISTER_PROXY_NODE_SQL: &str = r#"
UPDATE proxy_nodes
SET
status = 'offline'::proxynodestatus,
tunnel_connected = FALSE,
tunnel_connected_at = NOW(),
updated_at = NOW()
WHERE id = $1
"#;
const UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL: &str = r#"
UPDATE proxy_nodes
SET
name = COALESCE($2, name),
remote_config = $3,
config_version = config_version + 1,
updated_at = NOW()
WHERE id = $1
"#;
const RESET_STALE_TUNNEL_STATUSES_SQL: &str = r#"
UPDATE proxy_nodes
SET
tunnel_connected = FALSE,
status = 'offline'::proxynodestatus,
active_connections = 0,
tunnel_connected_at = NOW(),
updated_at = NOW()
WHERE is_manual = FALSE
AND tunnel_connected = TRUE
"#;
#[derive(Debug, Clone)]
pub struct SqlxProxyNodeRepository {
pool: PgPool,
@@ -199,6 +335,80 @@ impl SqlxProxyNodeRepository {
),
})
}
fn registration_lock_key(ip: &str, port: i32) -> i64 {
let mut hasher = Sha256::new();
hasher.update(ip.as_bytes());
hasher.update(b":");
hasher.update(port.to_string().as_bytes());
let digest = hasher.finalize();
let mut bytes = [0u8; 8];
bytes.copy_from_slice(&digest[..8]);
i64::from_be_bytes(bytes)
}
fn normalize_remote_config(
mutation: &ProxyNodeRemoteConfigMutation,
existing: Option<&serde_json::Value>,
) -> Option<serde_json::Value> {
let mut config = match existing {
Some(serde_json::Value::Object(map)) => map.clone(),
_ => serde_json::Map::new(),
};
if let Some(node_name) = mutation.node_name.as_ref() {
config.insert(
"node_name".to_string(),
serde_json::Value::String(node_name.clone()),
);
}
if let Some(allowed_ports) = mutation.allowed_ports.as_ref() {
config.insert(
"allowed_ports".to_string(),
serde_json::json!(allowed_ports),
);
}
if let Some(log_level) = mutation.log_level.as_ref() {
config.insert(
"log_level".to_string(),
serde_json::Value::String(log_level.clone()),
);
}
if let Some(heartbeat_interval) = mutation.heartbeat_interval {
config.insert(
"heartbeat_interval".to_string(),
serde_json::json!(heartbeat_interval),
);
}
if let Some(scheduling_state) = mutation.scheduling_state.as_ref() {
match scheduling_state {
Some(state) => {
config.insert(
"scheduling_state".to_string(),
serde_json::Value::String(state.clone()),
);
}
None => {
config.remove("scheduling_state");
}
}
}
if let Some(upgrade_to) = mutation.upgrade_to.as_ref() {
match upgrade_to {
Some(version) => {
config.insert(
"upgrade_to".to_string(),
serde_json::Value::String(version.clone()),
);
}
None => {
config.remove("upgrade_to");
}
}
}
(!config.is_empty()).then_some(serde_json::Value::Object(config))
}
}
#[async_trait]
@@ -240,6 +450,88 @@ impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
#[async_trait]
impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
async fn reset_stale_tunnel_statuses(&self) -> Result<usize, DataLayerError> {
let result = sqlx::query(RESET_STALE_TUNNEL_STATUSES_SQL)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(result.rows_affected() as usize)
}
async fn register_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
) -> Result<StoredProxyNode, DataLayerError> {
let normalized_proxy_metadata = normalize_proxy_metadata(
mutation.proxy_metadata.as_ref(),
mutation.proxy_version.as_deref(),
);
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()?;
let existing = sqlx::query(FIND_EXISTING_TUNNEL_NODE_SQL)
.bind(&mutation.ip)
.bind(mutation.port)
.fetch_optional(&mut *tx)
.await
.map_postgres_err()?;
let node_id = if let Some(row) = existing.as_ref() {
let existing = Self::row_to_stored(row)?;
sqlx::query(UPDATE_PROXY_NODE_REGISTRATION_SQL)
.bind(&existing.id)
.bind(&mutation.name)
.bind(&mutation.ip)
.bind(mutation.port)
.bind(mutation.region.as_deref())
.bind(mutation.registered_by.as_deref())
.bind(mutation.heartbeat_interval)
.bind(mutation.active_connections)
.bind(mutation.total_requests)
.bind(mutation.avg_latency_ms)
.bind(mutation.hardware_info.as_ref())
.bind(mutation.estimated_max_concurrency)
.bind(mutation.tunnel_mode)
.bind(normalized_proxy_metadata.as_ref())
.execute(&mut *tx)
.await
.map_postgres_err()?;
existing.id
} else {
let node_id = uuid::Uuid::new_v4().to_string();
sqlx::query(INSERT_PROXY_NODE_SQL)
.bind(&node_id)
.bind(&mutation.name)
.bind(&mutation.ip)
.bind(mutation.port)
.bind(mutation.region.as_deref())
.bind(mutation.registered_by.as_deref())
.bind(mutation.heartbeat_interval)
.bind(mutation.active_connections)
.bind(mutation.total_requests)
.bind(mutation.avg_latency_ms)
.bind(mutation.hardware_info.as_ref())
.bind(mutation.estimated_max_concurrency)
.bind(mutation.tunnel_mode)
.bind(normalized_proxy_metadata.as_ref())
.execute(&mut *tx)
.await
.map_postgres_err()?;
node_id
};
tx.commit().await.map_err(postgres_error)?;
self.find_proxy_node(&node_id).await?.ok_or_else(|| {
DataLayerError::UnexpectedValue("registered proxy node missing".to_string())
})
}
async fn apply_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,
@@ -274,7 +566,30 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
.await
.map_postgres_err()?;
self.find_proxy_node(&mutation.node_id).await
let updated = self.find_proxy_node(&mutation.node_id).await?;
let Some(updated) = updated else {
return Ok(None);
};
if reconcile_remote_config_after_heartbeat(
updated.remote_config.as_ref(),
mutation.proxy_version.as_deref(),
) != updated.remote_config
{
return self
.update_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: mutation.node_id.clone(),
node_name: None,
allowed_ports: None,
log_level: None,
heartbeat_interval: None,
scheduling_state: None,
upgrade_to: Some(None),
})
.await;
}
Ok(Some(updated))
}
async fn update_tunnel_status(
@@ -332,6 +647,10 @@ VALUES (
UPDATE proxy_nodes
SET
tunnel_connected = $2,
active_connections = CASE
WHEN $2 THEN active_connections
ELSE 0
END,
tunnel_connected_at = CASE
WHEN $3::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($3::double precision)
@@ -379,4 +698,49 @@ VALUES (
tx.commit().await.map_err(postgres_error)?;
self.find_proxy_node(&mutation.node_id).await
}
async fn unregister_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(UNREGISTER_PROXY_NODE_SQL)
.bind(node_id)
.execute(&self.pool)
.await
.map_postgres_err()?;
self.find_proxy_node(&existing.id).await
}
async fn update_remote_config(
&self,
mutation: &ProxyNodeRemoteConfigMutation,
) -> 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 remote_config =
Self::normalize_remote_config(mutation, existing.remote_config.as_ref());
sqlx::query(UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL)
.bind(&mutation.node_id)
.bind(mutation.node_name.as_deref())
.bind(remote_config.as_ref())
.execute(&self.pool)
.await
.map_postgres_err()?;
self.find_proxy_node(&mutation.node_id).await
}
}

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

@@ -162,6 +162,24 @@ pub struct ProxyNodeHeartbeatMutation {
pub proxy_version: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProxyNodeRegistrationMutation {
pub name: String,
pub ip: String,
pub port: i32,
pub region: Option<String>,
pub heartbeat_interval: i32,
pub active_connections: Option<i32>,
pub total_requests: Option<i64>,
pub avg_latency_ms: Option<f64>,
pub hardware_info: Option<serde_json::Value>,
pub estimated_max_concurrency: Option<i32>,
pub proxy_metadata: Option<serde_json::Value>,
pub proxy_version: Option<String>,
pub registered_by: Option<String>,
pub tunnel_mode: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProxyNodeTunnelStatusMutation {
pub node_id: String,
@@ -171,6 +189,17 @@ pub struct ProxyNodeTunnelStatusMutation {
pub observed_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProxyNodeRemoteConfigMutation {
pub node_id: String,
pub node_name: Option<String>,
pub allowed_ports: Option<Vec<u16>>,
pub log_level: Option<String>,
pub heartbeat_interval: Option<i32>,
pub scheduling_state: Option<Option<String>>,
pub upgrade_to: Option<Option<String>>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredProxyNodeEvent {
pub id: i64,
@@ -214,6 +243,91 @@ pub fn normalize_proxy_metadata(
}
}
fn normalize_proxy_version_label(value: &str) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
Some(
trimmed
.strip_prefix("proxy-v")
.unwrap_or(trimmed)
.to_ascii_lowercase(),
)
}
pub const PROXY_NODE_SCHEDULING_STATE_DRAINING: &str = "draining";
pub const PROXY_NODE_SCHEDULING_STATE_CORDONED: &str = "cordoned";
pub fn normalize_proxy_node_scheduling_state(value: &str) -> Option<&'static str> {
let trimmed = value.trim();
if trimmed.eq_ignore_ascii_case(PROXY_NODE_SCHEDULING_STATE_DRAINING) {
return Some(PROXY_NODE_SCHEDULING_STATE_DRAINING);
}
if trimmed.eq_ignore_ascii_case(PROXY_NODE_SCHEDULING_STATE_CORDONED) {
return Some(PROXY_NODE_SCHEDULING_STATE_CORDONED);
}
None
}
pub fn remote_config_scheduling_state(
remote_config: Option<&serde_json::Value>,
) -> Option<&'static str> {
remote_config
.and_then(serde_json::Value::as_object)
.and_then(|value| value.get("scheduling_state"))
.and_then(serde_json::Value::as_str)
.and_then(normalize_proxy_node_scheduling_state)
}
pub fn proxy_node_accepts_new_tunnels(node: &StoredProxyNode) -> bool {
remote_config_scheduling_state(node.remote_config.as_ref()).is_none()
}
pub fn proxy_reported_version(proxy_metadata: Option<&serde_json::Value>) -> Option<String> {
proxy_metadata
.and_then(serde_json::Value::as_object)
.and_then(|value| value.get("version"))
.and_then(serde_json::Value::as_str)
.and_then(normalize_proxy_version_label)
}
pub fn remote_config_upgrade_target(remote_config: Option<&serde_json::Value>) -> Option<String> {
remote_config
.and_then(serde_json::Value::as_object)
.and_then(|value| value.get("upgrade_to"))
.and_then(serde_json::Value::as_str)
.and_then(normalize_proxy_version_label)
}
pub fn reconcile_remote_config_after_heartbeat(
remote_config: Option<&serde_json::Value>,
proxy_version: Option<&str>,
) -> Option<serde_json::Value> {
let Some(mut config) = remote_config
.and_then(serde_json::Value::as_object)
.cloned()
else {
return remote_config.cloned();
};
let Some(target_version) = config
.get("upgrade_to")
.and_then(serde_json::Value::as_str)
.and_then(normalize_proxy_version_label)
else {
return Some(serde_json::Value::Object(config));
};
let Some(reported_version) = proxy_version.and_then(normalize_proxy_version_label) else {
return Some(serde_json::Value::Object(config));
};
if reported_version == target_version {
config.remove("upgrade_to");
}
(!config.is_empty()).then_some(serde_json::Value::Object(config))
}
#[async_trait]
pub trait ProxyNodeReadRepository: Send + Sync {
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, crate::DataLayerError>;
@@ -232,6 +346,13 @@ pub trait ProxyNodeReadRepository: Send + Sync {
#[async_trait]
pub trait ProxyNodeWriteRepository: Send + Sync {
async fn reset_stale_tunnel_statuses(&self) -> Result<usize, crate::DataLayerError>;
async fn register_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
) -> Result<StoredProxyNode, crate::DataLayerError>;
async fn apply_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,
@@ -241,4 +362,108 @@ pub trait ProxyNodeWriteRepository: Send + Sync {
&self,
mutation: &ProxyNodeTunnelStatusMutation,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
async fn unregister_node(
&self,
node_id: &str,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
async fn update_remote_config(
&self,
mutation: &ProxyNodeRemoteConfigMutation,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
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, StoredProxyNode,
};
#[test]
fn normalizes_reported_versions_and_clears_completed_upgrade_targets() {
let remote_config = json!({
"node_name": "edge-1",
"upgrade_to": "proxy-v2.0.0",
});
let proxy_metadata = json!({
"version": "2.0.0",
"arch": "arm64",
});
assert_eq!(
proxy_reported_version(Some(&proxy_metadata)).as_deref(),
Some("2.0.0")
);
assert_eq!(
remote_config_upgrade_target(Some(&remote_config)).as_deref(),
Some("2.0.0")
);
let reconciled =
reconcile_remote_config_after_heartbeat(Some(&remote_config), Some("proxy-v2.0.0"))
.expect("reconciled config should remain an object");
assert_eq!(reconciled.get("upgrade_to"), None);
assert_eq!(reconciled.get("node_name"), Some(&json!("edge-1")));
}
#[test]
fn normalizes_proxy_node_scheduling_state_and_detects_unschedulable_nodes() {
assert_eq!(
normalize_proxy_node_scheduling_state("draining"),
Some("draining")
);
assert_eq!(
normalize_proxy_node_scheduling_state(" CORDONED "),
Some("cordoned")
);
assert_eq!(normalize_proxy_node_scheduling_state("active"), None);
let remote_config = json!({
"node_name": "edge-1",
"scheduling_state": "draining",
});
assert_eq!(
remote_config_scheduling_state(Some(&remote_config)),
Some("draining")
);
let node = StoredProxyNode::new(
"node-1".to_string(),
"edge-1".to_string(),
"127.0.0.1".to_string(),
0,
false,
"online".to_string(),
30,
0,
0,
0,
0,
0,
true,
true,
0,
)
.expect("node should build")
.with_runtime_fields(
None,
None,
Some(1_800_000_000),
None,
None,
None,
None,
Some(1_800_000_001),
Some(remote_config),
Some(1_800_000_000),
Some(1_800_000_001),
);
assert!(!proxy_node_accepts_new_tunnels(&node));
}
}