feat: aether-proxy 远程配置下发、连通性测试与 setup TUI

- 后端新增远程配置管理 API (PUT /config) 和连通性测试 API (POST /test)
- 前端新增远程配置编辑对话框和节点连通性测试按钮
- aether-proxy 支持通过心跳接收并热加载远程配置 (端口白名单、日志级别、心跳间隔、时间戳容差)
- aether-proxy 新增 TOML 配置文件支持和交互式 setup TUI
- aether-proxy 心跳 404 时自动重注册节点
- plain proxy 响应改为流式传输,减少内存缓冲
- 新增 remote_config 和 config_version 数据库字段及迁移
This commit is contained in:
fawney19
2026-02-07 19:20:09 +08:00
parent 3b8398b2e5
commit 31bc452374
18 changed files with 2771 additions and 113 deletions

1127
aether-proxy/Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@@ -23,6 +23,9 @@ thiserror = "2"
bytes = "1"
hex = "0.4"
anyhow = "1"
toml = "0.8"
ratatui = "0.30"
crossterm = "0.28"
[profile.release]
lto = true

View File

@@ -36,10 +36,15 @@ impl std::fmt::Display for AuthError {
///
/// Expected format: `Basic base64(hmac:{timestamp}.{signature})`
/// where signature = hex(HMAC-SHA256(hmac_key, "{timestamp}\n{node_id}"))
///
/// `timestamp_tolerance` is accepted separately so the caller can supply
/// the value from [`DynamicConfig`](crate::runtime::DynamicConfig) (which
/// may be updated remotely).
pub fn validate_proxy_auth(
proxy_auth_header: Option<&str>,
config: &Config,
node_id: &str,
timestamp_tolerance: u64,
) -> Result<(), AuthError> {
let header = proxy_auth_header.ok_or(AuthError::MissingHeader)?;
@@ -83,7 +88,7 @@ pub fn validate_proxy_auth(
timestamp - now
};
if diff > config.timestamp_tolerance {
if diff > timestamp_tolerance {
return Err(AuthError::TimestampExpired);
}
@@ -146,7 +151,7 @@ mod tests {
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());
assert!(validate_proxy_auth(Some(&header), &config, "node-1", config.timestamp_tolerance).is_ok());
}
#[test]
@@ -154,7 +159,7 @@ mod tests {
let config = make_config();
let header = make_valid_auth(&config, "node-1");
assert!(matches!(
validate_proxy_auth(Some(&header), &config, "node-2"),
validate_proxy_auth(Some(&header), &config, "node-2", config.timestamp_tolerance),
Err(AuthError::SignatureMismatch)
));
}
@@ -163,7 +168,7 @@ mod tests {
fn test_missing_header() {
let config = make_config();
assert!(matches!(
validate_proxy_auth(None, &config, "node-1"),
validate_proxy_auth(None, &config, "node-1", config.timestamp_tolerance),
Err(AuthError::MissingHeader)
));
}
@@ -175,7 +180,7 @@ mod tests {
let header = format!("Basic {}", encoded);
let config = make_config();
assert!(matches!(
validate_proxy_auth(Some(&header), &config, "node-1"),
validate_proxy_auth(Some(&header), &config, "node-1", config.timestamp_tolerance),
Err(AuthError::InvalidUsername)
));
}

View File

@@ -1,4 +1,7 @@
use std::path::Path;
use clap::Parser;
use serde::{Deserialize, Serialize};
/// Aether forward proxy with HMAC authentication.
///
@@ -56,3 +59,91 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_LOG_JSON", default_value_t = false)]
pub log_json: bool,
}
// ---------------------------------------------------------------------------
// TOML config file support
// ---------------------------------------------------------------------------
/// Serializable config for TOML file persistence.
/// All fields are optional — only populated values are written.
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub management_token: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub hmac_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub listen_port: Option<u16>,
#[serde(skip_serializing_if = "Option::is_none")]
pub public_ip: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub node_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub node_region: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub heartbeat_interval: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub allowed_ports: Option<Vec<u16>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub timestamp_tolerance: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub log_level: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub log_json: Option<bool>,
}
impl ConfigFile {
/// Load from a TOML file.
pub fn load(path: &Path) -> anyhow::Result<Self> {
let content = std::fs::read_to_string(path)?;
Ok(toml::from_str(&content)?)
}
/// Save to a TOML file.
pub fn save(&self, path: &Path) -> anyhow::Result<()> {
let content = toml::to_string_pretty(self)?;
std::fs::write(path, content)?;
Ok(())
}
/// Inject values as environment variables so clap picks them up.
///
/// Only sets variables that are **not** already present in the
/// environment, preserving the precedence: CLI > env > config file.
pub fn inject_env(&self) {
macro_rules! set {
($env:expr, $val:expr) => {
if let Some(ref v) = $val {
if std::env::var($env).is_err() {
std::env::set_var($env, v.to_string());
}
}
};
}
set!("AETHER_PROXY_AETHER_URL", self.aether_url);
set!("AETHER_PROXY_MANAGEMENT_TOKEN", self.management_token);
set!("AETHER_PROXY_HMAC_KEY", self.hmac_key);
set!("AETHER_PROXY_LISTEN_PORT", self.listen_port);
set!("AETHER_PROXY_PUBLIC_IP", self.public_ip);
set!("AETHER_PROXY_NODE_NAME", self.node_name);
set!("AETHER_PROXY_NODE_REGION", self.node_region);
set!("AETHER_PROXY_HEARTBEAT_INTERVAL", self.heartbeat_interval);
set!("AETHER_PROXY_TIMESTAMP_TOLERANCE", self.timestamp_tolerance);
set!("AETHER_PROXY_LOG_LEVEL", self.log_level);
set!("AETHER_PROXY_LOG_JSON", self.log_json);
// allowed_ports needs special handling (comma-separated)
if let Some(ref ports) = self.allowed_ports {
if std::env::var("AETHER_PROXY_ALLOWED_PORTS").is_err() {
let s: String = ports
.iter()
.map(|p| p.to_string())
.collect::<Vec<_>>()
.join(",");
std::env::set_var("AETHER_PROXY_ALLOWED_PORTS", s);
}
}
}
}

View File

@@ -2,8 +2,11 @@ mod auth;
mod config;
mod proxy;
mod registration;
mod runtime;
mod setup;
use std::sync::Arc;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use clap::Parser;
use tokio::signal;
@@ -12,12 +15,46 @@ use tracing::{error, info};
use config::Config;
use registration::client::{detect_public_ip, AetherClient};
use runtime::DynamicConfig;
/// Default config file name.
const DEFAULT_CONFIG: &str = "aether-proxy.toml";
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let config = Config::parse();
let args: Vec<String> = std::env::args().collect();
// Initialize tracing
// ── Handle `setup` subcommand before clap parsing ────────────────────
if args.len() > 1 && args[1] == "setup" {
let path = args
.get(2)
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from(DEFAULT_CONFIG));
return setup::run(path);
}
// ── Load config file as env-var defaults (before clap) ───────────────
let config_file_path = std::env::var("AETHER_PROXY_CONFIG")
.unwrap_or_else(|_| DEFAULT_CONFIG.to_string());
if std::path::Path::new(&config_file_path).exists() {
if let Ok(file_cfg) = config::ConfigFile::load(std::path::Path::new(&config_file_path)) {
file_cfg.inject_env();
}
}
// ── Parse config; fall back to setup TUI if required args are missing ─
let config = match Config::try_parse() {
Ok(c) => c,
Err(e) => {
if e.kind() == clap::error::ErrorKind::MissingRequiredArgument {
eprintln!("缺少必要配置,启动交互式配置向导...\n");
return setup::run(PathBuf::from(&config_file_path));
}
e.exit();
}
};
// Initialize tracing (with hot-reload support)
init_tracing(&config);
info!(
@@ -37,10 +74,14 @@ async fn main() -> anyhow::Result<()> {
// 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");
let node_id = Arc::new(RwLock::new(node_id));
// Dynamic config (hot-reloadable via heartbeat)
let dynamic = Arc::new(RwLock::new(DynamicConfig::from_config(&config)));
// Shutdown signal channel
let (shutdown_tx, shutdown_rx) = watch::channel(false);
@@ -50,10 +91,12 @@ async fn main() -> anyhow::Result<()> {
let heartbeat_handle = {
let client = Arc::clone(&aether_client);
let node_id = Arc::clone(&node_id);
let interval = config.heartbeat_interval;
let config = Arc::clone(&config);
let dynamic = Arc::clone(&dynamic);
let public_ip = public_ip.clone();
let rx = shutdown_rx.clone();
tokio::spawn(async move {
registration::heartbeat::run(client, node_id, interval, rx).await;
registration::heartbeat::run(client, node_id, config, public_ip, dynamic, rx).await;
})
};
@@ -61,9 +104,10 @@ async fn main() -> anyhow::Result<()> {
let server_handle = {
let config = Arc::clone(&config);
let node_id = Arc::clone(&node_id);
let dynamic = Arc::clone(&dynamic);
let rx = shutdown_rx.clone();
tokio::spawn(async move {
if let Err(e) = proxy::server::run(config, node_id, rx).await {
if let Err(e) = proxy::server::run(config, node_id, dynamic, rx).await {
error!(error = %e, "proxy server error");
}
})
@@ -78,7 +122,8 @@ async fn main() -> anyhow::Result<()> {
let _ = shutdown_tx.send(true);
// Graceful unregister (best-effort)
if let Err(e) = aether_client.unregister(&node_id).await {
let current_node_id = node_id.read().unwrap().clone();
if let Err(e) = aether_client.unregister(&current_node_id).await {
error!(error = %e, "unregister failed during shutdown");
}
@@ -90,19 +135,30 @@ async fn main() -> anyhow::Result<()> {
}
fn init_tracing(config: &Config) {
use tracing_subscriber::EnvFilter;
use tracing_subscriber::prelude::*;
use tracing_subscriber::{reload, EnvFilter};
let filter = EnvFilter::try_new(&config.log_level)
.unwrap_or_else(|_| EnvFilter::new("info"));
let filter =
EnvFilter::try_new(&config.log_level).unwrap_or_else(|_| EnvFilter::new("info"));
let (filter_layer, reload_handle) = reload::Layer::new(filter);
// Register log-level hot-reloader
runtime::set_log_reloader(Box::new(move |level: &str| {
if let Ok(new_filter) = EnvFilter::try_new(level) {
let _ = reload_handle.modify(|f| *f = new_filter);
}
}));
if config.log_json {
tracing_subscriber::fmt()
.with_env_filter(filter)
.json()
tracing_subscriber::registry()
.with(filter_layer)
.with(tracing_subscriber::fmt::layer().json())
.init();
} else {
tracing_subscriber::fmt()
.with_env_filter(filter)
tracing_subscriber::registry()
.with(filter_layer)
.with(tracing_subscriber::fmt::layer())
.init();
}
}

View File

@@ -18,6 +18,7 @@ pub async fn handle_connect(
config: Arc<Config>,
node_id: &str,
allowed_ports: &HashSet<u16>,
timestamp_tolerance: u64,
) -> Response<http_body_util::Empty<bytes::Bytes>> {
// Extract Proxy-Authorization header
let proxy_auth = req
@@ -26,7 +27,7 @@ pub async fn handle_connect(
.and_then(|v| v.to_str().ok());
// HMAC authentication
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, node_id) {
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, node_id, timestamp_tolerance) {
warn!(error = %e, "CONNECT auth failed");
return proxy_auth_required(&e.to_string());
}

View File

@@ -10,15 +10,20 @@ use crate::auth;
use crate::config::Config;
use crate::proxy::target_filter;
/// Boxed body that unifies `Full` (error responses) and `Incoming` (streamed upstream).
pub type BoxBody =
http_body_util::combinators::BoxBody<bytes::Bytes, Box<dyn std::error::Error + Send + Sync>>;
/// Handle plain HTTP forward proxy requests (non-CONNECT).
///
/// Flow: validate auth -> check target filter -> forward request -> return response
/// Flow: validate auth -> check target filter -> forward request -> **stream** response
pub async fn handle_plain(
req: Request<Incoming>,
config: Arc<Config>,
node_id: &str,
allowed_ports: &HashSet<u16>,
) -> Response<Full<bytes::Bytes>> {
timestamp_tolerance: u64,
) -> Response<BoxBody> {
// Extract Proxy-Authorization header
let proxy_auth = req
.headers()
@@ -26,7 +31,7 @@ pub async fn handle_plain(
.and_then(|v| v.to_str().ok());
// HMAC authentication
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, node_id) {
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, node_id, timestamp_tolerance) {
warn!(error = %e, "HTTP proxy auth failed");
return proxy_auth_required(&e.to_string());
}
@@ -72,7 +77,7 @@ pub async fn handle_plain(
builder = builder.header(name, value);
}
// Collect the incoming body
// Collect the incoming request body (client payloads are small)
let body_bytes = match req.into_body().collect().await {
Ok(collected) => collected.to_bytes(),
Err(e) => {
@@ -111,15 +116,12 @@ pub async fn handle_plain(
match sender.send_request(outgoing).await {
Ok(resp) => {
// Stream the response body directly — no buffering
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))
let body: BoxBody = body
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })
.boxed();
Response::from_parts(parts, body)
}
Err(e) => {
warn!(error = %e, "HTTP proxy request failed");
@@ -128,35 +130,43 @@ pub async fn handle_plain(
}
}
fn proxy_auth_required(msg: &str) -> Response<Full<bytes::Bytes>> {
// ── Error response helpers ───────────────────────────────────────────────────
fn empty_box() -> BoxBody {
Full::new(bytes::Bytes::new())
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
.boxed()
}
fn proxy_auth_required(msg: &str) -> Response<BoxBody> {
Response::builder()
.status(407)
.header("Proxy-Authenticate", "HMAC-SHA256")
.header("X-Error", msg)
.body(Full::new(bytes::Bytes::new()))
.body(empty_box())
.unwrap()
}
fn forbidden(msg: &str) -> Response<Full<bytes::Bytes>> {
fn forbidden(msg: &str) -> Response<BoxBody> {
Response::builder()
.status(403)
.header("X-Error", msg)
.body(Full::new(bytes::Bytes::new()))
.body(empty_box())
.unwrap()
}
fn bad_request(msg: &str) -> Response<Full<bytes::Bytes>> {
fn bad_request(msg: &str) -> Response<BoxBody> {
Response::builder()
.status(400)
.header("X-Error", msg)
.body(Full::new(bytes::Bytes::new()))
.body(empty_box())
.unwrap()
}
fn bad_gateway(msg: &str) -> Response<Full<bytes::Bytes>> {
fn bad_gateway(msg: &str) -> Response<BoxBody> {
Response::builder()
.status(502)
.header("X-Error", msg)
.body(Full::new(bytes::Bytes::new()))
.body(empty_box())
.unwrap()
}

View File

@@ -1,6 +1,5 @@
use std::collections::HashSet;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::{Arc, RwLock};
use http_body_util::BodyExt;
use hyper::body::Incoming;
@@ -14,6 +13,7 @@ use tracing::{debug, info, warn};
use crate::config::Config;
use crate::proxy::{connect, plain};
use crate::runtime::SharedDynamicConfig;
/// Start the proxy server.
///
@@ -22,15 +22,14 @@ use crate::proxy::{connect, plain};
/// - Other HTTP requests -> plain forward proxy handler
pub async fn run(
config: Arc<Config>,
node_id: Arc<String>,
node_id: Arc<RwLock<String>>,
dynamic: SharedDynamicConfig,
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() => {
@@ -46,28 +45,33 @@ pub async fn run(
let config = Arc::clone(&config);
let node_id = Arc::clone(&node_id);
let allowed_ports = Arc::clone(&allowed_ports);
let dynamic = Arc::clone(&dynamic);
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);
let dynamic = Arc::clone(&dynamic);
async move {
type BoxBody = http_body_util::combinators::BoxBody<bytes::Bytes, Box<dyn std::error::Error + Send + Sync>>;
// Snapshot current dynamic values (may be updated by remote config)
let current_node_id = node_id.read().unwrap().clone();
let (allowed_ports, timestamp_tolerance) = {
let d = dynamic.read().unwrap();
(d.allowed_ports.clone(), d.timestamp_tolerance)
};
if req.method() == Method::CONNECT {
let resp = connect::handle_connect(
req,
config,
&node_id,
&current_node_id,
&allowed_ports,
timestamp_tolerance,
)
.await;
let resp = resp.map(|_| -> BoxBody {
@@ -80,14 +84,12 @@ pub async fn run(
let resp = plain::handle_plain(
req,
config,
&node_id,
&current_node_id,
&allowed_ports,
timestamp_tolerance,
)
.await;
let resp = resp.map(|body| -> BoxBody {
body.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
.boxed()
});
// plain::handle_plain already returns BoxBody (streaming)
Ok(resp)
}
}

View File

@@ -68,7 +68,6 @@ pub enum FilterError {
PrivateIp(IpAddr),
PortNotAllowed(u16),
DnsResolutionFailed(String),
AllAddressesPrivate(String),
}
impl std::fmt::Display for FilterError {
@@ -77,9 +76,6 @@ impl std::fmt::Display for FilterError {
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)
}
}
}
}

View File

@@ -1,9 +1,28 @@
use reqwest::Client;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use tracing::{debug, error, info, warn};
use crate::config::Config;
/// Heartbeat-specific error that distinguishes "node not found" (needs
/// re-registration) from transient / other failures.
#[derive(Debug)]
pub enum HeartbeatError {
/// HTTP 404 the node_id is no longer known to Aether.
NodeNotFound(String),
/// Any other failure (network, 5xx, etc.).
Other(anyhow::Error),
}
impl std::fmt::Display for HeartbeatError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NodeNotFound(msg) => write!(f, "node not found: {}", msg),
Self::Other(e) => write!(f, "{}", e),
}
}
}
#[derive(Debug, Serialize)]
struct RegisterRequest {
name: String,
@@ -30,6 +49,37 @@ struct HeartbeatRequest {
avg_latency_ms: Option<f64>,
}
/// Remote configuration pushed by the Aether management backend.
#[derive(Debug, Clone, Deserialize)]
pub struct RemoteConfig {
pub allowed_ports: Option<Vec<u16>>,
pub log_level: Option<String>,
pub heartbeat_interval: Option<u64>,
pub timestamp_tolerance: Option<u64>,
}
/// Parsed heartbeat response from Aether.
#[derive(Debug, Deserialize)]
struct HeartbeatResponseBody {
#[serde(default)]
node: Option<HeartbeatNodeInfo>,
}
#[derive(Debug, Deserialize)]
struct HeartbeatNodeInfo {
#[serde(default)]
remote_config: Option<RemoteConfig>,
#[serde(default)]
config_version: Option<u64>,
}
/// Heartbeat result returned to the caller.
#[derive(Debug)]
pub struct HeartbeatResult {
pub remote_config: Option<RemoteConfig>,
pub config_version: u64,
}
#[derive(Debug, Serialize)]
struct UnregisterRequest {
node_id: String,
@@ -101,13 +151,17 @@ impl AetherClient {
}
/// Send heartbeat to Aether.
///
/// On success, returns any remote config included in the response.
/// Returns [`HeartbeatError::NodeNotFound`] on HTTP 404 so the caller
/// can trigger re-registration.
pub async fn heartbeat(
&self,
node_id: &str,
active_connections: Option<i64>,
total_requests: Option<i64>,
avg_latency_ms: Option<f64>,
) -> anyhow::Result<()> {
) -> Result<HeartbeatResult, HeartbeatError> {
let url = format!("{}/api/admin/proxy-nodes/heartbeat", self.base_url);
let body = HeartbeatRequest {
node_id: node_id.to_string(),
@@ -124,17 +178,43 @@ impl AetherClient {
.header("Authorization", format!("Bearer {}", self.token))
.json(&body)
.send()
.await?;
.await
.map_err(|e| HeartbeatError::Other(e.into()))?;
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);
if status == StatusCode::NOT_FOUND {
return Err(HeartbeatError::NodeNotFound(text));
}
return Err(HeartbeatError::Other(anyhow::anyhow!(
"heartbeat failed (HTTP {}): {}",
status,
text
)));
}
debug!(node_id = %node_id, "heartbeat ok");
Ok(())
// Parse remote config from response (best-effort)
let result = match resp.json::<HeartbeatResponseBody>().await {
Ok(body) => {
let (remote_config, config_version) = match body.node {
Some(node) => (node.remote_config, node.config_version.unwrap_or(0)),
None => (None, 0),
};
HeartbeatResult {
remote_config,
config_version,
}
}
Err(_) => HeartbeatResult {
remote_config: None,
config_version: 0,
},
};
debug!(node_id = %node_id, config_version = result.config_version, "heartbeat ok");
Ok(result)
}
/// Unregister this node from Aether (graceful shutdown).

View File

@@ -1,46 +1,99 @@
use std::sync::Arc;
use std::sync::{Arc, RwLock};
use tokio::sync::watch;
use tracing::{debug, warn};
use tracing::{debug, error, info, warn};
use crate::registration::client::AetherClient;
use crate::config::Config;
use crate::registration::client::{AetherClient, HeartbeatError};
use crate::runtime::{self, SharedDynamicConfig};
/// Run periodic heartbeat task until shutdown signal.
///
/// When Aether responds with 404 (node not found), this task automatically
/// re-registers the node and updates the shared `node_id` so the proxy
/// server and future heartbeats use the new identity.
///
/// When the heartbeat response includes a `remote_config`, it is applied
/// to the [`DynamicConfig`](crate::runtime::DynamicConfig) so the proxy
/// picks up changes without a restart.
pub async fn run(
client: Arc<AetherClient>,
node_id: Arc<String>,
interval_secs: u64,
node_id: Arc<RwLock<String>>,
config: Arc<Config>,
public_ip: String,
dynamic: SharedDynamicConfig,
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;
// Skip the first tick (registration already acts as initial heartbeat)
let initial_interval = dynamic.read().unwrap().heartbeat_interval;
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(initial_interval)) => {}
_ = shutdown_rx.changed() => {
debug!("heartbeat task stopping (during initial wait)");
return;
}
}
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"
);
}
let current_node_id = node_id.read().unwrap().clone();
match client.heartbeat(&current_node_id, None, None, None).await {
Ok(result) => {
if consecutive_failures > 0 {
debug!(
previous_failures = consecutive_failures,
"heartbeat recovered"
);
}
consecutive_failures = 0;
// Apply remote config if present and version changed
if let Some(ref remote) = result.remote_config {
runtime::apply_remote_config(&dynamic, remote, result.config_version);
}
}
Err(HeartbeatError::NodeNotFound(_)) => {
warn!(
old_node_id = %current_node_id,
"node not found, re-registering"
);
match client.register(&config, &public_ip).await {
Ok(new_id) => {
info!(
old_node_id = %current_node_id,
new_node_id = %new_id,
"re-registered successfully"
);
*node_id.write().unwrap() = new_id;
consecutive_failures = 0;
}
Err(e) => {
consecutive_failures += 1;
warn!(
error!(
error = %e,
consecutive_failures,
"heartbeat failed"
"re-registration failed"
);
}
}
}
Err(HeartbeatError::Other(e)) => {
consecutive_failures += 1;
warn!(
error = %e,
consecutive_failures,
"heartbeat failed"
);
}
}
// Read interval from dynamic config (may have been updated remotely)
let interval_secs = dynamic.read().unwrap().heartbeat_interval;
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(interval_secs)) => {}
_ = shutdown_rx.changed() => {
debug!("heartbeat task stopping");
break;

112
aether-proxy/src/runtime.rs Normal file
View File

@@ -0,0 +1,112 @@
//! Runtime-mutable configuration that can be updated remotely via heartbeat.
//!
//! Fields in [`DynamicConfig`] are initially populated from the static
//! [`Config`](crate::config::Config) and may be overridden by the Aether
//! management backend through the heartbeat response.
use std::collections::HashSet;
use std::sync::{Arc, OnceLock, RwLock};
use tracing::info;
use crate::config::Config;
/// Configuration that can be changed at runtime without restart.
#[derive(Debug)]
pub struct DynamicConfig {
pub allowed_ports: HashSet<u16>,
pub timestamp_tolerance: u64,
pub log_level: String,
pub heartbeat_interval: u64,
/// Monotonically increasing version from the backend.
/// `0` means no remote config has ever been applied.
pub config_version: u64,
}
impl DynamicConfig {
/// Initialize from static config (startup defaults).
pub fn from_config(config: &Config) -> Self {
Self {
allowed_ports: config.allowed_ports.iter().copied().collect(),
timestamp_tolerance: config.timestamp_tolerance,
log_level: config.log_level.clone(),
heartbeat_interval: config.heartbeat_interval,
config_version: 0,
}
}
}
/// Shared dynamic config handle.
pub type SharedDynamicConfig = Arc<RwLock<DynamicConfig>>;
// ── Log-level hot-reload ─────────────────────────────────────────────────────
/// Global log-level reloader function, set during tracing init.
static LOG_RELOADER: OnceLock<Box<dyn Fn(&str) + Send + Sync>> = OnceLock::new();
/// Register the log-level reload function (called once from `init_tracing`).
pub fn set_log_reloader(f: Box<dyn Fn(&str) + Send + Sync>) {
let _ = LOG_RELOADER.set(f);
}
/// Apply a remote config update to the dynamic config.
///
/// Returns `true` if the config was actually changed.
pub fn apply_remote_config(
dynamic: &SharedDynamicConfig,
remote: &super::registration::client::RemoteConfig,
version: u64,
) -> bool {
let mut cfg = dynamic.write().unwrap();
if version <= cfg.config_version {
return false;
}
let mut changed = Vec::new();
if let Some(ref ports) = remote.allowed_ports {
let new_set: HashSet<u16> = ports.iter().copied().collect();
if new_set != cfg.allowed_ports {
changed.push(format!("allowed_ports → {:?}", ports));
cfg.allowed_ports = new_set;
}
}
if let Some(tol) = remote.timestamp_tolerance {
if tol != cfg.timestamp_tolerance {
changed.push(format!("timestamp_tolerance → {}", tol));
cfg.timestamp_tolerance = tol;
}
}
if let Some(interval) = remote.heartbeat_interval {
if interval != cfg.heartbeat_interval {
changed.push(format!("heartbeat_interval → {}s", interval));
cfg.heartbeat_interval = interval;
}
}
if let Some(ref level) = remote.log_level {
if *level != cfg.log_level {
changed.push(format!("log_level → {}", level));
cfg.log_level = level.clone();
// Hot-reload tracing filter
if let Some(reloader) = LOG_RELOADER.get() {
reloader(level);
}
}
}
cfg.config_version = version;
if !changed.is_empty() {
info!(
version,
changes = %changed.join(", "),
"remote config applied"
);
}
!changed.is_empty()
}

649
aether-proxy/src/setup.rs Normal file
View File

@@ -0,0 +1,649 @@
//! Interactive TUI for configuring aether-proxy.
//!
//! Launched via `aether-proxy setup [path]`. Presents a full-screen form
//! backed by ratatui where the user can navigate fields, edit values, and
//! save to a TOML config file.
use std::io;
use std::path::PathBuf;
use std::time::{Duration, Instant};
use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers};
use crossterm::execute;
use crossterm::terminal::{self, EnterAlternateScreen, LeaveAlternateScreen};
use ratatui::backend::CrosstermBackend;
use ratatui::layout::{Constraint, Layout, Rect};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span};
use ratatui::widgets::{Block, Borders, Paragraph};
use ratatui::Frame;
use ratatui::Terminal;
use crate::config::ConfigFile;
/// Column width reserved for the field label (chars).
const LABEL_WIDTH: usize = 22;
// ── Field types ──────────────────────────────────────────────────────────────
#[derive(Clone, Copy, PartialEq)]
enum FieldKind {
Text,
Secret,
Number,
Bool,
PortList,
LogLevel,
}
struct Field {
label: &'static str,
key: &'static str,
value: String,
kind: FieldKind,
required: bool,
help: &'static str,
}
// ── App state ────────────────────────────────────────────────────────────────
#[derive(PartialEq)]
enum Mode {
Normal,
Editing,
}
struct App {
fields: Vec<Field>,
selected: usize,
mode: Mode,
edit_buffer: String,
edit_cursor: usize, // char index
config_path: PathBuf,
modified: bool,
message: Option<(String, Instant, bool)>, // (text, when, is_error)
scroll_offset: usize,
saved_once: bool,
}
impl App {
fn new(config_path: PathBuf) -> Self {
Self {
fields: vec![
Field {
label: "Aether URL",
key: "aether_url",
value: String::new(),
kind: FieldKind::Text,
required: true,
help: "Aether 服务器 URL (如 https://aether.example.com)",
},
Field {
label: "Management Token",
key: "management_token",
value: String::new(),
kind: FieldKind::Secret,
required: true,
help: "Aether 管理 API Token (ae_xxx)",
},
Field {
label: "HMAC Key",
key: "hmac_key",
value: String::new(),
kind: FieldKind::Secret,
required: true,
help: "HMAC-SHA256 签名密钥,用于代理请求认证",
},
Field {
label: "Listen Port",
key: "listen_port",
value: "18080".into(),
kind: FieldKind::Number,
required: true,
help: "代理服务监听端口",
},
Field {
label: "Public IP",
key: "public_ip",
value: String::new(),
kind: FieldKind::Text,
required: false,
help: "节点公网 IP (留空则自动检测)",
},
Field {
label: "Node Name",
key: "node_name",
value: "proxy-01".into(),
kind: FieldKind::Text,
required: true,
help: "节点名称,用于在 Aether 后台识别",
},
Field {
label: "Node Region",
key: "node_region",
value: String::new(),
kind: FieldKind::Text,
required: false,
help: "节点区域标识 (如 ap-northeast-1)",
},
Field {
label: "Heartbeat Interval",
key: "heartbeat_interval",
value: "30".into(),
kind: FieldKind::Number,
required: true,
help: "心跳上报间隔 (秒)",
},
Field {
label: "Allowed Ports",
key: "allowed_ports",
value: "80, 443, 8080, 8443".into(),
kind: FieldKind::PortList,
required: true,
help: "允许代理的目标端口,逗号分隔",
},
Field {
label: "Timestamp Tolerance",
key: "timestamp_tolerance",
value: "300".into(),
kind: FieldKind::Number,
required: true,
help: "HMAC 时间戳容差窗口 (秒)",
},
Field {
label: "Log Level",
key: "log_level",
value: "info".into(),
kind: FieldKind::LogLevel,
required: true,
help: "日志级别 — Enter 切换: trace / debug / info / warn / error",
},
Field {
label: "Log JSON",
key: "log_json",
value: "false".into(),
kind: FieldKind::Bool,
required: true,
help: "是否以 JSON 格式输出日志 — Enter 切换",
},
],
selected: 0,
mode: Mode::Normal,
edit_buffer: String::new(),
edit_cursor: 0,
config_path,
modified: false,
message: None,
scroll_offset: 0,
saved_once: false,
}
}
// ── Config ↔ fields ──────────────────────────────────────────────────
fn load_from_file(&mut self) {
if let Ok(cfg) = ConfigFile::load(&self.config_path) {
self.apply_config(&cfg);
}
}
fn apply_config(&mut self, cfg: &ConfigFile) {
for field in &mut self.fields {
let val: Option<String> = match field.key {
"aether_url" => cfg.aether_url.clone(),
"management_token" => cfg.management_token.clone(),
"hmac_key" => cfg.hmac_key.clone(),
"listen_port" => cfg.listen_port.map(|v| v.to_string()),
"public_ip" => cfg.public_ip.clone(),
"node_name" => cfg.node_name.clone(),
"node_region" => cfg.node_region.clone(),
"heartbeat_interval" => cfg.heartbeat_interval.map(|v| v.to_string()),
"allowed_ports" => cfg.allowed_ports.as_ref().map(|p| {
p.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(", ")
}),
"timestamp_tolerance" => cfg.timestamp_tolerance.map(|v| v.to_string()),
"log_level" => cfg.log_level.clone(),
"log_json" => cfg.log_json.map(|v| v.to_string()),
_ => None,
};
if let Some(v) = val {
field.value = v;
}
}
}
fn to_config(&self) -> ConfigFile {
let get = |key: &str| -> Option<String> {
self.fields
.iter()
.find(|f| f.key == key)
.map(|f| f.value.clone())
.filter(|v| !v.is_empty())
};
ConfigFile {
aether_url: get("aether_url"),
management_token: get("management_token"),
hmac_key: get("hmac_key"),
listen_port: get("listen_port").and_then(|v| v.parse().ok()),
public_ip: get("public_ip"),
node_name: get("node_name"),
node_region: get("node_region"),
heartbeat_interval: get("heartbeat_interval").and_then(|v| v.parse().ok()),
allowed_ports: get("allowed_ports").map(|v| {
v.split(',')
.filter_map(|s| s.trim().parse().ok())
.collect()
}),
timestamp_tolerance: get("timestamp_tolerance").and_then(|v| v.parse().ok()),
log_level: get("log_level"),
log_json: get("log_json").and_then(|v| v.parse().ok()),
}
}
fn save(&mut self) -> anyhow::Result<()> {
let cfg = self.to_config();
cfg.save(&self.config_path)?;
self.modified = false;
self.saved_once = true;
self.message = Some((
format!("✓ 已保存到 {}", self.config_path.display()),
Instant::now(),
false,
));
Ok(())
}
// ── Scrolling ────────────────────────────────────────────────────────
fn ensure_visible(&mut self, visible_rows: usize) {
if visible_rows == 0 {
return;
}
if self.selected < self.scroll_offset {
self.scroll_offset = self.selected;
} else if self.selected >= self.scroll_offset + visible_rows {
self.scroll_offset = self.selected - visible_rows + 1;
}
}
// ── Key handling ─────────────────────────────────────────────────────
/// Returns `true` when the app should exit.
fn handle_key(&mut self, key: KeyEvent) -> bool {
// Expire old messages
if let Some((_, when, _)) = &self.message {
if when.elapsed() > Duration::from_secs(4) {
self.message = None;
}
}
match self.mode {
Mode::Normal => self.handle_normal(key),
Mode::Editing => {
self.handle_edit(key);
false
}
}
}
fn handle_normal(&mut self, key: KeyEvent) -> bool {
match key.code {
KeyCode::Char('q') | KeyCode::Esc => return true,
KeyCode::Char('s') if key.modifiers.contains(KeyModifiers::CONTROL) => {
if let Err(e) = self.save() {
self.message = Some((format!("{}", e), Instant::now(), true));
}
}
KeyCode::Up | KeyCode::Char('k') => {
self.selected = self.selected.saturating_sub(1);
}
KeyCode::Down | KeyCode::Char('j') => {
if self.selected + 1 < self.fields.len() {
self.selected += 1;
}
}
KeyCode::Home => self.selected = 0,
KeyCode::End => self.selected = self.fields.len() - 1,
KeyCode::Enter | KeyCode::Char(' ') => {
let field = &self.fields[self.selected];
match field.kind {
FieldKind::Bool => {
let toggled = if field.value == "true" { "false" } else { "true" };
self.fields[self.selected].value = toggled.into();
self.modified = true;
}
FieldKind::LogLevel => {
const LEVELS: &[&str] =
&["trace", "debug", "info", "warn", "error"];
let idx = LEVELS.iter().position(|l| *l == field.value).unwrap_or(2);
self.fields[self.selected].value =
LEVELS[(idx + 1) % LEVELS.len()].into();
self.modified = true;
}
_ => {
self.edit_buffer = field.value.clone();
self.edit_cursor = self.edit_buffer.chars().count();
self.mode = Mode::Editing;
}
}
}
KeyCode::Tab => {
// Quick save shortcut
if let Err(e) = self.save() {
self.message = Some((format!("{}", e), Instant::now(), true));
}
}
_ => {}
}
false
}
fn handle_edit(&mut self, key: KeyEvent) {
match key.code {
KeyCode::Esc => {
// Cancel — discard changes to this field
self.mode = Mode::Normal;
}
KeyCode::Enter => {
if self.validate_edit() {
self.fields[self.selected].value = self.edit_buffer.clone();
self.modified = true;
self.mode = Mode::Normal;
} else {
self.message =
Some(("✗ 格式无效".into(), Instant::now(), true));
}
}
KeyCode::Backspace => {
if self.edit_cursor > 0 {
self.edit_cursor -= 1;
let byte = self.char_byte_pos(self.edit_cursor);
self.edit_buffer.remove(byte);
}
}
KeyCode::Delete => {
if self.edit_cursor < self.edit_buffer.chars().count() {
let byte = self.char_byte_pos(self.edit_cursor);
self.edit_buffer.remove(byte);
}
}
KeyCode::Left => {
self.edit_cursor = self.edit_cursor.saturating_sub(1);
}
KeyCode::Right => {
let len = self.edit_buffer.chars().count();
if self.edit_cursor < len {
self.edit_cursor += 1;
}
}
KeyCode::Home => self.edit_cursor = 0,
KeyCode::End => self.edit_cursor = self.edit_buffer.chars().count(),
KeyCode::Char(c) => {
let byte = self.char_byte_pos(self.edit_cursor);
self.edit_buffer.insert(byte, c);
self.edit_cursor += 1;
}
_ => {}
}
}
fn validate_edit(&self) -> bool {
let kind = self.fields[self.selected].kind;
let buf = &self.edit_buffer;
match kind {
FieldKind::Number => buf.is_empty() || buf.parse::<u64>().is_ok(),
FieldKind::PortList => {
buf.is_empty()
|| buf
.split(',')
.all(|s| s.trim().is_empty() || s.trim().parse::<u16>().is_ok())
}
_ => true,
}
}
/// Byte offset of the char at `char_idx`.
fn char_byte_pos(&self, char_idx: usize) -> usize {
self.edit_buffer
.char_indices()
.nth(char_idx)
.map(|(i, _)| i)
.unwrap_or(self.edit_buffer.len())
}
}
// ── Rendering ────────────────────────────────────────────────────────────────
fn ui(f: &mut Frame, app: &mut App) {
let area = f.area();
// Outer block
let title = if app.modified {
" Aether Proxy Setup [*] "
} else {
" Aether Proxy Setup "
};
let outer = Block::default()
.borders(Borders::ALL)
.title(title)
.title_alignment(ratatui::layout::Alignment::Center)
.border_style(Style::default().fg(Color::Cyan));
let inner = outer.inner(area);
f.render_widget(outer, area);
// Split: fields | footer
let chunks = Layout::vertical([Constraint::Min(1), Constraint::Length(4)]).split(inner);
let fields_area = chunks[0];
let footer_area = chunks[1];
render_fields(f, app, fields_area);
render_footer(f, app, footer_area);
}
fn render_fields(f: &mut Frame, app: &mut App, area: Rect) {
let visible = area.height as usize;
app.ensure_visible(visible);
let mut lines: Vec<Line> = Vec::new();
for (i, field) in app.fields.iter().enumerate() {
if i < app.scroll_offset || i >= app.scroll_offset + visible {
continue;
}
let selected = i == app.selected;
let indicator = if selected { "" } else { " " };
let label_style = if selected {
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD)
} else {
Style::default().fg(Color::DarkGray)
};
let padded_label = format!("{:<width$}", field.label, width = LABEL_WIDTH);
// Value display
let (value_text, value_style) = if app.mode == Mode::Editing && selected {
(
app.edit_buffer.clone(),
Style::default().fg(Color::Yellow),
)
} else {
field_display(field)
};
lines.push(Line::from(vec![
Span::styled(indicator, label_style),
Span::styled(padded_label, label_style),
Span::raw(" "),
Span::styled(value_text, value_style),
]));
}
let paragraph = Paragraph::new(lines);
f.render_widget(paragraph, area);
// Cursor position while editing
if app.mode == Mode::Editing {
let row_in_view = app.selected - app.scroll_offset;
// prefix: 3 (indicator) + LABEL_WIDTH + 2 (gap) = 27
let prefix: u16 = 3 + LABEL_WIDTH as u16 + 2;
let cx = area.x + prefix + app.edit_cursor as u16;
let cy = area.y + row_in_view as u16;
if cx < area.x + area.width && cy < area.y + area.height {
f.set_cursor_position((cx, cy));
}
}
}
/// Returns (display_text, style) for a field in normal mode.
fn field_display(field: &Field) -> (String, Style) {
if field.value.is_empty() {
let text = if field.required {
"(必填)".into()
} else {
"".into()
};
let color = if field.required {
Color::Red
} else {
Color::DarkGray
};
return (text, Style::default().fg(color));
}
match field.kind {
FieldKind::Secret => (
"".repeat(field.value.len().min(20)),
Style::default().fg(Color::White),
),
FieldKind::Bool => {
if field.value == "true" {
("✓ 开启".into(), Style::default().fg(Color::Green))
} else {
("✗ 关闭".into(), Style::default().fg(Color::DarkGray))
}
}
FieldKind::LogLevel => {
let color = match field.value.as_str() {
"trace" => Color::Magenta,
"debug" => Color::Blue,
"info" => Color::Green,
"warn" => Color::Yellow,
"error" => Color::Red,
_ => Color::White,
};
(field.value.clone(), Style::default().fg(color))
}
_ => (field.value.clone(), Style::default().fg(Color::White)),
}
}
fn render_footer(f: &mut Frame, app: &App, area: Rect) {
let help = app.fields[app.selected].help;
let keybindings = if app.mode == Mode::Editing {
"Enter 确认 Esc 取消"
} else {
"↑↓ 选择 Enter 编辑 ^S 保存 q 退出"
};
let mut status_spans: Vec<Span> = vec![Span::styled(
keybindings,
Style::default().fg(Color::DarkGray),
)];
if let Some((msg, _, is_err)) = &app.message {
let color = if *is_err { Color::Red } else { Color::Green };
status_spans.push(Span::raw(" "));
status_spans.push(Span::styled(msg.clone(), Style::default().fg(color)));
}
let footer_text = vec![
Line::raw(""),
Line::from(Span::styled(
format!(" {}", help),
Style::default().fg(Color::DarkGray),
)),
Line::from(
status_spans
.into_iter()
.map(|mut s| {
// add left padding to first span
if s.content.as_ref() == keybindings {
s.content = format!(" {}", s.content).into();
}
s
})
.collect::<Vec<_>>(),
),
];
let footer = Paragraph::new(footer_text).block(
Block::default()
.borders(Borders::TOP)
.border_style(Style::default().fg(Color::DarkGray)),
);
f.render_widget(footer, area);
}
// ── Entry point ──────────────────────────────────────────────────────────────
pub fn run(config_path: PathBuf) -> anyhow::Result<()> {
// Setup terminal
terminal::enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(stdout, EnterAlternateScreen)?;
let backend = CrosstermBackend::new(stdout);
let mut terminal = Terminal::new(backend)?;
let mut app = App::new(config_path.clone());
app.load_from_file();
let result = event_loop(&mut terminal, &mut app);
// Restore terminal
terminal::disable_raw_mode()?;
execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
terminal.show_cursor()?;
result?;
// Post-TUI message
if app.saved_once {
eprintln!();
eprintln!(" 配置已保存到 {}", config_path.display());
eprintln!();
eprintln!(" 启动方式:");
eprintln!(" aether-proxy (自动读取 {})", config_path.display());
eprintln!();
}
Ok(())
}
fn event_loop(
terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
app: &mut App,
) -> anyhow::Result<()> {
loop {
terminal.draw(|f| ui(f, app))?;
if event::poll(Duration::from_millis(200))? {
if let Event::Key(key) = event::read()? {
// Only handle Press events (ignore Release on Windows)
if key.kind == KeyEventKind::Press && app.handle_key(key) {
break;
}
}
}
}
Ok(())
}