feat(tunnel): add tunnel IP family controls

This commit is contained in:
RWDai
2026-05-20 13:42:55 +08:00
parent ed75ae6d56
commit 7284165f39
4 changed files with 273 additions and 11 deletions

View File

@@ -571,6 +571,24 @@ pub struct Config {
)] )]
pub tunnel_connect_timeout_ms: u64, pub tunnel_connect_timeout_ms: u64,
/// Force direct WebSocket tunnel TCP connects, or Aether outbound proxy endpoint connects, to IPv4 addresses only.
#[arg(
long,
env = "AETHER_TUNNEL_IPV4_ONLY",
default_value_t = false,
conflicts_with = "tunnel_ipv6_only"
)]
pub tunnel_ipv4_only: bool,
/// Force direct WebSocket tunnel TCP connects, or Aether outbound proxy endpoint connects, to IPv6 addresses only.
#[arg(
long,
env = "AETHER_TUNNEL_IPV6_ONLY",
default_value_t = false,
conflicts_with = "tunnel_ipv4_only"
)]
pub tunnel_ipv6_only: bool,
/// WebSocket tunnel TCP keepalive in seconds (0 disables) /// WebSocket tunnel TCP keepalive in seconds (0 disables)
#[arg(long, env = "AETHER_TUNNEL_TCP_KEEPALIVE", default_value_t = 30)] #[arg(long, env = "AETHER_TUNNEL_TCP_KEEPALIVE", default_value_t = 30)]
pub tunnel_tcp_keepalive_secs: u64, pub tunnel_tcp_keepalive_secs: u64,
@@ -655,6 +673,9 @@ impl Config {
if tunnel_connect_timeout.is_zero() { if tunnel_connect_timeout.is_zero() {
anyhow::bail!("effective tunnel connect timeout must be > 0"); anyhow::bail!("effective tunnel connect timeout must be > 0");
} }
if self.tunnel_ipv4_only && self.tunnel_ipv6_only {
anyhow::bail!("tunnel_ipv4_only and tunnel_ipv6_only cannot both be enabled");
}
let tunnel_ping_interval = self.tunnel_ping_interval()?; let tunnel_ping_interval = self.tunnel_ping_interval()?;
if tunnel_ping_interval.is_zero() { if tunnel_ping_interval.is_zero() {
anyhow::bail!("effective tunnel ping interval must be > 0"); anyhow::bail!("effective tunnel ping interval must be > 0");
@@ -763,6 +784,16 @@ impl Config {
Ok(Duration::from_millis(self.tunnel_connect_timeout_ms)) Ok(Duration::from_millis(self.tunnel_connect_timeout_ms))
} }
pub fn tunnel_ip_family(&self) -> crate::egress_proxy::IpFamily {
if self.tunnel_ipv4_only {
crate::egress_proxy::IpFamily::Ipv4Only
} else if self.tunnel_ipv6_only {
crate::egress_proxy::IpFamily::Ipv6Only
} else {
crate::egress_proxy::IpFamily::Any
}
}
pub fn tunnel_stale_timeout(&self) -> anyhow::Result<Duration> { pub fn tunnel_stale_timeout(&self) -> anyhow::Result<Duration> {
Ok(Duration::from_millis(self.tunnel_stale_timeout_ms)) Ok(Duration::from_millis(self.tunnel_stale_timeout_ms))
} }
@@ -959,6 +990,10 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_connect_timeout_ms: Option<u64>, pub tunnel_connect_timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_ipv4_only: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_ipv6_only: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_tcp_keepalive_secs: Option<u64>, pub tunnel_tcp_keepalive_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_tcp_nodelay: Option<bool>, pub tunnel_tcp_nodelay: Option<bool>,
@@ -1147,6 +1182,8 @@ impl ConfigFile {
TUNNEL_CONNECT_TIMEOUT_MS_ENV, TUNNEL_CONNECT_TIMEOUT_MS_ENV,
self.tunnel_connect_timeout_ms self.tunnel_connect_timeout_ms
); );
set!("AETHER_TUNNEL_IPV4_ONLY", self.tunnel_ipv4_only);
set!("AETHER_TUNNEL_IPV6_ONLY", self.tunnel_ipv6_only);
set!( set!(
"AETHER_TUNNEL_TCP_KEEPALIVE", "AETHER_TUNNEL_TCP_KEEPALIVE",
self.tunnel_tcp_keepalive_secs self.tunnel_tcp_keepalive_secs
@@ -1338,6 +1375,20 @@ mod tests {
assert_eq!(cfg.allow_private_targets, Some(true)); assert_eq!(cfg.allow_private_targets, Some(true));
} }
#[test]
fn config_file_deserializes_tunnel_ip_family_flags() {
let cfg: ConfigFile = toml::from_str(
r#"
tunnel_ipv4_only = true
tunnel_ipv6_only = false
"#,
)
.expect("tunnel IP-family TOML");
assert_eq!(cfg.tunnel_ipv4_only, Some(true));
assert_eq!(cfg.tunnel_ipv6_only, Some(false));
}
#[test] #[test]
fn config_file_deserializes_upstream_proxy_url() { fn config_file_deserializes_upstream_proxy_url() {
let cfg: ConfigFile = toml::from_str("upstream_proxy_url = \"http://proxy.example:8080\"") let cfg: ConfigFile = toml::from_str("upstream_proxy_url = \"http://proxy.example:8080\"")
@@ -1531,6 +1582,106 @@ node_name = "tunnel-test"
assert!(config.allow_private_targets); assert!(config.allow_private_targets);
} }
#[test]
fn cli_defaults_tunnel_ip_family_to_any() {
let config = Config::parse_from([
"aether-tunnel",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"tunnel-test",
]);
assert!(!config.tunnel_ipv4_only);
assert!(!config.tunnel_ipv6_only);
assert_eq!(
config.tunnel_ip_family(),
crate::egress_proxy::IpFamily::Any
);
}
#[test]
fn cli_accepts_tunnel_ipv4_only() {
let config = Config::parse_from([
"aether-tunnel",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"tunnel-test",
"--tunnel-ipv4-only",
]);
assert!(config.tunnel_ipv4_only);
assert_eq!(
config.tunnel_ip_family(),
crate::egress_proxy::IpFamily::Ipv4Only
);
}
#[test]
fn cli_accepts_tunnel_ipv6_only() {
let config = Config::parse_from([
"aether-tunnel",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"tunnel-test",
"--tunnel-ipv6-only",
]);
assert!(config.tunnel_ipv6_only);
assert_eq!(
config.tunnel_ip_family(),
crate::egress_proxy::IpFamily::Ipv6Only
);
}
#[test]
fn cli_rejects_conflicting_tunnel_ip_family_flags() {
let error = Config::try_parse_from([
"aether-tunnel",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"tunnel-test",
"--tunnel-ipv4-only",
"--tunnel-ipv6-only",
])
.expect_err("conflicting tunnel IP-family flags should fail");
assert_eq!(error.kind(), clap::error::ErrorKind::ArgumentConflict);
}
#[test]
fn validate_rejects_conflicting_toml_tunnel_ip_family_flags() {
let config = Config {
tunnel_ipv4_only: true,
tunnel_ipv6_only: true,
..Config::parse_from([
"aether-tunnel",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"tunnel-test",
])
};
let error = config
.validate()
.expect_err("conflicting TOML-injected tunnel family flags should fail validation");
assert!(error.to_string().contains("tunnel_ipv4_only"));
}
#[test] #[test]
fn tunnel_fast_recovery_defaults_use_millisecond_values() { fn tunnel_fast_recovery_defaults_use_millisecond_values() {
let config = Config::parse_from([ let config = Config::parse_from([

View File

@@ -8,6 +8,31 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream; use tokio::net::TcpStream;
use url::Url; use url::Url;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum IpFamily {
Any,
Ipv4Only,
Ipv6Only,
}
impl IpFamily {
pub(crate) fn allows(self, addr: SocketAddr) -> bool {
match self {
Self::Any => true,
Self::Ipv4Only => addr.is_ipv4(),
Self::Ipv6Only => addr.is_ipv6(),
}
}
pub(crate) fn no_address_message(self, context: &str) -> String {
match self {
Self::Any => format!("{context} DNS returned no addresses"),
Self::Ipv4Only => format!("{context} DNS returned no IPv4 addresses"),
Self::Ipv6Only => format!("{context} DNS returned no IPv6 addresses"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum UpstreamProxyScheme { pub(crate) enum UpstreamProxyScheme {
Http, Http,
@@ -126,6 +151,7 @@ pub(crate) struct ProxyConnectOptions {
pub connect_timeout: Duration, pub connect_timeout: Duration,
pub tcp_nodelay: bool, pub tcp_nodelay: bool,
pub tcp_keepalive: Option<Duration>, pub tcp_keepalive: Option<Duration>,
pub ip_family: IpFamily,
} }
pub(crate) async fn connect_target_via_proxy( pub(crate) async fn connect_target_via_proxy(
@@ -139,6 +165,7 @@ pub(crate) async fn connect_target_via_proxy(
options.connect_timeout, options.connect_timeout,
options.tcp_nodelay, options.tcp_nodelay,
options.tcp_keepalive, options.tcp_keepalive,
options.ip_family,
) )
.await?; .await?;
@@ -159,6 +186,7 @@ pub(crate) async fn connect_proxy_tcp(
connect_timeout: Duration, connect_timeout: Duration,
tcp_nodelay: bool, tcp_nodelay: bool,
tcp_keepalive: Option<Duration>, tcp_keepalive: Option<Duration>,
ip_family: IpFamily,
) -> io::Result<TcpStream> { ) -> io::Result<TcpStream> {
let resolved = tokio::time::timeout( let resolved = tokio::time::timeout(
connect_timeout, connect_timeout,
@@ -169,7 +197,7 @@ pub(crate) async fn connect_proxy_tcp(
.map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?; .map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?;
let mut last_error = None; let mut last_error = None;
for addr in resolved { for addr in resolved.filter(|addr| ip_family.allows(*addr)) {
match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await { match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
Ok(Ok(stream)) => { Ok(Ok(stream)) => {
configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?; configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?;
@@ -185,7 +213,7 @@ pub(crate) async fn connect_proxy_tcp(
} }
} }
Err(last_error.unwrap_or_else(|| io::Error::other("proxy DNS returned no addresses"))) Err(last_error.unwrap_or_else(|| io::Error::other(ip_family.no_address_message("proxy"))))
} }
fn configure_tcp_stream( fn configure_tcp_stream(

View File

@@ -1,5 +1,7 @@
//! WebSocket tunnel client: connect, authenticate, and run the tunnel. //! WebSocket tunnel client: connect, authenticate, and run the tunnel.
use std::io;
use std::net::SocketAddr;
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
@@ -10,7 +12,9 @@ use tokio_tungstenite::tungstenite::http;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
use crate::egress_proxy::{connect_target_via_proxy, ProxyConnectOptions, UpstreamProxyConfig}; use crate::egress_proxy::{
connect_target_via_proxy, IpFamily, ProxyConnectOptions, UpstreamProxyConfig,
};
use crate::state::{AppState, ServerContext}; use crate::state::{AppState, ServerContext};
use aether_contracts::tunnel::{CURRENT_TUNNEL_PROTOCOL_VERSION, TUNNEL_PROTOCOL_VERSION_HEADER}; use aether_contracts::tunnel::{CURRENT_TUNNEL_PROTOCOL_VERSION, TUNNEL_PROTOCOL_VERSION_HEADER};
@@ -325,6 +329,7 @@ async fn connect_tunnel_tcp(
tcp_nodelay: state.config.tunnel_tcp_nodelay, tcp_nodelay: state.config.tunnel_tcp_nodelay,
tcp_keepalive: (state.config.tunnel_tcp_keepalive_secs > 0) tcp_keepalive: (state.config.tunnel_tcp_keepalive_secs > 0)
.then(|| Duration::from_secs(state.config.tunnel_tcp_keepalive_secs)), .then(|| Duration::from_secs(state.config.tunnel_tcp_keepalive_secs)),
ip_family: state.config.tunnel_ip_family(),
}, },
), ),
) )
@@ -338,15 +343,54 @@ async fn connect_tunnel_tcp(
.map_err(anyhow::Error::from); .map_err(anyhow::Error::from);
} }
tokio::time::timeout(connect_timeout, TcpStream::connect((host, port))) let ip_family = state.config.tunnel_ip_family();
tokio::time::timeout(
connect_timeout,
connect_direct_tunnel_tcp(host, port, ip_family),
)
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel TCP connect timeout ({}ms)",
connect_timeout.as_millis()
)
})?
.map_err(anyhow::Error::from)
}
async fn connect_direct_tunnel_tcp(
host: &str,
port: u16,
ip_family: IpFamily,
) -> io::Result<TcpStream> {
let resolved = tokio::net::lookup_host((host, port))
.await .await
.map_err(|_| { .map_err(|err| io::Error::other(format!("tunnel DNS failed: {err}")))?;
anyhow::anyhow!( let addrs = filter_socket_addrs(resolved, ip_family);
"tunnel TCP connect timeout ({}ms)",
connect_timeout.as_millis() if addrs.is_empty() {
) return Err(io::Error::other(ip_family.no_address_message("tunnel")));
})? }
.map_err(anyhow::Error::from)
let mut last_error = None;
for addr in addrs {
match TcpStream::connect(addr).await {
Ok(stream) => return Ok(stream),
Err(error) => last_error = Some(error),
}
}
Err(last_error.unwrap_or_else(|| io::Error::other("tunnel DNS returned no addresses")))
}
fn filter_socket_addrs(
addrs: impl IntoIterator<Item = SocketAddr>,
ip_family: IpFamily,
) -> Vec<SocketAddr> {
addrs
.into_iter()
.filter(|addr| ip_family.allows(*addr))
.collect()
} }
/// Configure TCP keepalive and NODELAY on an established socket. /// Configure TCP keepalive and NODELAY on an established socket.
@@ -392,3 +436,40 @@ fn build_tunnel_url(server: &ServerContext) -> String {
}; };
format!("{}/api/internal/proxy-tunnel", ws_base) format!("{}/api/internal/proxy-tunnel", ws_base)
} }
#[cfg(test)]
mod tests {
use std::net::{Ipv4Addr, Ipv6Addr};
use super::*;
fn mixed_addrs() -> Vec<SocketAddr> {
vec![
SocketAddr::from((Ipv6Addr::LOCALHOST, 443)),
SocketAddr::from((Ipv4Addr::LOCALHOST, 443)),
]
}
#[test]
fn filter_socket_addrs_keeps_all_addresses_by_default() {
let addrs = filter_socket_addrs(mixed_addrs(), IpFamily::Any);
assert_eq!(addrs.len(), 2);
assert!(addrs[0].is_ipv6());
assert!(addrs[1].is_ipv4());
}
#[test]
fn filter_socket_addrs_keeps_only_ipv4_addresses() {
let addrs = filter_socket_addrs(mixed_addrs(), IpFamily::Ipv4Only);
assert_eq!(addrs, vec![SocketAddr::from((Ipv4Addr::LOCALHOST, 443))]);
}
#[test]
fn filter_socket_addrs_keeps_only_ipv6_addresses() {
let addrs = filter_socket_addrs(mixed_addrs(), IpFamily::Ipv6Only);
assert_eq!(addrs, vec![SocketAddr::from((Ipv6Addr::LOCALHOST, 443))]);
}
}

View File

@@ -281,6 +281,7 @@ impl Service<Uri> for InstrumentedConnector {
connect_timeout: self.connect_timeout, connect_timeout: self.connect_timeout,
tcp_nodelay: self.tcp_nodelay, tcp_nodelay: self.tcp_nodelay,
tcp_keepalive: self.tcp_keepalive, tcp_keepalive: self.tcp_keepalive,
ip_family: crate::egress_proxy::IpFamily::Any,
}; };
let connect_start = std::time::Instant::now(); let connect_start = std::time::Instant::now();
return Box::pin(async move { return Box::pin(async move {
@@ -347,6 +348,7 @@ async fn connect_via_proxy(
options.connect_timeout, options.connect_timeout,
options.tcp_nodelay, options.tcp_nodelay,
options.tcp_keepalive, options.tcp_keepalive,
options.ip_family,
) )
.await?; .await?;