mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +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:
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;
|
||||
Reference in New Issue
Block a user