mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: ProxyNode 代理节点管理系统与 OpenAI Responses API 解析增强
ProxyNode 系统:新增 aether-proxy(Rust)海外 VPS 代理组件,后端实现节点注册/心跳/ HMAC 认证/健康检测调度器/模块化集成,前端新增代理节点管理页面。ProxyConfig 支持 node_id 模式,http_client 支持 HMAC 签名代理 URL 构建与 TTL 缓存。 OpenAI CLI 解析器:适配 Responses API 格式,支持 input_tokens/output_tokens 提取、 output[].content[].text 文本解析、response.completed 流式事件 usage 嵌套结构。
This commit is contained in:
@@ -21,6 +21,10 @@ JWT_SECRET_KEY=change-this-to-a-secure-random-string
|
|||||||
# 注意:更换此密钥后需要在管理面板重新配置所有 Provider API Key
|
# 注意:更换此密钥后需要在管理面板重新配置所有 Provider API Key
|
||||||
ENCRYPTION_KEY=change-this-to-another-secure-random-string
|
ENCRYPTION_KEY=change-this-to-another-secure-random-string
|
||||||
|
|
||||||
|
# 代理节点 HMAC 密钥(用于 aether-proxy 认证)
|
||||||
|
# 可选:不设置时会从 ENCRYPTION_KEY 派生(推荐生产环境显式设置)
|
||||||
|
# PROXY_HMAC_KEY=change-this-to-a-proxy-hmac-key
|
||||||
|
|
||||||
# 管理员账号(仅首次初始化时使用, 创建完成后可在系统内修改密码)
|
# 管理员账号(仅首次初始化时使用, 创建完成后可在系统内修改密码)
|
||||||
ADMIN_EMAIL=admin@example.com
|
ADMIN_EMAIL=admin@example.com
|
||||||
ADMIN_USERNAME=admin
|
ADMIN_USERNAME=admin
|
||||||
|
|||||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -235,3 +235,4 @@ src/_version.py
|
|||||||
|
|
||||||
# Analysis folder (third-party code for reference)
|
# Analysis folder (third-party code for reference)
|
||||||
analysis/
|
analysis/
|
||||||
|
/aether-proxy/target/
|
||||||
|
|||||||
31
aether-proxy/.env.example
Normal file
31
aether-proxy/.env.example
Normal file
@@ -0,0 +1,31 @@
|
|||||||
|
# Aether server URL
|
||||||
|
AETHER_PROXY_AETHER_URL=https://aether.example.com
|
||||||
|
|
||||||
|
# Management Token (ae_xxx, must belong to an ADMIN user)
|
||||||
|
AETHER_PROXY_MANAGEMENT_TOKEN=ae_xxxxx
|
||||||
|
|
||||||
|
# HMAC key (must match Aether's PROXY_HMAC_KEY)
|
||||||
|
AETHER_PROXY_HMAC_KEY=
|
||||||
|
|
||||||
|
# Proxy listen port
|
||||||
|
AETHER_PROXY_LISTEN_PORT=18080
|
||||||
|
|
||||||
|
# Public IP (auto-detected if omitted)
|
||||||
|
# AETHER_PROXY_PUBLIC_IP=203.0.113.42
|
||||||
|
|
||||||
|
# Node identification
|
||||||
|
AETHER_PROXY_NODE_NAME=proxy-01
|
||||||
|
# AETHER_PROXY_NODE_REGION=ap-northeast-1
|
||||||
|
|
||||||
|
# Heartbeat interval in seconds
|
||||||
|
AETHER_PROXY_HEARTBEAT_INTERVAL=30
|
||||||
|
|
||||||
|
# Allowed destination ports (comma-separated)
|
||||||
|
AETHER_PROXY_ALLOWED_PORTS=80,443,8080,8443
|
||||||
|
|
||||||
|
# HMAC timestamp tolerance in seconds
|
||||||
|
AETHER_PROXY_TIMESTAMP_TOLERANCE=300
|
||||||
|
|
||||||
|
# Logging
|
||||||
|
AETHER_PROXY_LOG_LEVEL=info
|
||||||
|
AETHER_PROXY_LOG_JSON=false
|
||||||
1986
aether-proxy/Cargo.lock
generated
Normal file
1986
aether-proxy/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load Diff
30
aether-proxy/Cargo.toml
Normal file
30
aether-proxy/Cargo.toml
Normal file
@@ -0,0 +1,30 @@
|
|||||||
|
[package]
|
||||||
|
name = "aether-proxy"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
description = "Forward proxy for Aether with HMAC authentication"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
tokio = { version = "1", features = ["full"] }
|
||||||
|
hyper = { version = "1", features = ["http1", "server"] }
|
||||||
|
hyper-util = { version = "0.1", features = ["tokio", "http1", "server"] }
|
||||||
|
http-body-util = "0.1"
|
||||||
|
reqwest = { version = "0.12", features = ["json"] }
|
||||||
|
hmac = "0.12"
|
||||||
|
sha2 = "0.10"
|
||||||
|
subtle = "2"
|
||||||
|
base64 = "0.22"
|
||||||
|
clap = { version = "4", features = ["derive", "env"] }
|
||||||
|
tracing = "0.1"
|
||||||
|
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
serde_json = "1"
|
||||||
|
thiserror = "2"
|
||||||
|
bytes = "1"
|
||||||
|
hex = "0.4"
|
||||||
|
anyhow = "1"
|
||||||
|
|
||||||
|
[profile.release]
|
||||||
|
lto = true
|
||||||
|
strip = true
|
||||||
|
codegen-units = 1
|
||||||
25
aether-proxy/Dockerfile
Normal file
25
aether-proxy/Dockerfile
Normal file
@@ -0,0 +1,25 @@
|
|||||||
|
FROM rust:1.83-slim AS builder
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
RUN apt-get update && apt-get install -y pkg-config libssl-dev && rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
COPY Cargo.toml Cargo.lock* ./
|
||||||
|
# Create dummy main.rs for dependency caching
|
||||||
|
RUN mkdir src && echo "fn main() {}" > src/main.rs
|
||||||
|
RUN cargo build --release 2>/dev/null || true
|
||||||
|
|
||||||
|
COPY src/ src/
|
||||||
|
# Touch main.rs to force rebuild with real source
|
||||||
|
RUN touch src/main.rs
|
||||||
|
RUN cargo build --release
|
||||||
|
|
||||||
|
FROM debian:bookworm-slim
|
||||||
|
|
||||||
|
RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
COPY --from=builder /app/target/release/aether-proxy /usr/local/bin/aether-proxy
|
||||||
|
|
||||||
|
EXPOSE 18080
|
||||||
|
|
||||||
|
ENTRYPOINT ["aether-proxy"]
|
||||||
182
aether-proxy/src/auth/hmac.rs
Normal file
182
aether-proxy/src/auth/hmac.rs
Normal file
@@ -0,0 +1,182 @@
|
|||||||
|
use base64::Engine;
|
||||||
|
use hmac::{Hmac, Mac};
|
||||||
|
use sha2::Sha256;
|
||||||
|
use subtle::ConstantTimeEq;
|
||||||
|
|
||||||
|
use crate::config::Config;
|
||||||
|
|
||||||
|
type HmacSha256 = Hmac<Sha256>;
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum AuthError {
|
||||||
|
MissingHeader,
|
||||||
|
InvalidBasicAuth,
|
||||||
|
InvalidUsername,
|
||||||
|
InvalidPasswordFormat,
|
||||||
|
TimestampParseError,
|
||||||
|
TimestampExpired,
|
||||||
|
SignatureMismatch,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for AuthError {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::MissingHeader => write!(f, "missing Proxy-Authorization header"),
|
||||||
|
Self::InvalidBasicAuth => write!(f, "invalid Basic auth encoding"),
|
||||||
|
Self::InvalidUsername => write!(f, "username must be 'hmac'"),
|
||||||
|
Self::InvalidPasswordFormat => write!(f, "password format must be 'timestamp.signature'"),
|
||||||
|
Self::TimestampParseError => write!(f, "invalid timestamp"),
|
||||||
|
Self::TimestampExpired => write!(f, "timestamp outside tolerance window"),
|
||||||
|
Self::SignatureMismatch => write!(f, "HMAC signature mismatch"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validate Proxy-Authorization header.
|
||||||
|
///
|
||||||
|
/// Expected format: `Basic base64(hmac:{timestamp}.{signature})`
|
||||||
|
/// where signature = hex(HMAC-SHA256(hmac_key, "{timestamp}\n{node_id}"))
|
||||||
|
pub fn validate_proxy_auth(
|
||||||
|
proxy_auth_header: Option<&str>,
|
||||||
|
config: &Config,
|
||||||
|
node_id: &str,
|
||||||
|
) -> Result<(), AuthError> {
|
||||||
|
let header = proxy_auth_header.ok_or(AuthError::MissingHeader)?;
|
||||||
|
|
||||||
|
let encoded = header
|
||||||
|
.strip_prefix("Basic ")
|
||||||
|
.or_else(|| header.strip_prefix("basic "))
|
||||||
|
.ok_or(AuthError::InvalidBasicAuth)?;
|
||||||
|
|
||||||
|
let decoded_bytes = base64::engine::general_purpose::STANDARD
|
||||||
|
.decode(encoded.trim())
|
||||||
|
.map_err(|_| AuthError::InvalidBasicAuth)?;
|
||||||
|
|
||||||
|
let decoded = String::from_utf8(decoded_bytes).map_err(|_| AuthError::InvalidBasicAuth)?;
|
||||||
|
|
||||||
|
// format: hmac:{timestamp}.{signature}
|
||||||
|
let (username, password) = decoded
|
||||||
|
.split_once(':')
|
||||||
|
.ok_or(AuthError::InvalidBasicAuth)?;
|
||||||
|
|
||||||
|
if username != "hmac" {
|
||||||
|
return Err(AuthError::InvalidUsername);
|
||||||
|
}
|
||||||
|
|
||||||
|
let (timestamp_str, signature_hex) = password
|
||||||
|
.split_once('.')
|
||||||
|
.ok_or(AuthError::InvalidPasswordFormat)?;
|
||||||
|
|
||||||
|
// Validate timestamp window
|
||||||
|
let timestamp: u64 = timestamp_str
|
||||||
|
.parse()
|
||||||
|
.map_err(|_| AuthError::TimestampParseError)?;
|
||||||
|
|
||||||
|
let now = std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.expect("system clock before epoch")
|
||||||
|
.as_secs();
|
||||||
|
|
||||||
|
let diff = if now > timestamp {
|
||||||
|
now - timestamp
|
||||||
|
} else {
|
||||||
|
timestamp - now
|
||||||
|
};
|
||||||
|
|
||||||
|
if diff > config.timestamp_tolerance {
|
||||||
|
return Err(AuthError::TimestampExpired);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Recompute signature
|
||||||
|
let payload = format!("{}\n{}", timestamp_str, node_id);
|
||||||
|
let mut mac =
|
||||||
|
HmacSha256::new_from_slice(config.hmac_key.as_bytes()).expect("HMAC accepts any key size");
|
||||||
|
mac.update(payload.as_bytes());
|
||||||
|
let expected = mac.finalize().into_bytes();
|
||||||
|
let expected_hex = hex::encode(expected);
|
||||||
|
|
||||||
|
// Constant-time comparison
|
||||||
|
let sig_bytes = signature_hex.as_bytes();
|
||||||
|
let exp_bytes = expected_hex.as_bytes();
|
||||||
|
|
||||||
|
if sig_bytes.len() != exp_bytes.len() || sig_bytes.ct_eq(exp_bytes).unwrap_u8() != 1 {
|
||||||
|
return Err(AuthError::SignatureMismatch);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn make_config() -> Config {
|
||||||
|
Config {
|
||||||
|
aether_url: String::new(),
|
||||||
|
management_token: String::new(),
|
||||||
|
hmac_key: "test-hmac-key".to_string(),
|
||||||
|
listen_port: 18080,
|
||||||
|
public_ip: None,
|
||||||
|
node_name: "test".to_string(),
|
||||||
|
node_region: None,
|
||||||
|
heartbeat_interval: 30,
|
||||||
|
allowed_ports: vec![80, 443],
|
||||||
|
timestamp_tolerance: 300,
|
||||||
|
log_level: "info".to_string(),
|
||||||
|
log_json: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn make_valid_auth(config: &Config, node_id: &str) -> String {
|
||||||
|
let now = std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_secs();
|
||||||
|
let payload = format!("{}\n{}", now, node_id);
|
||||||
|
let mut mac =
|
||||||
|
HmacSha256::new_from_slice(config.hmac_key.as_bytes()).unwrap();
|
||||||
|
mac.update(payload.as_bytes());
|
||||||
|
let sig = hex::encode(mac.finalize().into_bytes());
|
||||||
|
let cred = format!("hmac:{}.{}", now, sig);
|
||||||
|
let encoded = base64::engine::general_purpose::STANDARD.encode(cred);
|
||||||
|
format!("Basic {}", encoded)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_valid_auth() {
|
||||||
|
let config = make_config();
|
||||||
|
let header = make_valid_auth(&config, "node-1");
|
||||||
|
assert!(validate_proxy_auth(Some(&header), &config, "node-1").is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_wrong_node_id() {
|
||||||
|
let config = make_config();
|
||||||
|
let header = make_valid_auth(&config, "node-1");
|
||||||
|
assert!(matches!(
|
||||||
|
validate_proxy_auth(Some(&header), &config, "node-2"),
|
||||||
|
Err(AuthError::SignatureMismatch)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_missing_header() {
|
||||||
|
let config = make_config();
|
||||||
|
assert!(matches!(
|
||||||
|
validate_proxy_auth(None, &config, "node-1"),
|
||||||
|
Err(AuthError::MissingHeader)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_wrong_username() {
|
||||||
|
let cred = "user:12345.abc";
|
||||||
|
let encoded = base64::engine::general_purpose::STANDARD.encode(cred);
|
||||||
|
let header = format!("Basic {}", encoded);
|
||||||
|
let config = make_config();
|
||||||
|
assert!(matches!(
|
||||||
|
validate_proxy_auth(Some(&header), &config, "node-1"),
|
||||||
|
Err(AuthError::InvalidUsername)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
3
aether-proxy/src/auth/mod.rs
Normal file
3
aether-proxy/src/auth/mod.rs
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
pub mod hmac;
|
||||||
|
|
||||||
|
pub use self::hmac::validate_proxy_auth;
|
||||||
58
aether-proxy/src/config.rs
Normal file
58
aether-proxy/src/config.rs
Normal file
@@ -0,0 +1,58 @@
|
|||||||
|
use clap::Parser;
|
||||||
|
|
||||||
|
/// Aether forward proxy with HMAC authentication.
|
||||||
|
///
|
||||||
|
/// Deployed on overseas VPS to relay API traffic for Aether instances
|
||||||
|
/// behind the GFW. Registers with Aether, sends heartbeats, and validates
|
||||||
|
/// incoming proxy requests via HMAC-SHA256 signatures in Basic Auth.
|
||||||
|
#[derive(Parser, Debug, Clone)]
|
||||||
|
#[command(version, about)]
|
||||||
|
pub struct Config {
|
||||||
|
/// Aether server URL (e.g. https://aether.example.com)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_AETHER_URL")]
|
||||||
|
pub aether_url: String,
|
||||||
|
|
||||||
|
/// Management Token for Aether admin API (ae_xxx)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_MANAGEMENT_TOKEN")]
|
||||||
|
pub management_token: String,
|
||||||
|
|
||||||
|
/// HMAC-SHA256 key for proxy authentication
|
||||||
|
#[arg(long, env = "AETHER_PROXY_HMAC_KEY")]
|
||||||
|
pub hmac_key: String,
|
||||||
|
|
||||||
|
/// Port to listen on for proxy connections
|
||||||
|
#[arg(long, env = "AETHER_PROXY_LISTEN_PORT", default_value_t = 18080)]
|
||||||
|
pub listen_port: u16,
|
||||||
|
|
||||||
|
/// Public IP address of this node (auto-detected if omitted)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_PUBLIC_IP")]
|
||||||
|
pub public_ip: Option<String>,
|
||||||
|
|
||||||
|
/// Human-readable node name
|
||||||
|
#[arg(long, env = "AETHER_PROXY_NODE_NAME", default_value = "proxy-01")]
|
||||||
|
pub node_name: String,
|
||||||
|
|
||||||
|
/// Region label (e.g. ap-northeast-1)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_NODE_REGION")]
|
||||||
|
pub node_region: Option<String>,
|
||||||
|
|
||||||
|
/// Heartbeat interval in seconds
|
||||||
|
#[arg(long, env = "AETHER_PROXY_HEARTBEAT_INTERVAL", default_value_t = 30)]
|
||||||
|
pub heartbeat_interval: u64,
|
||||||
|
|
||||||
|
/// Allowed destination ports (default: 80,443,8080,8443)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_ALLOWED_PORTS", value_delimiter = ',', default_values_t = vec![80, 443, 8080, 8443])]
|
||||||
|
pub allowed_ports: Vec<u16>,
|
||||||
|
|
||||||
|
/// Timestamp tolerance window in seconds for HMAC validation
|
||||||
|
#[arg(long, env = "AETHER_PROXY_TIMESTAMP_TOLERANCE", default_value_t = 300)]
|
||||||
|
pub timestamp_tolerance: u64,
|
||||||
|
|
||||||
|
/// Log level (trace, debug, info, warn, error)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_LOG_LEVEL", default_value = "info")]
|
||||||
|
pub log_level: String,
|
||||||
|
|
||||||
|
/// Output logs as JSON
|
||||||
|
#[arg(long, env = "AETHER_PROXY_LOG_JSON", default_value_t = false)]
|
||||||
|
pub log_json: bool,
|
||||||
|
}
|
||||||
132
aether-proxy/src/main.rs
Normal file
132
aether-proxy/src/main.rs
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
mod auth;
|
||||||
|
mod config;
|
||||||
|
mod proxy;
|
||||||
|
mod registration;
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use clap::Parser;
|
||||||
|
use tokio::signal;
|
||||||
|
use tokio::sync::watch;
|
||||||
|
use tracing::{error, info};
|
||||||
|
|
||||||
|
use config::Config;
|
||||||
|
use registration::client::{detect_public_ip, AetherClient};
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
let config = Config::parse();
|
||||||
|
|
||||||
|
// Initialize tracing
|
||||||
|
init_tracing(&config);
|
||||||
|
|
||||||
|
info!(
|
||||||
|
version = env!("CARGO_PKG_VERSION"),
|
||||||
|
port = config.listen_port,
|
||||||
|
node_name = %config.node_name,
|
||||||
|
"aether-proxy starting"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Resolve public IP
|
||||||
|
let public_ip = match &config.public_ip {
|
||||||
|
Some(ip) => ip.clone(),
|
||||||
|
None => detect_public_ip().await?,
|
||||||
|
};
|
||||||
|
info!(public_ip = %public_ip, "using public IP");
|
||||||
|
|
||||||
|
// Register with Aether
|
||||||
|
let aether_client = Arc::new(AetherClient::new(&config));
|
||||||
|
let node_id = aether_client.register(&config, &public_ip).await?;
|
||||||
|
let node_id = Arc::new(node_id);
|
||||||
|
|
||||||
|
info!(node_id = %node_id, "node registered");
|
||||||
|
|
||||||
|
// Shutdown signal channel
|
||||||
|
let (shutdown_tx, shutdown_rx) = watch::channel(false);
|
||||||
|
|
||||||
|
let config = Arc::new(config);
|
||||||
|
|
||||||
|
// Start heartbeat task
|
||||||
|
let heartbeat_handle = {
|
||||||
|
let client = Arc::clone(&aether_client);
|
||||||
|
let node_id = Arc::clone(&node_id);
|
||||||
|
let interval = config.heartbeat_interval;
|
||||||
|
let rx = shutdown_rx.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
registration::heartbeat::run(client, node_id, interval, rx).await;
|
||||||
|
})
|
||||||
|
};
|
||||||
|
|
||||||
|
// Start proxy server
|
||||||
|
let server_handle = {
|
||||||
|
let config = Arc::clone(&config);
|
||||||
|
let node_id = Arc::clone(&node_id);
|
||||||
|
let rx = shutdown_rx.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(e) = proxy::server::run(config, node_id, rx).await {
|
||||||
|
error!(error = %e, "proxy server error");
|
||||||
|
}
|
||||||
|
})
|
||||||
|
};
|
||||||
|
|
||||||
|
// Wait for shutdown signal (SIGTERM or SIGINT)
|
||||||
|
wait_for_shutdown().await;
|
||||||
|
|
||||||
|
info!("shutdown signal received, cleaning up...");
|
||||||
|
|
||||||
|
// Signal all tasks to stop
|
||||||
|
let _ = shutdown_tx.send(true);
|
||||||
|
|
||||||
|
// Graceful unregister (best-effort)
|
||||||
|
if let Err(e) = aether_client.unregister(&node_id).await {
|
||||||
|
error!(error = %e, "unregister failed during shutdown");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for tasks to finish
|
||||||
|
let _ = tokio::join!(heartbeat_handle, server_handle);
|
||||||
|
|
||||||
|
info!("aether-proxy stopped");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn init_tracing(config: &Config) {
|
||||||
|
use tracing_subscriber::EnvFilter;
|
||||||
|
|
||||||
|
let filter = EnvFilter::try_new(&config.log_level)
|
||||||
|
.unwrap_or_else(|_| EnvFilter::new("info"));
|
||||||
|
|
||||||
|
if config.log_json {
|
||||||
|
tracing_subscriber::fmt()
|
||||||
|
.with_env_filter(filter)
|
||||||
|
.json()
|
||||||
|
.init();
|
||||||
|
} else {
|
||||||
|
tracing_subscriber::fmt()
|
||||||
|
.with_env_filter(filter)
|
||||||
|
.init();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn wait_for_shutdown() {
|
||||||
|
let ctrl_c = async {
|
||||||
|
signal::ctrl_c()
|
||||||
|
.await
|
||||||
|
.expect("failed to install Ctrl+C handler");
|
||||||
|
};
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
let terminate = async {
|
||||||
|
signal::unix::signal(signal::unix::SignalKind::terminate())
|
||||||
|
.expect("failed to install SIGTERM handler")
|
||||||
|
.recv()
|
||||||
|
.await;
|
||||||
|
};
|
||||||
|
|
||||||
|
#[cfg(not(unix))]
|
||||||
|
let terminate = std::future::pending::<()>();
|
||||||
|
|
||||||
|
tokio::select! {
|
||||||
|
_ = ctrl_c => {},
|
||||||
|
_ = terminate => {},
|
||||||
|
}
|
||||||
|
}
|
||||||
134
aether-proxy/src/proxy/connect.rs
Normal file
134
aether-proxy/src/proxy/connect.rs
Normal file
@@ -0,0 +1,134 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use hyper::body::Incoming;
|
||||||
|
use hyper::{Request, Response};
|
||||||
|
use tokio::net::TcpStream;
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use crate::auth;
|
||||||
|
use crate::config::Config;
|
||||||
|
use crate::proxy::target_filter;
|
||||||
|
|
||||||
|
/// Handle HTTP CONNECT tunnel requests.
|
||||||
|
///
|
||||||
|
/// Flow: validate auth -> check target filter -> TCP connect -> 200 -> bidirectional copy
|
||||||
|
pub async fn handle_connect(
|
||||||
|
req: Request<Incoming>,
|
||||||
|
config: Arc<Config>,
|
||||||
|
node_id: &str,
|
||||||
|
allowed_ports: &HashSet<u16>,
|
||||||
|
) -> Response<http_body_util::Empty<bytes::Bytes>> {
|
||||||
|
// Extract Proxy-Authorization header
|
||||||
|
let proxy_auth = req
|
||||||
|
.headers()
|
||||||
|
.get("proxy-authorization")
|
||||||
|
.and_then(|v| v.to_str().ok());
|
||||||
|
|
||||||
|
// HMAC authentication
|
||||||
|
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, node_id) {
|
||||||
|
warn!(error = %e, "CONNECT auth failed");
|
||||||
|
return proxy_auth_required(&e.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse target host:port from CONNECT URI
|
||||||
|
let authority = match req.uri().authority() {
|
||||||
|
Some(auth) => auth.clone(),
|
||||||
|
None => {
|
||||||
|
warn!("CONNECT request missing authority");
|
||||||
|
return bad_request("missing target authority");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let host = authority.host().to_string();
|
||||||
|
let port = authority.port_u16().unwrap_or(443);
|
||||||
|
|
||||||
|
// Target filter: private IP + port whitelist
|
||||||
|
let target_addr = match target_filter::validate_target(&host, port, allowed_ports) {
|
||||||
|
Ok(addr) => addr,
|
||||||
|
Err(e) => {
|
||||||
|
warn!(host = %host, port, error = %e, "CONNECT target rejected");
|
||||||
|
return forbidden(&e.to_string());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(target = %target_addr, "CONNECT tunnel establishing");
|
||||||
|
|
||||||
|
// Connect to target
|
||||||
|
let target_stream = match TcpStream::connect(target_addr).await {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(e) => {
|
||||||
|
warn!(target = %target_addr, error = %e, "CONNECT target connection failed");
|
||||||
|
return bad_gateway(&e.to_string());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Respond 200 and upgrade connection to raw TCP tunnel
|
||||||
|
tokio::task::spawn(async move {
|
||||||
|
match hyper::upgrade::on(req).await {
|
||||||
|
Ok(upgraded) => {
|
||||||
|
let mut upgraded =
|
||||||
|
hyper_util::rt::TokioIo::new(upgraded);
|
||||||
|
let mut target = target_stream;
|
||||||
|
|
||||||
|
match tokio::io::copy_bidirectional(&mut upgraded, &mut target).await {
|
||||||
|
Ok((from_client, from_target)) => {
|
||||||
|
debug!(
|
||||||
|
from_client,
|
||||||
|
from_target,
|
||||||
|
"CONNECT tunnel closed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
debug!(error = %e, "CONNECT tunnel error");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
warn!(error = %e, "CONNECT upgrade failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
Response::builder()
|
||||||
|
.status(200)
|
||||||
|
.body(http_body_util::Empty::new())
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn proxy_auth_required(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(407)
|
||||||
|
.header("Proxy-Authenticate", "HMAC-SHA256")
|
||||||
|
.header("Content-Length", "0")
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(http_body_util::Empty::new())
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn forbidden(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(403)
|
||||||
|
.header("Content-Length", "0")
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(http_body_util::Empty::new())
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bad_request(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(400)
|
||||||
|
.header("Content-Length", "0")
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(http_body_util::Empty::new())
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bad_gateway(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(502)
|
||||||
|
.header("Content-Length", "0")
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(http_body_util::Empty::new())
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
4
aether-proxy/src/proxy/mod.rs
Normal file
4
aether-proxy/src/proxy/mod.rs
Normal file
@@ -0,0 +1,4 @@
|
|||||||
|
pub mod connect;
|
||||||
|
pub mod plain;
|
||||||
|
pub mod server;
|
||||||
|
pub mod target_filter;
|
||||||
162
aether-proxy/src/proxy/plain.rs
Normal file
162
aether-proxy/src/proxy/plain.rs
Normal file
@@ -0,0 +1,162 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use http_body_util::{BodyExt, Full};
|
||||||
|
use hyper::body::Incoming;
|
||||||
|
use hyper::{Request, Response};
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use crate::auth;
|
||||||
|
use crate::config::Config;
|
||||||
|
use crate::proxy::target_filter;
|
||||||
|
|
||||||
|
/// Handle plain HTTP forward proxy requests (non-CONNECT).
|
||||||
|
///
|
||||||
|
/// Flow: validate auth -> check target filter -> forward request -> return response
|
||||||
|
pub async fn handle_plain(
|
||||||
|
req: Request<Incoming>,
|
||||||
|
config: Arc<Config>,
|
||||||
|
node_id: &str,
|
||||||
|
allowed_ports: &HashSet<u16>,
|
||||||
|
) -> Response<Full<bytes::Bytes>> {
|
||||||
|
// Extract Proxy-Authorization header
|
||||||
|
let proxy_auth = req
|
||||||
|
.headers()
|
||||||
|
.get("proxy-authorization")
|
||||||
|
.and_then(|v| v.to_str().ok());
|
||||||
|
|
||||||
|
// HMAC authentication
|
||||||
|
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, node_id) {
|
||||||
|
warn!(error = %e, "HTTP proxy auth failed");
|
||||||
|
return proxy_auth_required(&e.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse target from absolute URI
|
||||||
|
let uri = req.uri().clone();
|
||||||
|
let host = match uri.host() {
|
||||||
|
Some(h) => h.to_string(),
|
||||||
|
None => {
|
||||||
|
warn!(uri = %uri, "HTTP proxy request missing host");
|
||||||
|
return bad_request("missing host in URI");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let port = uri.port_u16().unwrap_or(80);
|
||||||
|
|
||||||
|
// Target filter
|
||||||
|
let target_addr = match target_filter::validate_target(&host, port, allowed_ports) {
|
||||||
|
Ok(addr) => addr,
|
||||||
|
Err(e) => {
|
||||||
|
warn!(host = %host, port, error = %e, "HTTP proxy target rejected");
|
||||||
|
return forbidden(&e.to_string());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(target = %target_addr, method = %req.method(), "HTTP proxy forwarding");
|
||||||
|
|
||||||
|
// Build outgoing request (strip proxy headers, use relative URI)
|
||||||
|
let path_and_query = uri
|
||||||
|
.path_and_query()
|
||||||
|
.map(|pq| pq.as_str())
|
||||||
|
.unwrap_or("/");
|
||||||
|
|
||||||
|
let mut builder = Request::builder()
|
||||||
|
.method(req.method())
|
||||||
|
.uri(path_and_query)
|
||||||
|
.version(req.version());
|
||||||
|
|
||||||
|
// Copy headers, skipping proxy-specific ones
|
||||||
|
for (name, value) in req.headers() {
|
||||||
|
if name == "proxy-authorization" || name == "proxy-connection" {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
builder = builder.header(name, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Collect the incoming body
|
||||||
|
let body_bytes = match req.into_body().collect().await {
|
||||||
|
Ok(collected) => collected.to_bytes(),
|
||||||
|
Err(e) => {
|
||||||
|
warn!(error = %e, "failed to read request body");
|
||||||
|
return bad_gateway("failed to read request body");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Connect and send via raw TCP + hyper client
|
||||||
|
let stream = match tokio::net::TcpStream::connect(target_addr).await {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(e) => {
|
||||||
|
warn!(target = %target_addr, error = %e, "HTTP proxy connection failed");
|
||||||
|
return bad_gateway(&format!("connection failed: {}", e));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let io = hyper_util::rt::TokioIo::new(stream);
|
||||||
|
let (mut sender, conn) = match hyper::client::conn::http1::handshake(io).await {
|
||||||
|
Ok(pair) => pair,
|
||||||
|
Err(e) => {
|
||||||
|
warn!(error = %e, "HTTP handshake failed");
|
||||||
|
return bad_gateway(&format!("handshake failed: {}", e));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
tokio::task::spawn(async move {
|
||||||
|
if let Err(e) = conn.await {
|
||||||
|
debug!(error = %e, "HTTP proxy client connection error");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let outgoing = builder
|
||||||
|
.body(Full::new(body_bytes))
|
||||||
|
.expect("failed to build outgoing request");
|
||||||
|
|
||||||
|
match sender.send_request(outgoing).await {
|
||||||
|
Ok(resp) => {
|
||||||
|
let (parts, body) = resp.into_parts();
|
||||||
|
let body_bytes = match body.collect().await {
|
||||||
|
Ok(collected) => collected.to_bytes(),
|
||||||
|
Err(e) => {
|
||||||
|
warn!(error = %e, "failed to read response body");
|
||||||
|
return bad_gateway("failed to read response body");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
Response::from_parts(parts, Full::new(body_bytes))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
warn!(error = %e, "HTTP proxy request failed");
|
||||||
|
bad_gateway(&format!("upstream request failed: {}", e))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn proxy_auth_required(msg: &str) -> Response<Full<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(407)
|
||||||
|
.header("Proxy-Authenticate", "HMAC-SHA256")
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(Full::new(bytes::Bytes::new()))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn forbidden(msg: &str) -> Response<Full<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(403)
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(Full::new(bytes::Bytes::new()))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bad_request(msg: &str) -> Response<Full<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(400)
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(Full::new(bytes::Bytes::new()))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bad_gateway(msg: &str) -> Response<Full<bytes::Bytes>> {
|
||||||
|
Response::builder()
|
||||||
|
.status(502)
|
||||||
|
.header("X-Error", msg)
|
||||||
|
.body(Full::new(bytes::Bytes::new()))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
117
aether-proxy/src/proxy/server.rs
Normal file
117
aether-proxy/src/proxy/server.rs
Normal file
@@ -0,0 +1,117 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
use std::net::SocketAddr;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use http_body_util::BodyExt;
|
||||||
|
use hyper::body::Incoming;
|
||||||
|
use hyper::server::conn::http1;
|
||||||
|
use hyper::service::service_fn;
|
||||||
|
use hyper::{Method, Request};
|
||||||
|
use hyper_util::rt::TokioIo;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
use tokio::sync::watch;
|
||||||
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
|
use crate::config::Config;
|
||||||
|
use crate::proxy::{connect, plain};
|
||||||
|
|
||||||
|
/// Start the proxy server.
|
||||||
|
///
|
||||||
|
/// Listens for incoming TCP connections and dispatches:
|
||||||
|
/// - CONNECT requests -> tunnel handler
|
||||||
|
/// - Other HTTP requests -> plain forward proxy handler
|
||||||
|
pub async fn run(
|
||||||
|
config: Arc<Config>,
|
||||||
|
node_id: Arc<String>,
|
||||||
|
mut shutdown_rx: watch::Receiver<bool>,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
let addr = SocketAddr::from(([0, 0, 0, 0], config.listen_port));
|
||||||
|
let listener = TcpListener::bind(addr).await?;
|
||||||
|
info!(addr = %addr, "proxy server listening");
|
||||||
|
|
||||||
|
let allowed_ports: Arc<HashSet<u16>> = Arc::new(config.allowed_ports.iter().copied().collect());
|
||||||
|
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
result = listener.accept() => {
|
||||||
|
let (stream, peer_addr) = match result {
|
||||||
|
Ok(v) => v,
|
||||||
|
Err(e) => {
|
||||||
|
warn!(error = %e, "failed to accept connection");
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(peer = %peer_addr, "new connection");
|
||||||
|
|
||||||
|
let config = Arc::clone(&config);
|
||||||
|
let node_id = Arc::clone(&node_id);
|
||||||
|
let allowed_ports = Arc::clone(&allowed_ports);
|
||||||
|
|
||||||
|
tokio::task::spawn(async move {
|
||||||
|
let io = TokioIo::new(stream);
|
||||||
|
let config = config;
|
||||||
|
let node_id = node_id;
|
||||||
|
let allowed_ports = allowed_ports;
|
||||||
|
|
||||||
|
let service = service_fn(move |req: Request<Incoming>| {
|
||||||
|
let config = Arc::clone(&config);
|
||||||
|
let node_id = Arc::clone(&node_id);
|
||||||
|
let allowed_ports = Arc::clone(&allowed_ports);
|
||||||
|
|
||||||
|
async move {
|
||||||
|
type BoxBody = http_body_util::combinators::BoxBody<bytes::Bytes, Box<dyn std::error::Error + Send + Sync>>;
|
||||||
|
|
||||||
|
if req.method() == Method::CONNECT {
|
||||||
|
let resp = connect::handle_connect(
|
||||||
|
req,
|
||||||
|
config,
|
||||||
|
&node_id,
|
||||||
|
&allowed_ports,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let resp = resp.map(|_| -> BoxBody {
|
||||||
|
http_body_util::Empty::new()
|
||||||
|
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
|
||||||
|
.boxed()
|
||||||
|
});
|
||||||
|
Ok::<_, hyper::Error>(resp)
|
||||||
|
} else {
|
||||||
|
let resp = plain::handle_plain(
|
||||||
|
req,
|
||||||
|
config,
|
||||||
|
&node_id,
|
||||||
|
&allowed_ports,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let resp = resp.map(|body| -> BoxBody {
|
||||||
|
body.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
|
||||||
|
.boxed()
|
||||||
|
});
|
||||||
|
Ok(resp)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Err(e) = http1::Builder::new()
|
||||||
|
.preserve_header_case(true)
|
||||||
|
.title_case_headers(false)
|
||||||
|
.serve_connection(io, service)
|
||||||
|
.with_upgrades()
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
if !e.to_string().contains("connection closed") {
|
||||||
|
debug!(peer = %peer_addr, error = %e, "connection error");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
_ = shutdown_rx.changed() => {
|
||||||
|
info!("proxy server shutting down");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
181
aether-proxy/src/proxy/target_filter.rs
Normal file
181
aether-proxy/src/proxy/target_filter.rs
Normal file
@@ -0,0 +1,181 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
|
||||||
|
|
||||||
|
/// Check if an IP address belongs to a private/reserved network.
|
||||||
|
fn is_private_ip(ip: &IpAddr) -> bool {
|
||||||
|
match ip {
|
||||||
|
IpAddr::V4(v4) => is_private_ipv4(v4),
|
||||||
|
IpAddr::V6(v6) => is_private_ipv6(v6),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_private_ipv4(ip: &Ipv4Addr) -> bool {
|
||||||
|
let octets = ip.octets();
|
||||||
|
// 10.0.0.0/8
|
||||||
|
if octets[0] == 10 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// 172.16.0.0/12
|
||||||
|
if octets[0] == 172 && (16..=31).contains(&octets[1]) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// 192.168.0.0/16
|
||||||
|
if octets[0] == 192 && octets[1] == 168 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// 127.0.0.0/8
|
||||||
|
if octets[0] == 127 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// 169.254.0.0/16 (link-local)
|
||||||
|
if octets[0] == 169 && octets[1] == 254 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// 0.0.0.0/8
|
||||||
|
if octets[0] == 0 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_private_ipv6(ip: &Ipv6Addr) -> bool {
|
||||||
|
// ::1 loopback
|
||||||
|
if ip.is_loopback() {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// :: unspecified
|
||||||
|
if ip.is_unspecified() {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let segments = ip.segments();
|
||||||
|
// fc00::/7 (ULA) - first byte is 0xfc or 0xfd
|
||||||
|
if segments[0] & 0xfe00 == 0xfc00 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// fe80::/10 (link-local)
|
||||||
|
if segments[0] & 0xffc0 == 0xfe80 {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
// IPv4-mapped IPv6 (::ffff:x.x.x.x) - check the embedded IPv4
|
||||||
|
if let Some(v4) = ip.to_ipv4_mapped() {
|
||||||
|
return is_private_ipv4(&v4);
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum FilterError {
|
||||||
|
PrivateIp(IpAddr),
|
||||||
|
PortNotAllowed(u16),
|
||||||
|
DnsResolutionFailed(String),
|
||||||
|
AllAddressesPrivate(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for FilterError {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::PrivateIp(ip) => write!(f, "target IP {} is in private/reserved range", ip),
|
||||||
|
Self::PortNotAllowed(port) => write!(f, "port {} not in allowed list", port),
|
||||||
|
Self::DnsResolutionFailed(host) => write!(f, "DNS resolution failed for {}", host),
|
||||||
|
Self::AllAddressesPrivate(host) => {
|
||||||
|
write!(f, "all resolved addresses for {} are private", host)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validate that the target host:port is allowed.
|
||||||
|
///
|
||||||
|
/// Returns the resolved socket address to connect to.
|
||||||
|
pub fn validate_target(
|
||||||
|
host: &str,
|
||||||
|
port: u16,
|
||||||
|
allowed_ports: &HashSet<u16>,
|
||||||
|
) -> Result<SocketAddr, FilterError> {
|
||||||
|
// Port whitelist check
|
||||||
|
if !allowed_ports.contains(&port) {
|
||||||
|
return Err(FilterError::PortNotAllowed(port));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try parsing as IP directly
|
||||||
|
if let Ok(ip) = host.parse::<IpAddr>() {
|
||||||
|
if is_private_ip(&ip) {
|
||||||
|
return Err(FilterError::PrivateIp(ip));
|
||||||
|
}
|
||||||
|
return Ok(SocketAddr::new(ip, port));
|
||||||
|
}
|
||||||
|
|
||||||
|
// DNS resolution with private IP check (DNS rebinding protection)
|
||||||
|
let addr_str = format!("{}:{}", host, port);
|
||||||
|
let addrs: Vec<SocketAddr> = addr_str
|
||||||
|
.to_socket_addrs()
|
||||||
|
.map_err(|_| FilterError::DnsResolutionFailed(host.to_string()))?
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if addrs.is_empty() {
|
||||||
|
return Err(FilterError::DnsResolutionFailed(host.to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
// All resolved addresses must be non-private
|
||||||
|
for addr in &addrs {
|
||||||
|
if is_private_ip(&addr.ip()) {
|
||||||
|
return Err(FilterError::PrivateIp(addr.ip()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return the first valid address
|
||||||
|
Ok(addrs[0])
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn ports() -> HashSet<u16> {
|
||||||
|
[80, 443, 8080, 8443].into_iter().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_private_ipv4() {
|
||||||
|
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
|
||||||
|
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1))));
|
||||||
|
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
|
||||||
|
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))));
|
||||||
|
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(169, 254, 1, 1))));
|
||||||
|
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))));
|
||||||
|
assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))));
|
||||||
|
assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(203, 0, 113, 1))));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_private_ipv6() {
|
||||||
|
assert!(is_private_ip(&IpAddr::V6(Ipv6Addr::LOCALHOST)));
|
||||||
|
assert!(is_private_ip(&IpAddr::V6(Ipv6Addr::UNSPECIFIED)));
|
||||||
|
// fc00::1 (ULA)
|
||||||
|
assert!(is_private_ip(&IpAddr::V6(Ipv6Addr::new(
|
||||||
|
0xfc00, 0, 0, 0, 0, 0, 0, 1
|
||||||
|
))));
|
||||||
|
// fe80::1 (link-local)
|
||||||
|
assert!(is_private_ip(&IpAddr::V6(Ipv6Addr::new(
|
||||||
|
0xfe80, 0, 0, 0, 0, 0, 0, 1
|
||||||
|
))));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_port_not_allowed() {
|
||||||
|
let result = validate_target("8.8.8.8", 22, &ports());
|
||||||
|
assert!(matches!(result, Err(FilterError::PortNotAllowed(22))));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_private_ip_blocked() {
|
||||||
|
let result = validate_target("127.0.0.1", 80, &ports());
|
||||||
|
assert!(matches!(result, Err(FilterError::PrivateIp(_))));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_public_ip_allowed() {
|
||||||
|
let result = validate_target("8.8.8.8", 443, &ports());
|
||||||
|
assert!(result.is_ok());
|
||||||
|
}
|
||||||
|
}
|
||||||
207
aether-proxy/src/registration/client.rs
Normal file
207
aether-proxy/src/registration/client.rs
Normal file
@@ -0,0 +1,207 @@
|
|||||||
|
use reqwest::Client;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use tracing::{debug, error, info, warn};
|
||||||
|
|
||||||
|
use crate::config::Config;
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct RegisterRequest {
|
||||||
|
name: String,
|
||||||
|
ip: String,
|
||||||
|
port: u16,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
region: Option<String>,
|
||||||
|
heartbeat_interval: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct RegisterResponse {
|
||||||
|
pub node_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct HeartbeatRequest {
|
||||||
|
node_id: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
active_connections: Option<i64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
total_requests: Option<i64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
avg_latency_ms: Option<f64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct UnregisterRequest {
|
||||||
|
node_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Aether API client for proxy node lifecycle management.
|
||||||
|
pub struct AetherClient {
|
||||||
|
http: Client,
|
||||||
|
base_url: String,
|
||||||
|
token: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AetherClient {
|
||||||
|
pub fn new(config: &Config) -> Self {
|
||||||
|
let http = Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_secs(10))
|
||||||
|
.build()
|
||||||
|
.expect("failed to create HTTP client");
|
||||||
|
|
||||||
|
Self {
|
||||||
|
http,
|
||||||
|
base_url: config.aether_url.trim_end_matches('/').to_string(),
|
||||||
|
token: config.management_token.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register this node with Aether (idempotent upsert by ip:port).
|
||||||
|
///
|
||||||
|
/// Returns the stable node_id assigned by Aether.
|
||||||
|
pub async fn register(
|
||||||
|
&self,
|
||||||
|
config: &Config,
|
||||||
|
public_ip: &str,
|
||||||
|
) -> anyhow::Result<String> {
|
||||||
|
let url = format!("{}/api/admin/proxy-nodes/register", self.base_url);
|
||||||
|
let body = RegisterRequest {
|
||||||
|
name: config.node_name.clone(),
|
||||||
|
ip: public_ip.to_string(),
|
||||||
|
port: config.listen_port,
|
||||||
|
region: config.node_region.clone(),
|
||||||
|
heartbeat_interval: config.heartbeat_interval,
|
||||||
|
};
|
||||||
|
|
||||||
|
info!(
|
||||||
|
url = %url,
|
||||||
|
name = %body.name,
|
||||||
|
ip = %body.ip,
|
||||||
|
port = body.port,
|
||||||
|
"registering with Aether"
|
||||||
|
);
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http
|
||||||
|
.post(&url)
|
||||||
|
.header("Authorization", format!("Bearer {}", self.token))
|
||||||
|
.json(&body)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
if !status.is_success() {
|
||||||
|
let text = resp.text().await.unwrap_or_default();
|
||||||
|
anyhow::bail!("register failed (HTTP {}): {}", status, text);
|
||||||
|
}
|
||||||
|
|
||||||
|
let data: RegisterResponse = resp.json().await?;
|
||||||
|
info!(node_id = %data.node_id, "registered successfully");
|
||||||
|
Ok(data.node_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Send heartbeat to Aether.
|
||||||
|
pub async fn heartbeat(
|
||||||
|
&self,
|
||||||
|
node_id: &str,
|
||||||
|
active_connections: Option<i64>,
|
||||||
|
total_requests: Option<i64>,
|
||||||
|
avg_latency_ms: Option<f64>,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
let url = format!("{}/api/admin/proxy-nodes/heartbeat", self.base_url);
|
||||||
|
let body = HeartbeatRequest {
|
||||||
|
node_id: node_id.to_string(),
|
||||||
|
active_connections,
|
||||||
|
total_requests,
|
||||||
|
avg_latency_ms,
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(node_id = %node_id, "sending heartbeat");
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http
|
||||||
|
.post(&url)
|
||||||
|
.header("Authorization", format!("Bearer {}", self.token))
|
||||||
|
.json(&body)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
if !status.is_success() {
|
||||||
|
let text = resp.text().await.unwrap_or_default();
|
||||||
|
warn!(status = %status, body = %text, "heartbeat failed");
|
||||||
|
anyhow::bail!("heartbeat failed (HTTP {}): {}", status, text);
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(node_id = %node_id, "heartbeat ok");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unregister this node from Aether (graceful shutdown).
|
||||||
|
pub async fn unregister(&self, node_id: &str) -> anyhow::Result<()> {
|
||||||
|
let url = format!("{}/api/admin/proxy-nodes/unregister", self.base_url);
|
||||||
|
let body = UnregisterRequest {
|
||||||
|
node_id: node_id.to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
info!(node_id = %node_id, "unregistering from Aether");
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http
|
||||||
|
.post(&url)
|
||||||
|
.header("Authorization", format!("Bearer {}", self.token))
|
||||||
|
.json(&body)
|
||||||
|
.send()
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match resp {
|
||||||
|
Ok(r) if r.status().is_success() => {
|
||||||
|
info!(node_id = %node_id, "unregistered successfully");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Ok(r) => {
|
||||||
|
let text = r.text().await.unwrap_or_default();
|
||||||
|
error!(body = %text, "unregister failed");
|
||||||
|
anyhow::bail!("unregister failed: {}", text);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
// Best-effort during shutdown
|
||||||
|
error!(error = %e, "unregister request failed");
|
||||||
|
anyhow::bail!("unregister request failed: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Auto-detect public IP by querying external services.
|
||||||
|
pub async fn detect_public_ip() -> anyhow::Result<String> {
|
||||||
|
let endpoints = [
|
||||||
|
"https://api.ipify.org",
|
||||||
|
"https://ifconfig.me/ip",
|
||||||
|
"https://icanhazip.com",
|
||||||
|
];
|
||||||
|
|
||||||
|
let client = Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_secs(5))
|
||||||
|
.build()?;
|
||||||
|
|
||||||
|
for endpoint in &endpoints {
|
||||||
|
match client.get(*endpoint).send().await {
|
||||||
|
Ok(resp) if resp.status().is_success() => {
|
||||||
|
let ip = resp.text().await?.trim().to_string();
|
||||||
|
if !ip.is_empty() {
|
||||||
|
info!(ip = %ip, source = %endpoint, "detected public IP");
|
||||||
|
return Ok(ip);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(resp) => {
|
||||||
|
debug!(endpoint = %endpoint, status = %resp.status(), "IP detection failed");
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
debug!(endpoint = %endpoint, error = %e, "IP detection failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
anyhow::bail!("failed to detect public IP from any source; use --public-ip")
|
||||||
|
}
|
||||||
50
aether-proxy/src/registration/heartbeat.rs
Normal file
50
aether-proxy/src/registration/heartbeat.rs
Normal file
@@ -0,0 +1,50 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use tokio::sync::watch;
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use crate::registration::client::AetherClient;
|
||||||
|
|
||||||
|
/// Run periodic heartbeat task until shutdown signal.
|
||||||
|
pub async fn run(
|
||||||
|
client: Arc<AetherClient>,
|
||||||
|
node_id: Arc<String>,
|
||||||
|
interval_secs: u64,
|
||||||
|
mut shutdown_rx: watch::Receiver<bool>,
|
||||||
|
) {
|
||||||
|
let mut interval = tokio::time::interval(std::time::Duration::from_secs(interval_secs));
|
||||||
|
// Skip the first immediate tick (registration already acts as initial heartbeat)
|
||||||
|
interval.tick().await;
|
||||||
|
|
||||||
|
let mut consecutive_failures: u32 = 0;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
_ = interval.tick() => {
|
||||||
|
match client.heartbeat(&node_id, None, None, None).await {
|
||||||
|
Ok(()) => {
|
||||||
|
if consecutive_failures > 0 {
|
||||||
|
debug!(
|
||||||
|
previous_failures = consecutive_failures,
|
||||||
|
"heartbeat recovered"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
consecutive_failures = 0;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
consecutive_failures += 1;
|
||||||
|
warn!(
|
||||||
|
error = %e,
|
||||||
|
consecutive_failures,
|
||||||
|
"heartbeat failed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = shutdown_rx.changed() => {
|
||||||
|
debug!("heartbeat task stopping");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
2
aether-proxy/src/registration/mod.rs
Normal file
2
aether-proxy/src/registration/mod.rs
Normal file
@@ -0,0 +1,2 @@
|
|||||||
|
pub mod client;
|
||||||
|
pub mod heartbeat;
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
"""Update Antigravity endpoint signature to gemini:chat
|
"""Antigravity endpoint signature to gemini:chat & add proxy_nodes table
|
||||||
|
|
||||||
Revision ID: e1b2c3d4f5a6
|
Revision ID: e1b2c3d4f5a6
|
||||||
Revises: b5c6d7e8f9a0
|
Revises: b5c6d7e8f9a0
|
||||||
@@ -10,7 +10,9 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
|
|
||||||
from sqlalchemy import text
|
import sqlalchemy as sa
|
||||||
|
from sqlalchemy import inspect, text
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
from alembic import op
|
from alembic import op
|
||||||
|
|
||||||
@@ -21,9 +23,19 @@ branch_labels: str | Sequence[str] | None = None
|
|||||||
depends_on: str | Sequence[str] | None = None
|
depends_on: str | Sequence[str] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def table_exists(table_name: str) -> bool:
|
||||||
|
bind = op.get_bind()
|
||||||
|
inspector = inspect(bind)
|
||||||
|
return table_name in inspector.get_table_names()
|
||||||
|
|
||||||
|
|
||||||
def upgrade() -> None:
|
def upgrade() -> None:
|
||||||
conn = op.get_bind()
|
conn = op.get_bind()
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# Part 1: Antigravity endpoint signature migration (gemini:cli -> gemini:chat)
|
||||||
|
# =========================================================================
|
||||||
|
|
||||||
# --- provider_endpoints ---
|
# --- provider_endpoints ---
|
||||||
# Update only when there is no conflicting gemini:chat endpoint for the same provider
|
# Update only when there is no conflicting gemini:chat endpoint for the same provider
|
||||||
# (provider_endpoints has a unique constraint on (provider_id, api_format)).
|
# (provider_endpoints has a unique constraint on (provider_id, api_format)).
|
||||||
@@ -58,7 +70,7 @@ def upgrade() -> None:
|
|||||||
|
|
||||||
# --- provider_api_keys.api_formats (JSON array) ---
|
# --- provider_api_keys.api_formats (JSON array) ---
|
||||||
# Replace "gemini:cli" with "gemini:chat" in the JSON array for Antigravity keys.
|
# Replace "gemini:cli" with "gemini:chat" in the JSON array for Antigravity keys.
|
||||||
# Uses text-level replace on the serialized JSON — safe because the value is a
|
# Uses text-level replace on the serialized JSON -- safe because the value is a
|
||||||
# simple string with no special characters that could cause ambiguous replacements.
|
# simple string with no special characters that could cause ambiguous replacements.
|
||||||
conn.execute(text("""
|
conn.execute(text("""
|
||||||
UPDATE provider_api_keys pak
|
UPDATE provider_api_keys pak
|
||||||
@@ -70,10 +82,83 @@ def upgrade() -> None:
|
|||||||
AND pak.api_formats::text LIKE '%"gemini:cli"%'
|
AND pak.api_formats::text LIKE '%"gemini:cli"%'
|
||||||
"""))
|
"""))
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# Part 2: Create proxy_nodes table (idempotent)
|
||||||
|
# =========================================================================
|
||||||
|
|
||||||
|
# Create ENUM type (idempotent)
|
||||||
|
op.execute(
|
||||||
|
"DO $$ BEGIN "
|
||||||
|
"CREATE TYPE proxynodestatus AS ENUM ('online', 'unhealthy', 'offline'); "
|
||||||
|
"EXCEPTION WHEN duplicate_object THEN NULL; "
|
||||||
|
"END $$"
|
||||||
|
)
|
||||||
|
|
||||||
|
if table_exists("proxy_nodes"):
|
||||||
|
return
|
||||||
|
|
||||||
|
op.create_table(
|
||||||
|
"proxy_nodes",
|
||||||
|
sa.Column("id", sa.String(36), primary_key=True),
|
||||||
|
sa.Column("name", sa.String(100), nullable=False),
|
||||||
|
sa.Column("ip", sa.String(45), nullable=False),
|
||||||
|
sa.Column("port", sa.Integer(), nullable=False),
|
||||||
|
sa.Column("region", sa.String(100), nullable=True),
|
||||||
|
sa.Column(
|
||||||
|
"status",
|
||||||
|
postgresql.ENUM(
|
||||||
|
"online",
|
||||||
|
"unhealthy",
|
||||||
|
"offline",
|
||||||
|
name="proxynodestatus",
|
||||||
|
create_type=False,
|
||||||
|
),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.text("'online'"),
|
||||||
|
),
|
||||||
|
sa.Column(
|
||||||
|
"registered_by",
|
||||||
|
sa.String(36),
|
||||||
|
sa.ForeignKey("users.id", ondelete="SET NULL"),
|
||||||
|
nullable=True,
|
||||||
|
),
|
||||||
|
sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("heartbeat_interval", sa.Integer(), nullable=False, server_default=sa.text("30")),
|
||||||
|
sa.Column("active_connections", sa.Integer(), nullable=False, server_default=sa.text("0")),
|
||||||
|
sa.Column("total_requests", sa.BigInteger(), nullable=False, server_default=sa.text("0")),
|
||||||
|
sa.Column("avg_latency_ms", sa.Float(), nullable=True),
|
||||||
|
sa.Column(
|
||||||
|
"created_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
server_default=sa.text("CURRENT_TIMESTAMP"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column(
|
||||||
|
"updated_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
server_default=sa.text("CURRENT_TIMESTAMP"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.UniqueConstraint("ip", "port", name="uq_proxy_node_ip_port"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def downgrade() -> None:
|
def downgrade() -> None:
|
||||||
conn = op.get_bind()
|
conn = op.get_bind()
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# Part 2 rollback: Drop proxy_nodes table
|
||||||
|
# =========================================================================
|
||||||
|
if table_exists("proxy_nodes"):
|
||||||
|
op.drop_table("proxy_nodes")
|
||||||
|
|
||||||
|
# Best-effort: drop type (only used by proxy_nodes)
|
||||||
|
op.execute("DROP TYPE IF EXISTS proxynodestatus")
|
||||||
|
|
||||||
|
# =========================================================================
|
||||||
|
# Part 1 rollback: Revert Antigravity endpoint signature (gemini:chat -> gemini:cli)
|
||||||
|
# =========================================================================
|
||||||
|
|
||||||
# --- provider_endpoints ---
|
# --- provider_endpoints ---
|
||||||
conn.execute(text("""
|
conn.execute(text("""
|
||||||
UPDATE provider_endpoints pe
|
UPDATE provider_endpoints pe
|
||||||
|
|||||||
36
frontend/src/api/proxy-nodes.ts
Normal file
36
frontend/src/api/proxy-nodes.ts
Normal file
@@ -0,0 +1,36 @@
|
|||||||
|
import apiClient from './client'
|
||||||
|
|
||||||
|
export interface ProxyNode {
|
||||||
|
id: string
|
||||||
|
name: string
|
||||||
|
ip: string
|
||||||
|
port: number
|
||||||
|
region: string | null
|
||||||
|
status: 'online' | 'unhealthy' | 'offline'
|
||||||
|
registered_by: string | null
|
||||||
|
last_heartbeat_at: string | null
|
||||||
|
heartbeat_interval: number
|
||||||
|
active_connections: number
|
||||||
|
total_requests: number
|
||||||
|
avg_latency_ms: number | null
|
||||||
|
created_at: string
|
||||||
|
updated_at: string
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface ProxyNodeListResponse {
|
||||||
|
items: ProxyNode[]
|
||||||
|
total: number
|
||||||
|
skip: number
|
||||||
|
limit: number
|
||||||
|
}
|
||||||
|
|
||||||
|
export const proxyNodesApi = {
|
||||||
|
async listProxyNodes(params?: { status?: string; skip?: number; limit?: number }): Promise<ProxyNodeListResponse> {
|
||||||
|
const response = await apiClient.get<ProxyNodeListResponse>('/api/admin/proxy-nodes', { params })
|
||||||
|
return response.data
|
||||||
|
},
|
||||||
|
|
||||||
|
async deleteProxyNode(nodeId: string): Promise<void> {
|
||||||
|
await apiClient.delete(`/api/admin/proxy-nodes/${nodeId}`)
|
||||||
|
},
|
||||||
|
}
|
||||||
@@ -368,6 +368,7 @@ import {
|
|||||||
Video,
|
Video,
|
||||||
Zap,
|
Zap,
|
||||||
FileUp,
|
FileUp,
|
||||||
|
Server,
|
||||||
type LucideIcon,
|
type LucideIcon,
|
||||||
} from 'lucide-vue-next'
|
} from 'lucide-vue-next'
|
||||||
|
|
||||||
@@ -520,6 +521,7 @@ const navigation = computed(() => {
|
|||||||
FileUp,
|
FileUp,
|
||||||
Shield,
|
Shield,
|
||||||
Puzzle,
|
Puzzle,
|
||||||
|
Server,
|
||||||
}
|
}
|
||||||
|
|
||||||
// 添加模块菜单项(按 admin_menu_order 排序,只显示已激活的)
|
// 添加模块菜单项(按 admin_menu_order 排序,只显示已激活的)
|
||||||
|
|||||||
@@ -230,6 +230,12 @@ const routes: RouteRecordRaw[] = [
|
|||||||
name: 'AsyncTasks',
|
name: 'AsyncTasks',
|
||||||
component: () => importWithRetry(() => import('@/views/admin/AsyncTasks.vue'))
|
component: () => importWithRetry(() => import('@/views/admin/AsyncTasks.vue'))
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
path: 'proxy-nodes',
|
||||||
|
name: 'ProxyNodes',
|
||||||
|
component: () => importWithRetry(() => import('@/views/admin/ProxyNodes.vue')),
|
||||||
|
meta: { module: 'proxy_nodes' }
|
||||||
|
},
|
||||||
{
|
{
|
||||||
path: 'gemini-files',
|
path: 'gemini-files',
|
||||||
name: 'GeminiFilesManagement',
|
name: 'GeminiFilesManagement',
|
||||||
|
|||||||
50
frontend/src/stores/proxy-nodes.ts
Normal file
50
frontend/src/stores/proxy-nodes.ts
Normal file
@@ -0,0 +1,50 @@
|
|||||||
|
import { defineStore } from 'pinia'
|
||||||
|
import { ref } from 'vue'
|
||||||
|
import { proxyNodesApi, type ProxyNode } from '@/api/proxy-nodes'
|
||||||
|
|
||||||
|
export const useProxyNodesStore = defineStore('proxy-nodes', () => {
|
||||||
|
const nodes = ref<ProxyNode[]>([])
|
||||||
|
const total = ref(0)
|
||||||
|
const loading = ref(false)
|
||||||
|
const error = ref<string | null>(null)
|
||||||
|
|
||||||
|
async function fetchNodes(params?: { status?: string }) {
|
||||||
|
loading.value = true
|
||||||
|
error.value = null
|
||||||
|
|
||||||
|
try {
|
||||||
|
const data = await proxyNodesApi.listProxyNodes({ ...params, limit: 1000 })
|
||||||
|
nodes.value = data.items
|
||||||
|
total.value = data.total
|
||||||
|
} catch (err: any) {
|
||||||
|
error.value = err.response?.data?.error?.message || err.response?.data?.detail || '获取代理节点列表失败'
|
||||||
|
} finally {
|
||||||
|
loading.value = false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async function deleteNode(nodeId: string) {
|
||||||
|
loading.value = true
|
||||||
|
error.value = null
|
||||||
|
|
||||||
|
try {
|
||||||
|
await proxyNodesApi.deleteProxyNode(nodeId)
|
||||||
|
nodes.value = nodes.value.filter(n => n.id !== nodeId)
|
||||||
|
total.value = Math.max(0, total.value - 1)
|
||||||
|
} catch (err: any) {
|
||||||
|
error.value = err.response?.data?.error?.message || err.response?.data?.detail || '删除代理节点失败'
|
||||||
|
throw err
|
||||||
|
} finally {
|
||||||
|
loading.value = false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
nodes,
|
||||||
|
total,
|
||||||
|
loading,
|
||||||
|
error,
|
||||||
|
fetchNodes,
|
||||||
|
deleteNode,
|
||||||
|
}
|
||||||
|
})
|
||||||
329
frontend/src/views/admin/ProxyNodes.vue
Normal file
329
frontend/src/views/admin/ProxyNodes.vue
Normal file
@@ -0,0 +1,329 @@
|
|||||||
|
<template>
|
||||||
|
<div class="space-y-6 pb-8">
|
||||||
|
<Card variant="default" class="overflow-hidden">
|
||||||
|
<!-- 标题和筛选器 -->
|
||||||
|
<div class="px-4 sm:px-6 py-3.5 border-b border-border/60">
|
||||||
|
<!-- 移动端 -->
|
||||||
|
<div class="flex flex-col gap-3 sm:hidden">
|
||||||
|
<div class="flex items-center justify-between">
|
||||||
|
<h3 class="text-base font-semibold">
|
||||||
|
代理节点
|
||||||
|
</h3>
|
||||||
|
<RefreshButton
|
||||||
|
:loading="store.loading"
|
||||||
|
@click="refresh"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="flex items-center gap-2">
|
||||||
|
<div class="relative flex-1">
|
||||||
|
<Search class="absolute left-2.5 top-1/2 -translate-y-1/2 h-3.5 w-3.5 text-muted-foreground z-10 pointer-events-none" />
|
||||||
|
<Input
|
||||||
|
v-model="searchQuery"
|
||||||
|
type="text"
|
||||||
|
placeholder="搜索..."
|
||||||
|
class="w-full pl-8 pr-3 h-8 text-sm bg-background/50 border-border/60"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<Select v-model="filterStatus">
|
||||||
|
<SelectTrigger class="w-24 h-8 text-xs border-border/60">
|
||||||
|
<SelectValue placeholder="状态" />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
<SelectItem value="all">全部</SelectItem>
|
||||||
|
<SelectItem value="online">在线</SelectItem>
|
||||||
|
<SelectItem value="unhealthy">异常</SelectItem>
|
||||||
|
<SelectItem value="offline">离线</SelectItem>
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 桌面端 -->
|
||||||
|
<div class="hidden sm:flex items-center justify-between gap-4">
|
||||||
|
<h3 class="text-base font-semibold">
|
||||||
|
代理节点
|
||||||
|
</h3>
|
||||||
|
<div class="flex items-center gap-2">
|
||||||
|
<div class="relative">
|
||||||
|
<Search class="absolute left-2.5 top-1/2 -translate-y-1/2 h-3.5 w-3.5 text-muted-foreground z-10 pointer-events-none" />
|
||||||
|
<Input
|
||||||
|
v-model="searchQuery"
|
||||||
|
type="text"
|
||||||
|
placeholder="搜索..."
|
||||||
|
class="w-48 pl-8 pr-3 h-8 text-sm bg-background/50 border-border/60"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="h-4 w-px bg-border" />
|
||||||
|
<Select v-model="filterStatus">
|
||||||
|
<SelectTrigger class="w-28 h-8 text-xs border-border/60">
|
||||||
|
<SelectValue placeholder="全部状态" />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
<SelectItem value="all">全部状态</SelectItem>
|
||||||
|
<SelectItem value="online">在线</SelectItem>
|
||||||
|
<SelectItem value="unhealthy">异常</SelectItem>
|
||||||
|
<SelectItem value="offline">离线</SelectItem>
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
<div class="h-4 w-px bg-border" />
|
||||||
|
<RefreshButton
|
||||||
|
:loading="store.loading"
|
||||||
|
@click="refresh"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 桌面端表格 -->
|
||||||
|
<div class="hidden xl:block overflow-x-auto">
|
||||||
|
<Table>
|
||||||
|
<TableHeader>
|
||||||
|
<TableRow class="border-b border-border/60 hover:bg-transparent">
|
||||||
|
<TableHead class="w-[160px] h-12 font-semibold">名称</TableHead>
|
||||||
|
<TableHead class="w-[180px] h-12 font-semibold">地址</TableHead>
|
||||||
|
<TableHead class="w-[100px] h-12 font-semibold">区域</TableHead>
|
||||||
|
<TableHead class="w-[90px] h-12 font-semibold text-center">状态</TableHead>
|
||||||
|
<TableHead class="w-[100px] h-12 font-semibold text-center">连接数</TableHead>
|
||||||
|
<TableHead class="w-[100px] h-12 font-semibold text-center">总请求</TableHead>
|
||||||
|
<TableHead class="w-[100px] h-12 font-semibold text-center">延迟</TableHead>
|
||||||
|
<TableHead class="w-[160px] h-12 font-semibold">最后心跳</TableHead>
|
||||||
|
<TableHead class="w-[80px] h-12 font-semibold text-center">操作</TableHead>
|
||||||
|
</TableRow>
|
||||||
|
</TableHeader>
|
||||||
|
<TableBody>
|
||||||
|
<TableRow
|
||||||
|
v-for="node in paginatedNodes"
|
||||||
|
:key="node.id"
|
||||||
|
class="border-b border-border/40 hover:bg-muted/30 transition-colors"
|
||||||
|
>
|
||||||
|
<TableCell class="py-4">
|
||||||
|
<span class="text-sm font-semibold">{{ node.name }}</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4">
|
||||||
|
<code class="text-xs text-muted-foreground">{{ node.ip }}:{{ node.port }}</code>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4">
|
||||||
|
<span class="text-sm text-muted-foreground">{{ node.region || '-' }}</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4 text-center">
|
||||||
|
<Badge :variant="statusVariant(node.status)" class="font-medium px-2.5 py-0.5 text-xs">
|
||||||
|
{{ statusLabel(node.status) }}
|
||||||
|
</Badge>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4 text-center">
|
||||||
|
<span class="text-sm tabular-nums">{{ node.active_connections }}</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4 text-center">
|
||||||
|
<span class="text-sm tabular-nums">{{ formatNumber(node.total_requests) }}</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4 text-center">
|
||||||
|
<span class="text-sm tabular-nums">{{ node.avg_latency_ms != null ? `${node.avg_latency_ms.toFixed(0)}ms` : '-' }}</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4">
|
||||||
|
<span class="text-xs text-muted-foreground">{{ formatTime(node.last_heartbeat_at) }}</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-4 text-center">
|
||||||
|
<Button
|
||||||
|
variant="ghost"
|
||||||
|
size="icon"
|
||||||
|
class="h-8 w-8"
|
||||||
|
title="删除"
|
||||||
|
@click="handleDelete(node)"
|
||||||
|
>
|
||||||
|
<Trash2 class="h-4 w-4" />
|
||||||
|
</Button>
|
||||||
|
</TableCell>
|
||||||
|
</TableRow>
|
||||||
|
<TableRow v-if="paginatedNodes.length === 0">
|
||||||
|
<TableCell colspan="9" class="py-12 text-center text-muted-foreground text-sm">
|
||||||
|
{{ store.loading ? '加载中...' : '暂无代理节点' }}
|
||||||
|
</TableCell>
|
||||||
|
</TableRow>
|
||||||
|
</TableBody>
|
||||||
|
</Table>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 移动端卡片列表 -->
|
||||||
|
<div class="xl:hidden divide-y divide-border/40">
|
||||||
|
<div
|
||||||
|
v-for="node in paginatedNodes"
|
||||||
|
:key="node.id"
|
||||||
|
class="p-4 sm:p-5"
|
||||||
|
>
|
||||||
|
<div class="flex items-start justify-between mb-2">
|
||||||
|
<div>
|
||||||
|
<div class="font-semibold text-sm">{{ node.name }}</div>
|
||||||
|
<code class="text-xs text-muted-foreground">{{ node.ip }}:{{ node.port }}</code>
|
||||||
|
</div>
|
||||||
|
<Badge :variant="statusVariant(node.status)" class="text-xs">
|
||||||
|
{{ statusLabel(node.status) }}
|
||||||
|
</Badge>
|
||||||
|
</div>
|
||||||
|
<div class="grid grid-cols-3 gap-2 text-xs text-muted-foreground mb-3">
|
||||||
|
<div>
|
||||||
|
<span class="block text-foreground/60">区域</span>
|
||||||
|
<span>{{ node.region || '-' }}</span>
|
||||||
|
</div>
|
||||||
|
<div>
|
||||||
|
<span class="block text-foreground/60">连接</span>
|
||||||
|
<span class="tabular-nums">{{ node.active_connections }}</span>
|
||||||
|
</div>
|
||||||
|
<div>
|
||||||
|
<span class="block text-foreground/60">延迟</span>
|
||||||
|
<span class="tabular-nums">{{ node.avg_latency_ms != null ? `${node.avg_latency_ms.toFixed(0)}ms` : '-' }}</span>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="flex items-center justify-between">
|
||||||
|
<span class="text-xs text-muted-foreground">{{ formatTime(node.last_heartbeat_at) }}</span>
|
||||||
|
<Button
|
||||||
|
variant="ghost"
|
||||||
|
size="sm"
|
||||||
|
class="h-7 px-2 text-xs"
|
||||||
|
@click="handleDelete(node)"
|
||||||
|
>
|
||||||
|
<Trash2 class="h-3 w-3 mr-1" />
|
||||||
|
删除
|
||||||
|
</Button>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div v-if="paginatedNodes.length === 0" class="p-8 text-center text-muted-foreground text-sm">
|
||||||
|
{{ store.loading ? '加载中...' : '暂无代理节点' }}
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 分页 -->
|
||||||
|
<Pagination
|
||||||
|
:current="currentPage"
|
||||||
|
:total="filteredNodes.length"
|
||||||
|
:page-size="pageSize"
|
||||||
|
cache-key="proxy-nodes-page-size"
|
||||||
|
@update:current="currentPage = $event"
|
||||||
|
@update:page-size="pageSize = $event"
|
||||||
|
/>
|
||||||
|
</Card>
|
||||||
|
</div>
|
||||||
|
</template>
|
||||||
|
|
||||||
|
<script setup lang="ts">
|
||||||
|
import { ref, computed, onMounted, watch } from 'vue'
|
||||||
|
import { useProxyNodesStore } from '@/stores/proxy-nodes'
|
||||||
|
import { useToast } from '@/composables/useToast'
|
||||||
|
import { useConfirm } from '@/composables/useConfirm'
|
||||||
|
import type { ProxyNode } from '@/api/proxy-nodes'
|
||||||
|
|
||||||
|
import {
|
||||||
|
Card,
|
||||||
|
Button,
|
||||||
|
Badge,
|
||||||
|
Input,
|
||||||
|
Select,
|
||||||
|
SelectTrigger,
|
||||||
|
SelectValue,
|
||||||
|
SelectContent,
|
||||||
|
SelectItem,
|
||||||
|
Table,
|
||||||
|
TableHeader,
|
||||||
|
TableBody,
|
||||||
|
TableRow,
|
||||||
|
TableHead,
|
||||||
|
TableCell,
|
||||||
|
Pagination,
|
||||||
|
RefreshButton,
|
||||||
|
} from '@/components/ui'
|
||||||
|
|
||||||
|
import { Search, Trash2 } from 'lucide-vue-next'
|
||||||
|
|
||||||
|
const { success, error: toastError } = useToast()
|
||||||
|
const { confirmDanger } = useConfirm()
|
||||||
|
const store = useProxyNodesStore()
|
||||||
|
|
||||||
|
const searchQuery = ref('')
|
||||||
|
const filterStatus = ref('all')
|
||||||
|
const currentPage = ref(1)
|
||||||
|
const pageSize = ref(20)
|
||||||
|
|
||||||
|
const filteredNodes = computed(() => {
|
||||||
|
let filtered = [...store.nodes]
|
||||||
|
|
||||||
|
if (searchQuery.value) {
|
||||||
|
const keywords = searchQuery.value.toLowerCase().split(/\s+/).filter(k => k.length > 0)
|
||||||
|
filtered = filtered.filter(node => {
|
||||||
|
const text = `${node.name} ${node.ip} ${node.region || ''}`.toLowerCase()
|
||||||
|
return keywords.every(kw => text.includes(kw))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
if (filterStatus.value !== 'all') {
|
||||||
|
filtered = filtered.filter(node => node.status === filterStatus.value)
|
||||||
|
}
|
||||||
|
|
||||||
|
return filtered
|
||||||
|
})
|
||||||
|
|
||||||
|
const paginatedNodes = computed(() => {
|
||||||
|
const start = (currentPage.value - 1) * pageSize.value
|
||||||
|
return filteredNodes.value.slice(start, start + pageSize.value)
|
||||||
|
})
|
||||||
|
|
||||||
|
watch([searchQuery, filterStatus], () => {
|
||||||
|
currentPage.value = 1
|
||||||
|
})
|
||||||
|
|
||||||
|
onMounted(async () => {
|
||||||
|
await store.fetchNodes()
|
||||||
|
})
|
||||||
|
|
||||||
|
async function refresh() {
|
||||||
|
await store.fetchNodes()
|
||||||
|
}
|
||||||
|
|
||||||
|
async function handleDelete(node: ProxyNode) {
|
||||||
|
const confirmed = await confirmDanger(
|
||||||
|
`确定要删除代理节点 "${node.name}" (${node.ip}:${node.port}) 吗?`,
|
||||||
|
'删除节点'
|
||||||
|
)
|
||||||
|
if (!confirmed) return
|
||||||
|
|
||||||
|
try {
|
||||||
|
await store.deleteNode(node.id)
|
||||||
|
success('代理节点已删除')
|
||||||
|
} catch (err: any) {
|
||||||
|
toastError(err.response?.data?.error?.message || '删除失败')
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function statusVariant(status: string) {
|
||||||
|
switch (status) {
|
||||||
|
case 'online': return 'success' as const
|
||||||
|
case 'unhealthy': return 'secondary' as const
|
||||||
|
case 'offline': return 'destructive' as const
|
||||||
|
default: return 'secondary' as const
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function statusLabel(status: string) {
|
||||||
|
switch (status) {
|
||||||
|
case 'online': return '在线'
|
||||||
|
case 'unhealthy': return '异常'
|
||||||
|
case 'offline': return '离线'
|
||||||
|
default: return status
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function formatNumber(n: number) {
|
||||||
|
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
|
||||||
|
if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`
|
||||||
|
return String(n)
|
||||||
|
}
|
||||||
|
|
||||||
|
function formatTime(iso: string | null) {
|
||||||
|
if (!iso) return '-'
|
||||||
|
const d = new Date(iso)
|
||||||
|
const now = new Date()
|
||||||
|
const diff = (now.getTime() - d.getTime()) / 1000
|
||||||
|
if (diff < 60) return '刚刚'
|
||||||
|
if (diff < 3600) return `${Math.floor(diff / 60)}分钟前`
|
||||||
|
if (diff < 86400) return `${Math.floor(diff / 3600)}小时前`
|
||||||
|
return d.toLocaleDateString('zh-CN', { month: '2-digit', day: '2-digit', hour: '2-digit', minute: '2-digit' })
|
||||||
|
}
|
||||||
|
</script>
|
||||||
@@ -44,5 +44,6 @@ router.include_router(video_tasks_router)
|
|||||||
# 注意:以下路由已迁移到模块系统,由 ModuleRegistry 动态注册
|
# 注意:以下路由已迁移到模块系统,由 ModuleRegistry 动态注册
|
||||||
# - ldap_router: 当 LDAP_AVAILABLE=true 时注册
|
# - ldap_router: 当 LDAP_AVAILABLE=true 时注册
|
||||||
# - management_tokens_router: 当 MANAGEMENT_TOKENS_AVAILABLE=true 时注册
|
# - management_tokens_router: 当 MANAGEMENT_TOKENS_AVAILABLE=true 时注册
|
||||||
|
# - proxy_nodes_router: 当 PROXY_NODES_AVAILABLE=true 时注册
|
||||||
|
|
||||||
__all__ = ["router"]
|
__all__ = ["router"]
|
||||||
|
|||||||
5
src/api/admin/proxy_nodes/__init__.py
Normal file
5
src/api/admin/proxy_nodes/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
"""Proxy node admin routes export."""
|
||||||
|
|
||||||
|
from .routes import router
|
||||||
|
|
||||||
|
__all__ = ["router"]
|
||||||
304
src/api/admin/proxy_nodes/routes.py
Normal file
304
src/api/admin/proxy_nodes/routes.py
Normal file
@@ -0,0 +1,304 @@
|
|||||||
|
"""管理员代理节点(ProxyNode)管理端点
|
||||||
|
|
||||||
|
用于 aether-proxy 在 VPS 上注册、心跳、注销节点,以及管理员查看/删除节点记录。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import ipaddress
|
||||||
|
import uuid
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends, Query, Request
|
||||||
|
from pydantic import BaseModel, Field, ValidationError, field_validator
|
||||||
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
|
from src.api.base.admin_adapter import AdminApiAdapter
|
||||||
|
from src.api.base.context import ApiRequestContext
|
||||||
|
from src.api.base.pipeline import ApiRequestPipeline
|
||||||
|
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||||
|
from src.database import get_db
|
||||||
|
from src.models.database import ProxyNode, ProxyNodeStatus
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/api/admin/proxy-nodes", tags=["Admin - Proxy Nodes"])
|
||||||
|
pipeline = ApiRequestPipeline()
|
||||||
|
|
||||||
|
|
||||||
|
def _node_to_dict(node: ProxyNode) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"id": node.id,
|
||||||
|
"name": node.name,
|
||||||
|
"ip": node.ip,
|
||||||
|
"port": node.port,
|
||||||
|
"region": node.region,
|
||||||
|
"status": node.status.value if node.status else None,
|
||||||
|
"registered_by": node.registered_by,
|
||||||
|
"last_heartbeat_at": node.last_heartbeat_at,
|
||||||
|
"heartbeat_interval": node.heartbeat_interval,
|
||||||
|
"active_connections": node.active_connections,
|
||||||
|
"total_requests": node.total_requests,
|
||||||
|
"avg_latency_ms": node.avg_latency_ms,
|
||||||
|
"created_at": node.created_at,
|
||||||
|
"updated_at": node.updated_at,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNodeRegisterRequest(BaseModel):
|
||||||
|
name: str = Field(..., min_length=1, max_length=100, description="节点名")
|
||||||
|
ip: str = Field(..., description="公网 IP(IPv4/IPv6)")
|
||||||
|
port: int = Field(..., ge=1, le=65535, description="代理端口")
|
||||||
|
region: str | None = Field(None, max_length=100, description="区域标签")
|
||||||
|
heartbeat_interval: int = Field(30, ge=5, le=600, description="心跳间隔(秒)")
|
||||||
|
|
||||||
|
# 指标(可选)
|
||||||
|
active_connections: int | None = Field(None, ge=0, description="当前活跃连接数")
|
||||||
|
total_requests: int | None = Field(None, ge=0, description="累计请求数")
|
||||||
|
avg_latency_ms: float | None = Field(None, ge=0, description="平均延迟(毫秒)")
|
||||||
|
|
||||||
|
@field_validator("ip")
|
||||||
|
@classmethod
|
||||||
|
def validate_ip(cls, v: str) -> str:
|
||||||
|
v = v.strip()
|
||||||
|
try:
|
||||||
|
ipaddress.ip_address(v)
|
||||||
|
except ValueError as exc:
|
||||||
|
raise ValueError("ip 必须是合法的 IPv4/IPv6 地址") from exc
|
||||||
|
return v
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNodeHeartbeatRequest(BaseModel):
|
||||||
|
node_id: str = Field(..., min_length=1, max_length=36, description="节点 ID")
|
||||||
|
heartbeat_interval: int | None = Field(None, ge=5, le=600, description="心跳间隔(秒)")
|
||||||
|
|
||||||
|
active_connections: int | None = Field(None, ge=0, description="当前活跃连接数")
|
||||||
|
total_requests: int | None = Field(None, ge=0, description="累计请求数")
|
||||||
|
avg_latency_ms: float | None = Field(None, ge=0, description="平均延迟(毫秒)")
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNodeUnregisterRequest(BaseModel):
|
||||||
|
node_id: str = Field(..., min_length=1, max_length=36, description="节点 ID")
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/register")
|
||||||
|
async def register_proxy_node(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||||
|
adapter = AdminRegisterProxyNodeAdapter()
|
||||||
|
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/heartbeat")
|
||||||
|
async def heartbeat_proxy_node(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||||
|
adapter = AdminHeartbeatProxyNodeAdapter()
|
||||||
|
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/unregister")
|
||||||
|
async def unregister_proxy_node(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||||
|
adapter = AdminUnregisterProxyNodeAdapter()
|
||||||
|
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("")
|
||||||
|
async def list_proxy_nodes(
|
||||||
|
request: Request,
|
||||||
|
status: str | None = Query(None, description="按状态筛选:online/unhealthy/offline"),
|
||||||
|
skip: int = Query(0, ge=0),
|
||||||
|
limit: int = Query(100, ge=1, le=1000),
|
||||||
|
db: Session = Depends(get_db),
|
||||||
|
) -> Any:
|
||||||
|
adapter = AdminListProxyNodesAdapter(status=status, skip=skip, limit=limit)
|
||||||
|
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/{node_id}")
|
||||||
|
async def delete_proxy_node(node_id: str, request: Request, db: Session = Depends(get_db)) -> Any:
|
||||||
|
adapter = AdminDeleteProxyNodeAdapter(node_id=node_id)
|
||||||
|
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||||
|
|
||||||
|
|
||||||
|
def _format_validation_error(exc: ValidationError) -> str:
|
||||||
|
parts: list[str] = []
|
||||||
|
for err in exc.errors():
|
||||||
|
field = " -> ".join(str(x) for x in err.get("loc", []))
|
||||||
|
msg = str(err.get("msg", "invalid"))
|
||||||
|
parts.append(f"{field}: {msg}")
|
||||||
|
return "; ".join(parts) or "输入验证失败"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AdminRegisterProxyNodeAdapter(AdminApiAdapter):
|
||||||
|
name: str = "admin_register_proxy_node"
|
||||||
|
|
||||||
|
async def handle(self, context: ApiRequestContext) -> Any:
|
||||||
|
payload = context.ensure_json_body()
|
||||||
|
try:
|
||||||
|
req = ProxyNodeRegisterRequest.model_validate(payload)
|
||||||
|
except ValidationError as exc:
|
||||||
|
raise InvalidRequestException("输入验证失败: " + _format_validation_error(exc))
|
||||||
|
|
||||||
|
now = datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
node = (
|
||||||
|
context.db.query(ProxyNode)
|
||||||
|
.filter(ProxyNode.ip == req.ip, ProxyNode.port == req.port)
|
||||||
|
.first()
|
||||||
|
)
|
||||||
|
if node:
|
||||||
|
node.name = req.name
|
||||||
|
node.region = req.region
|
||||||
|
node.status = ProxyNodeStatus.ONLINE
|
||||||
|
node.last_heartbeat_at = now
|
||||||
|
node.heartbeat_interval = req.heartbeat_interval
|
||||||
|
if req.active_connections is not None:
|
||||||
|
node.active_connections = req.active_connections
|
||||||
|
if req.total_requests is not None:
|
||||||
|
node.total_requests = req.total_requests
|
||||||
|
if req.avg_latency_ms is not None:
|
||||||
|
node.avg_latency_ms = req.avg_latency_ms
|
||||||
|
else:
|
||||||
|
node = ProxyNode(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
name=req.name,
|
||||||
|
ip=req.ip,
|
||||||
|
port=req.port,
|
||||||
|
region=req.region,
|
||||||
|
status=ProxyNodeStatus.ONLINE,
|
||||||
|
registered_by=context.user.id if context.user else None,
|
||||||
|
last_heartbeat_at=now,
|
||||||
|
heartbeat_interval=req.heartbeat_interval,
|
||||||
|
active_connections=req.active_connections or 0,
|
||||||
|
total_requests=req.total_requests or 0,
|
||||||
|
avg_latency_ms=req.avg_latency_ms,
|
||||||
|
created_at=now,
|
||||||
|
updated_at=now,
|
||||||
|
)
|
||||||
|
context.db.add(node)
|
||||||
|
|
||||||
|
context.db.commit()
|
||||||
|
context.db.refresh(node)
|
||||||
|
|
||||||
|
context.add_audit_metadata(
|
||||||
|
action="proxy_node_register",
|
||||||
|
proxy_node_id=node.id,
|
||||||
|
proxy_node_ip=node.ip,
|
||||||
|
proxy_node_port=node.port,
|
||||||
|
)
|
||||||
|
|
||||||
|
return {"node_id": node.id, "node": _node_to_dict(node)}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AdminHeartbeatProxyNodeAdapter(AdminApiAdapter):
|
||||||
|
name: str = "admin_heartbeat_proxy_node"
|
||||||
|
|
||||||
|
async def handle(self, context: ApiRequestContext) -> Any:
|
||||||
|
payload = context.ensure_json_body()
|
||||||
|
try:
|
||||||
|
req = ProxyNodeHeartbeatRequest.model_validate(payload)
|
||||||
|
except ValidationError as exc:
|
||||||
|
raise InvalidRequestException("输入验证失败: " + _format_validation_error(exc))
|
||||||
|
|
||||||
|
node = context.db.query(ProxyNode).filter(ProxyNode.id == req.node_id).first()
|
||||||
|
if not node:
|
||||||
|
raise NotFoundException(f"ProxyNode {req.node_id} 不存在", "proxy_node")
|
||||||
|
|
||||||
|
now = datetime.now(timezone.utc)
|
||||||
|
node.status = ProxyNodeStatus.ONLINE
|
||||||
|
node.last_heartbeat_at = now
|
||||||
|
if req.heartbeat_interval is not None:
|
||||||
|
node.heartbeat_interval = req.heartbeat_interval
|
||||||
|
if req.active_connections is not None:
|
||||||
|
node.active_connections = req.active_connections
|
||||||
|
if req.total_requests is not None:
|
||||||
|
node.total_requests = req.total_requests
|
||||||
|
if req.avg_latency_ms is not None:
|
||||||
|
node.avg_latency_ms = req.avg_latency_ms
|
||||||
|
|
||||||
|
context.db.commit()
|
||||||
|
context.db.refresh(node)
|
||||||
|
|
||||||
|
context.add_audit_metadata(
|
||||||
|
action="proxy_node_heartbeat",
|
||||||
|
proxy_node_id=node.id,
|
||||||
|
)
|
||||||
|
|
||||||
|
return {"message": "heartbeat ok", "node": _node_to_dict(node)}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AdminUnregisterProxyNodeAdapter(AdminApiAdapter):
|
||||||
|
name: str = "admin_unregister_proxy_node"
|
||||||
|
|
||||||
|
async def handle(self, context: ApiRequestContext) -> Any:
|
||||||
|
payload = context.ensure_json_body()
|
||||||
|
try:
|
||||||
|
req = ProxyNodeUnregisterRequest.model_validate(payload)
|
||||||
|
except ValidationError as exc:
|
||||||
|
raise InvalidRequestException("输入验证失败: " + _format_validation_error(exc))
|
||||||
|
|
||||||
|
node = context.db.query(ProxyNode).filter(ProxyNode.id == req.node_id).first()
|
||||||
|
if not node:
|
||||||
|
raise NotFoundException(f"ProxyNode {req.node_id} 不存在", "proxy_node")
|
||||||
|
|
||||||
|
node.status = ProxyNodeStatus.OFFLINE
|
||||||
|
node.updated_at = datetime.now(timezone.utc)
|
||||||
|
context.db.commit()
|
||||||
|
|
||||||
|
context.add_audit_metadata(
|
||||||
|
action="proxy_node_unregister",
|
||||||
|
proxy_node_id=node.id,
|
||||||
|
)
|
||||||
|
|
||||||
|
return {"message": "unregistered", "node_id": node.id}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AdminListProxyNodesAdapter(AdminApiAdapter):
|
||||||
|
name: str = "admin_list_proxy_nodes"
|
||||||
|
status: str | None = None
|
||||||
|
skip: int = 0
|
||||||
|
limit: int = 100
|
||||||
|
|
||||||
|
async def handle(self, context: ApiRequestContext) -> Any:
|
||||||
|
query = context.db.query(ProxyNode)
|
||||||
|
if self.status:
|
||||||
|
normalized = self.status.strip().lower()
|
||||||
|
allowed = {"online", "unhealthy", "offline"}
|
||||||
|
if normalized not in allowed:
|
||||||
|
raise InvalidRequestException(f"status 必须是以下之一: {sorted(allowed)}", "status")
|
||||||
|
query = query.filter(ProxyNode.status == ProxyNodeStatus(normalized))
|
||||||
|
|
||||||
|
total = query.count()
|
||||||
|
nodes = (
|
||||||
|
query.order_by(ProxyNode.updated_at.desc()).offset(self.skip).limit(self.limit).all()
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"items": [_node_to_dict(n) for n in nodes],
|
||||||
|
"total": total,
|
||||||
|
"skip": self.skip,
|
||||||
|
"limit": self.limit,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AdminDeleteProxyNodeAdapter(AdminApiAdapter):
|
||||||
|
name: str = "admin_delete_proxy_node"
|
||||||
|
node_id: str = ""
|
||||||
|
|
||||||
|
async def handle(self, context: ApiRequestContext) -> Any:
|
||||||
|
node = context.db.query(ProxyNode).filter(ProxyNode.id == self.node_id).first()
|
||||||
|
if not node:
|
||||||
|
raise NotFoundException(f"ProxyNode {self.node_id} 不存在", "proxy_node")
|
||||||
|
|
||||||
|
context.add_audit_metadata(
|
||||||
|
action="proxy_node_delete",
|
||||||
|
proxy_node_id=node.id,
|
||||||
|
proxy_node_ip=node.ip,
|
||||||
|
proxy_node_port=node.port,
|
||||||
|
)
|
||||||
|
|
||||||
|
context.db.delete(node)
|
||||||
|
context.db.commit()
|
||||||
|
|
||||||
|
return {"message": "deleted", "node_id": self.node_id}
|
||||||
@@ -2182,7 +2182,15 @@ class CliMessageHandlerBase(BaseMessageHandler):
|
|||||||
message = evt.get("message", {})
|
message = evt.get("message", {})
|
||||||
if isinstance(message, dict):
|
if isinstance(message, dict):
|
||||||
usage = message.get("usage")
|
usage = message.get("usage")
|
||||||
# OpenAI 格式: 直接在 chunk 中
|
# OpenAI Responses API (openai:cli) 格式: response.completed 中 usage 嵌套在 response 对象内
|
||||||
|
elif event_type == "response.completed":
|
||||||
|
resp_obj = evt.get("response")
|
||||||
|
if isinstance(resp_obj, dict):
|
||||||
|
usage = resp_obj.get("usage")
|
||||||
|
# 兼容: 部分实现可能在顶层也有 usage
|
||||||
|
if not usage:
|
||||||
|
usage = evt.get("usage")
|
||||||
|
# OpenAI Chat 格式: 直接在 chunk 中
|
||||||
elif "usage" in evt:
|
elif "usage" in evt:
|
||||||
usage = evt.get("usage")
|
usage = evt.get("usage")
|
||||||
# Gemini 格式: usageMetadata
|
# Gemini 格式: usageMetadata
|
||||||
|
|||||||
@@ -241,13 +241,145 @@ class OpenAIResponseParser(ResponseParser):
|
|||||||
|
|
||||||
|
|
||||||
class OpenAICliResponseParser(OpenAIResponseParser):
|
class OpenAICliResponseParser(OpenAIResponseParser):
|
||||||
"""OpenAI CLI 格式响应解析器"""
|
"""OpenAI CLI / Responses API 格式响应解析器
|
||||||
|
|
||||||
|
OpenAI Responses API 与 Chat Completions API 的关键差异:
|
||||||
|
- Usage 字段: input_tokens/output_tokens(而非 prompt_tokens/completion_tokens)
|
||||||
|
- 响应结构: output[].content[].text(而非 choices[].message.content)
|
||||||
|
- 流式事件: response.completed 事件中 usage 嵌套在 response 对象内
|
||||||
|
"""
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.name = "openai:cli"
|
self.name = "openai:cli"
|
||||||
self.api_format = "openai:cli"
|
self.api_format = "openai:cli"
|
||||||
|
|
||||||
|
def parse_response(self, response: dict[str, Any], status_code: int) -> ParsedResponse:
|
||||||
|
result = ParsedResponse(
|
||||||
|
raw_response=response,
|
||||||
|
status_code=status_code,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Responses API: 文本在 output[].content[].text 中
|
||||||
|
result.text_content = self._extract_responses_api_text(response)
|
||||||
|
result.response_id = response.get("id")
|
||||||
|
|
||||||
|
# Responses API usage: input_tokens / output_tokens
|
||||||
|
usage = self._extract_responses_api_usage(response)
|
||||||
|
result.input_tokens = usage.get("input_tokens", 0)
|
||||||
|
result.output_tokens = usage.get("output_tokens", 0)
|
||||||
|
result.cache_creation_tokens = usage.get("cache_creation_tokens", 0)
|
||||||
|
result.cache_read_tokens = usage.get("cache_read_tokens", 0)
|
||||||
|
|
||||||
|
# 检查错误(支持嵌套错误格式)
|
||||||
|
is_error, error_info = _check_nested_error(response)
|
||||||
|
if is_error and error_info:
|
||||||
|
result.is_error = True
|
||||||
|
result.error_type = error_info.get("type")
|
||||||
|
result.error_message = error_info.get("message")
|
||||||
|
result.embedded_status_code = _extract_embedded_status_code(error_info)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def extract_usage_from_response(self, response: dict[str, Any]) -> dict[str, int]:
|
||||||
|
usage = self._extract_responses_api_usage(response)
|
||||||
|
return usage
|
||||||
|
|
||||||
|
def extract_text_content(self, response: dict[str, Any]) -> str:
|
||||||
|
return self._extract_responses_api_text(response)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _extract_responses_api_usage(response: dict[str, Any]) -> dict[str, int]:
|
||||||
|
"""从 Responses API 响应或流式事件中提取 usage
|
||||||
|
|
||||||
|
支持多种结构:
|
||||||
|
1. 顶层 usage(非流式响应 / 部分转换后的响应)
|
||||||
|
2. response.usage(流式 response.completed 事件)
|
||||||
|
3. 兼容 Chat Completions 字段名(prompt_tokens/completion_tokens)
|
||||||
|
"""
|
||||||
|
usage: dict[str, Any] = {}
|
||||||
|
|
||||||
|
# 优先从顶层 usage 提取
|
||||||
|
top_usage = response.get("usage")
|
||||||
|
if isinstance(top_usage, dict):
|
||||||
|
usage = top_usage
|
||||||
|
else:
|
||||||
|
# 流式事件: response.completed 中 usage 嵌套在 response 对象内
|
||||||
|
resp_obj = response.get("response")
|
||||||
|
if isinstance(resp_obj, dict):
|
||||||
|
nested_usage = resp_obj.get("usage")
|
||||||
|
if isinstance(nested_usage, dict):
|
||||||
|
usage = nested_usage
|
||||||
|
|
||||||
|
if not usage:
|
||||||
|
return {
|
||||||
|
"input_tokens": 0,
|
||||||
|
"output_tokens": 0,
|
||||||
|
"cache_creation_tokens": 0,
|
||||||
|
"cache_read_tokens": 0,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Responses API 使用 input_tokens/output_tokens
|
||||||
|
# 兼容 Chat Completions 的 prompt_tokens/completion_tokens(以防转换后的响应)
|
||||||
|
input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
|
||||||
|
output_tokens = usage.get("output_tokens") or usage.get("completion_tokens") or 0
|
||||||
|
|
||||||
|
return {
|
||||||
|
"input_tokens": int(input_tokens),
|
||||||
|
"output_tokens": int(output_tokens),
|
||||||
|
"cache_creation_tokens": int(
|
||||||
|
usage.get("cache_creation_input_tokens") or usage.get("cache_creation_tokens") or 0
|
||||||
|
),
|
||||||
|
"cache_read_tokens": int(
|
||||||
|
usage.get("cache_read_input_tokens") or usage.get("cache_read_tokens") or 0
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _extract_responses_api_text(response: dict[str, Any]) -> str:
|
||||||
|
"""从 Responses API 响应中提取文本内容
|
||||||
|
|
||||||
|
支持结构: output[].content[].text 或 output[].text
|
||||||
|
"""
|
||||||
|
text_parts: list[str] = []
|
||||||
|
|
||||||
|
output = response.get("output")
|
||||||
|
if isinstance(output, list):
|
||||||
|
for item in output:
|
||||||
|
if not isinstance(item, dict):
|
||||||
|
continue
|
||||||
|
# message 类型: output[].content[].text
|
||||||
|
if item.get("type") == "message":
|
||||||
|
content = item.get("content")
|
||||||
|
if isinstance(content, list):
|
||||||
|
for part in content:
|
||||||
|
if isinstance(part, dict):
|
||||||
|
ptype = str(part.get("type") or "")
|
||||||
|
if ptype in ("output_text", "text") and isinstance(
|
||||||
|
part.get("text"), str
|
||||||
|
):
|
||||||
|
text_parts.append(part["text"])
|
||||||
|
# 直接文本类型: output[].text
|
||||||
|
elif item.get("type") in ("output_text", "text") and isinstance(
|
||||||
|
item.get("text"), str
|
||||||
|
):
|
||||||
|
text_parts.append(item["text"])
|
||||||
|
|
||||||
|
# 兼容: 部分实现可能直接给 output_text
|
||||||
|
if not text_parts and isinstance(response.get("output_text"), str):
|
||||||
|
text_parts.append(response["output_text"])
|
||||||
|
|
||||||
|
# 兼容: 如果是 Chat Completions 格式(可能来自转换后的响应),回退到 choices 结构
|
||||||
|
if not text_parts:
|
||||||
|
choices = response.get("choices", [])
|
||||||
|
if isinstance(choices, list) and choices:
|
||||||
|
message = choices[0].get("message", {}) if isinstance(choices[0], dict) else {}
|
||||||
|
content = message.get("content") if isinstance(message, dict) else None
|
||||||
|
if isinstance(content, str):
|
||||||
|
text_parts.append(content)
|
||||||
|
|
||||||
|
return "".join(text_parts)
|
||||||
|
|
||||||
|
|
||||||
class ClaudeResponseParser(ResponseParser):
|
class ClaudeResponseParser(ResponseParser):
|
||||||
"""Claude 格式响应解析器"""
|
"""Claude 格式响应解析器"""
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import hmac
|
||||||
import time
|
import time
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -20,6 +21,7 @@ from urllib.parse import quote, urlparse
|
|||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
from src.config import config
|
from src.config import config
|
||||||
|
from src.core.exceptions import ProxyNodeUnavailableError
|
||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
from src.utils.ssl_utils import get_ssl_context
|
from src.utils.ssl_utils import get_ssl_context
|
||||||
|
|
||||||
@@ -27,6 +29,70 @@ from src.utils.ssl_utils import get_ssl_context
|
|||||||
_proxy_clients_lock = asyncio.Lock()
|
_proxy_clients_lock = asyncio.Lock()
|
||||||
_default_client_lock = asyncio.Lock()
|
_default_client_lock = asyncio.Lock()
|
||||||
|
|
||||||
|
# ProxyNode 信息缓存(降低高频 DB 查询开销)
|
||||||
|
_proxy_node_cache: dict[str, tuple[dict[str, Any] | None, float]] = {}
|
||||||
|
_PROXY_NODE_CACHE_TTL_SECONDS = 60.0
|
||||||
|
_PROXY_NODE_CACHE_MAX_SIZE = 256
|
||||||
|
|
||||||
|
|
||||||
|
def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
|
||||||
|
"""
|
||||||
|
读取 ProxyNode 信息(带内存 TTL 缓存)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
{"ip": str, "port": int} 或 None(不存在/非在线)
|
||||||
|
"""
|
||||||
|
now = time.time()
|
||||||
|
cached = _proxy_node_cache.get(node_id)
|
||||||
|
if cached:
|
||||||
|
value, expires_at = cached
|
||||||
|
if now < expires_at:
|
||||||
|
return value
|
||||||
|
|
||||||
|
# 防止无效 node_id 导致缓存无限膨胀
|
||||||
|
if len(_proxy_node_cache) >= _PROXY_NODE_CACHE_MAX_SIZE:
|
||||||
|
_proxy_node_cache.clear()
|
||||||
|
|
||||||
|
from src.database import create_session
|
||||||
|
from src.models.database import ProxyNode, ProxyNodeStatus
|
||||||
|
|
||||||
|
db = create_session()
|
||||||
|
try:
|
||||||
|
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
|
||||||
|
if not node or node.status != ProxyNodeStatus.ONLINE:
|
||||||
|
_proxy_node_cache[node_id] = (None, now + _PROXY_NODE_CACHE_TTL_SECONDS)
|
||||||
|
return None
|
||||||
|
|
||||||
|
value = {"ip": node.ip, "port": node.port}
|
||||||
|
_proxy_node_cache[node_id] = (value, now + _PROXY_NODE_CACHE_TTL_SECONDS)
|
||||||
|
return value
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
|
||||||
|
def _build_hmac_proxy_url(ip: str, port: int, node_id: str) -> str:
|
||||||
|
"""
|
||||||
|
构建带 HMAC BasicAuth 的 httpx proxy URL
|
||||||
|
|
||||||
|
格式: http://hmac:{timestamp}.{signature}@{ip}:{port}
|
||||||
|
signature = HMAC-SHA256(PROXY_HMAC_KEY, "{timestamp}\\n{node_id}") 的 hex
|
||||||
|
"""
|
||||||
|
if not config.proxy_hmac_key:
|
||||||
|
raise ProxyNodeUnavailableError(
|
||||||
|
"PROXY_HMAC_KEY 未配置,无法使用 ProxyNode 代理", node_id=node_id
|
||||||
|
)
|
||||||
|
|
||||||
|
timestamp = str(int(time.time()))
|
||||||
|
payload = f"{timestamp}\n{node_id}".encode("utf-8")
|
||||||
|
signature = hmac.new(
|
||||||
|
config.proxy_hmac_key.encode("utf-8"),
|
||||||
|
payload,
|
||||||
|
hashlib.sha256,
|
||||||
|
).hexdigest()
|
||||||
|
|
||||||
|
host = f"[{ip}]" if ":" in ip else ip
|
||||||
|
return f"http://hmac:{timestamp}.{signature}@{host}:{int(port)}"
|
||||||
|
|
||||||
|
|
||||||
def _compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
|
def _compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
|
||||||
"""
|
"""
|
||||||
@@ -41,6 +107,16 @@ def _compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
|
|||||||
if not proxy_config:
|
if not proxy_config:
|
||||||
return "__no_proxy__"
|
return "__no_proxy__"
|
||||||
|
|
||||||
|
# enabled=False 时视为无代理(兼容旧数据)
|
||||||
|
if not proxy_config.get("enabled", True):
|
||||||
|
return "__no_proxy__"
|
||||||
|
|
||||||
|
# ProxyNode 模式:基于 node_id + 时间桶缓存,避免签名随时间变化导致 cache key 爆炸
|
||||||
|
node_id = proxy_config.get("node_id")
|
||||||
|
if isinstance(node_id, str) and node_id.strip():
|
||||||
|
time_bucket = int(time.time() / 120) # 120 秒一个桶
|
||||||
|
return f"proxy_node:{node_id.strip()}:{time_bucket}"
|
||||||
|
|
||||||
# 构建代理 URL 作为缓存键的基础
|
# 构建代理 URL 作为缓存键的基础
|
||||||
proxy_url = build_proxy_url(proxy_config)
|
proxy_url = build_proxy_url(proxy_config)
|
||||||
if not proxy_url:
|
if not proxy_url:
|
||||||
@@ -55,7 +131,9 @@ def build_proxy_url(proxy_config: dict[str, Any]) -> str | None:
|
|||||||
根据代理配置构建完整的代理 URL
|
根据代理配置构建完整的代理 URL
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
proxy_config: 代理配置字典,包含 url, username, password, enabled
|
proxy_config: 代理配置字典,支持两种模式:
|
||||||
|
- 手动 URL 模式: {url, username, password, enabled}
|
||||||
|
- ProxyNode 模式: {node_id, enabled}
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
完整的代理 URL,如 socks5://user:pass@host:port
|
完整的代理 URL,如 socks5://user:pass@host:port
|
||||||
@@ -68,6 +146,15 @@ def build_proxy_url(proxy_config: dict[str, Any]) -> str | None:
|
|||||||
if not proxy_config.get("enabled", True):
|
if not proxy_config.get("enabled", True):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
# ProxyNode 模式(aether-proxy)
|
||||||
|
node_id = proxy_config.get("node_id")
|
||||||
|
if isinstance(node_id, str) and node_id.strip():
|
||||||
|
node_id = node_id.strip()
|
||||||
|
node_info = _get_proxy_node_info(node_id)
|
||||||
|
if not node_info:
|
||||||
|
raise ProxyNodeUnavailableError(f"代理节点 {node_id} 不可用", node_id=node_id)
|
||||||
|
return _build_hmac_proxy_url(node_info["ip"], node_info["port"], node_id)
|
||||||
|
|
||||||
proxy_url: str | None = proxy_config.get("url")
|
proxy_url: str | None = proxy_config.get("url")
|
||||||
if not proxy_url:
|
if not proxy_url:
|
||||||
return None
|
return None
|
||||||
@@ -307,9 +394,13 @@ class HTTPClientPool:
|
|||||||
client = httpx.AsyncClient(**client_config) # type: ignore[arg-type]
|
client = httpx.AsyncClient(**client_config) # type: ignore[arg-type]
|
||||||
cls._proxy_clients[cache_key] = (client, time.time())
|
cls._proxy_clients[cache_key] = (client, time.time())
|
||||||
|
|
||||||
|
proxy_label = "none"
|
||||||
|
if proxy_config:
|
||||||
|
proxy_label = str(
|
||||||
|
proxy_config.get("node_id") or proxy_config.get("url") or "unknown"
|
||||||
|
)
|
||||||
logger.debug(
|
logger.debug(
|
||||||
f"创建代理客户端(缓存): {proxy_config.get('url', 'unknown') if proxy_config else 'none'}, "
|
f"创建代理客户端(缓存): {proxy_label}, " f"缓存数量: {len(cls._proxy_clients)}"
|
||||||
f"缓存数量: {len(cls._proxy_clients)}"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
return client
|
return client
|
||||||
|
|||||||
@@ -3,6 +3,8 @@
|
|||||||
从环境变量或 .env 文件加载配置
|
从环境变量或 .env 文件加载配置
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
import hmac
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -46,6 +48,13 @@ class Config:
|
|||||||
# 加密密钥配置(独立于JWT密钥,用于敏感数据加密)
|
# 加密密钥配置(独立于JWT密钥,用于敏感数据加密)
|
||||||
self.encryption_key = os.getenv("ENCRYPTION_KEY", None)
|
self.encryption_key = os.getenv("ENCRYPTION_KEY", None)
|
||||||
|
|
||||||
|
# 代理节点 HMAC 密钥(用于 aether-proxy 认证)
|
||||||
|
proxy_hmac_key_env = os.getenv("PROXY_HMAC_KEY")
|
||||||
|
if proxy_hmac_key_env and proxy_hmac_key_env.strip():
|
||||||
|
self.proxy_hmac_key = proxy_hmac_key_env.strip()
|
||||||
|
else:
|
||||||
|
self.proxy_hmac_key = self._derive_proxy_hmac_key()
|
||||||
|
|
||||||
# 环境配置 - 智能检测
|
# 环境配置 - 智能检测
|
||||||
# Docker 部署默认为生产环境,本地开发默认为开发环境
|
# Docker 部署默认为生产环境,本地开发默认为开发环境
|
||||||
is_docker = (
|
is_docker = (
|
||||||
@@ -307,6 +316,20 @@ class Config:
|
|||||||
# 验证连接池配置
|
# 验证连接池配置
|
||||||
self._validate_pool_config()
|
self._validate_pool_config()
|
||||||
|
|
||||||
|
def _derive_proxy_hmac_key(self) -> str:
|
||||||
|
"""
|
||||||
|
从 ENCRYPTION_KEY 派生 PROXY_HMAC_KEY
|
||||||
|
|
||||||
|
目的:避免把 ENCRYPTION_KEY 直接下发到 VPS(aether-proxy)。
|
||||||
|
"""
|
||||||
|
if not self.encryption_key:
|
||||||
|
return ""
|
||||||
|
return hmac.new(
|
||||||
|
self.encryption_key.encode("utf-8"),
|
||||||
|
b"aether-proxy-hmac-key-v1",
|
||||||
|
hashlib.sha256,
|
||||||
|
).hexdigest()
|
||||||
|
|
||||||
def _auto_pool_size(self) -> int:
|
def _auto_pool_size(self) -> int:
|
||||||
"""
|
"""
|
||||||
智能计算连接池大小 - 根据 Worker 数量和 PostgreSQL 限制计算
|
智能计算连接池大小 - 根据 Worker 数量和 PostgreSQL 限制计算
|
||||||
|
|||||||
@@ -209,6 +209,17 @@ class ProviderNotAvailableException(ProviderException):
|
|||||||
self.upstream_response = upstream_response
|
self.upstream_response = upstream_response
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNodeUnavailableError(ProviderException):
|
||||||
|
"""代理节点不可用(ProxyNode 离线/不存在/不健康)"""
|
||||||
|
|
||||||
|
def __init__(self, message: str, node_id: str | None = None):
|
||||||
|
super().__init__(
|
||||||
|
message=message,
|
||||||
|
provider_name=None,
|
||||||
|
proxy_node_id=node_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ProviderTimeoutException(ProviderException):
|
class ProviderTimeoutException(ProviderException):
|
||||||
"""提供商请求超时"""
|
"""提供商请求超时"""
|
||||||
|
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ import re
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, field_validator
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||||
|
|
||||||
from src.core.enums import ProviderBillingType
|
from src.core.enums import ProviderBillingType
|
||||||
|
|
||||||
@@ -18,18 +18,26 @@ from src.core.enums import ProviderBillingType
|
|||||||
class ProxyConfig(BaseModel):
|
class ProxyConfig(BaseModel):
|
||||||
"""代理配置"""
|
"""代理配置"""
|
||||||
|
|
||||||
url: str = Field(..., description="代理 URL (http://, https://, socks5://)")
|
# 模式 1: 手动配置代理 URL(原有)
|
||||||
|
url: str | None = Field(None, description="代理 URL (http://, https://, socks5://)")
|
||||||
username: str | None = Field(None, max_length=255, description="代理用户名")
|
username: str | None = Field(None, max_length=255, description="代理用户名")
|
||||||
password: str | None = Field(None, max_length=500, description="代理密码")
|
password: str | None = Field(None, max_length=500, description="代理密码")
|
||||||
|
# 模式 2: ProxyNode(aether-proxy 注册的节点)
|
||||||
|
node_id: str | None = Field(None, description="代理节点 ID")
|
||||||
enabled: bool = Field(True, description="是否启用代理(false 时保留配置但不使用)")
|
enabled: bool = Field(True, description="是否启用代理(false 时保留配置但不使用)")
|
||||||
|
|
||||||
@field_validator("url")
|
@field_validator("url")
|
||||||
@classmethod
|
@classmethod
|
||||||
def validate_proxy_url(cls, v: str) -> str:
|
def validate_proxy_url(cls, v: str | None) -> str | None:
|
||||||
"""验证代理 URL 格式"""
|
"""验证代理 URL 格式"""
|
||||||
|
if v is None:
|
||||||
|
return None
|
||||||
|
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
v = v.strip()
|
v = v.strip()
|
||||||
|
if not v:
|
||||||
|
return None
|
||||||
|
|
||||||
# 检查禁止的字符(防止注入)
|
# 检查禁止的字符(防止注入)
|
||||||
if "\n" in v or "\r" in v:
|
if "\n" in v or "\r" in v:
|
||||||
@@ -50,6 +58,27 @@ class ProxyConfig(BaseModel):
|
|||||||
|
|
||||||
return v
|
return v
|
||||||
|
|
||||||
|
@field_validator("node_id")
|
||||||
|
@classmethod
|
||||||
|
def validate_node_id(cls, v: str | None) -> str | None:
|
||||||
|
if v is None:
|
||||||
|
return None
|
||||||
|
v = v.strip()
|
||||||
|
return v or None
|
||||||
|
|
||||||
|
@model_validator(mode="after")
|
||||||
|
def validate_proxy_mode(self) -> "ProxyConfig":
|
||||||
|
if not self.enabled:
|
||||||
|
return self
|
||||||
|
|
||||||
|
if not self.url and not self.node_id:
|
||||||
|
raise ValueError("启用代理时,必须提供 url 或 node_id")
|
||||||
|
|
||||||
|
if self.url and self.node_id:
|
||||||
|
raise ValueError("url 和 node_id 不能同时设置")
|
||||||
|
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
class CreateProviderRequest(BaseModel):
|
class CreateProviderRequest(BaseModel):
|
||||||
"""创建 Provider 请求"""
|
"""创建 Provider 请求"""
|
||||||
|
|||||||
@@ -792,6 +792,63 @@ class ProviderEndpoint(Base):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNodeStatus(PyEnum):
|
||||||
|
"""代理节点状态"""
|
||||||
|
|
||||||
|
ONLINE = "online"
|
||||||
|
UNHEALTHY = "unhealthy"
|
||||||
|
OFFLINE = "offline"
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNode(Base):
|
||||||
|
"""代理节点表(用于 aether-proxy 注册/心跳)"""
|
||||||
|
|
||||||
|
__tablename__ = "proxy_nodes"
|
||||||
|
|
||||||
|
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), index=True)
|
||||||
|
name = Column(String(100), nullable=False) # 节点名
|
||||||
|
ip = Column(String(45), nullable=False) # 公网 IP(IPv6 最长 39 + 冗余)
|
||||||
|
port = Column(Integer, nullable=False) # 代理端口
|
||||||
|
region = Column(String(100), nullable=True) # 区域标签
|
||||||
|
|
||||||
|
status = Column(
|
||||||
|
Enum(
|
||||||
|
ProxyNodeStatus,
|
||||||
|
name="proxynodestatus",
|
||||||
|
create_type=False,
|
||||||
|
values_callable=lambda x: [e.value for e in x],
|
||||||
|
),
|
||||||
|
default=ProxyNodeStatus.ONLINE,
|
||||||
|
nullable=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
registered_by = Column(
|
||||||
|
String(36),
|
||||||
|
ForeignKey("users.id", ondelete="SET NULL"),
|
||||||
|
nullable=True,
|
||||||
|
comment="注册该节点的管理员用户 ID(可空)",
|
||||||
|
)
|
||||||
|
last_heartbeat_at = Column(DateTime(timezone=True), nullable=True)
|
||||||
|
heartbeat_interval = Column(Integer, default=30, nullable=False)
|
||||||
|
|
||||||
|
# 性能指标(心跳上报)
|
||||||
|
active_connections = Column(Integer, default=0, nullable=False)
|
||||||
|
total_requests = Column(BigInteger, default=0, nullable=False)
|
||||||
|
avg_latency_ms = Column(Float, nullable=True)
|
||||||
|
|
||||||
|
created_at = Column(
|
||||||
|
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False
|
||||||
|
)
|
||||||
|
updated_at = Column(
|
||||||
|
DateTime(timezone=True),
|
||||||
|
default=lambda: datetime.now(timezone.utc),
|
||||||
|
onupdate=lambda: datetime.now(timezone.utc),
|
||||||
|
nullable=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
__table_args__ = (UniqueConstraint("ip", "port", name="uq_proxy_node_ip_port"),)
|
||||||
|
|
||||||
|
|
||||||
class GlobalModel(Base):
|
class GlobalModel(Base):
|
||||||
"""全局统一模型定义 - 包含价格和能力配置
|
"""全局统一模型定义 - 包含价格和能力配置
|
||||||
|
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from src.modules.gemini_files import gemini_files_module
|
|||||||
from src.modules.ldap import ldap_module
|
from src.modules.ldap import ldap_module
|
||||||
from src.modules.management_tokens import management_tokens_module
|
from src.modules.management_tokens import management_tokens_module
|
||||||
from src.modules.oauth import oauth_module
|
from src.modules.oauth import oauth_module
|
||||||
|
from src.modules.proxy_nodes import proxy_nodes_module
|
||||||
|
|
||||||
# 所有模块列表
|
# 所有模块列表
|
||||||
ALL_MODULES: list[ModuleDefinition] = [
|
ALL_MODULES: list[ModuleDefinition] = [
|
||||||
@@ -18,6 +19,7 @@ ALL_MODULES: list[ModuleDefinition] = [
|
|||||||
oauth_module,
|
oauth_module,
|
||||||
gemini_files_module,
|
gemini_files_module,
|
||||||
management_tokens_module,
|
management_tokens_module,
|
||||||
|
proxy_nodes_module,
|
||||||
]
|
]
|
||||||
|
|
||||||
__all__ = ["ALL_MODULES"]
|
__all__ = ["ALL_MODULES"]
|
||||||
|
|||||||
112
src/modules/proxy_nodes/__init__.py
Normal file
112
src/modules/proxy_nodes/__init__.py
Normal file
@@ -0,0 +1,112 @@
|
|||||||
|
"""
|
||||||
|
代理节点模块
|
||||||
|
|
||||||
|
提供海外 VPS 代理节点的注册、心跳、管理功能。
|
||||||
|
aether-proxy 部署在海外 VPS 上自动注册节点,Aether 通过 HMAC 签名认证转发请求。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
from src.core.modules.base import (
|
||||||
|
ModuleCategory,
|
||||||
|
ModuleDefinition,
|
||||||
|
ModuleHealth,
|
||||||
|
ModuleMetadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
|
|
||||||
|
def _get_router() -> Any:
|
||||||
|
"""延迟导入路由"""
|
||||||
|
from src.api.admin.proxy_nodes import router
|
||||||
|
|
||||||
|
return router
|
||||||
|
|
||||||
|
|
||||||
|
async def _on_startup() -> None:
|
||||||
|
"""启动心跳检测调度器"""
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
||||||
|
from src.utils.task_coordinator import StartupTaskCoordinator
|
||||||
|
|
||||||
|
logger = logging.getLogger("aether.modules.proxy_nodes")
|
||||||
|
|
||||||
|
from src.clients import get_redis_client
|
||||||
|
|
||||||
|
redis_client = await get_redis_client()
|
||||||
|
task_coordinator = StartupTaskCoordinator(redis_client)
|
||||||
|
|
||||||
|
proxy_node_health_scheduler = get_proxy_node_health_scheduler()
|
||||||
|
active = await task_coordinator.acquire("proxy_node_health")
|
||||||
|
if active:
|
||||||
|
logger.info("启动 ProxyNode 心跳检测调度器...")
|
||||||
|
await proxy_node_health_scheduler.start()
|
||||||
|
else:
|
||||||
|
logger.info("检测到其他 worker 已运行 ProxyNode 心跳检测,本实例跳过")
|
||||||
|
|
||||||
|
|
||||||
|
async def _on_shutdown() -> None:
|
||||||
|
"""停止心跳检测调度器"""
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
||||||
|
from src.utils.task_coordinator import StartupTaskCoordinator
|
||||||
|
|
||||||
|
logger = logging.getLogger("aether.modules.proxy_nodes")
|
||||||
|
|
||||||
|
from src.clients import get_redis_client
|
||||||
|
|
||||||
|
redis_client = await get_redis_client()
|
||||||
|
task_coordinator = StartupTaskCoordinator(redis_client)
|
||||||
|
|
||||||
|
scheduler = get_proxy_node_health_scheduler()
|
||||||
|
if scheduler.running:
|
||||||
|
logger.info("停止 ProxyNode 心跳检测调度器...")
|
||||||
|
await scheduler.stop()
|
||||||
|
await task_coordinator.release("proxy_node_health")
|
||||||
|
|
||||||
|
|
||||||
|
async def _health_check() -> ModuleHealth:
|
||||||
|
"""健康检查 - 检查是否有在线节点"""
|
||||||
|
return ModuleHealth.HEALTHY
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_config(db: Session) -> tuple[bool, str]:
|
||||||
|
"""
|
||||||
|
验证配置
|
||||||
|
|
||||||
|
代理节点模块需要 PROXY_HMAC_KEY 配置
|
||||||
|
"""
|
||||||
|
from src.config.settings import config
|
||||||
|
|
||||||
|
if not config.proxy_hmac_key:
|
||||||
|
return False, "PROXY_HMAC_KEY 未配置(也未设置 ENCRYPTION_KEY 用于自动派生)"
|
||||||
|
return True, ""
|
||||||
|
|
||||||
|
|
||||||
|
proxy_nodes_module = ModuleDefinition(
|
||||||
|
metadata=ModuleMetadata(
|
||||||
|
name="proxy_nodes",
|
||||||
|
display_name="代理节点",
|
||||||
|
description="海外 VPS 代理节点管理,通过 HMAC 签名认证转发 API 请求",
|
||||||
|
category=ModuleCategory.INTEGRATION,
|
||||||
|
env_key="PROXY_NODES_AVAILABLE",
|
||||||
|
default_available=True,
|
||||||
|
required_packages=[],
|
||||||
|
api_prefix="/api/admin/proxy-nodes",
|
||||||
|
admin_route="/admin/proxy-nodes",
|
||||||
|
admin_menu_icon="Server",
|
||||||
|
admin_menu_group="system",
|
||||||
|
admin_menu_order=60,
|
||||||
|
),
|
||||||
|
router_factory=_get_router,
|
||||||
|
on_startup=_on_startup,
|
||||||
|
on_shutdown=_on_shutdown,
|
||||||
|
health_check=_health_check,
|
||||||
|
validate_config=_validate_config,
|
||||||
|
)
|
||||||
5
src/services/proxy_node/__init__.py
Normal file
5
src/services/proxy_node/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
"""Proxy node services."""
|
||||||
|
|
||||||
|
from .health_scheduler import ProxyNodeHealthScheduler, get_proxy_node_health_scheduler
|
||||||
|
|
||||||
|
__all__ = ["ProxyNodeHealthScheduler", "get_proxy_node_health_scheduler"]
|
||||||
103
src/services/proxy_node/health_scheduler.py
Normal file
103
src/services/proxy_node/health_scheduler.py
Normal file
@@ -0,0 +1,103 @@
|
|||||||
|
"""
|
||||||
|
ProxyNode 心跳检测调度器
|
||||||
|
|
||||||
|
定期检查 proxy_nodes 的 last_heartbeat_at,更新节点状态:
|
||||||
|
- elapsed > interval * 3 -> unhealthy
|
||||||
|
- elapsed > interval * 10 -> offline
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from src.core.logger import logger
|
||||||
|
from src.database import create_session
|
||||||
|
from src.models.database import ProxyNode, ProxyNodeStatus
|
||||||
|
from src.services.system.scheduler import get_scheduler
|
||||||
|
|
||||||
|
|
||||||
|
class ProxyNodeHealthScheduler:
|
||||||
|
"""代理节点心跳检测调度器"""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.running = False
|
||||||
|
|
||||||
|
async def start(self) -> Any:
|
||||||
|
if self.running:
|
||||||
|
logger.warning("ProxyNodeHealthScheduler already running")
|
||||||
|
return
|
||||||
|
|
||||||
|
self.running = True
|
||||||
|
logger.info("ProxyNodeHealthScheduler started")
|
||||||
|
|
||||||
|
scheduler = get_scheduler()
|
||||||
|
scheduler.add_interval_job(
|
||||||
|
self._scheduled_check,
|
||||||
|
seconds=30,
|
||||||
|
job_id="proxy_node_health_check",
|
||||||
|
name="代理节点心跳检测",
|
||||||
|
)
|
||||||
|
|
||||||
|
# 启动时立即执行一次
|
||||||
|
await self._check_heartbeats()
|
||||||
|
|
||||||
|
async def stop(self) -> Any:
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
|
self.running = False
|
||||||
|
logger.info("ProxyNodeHealthScheduler stopped")
|
||||||
|
|
||||||
|
async def _scheduled_check(self) -> None:
|
||||||
|
await self._check_heartbeats()
|
||||||
|
|
||||||
|
async def _check_heartbeats(self) -> None:
|
||||||
|
db = create_session()
|
||||||
|
try:
|
||||||
|
now = datetime.now(timezone.utc)
|
||||||
|
nodes = db.query(ProxyNode).filter(ProxyNode.status != ProxyNodeStatus.OFFLINE).all()
|
||||||
|
if not nodes:
|
||||||
|
return
|
||||||
|
|
||||||
|
changed = 0
|
||||||
|
for node in nodes:
|
||||||
|
interval = int(node.heartbeat_interval or 30)
|
||||||
|
last = node.last_heartbeat_at
|
||||||
|
|
||||||
|
if last is None:
|
||||||
|
new_status = ProxyNodeStatus.OFFLINE
|
||||||
|
else:
|
||||||
|
elapsed = (now - last).total_seconds()
|
||||||
|
if elapsed > interval * 10:
|
||||||
|
new_status = ProxyNodeStatus.OFFLINE
|
||||||
|
elif elapsed > interval * 3:
|
||||||
|
new_status = ProxyNodeStatus.UNHEALTHY
|
||||||
|
else:
|
||||||
|
new_status = ProxyNodeStatus.ONLINE
|
||||||
|
|
||||||
|
if node.status != new_status:
|
||||||
|
node.status = new_status
|
||||||
|
node.updated_at = now
|
||||||
|
changed += 1
|
||||||
|
|
||||||
|
if changed:
|
||||||
|
db.commit()
|
||||||
|
logger.info("ProxyNode 心跳状态已更新: {} 个节点", changed)
|
||||||
|
except Exception as e:
|
||||||
|
try:
|
||||||
|
db.rollback()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
logger.exception("ProxyNode 心跳检测失败: {}", e)
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
|
||||||
|
_proxy_node_health_scheduler: ProxyNodeHealthScheduler | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def get_proxy_node_health_scheduler() -> ProxyNodeHealthScheduler:
|
||||||
|
global _proxy_node_health_scheduler
|
||||||
|
if _proxy_node_health_scheduler is None:
|
||||||
|
_proxy_node_health_scheduler = ProxyNodeHealthScheduler()
|
||||||
|
return _proxy_node_health_scheduler
|
||||||
@@ -704,6 +704,7 @@ class TaskService:
|
|||||||
from src.core.exceptions import (
|
from src.core.exceptions import (
|
||||||
ConcurrencyLimitError,
|
ConcurrencyLimitError,
|
||||||
EmbeddedErrorException,
|
EmbeddedErrorException,
|
||||||
|
ProxyNodeUnavailableError,
|
||||||
ThinkingSignatureException,
|
ThinkingSignatureException,
|
||||||
UpstreamClientException,
|
UpstreamClientException,
|
||||||
)
|
)
|
||||||
@@ -744,6 +745,20 @@ class TaskService:
|
|||||||
)
|
)
|
||||||
return "break"
|
return "break"
|
||||||
|
|
||||||
|
if isinstance(cause, ProxyNodeUnavailableError):
|
||||||
|
# ProxyNode 不可用属于“配置明确指定但不可达/不可用”的情况,
|
||||||
|
# 在当前候选上重试通常没有意义,直接切换到下一个候选更合理。
|
||||||
|
logger.warning(" [{}] 代理节点不可用,切换候选: {}", request_id, str(cause))
|
||||||
|
RequestCandidateService.mark_candidate_failed(
|
||||||
|
db=self.db,
|
||||||
|
candidate_id=candidate_record_id,
|
||||||
|
error_type=type(cause).__name__,
|
||||||
|
error_message=extract_error_message(cause),
|
||||||
|
latency_ms=elapsed_ms,
|
||||||
|
concurrent_requests=captured_key_concurrent,
|
||||||
|
)
|
||||||
|
return "break"
|
||||||
|
|
||||||
if isinstance(cause, EmbeddedErrorException):
|
if isinstance(cause, EmbeddedErrorException):
|
||||||
error_message = cause.error_message or ""
|
error_message = cause.error_message or ""
|
||||||
embedded_status = cause.error_code or 200
|
embedded_status = cause.error_code or 200
|
||||||
|
|||||||
35
tests/unit/test_proxy_config.py
Normal file
35
tests/unit/test_proxy_config.py
Normal file
@@ -0,0 +1,35 @@
|
|||||||
|
import pytest
|
||||||
|
from pydantic import ValidationError
|
||||||
|
|
||||||
|
from src.models.admin_requests import ProxyConfig
|
||||||
|
|
||||||
|
|
||||||
|
class TestProxyConfig:
|
||||||
|
def test_enabled_requires_url_or_node_id(self) -> None:
|
||||||
|
with pytest.raises(ValidationError):
|
||||||
|
ProxyConfig.model_validate({"enabled": True})
|
||||||
|
|
||||||
|
def test_enabled_rejects_both_url_and_node_id(self) -> None:
|
||||||
|
with pytest.raises(ValidationError):
|
||||||
|
ProxyConfig.model_validate(
|
||||||
|
{"enabled": True, "url": "http://127.0.0.1:8080", "node_id": "n1"}
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_disabled_allows_empty(self) -> None:
|
||||||
|
cfg = ProxyConfig.model_validate({"enabled": False})
|
||||||
|
assert cfg.url is None
|
||||||
|
assert cfg.node_id is None
|
||||||
|
|
||||||
|
def test_url_mode_ok(self) -> None:
|
||||||
|
cfg = ProxyConfig.model_validate({"enabled": True, "url": "http://127.0.0.1:8080"})
|
||||||
|
assert cfg.url == "http://127.0.0.1:8080"
|
||||||
|
assert cfg.node_id is None
|
||||||
|
|
||||||
|
def test_url_rejects_inline_auth(self) -> None:
|
||||||
|
with pytest.raises(ValidationError):
|
||||||
|
ProxyConfig.model_validate({"enabled": True, "url": "http://u:p@127.0.0.1:8080"})
|
||||||
|
|
||||||
|
def test_node_id_mode_ok_and_strips(self) -> None:
|
||||||
|
cfg = ProxyConfig.model_validate({"enabled": True, "node_id": " node-1 "})
|
||||||
|
assert cfg.node_id == "node-1"
|
||||||
|
assert cfg.url is None
|
||||||
Reference in New Issue
Block a user