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:
fawney19
2026-02-07 12:24:42 +08:00
parent 62f852b851
commit 1180634269
40 changed files with 4761 additions and 11 deletions

31
aether-proxy/.env.example Normal file
View 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

File diff suppressed because it is too large Load Diff

30
aether-proxy/Cargo.toml Normal file
View 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
View 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"]

View 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)
));
}
}

View File

@@ -0,0 +1,3 @@
pub mod hmac;
pub use self::hmac::validate_proxy_auth;

View 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
View 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 => {},
}
}

View 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()
}

View File

@@ -0,0 +1,4 @@
pub mod connect;
pub mod plain;
pub mod server;
pub mod target_filter;

View 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()
}

View 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(())
}

View 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());
}
}

View 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")
}

View 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;
}
}
}
}

View File

@@ -0,0 +1,2 @@
pub mod client;
pub mod heartbeat;