feat(smart-proxy): add UDP transport support with QUIC/HTTP3 routing and datagram handler relay
This commit is contained in:
@@ -32,6 +32,7 @@ arc-swap = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
rustls = { workspace = true }
|
||||
rustls-pemfile = { workspace = true }
|
||||
tokio-rustls = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
dashmap = { workspace = true }
|
||||
|
||||
@@ -47,7 +47,7 @@ pub use rustproxy_security;
|
||||
|
||||
use rustproxy_config::{RouteConfig, RustProxyOptions, TlsMode, CertificateSpec, ForwardingEngine};
|
||||
use rustproxy_routing::RouteManager;
|
||||
use rustproxy_passthrough::{TcpListenerManager, TlsCertConfig, ConnectionConfig};
|
||||
use rustproxy_passthrough::{TcpListenerManager, UdpListenerManager, TlsCertConfig, ConnectionConfig};
|
||||
use rustproxy_metrics::{MetricsCollector, Metrics, Statistics};
|
||||
use rustproxy_tls::{CertManager, CertStore, CertBundle, CertMetadata, CertSource};
|
||||
use rustproxy_nftables::{NftManager, rule_builder};
|
||||
@@ -68,6 +68,7 @@ pub struct RustProxy {
|
||||
options: RustProxyOptions,
|
||||
route_table: ArcSwap<RouteManager>,
|
||||
listener_manager: Option<TcpListenerManager>,
|
||||
udp_listener_manager: Option<UdpListenerManager>,
|
||||
metrics: Arc<MetricsCollector>,
|
||||
cert_manager: Option<Arc<tokio::sync::Mutex<CertManager>>>,
|
||||
challenge_server: Option<challenge_server::ChallengeServer>,
|
||||
@@ -114,6 +115,7 @@ impl RustProxy {
|
||||
options,
|
||||
route_table: ArcSwap::from(Arc::new(route_manager)),
|
||||
listener_manager: None,
|
||||
udp_listener_manager: None,
|
||||
metrics: Arc::new(MetricsCollector::with_retention(retention)),
|
||||
cert_manager,
|
||||
challenge_server: None,
|
||||
@@ -153,6 +155,7 @@ impl RustProxy {
|
||||
send_proxy_protocol: None,
|
||||
headers: None,
|
||||
advanced: None,
|
||||
backend_transport: None,
|
||||
priority: None,
|
||||
}
|
||||
]);
|
||||
@@ -289,17 +292,62 @@ impl RustProxy {
|
||||
}
|
||||
}
|
||||
|
||||
// Build QUIC TLS config before set_tls_configs consumes the map
|
||||
let quic_tls_config = Self::build_quic_tls_config(&tls_configs);
|
||||
|
||||
if !tls_configs.is_empty() {
|
||||
debug!("Loaded TLS certificates for {} domains", tls_configs.len());
|
||||
listener.set_tls_configs(tls_configs);
|
||||
}
|
||||
|
||||
// Bind all ports
|
||||
for port in &ports {
|
||||
// Determine which ports need TCP vs UDP based on route transport config
|
||||
let mut tcp_ports = std::collections::HashSet::new();
|
||||
let mut udp_ports = std::collections::HashSet::new();
|
||||
for route in &self.options.routes {
|
||||
if !route.is_enabled() { continue; }
|
||||
let transport = route.route_match.transport.as_ref();
|
||||
let route_ports = route.route_match.ports.to_ports();
|
||||
for port in route_ports {
|
||||
match transport {
|
||||
Some(rustproxy_config::TransportProtocol::Udp) => {
|
||||
udp_ports.insert(port);
|
||||
}
|
||||
Some(rustproxy_config::TransportProtocol::All) => {
|
||||
tcp_ports.insert(port);
|
||||
udp_ports.insert(port);
|
||||
}
|
||||
Some(rustproxy_config::TransportProtocol::Tcp) | None => {
|
||||
tcp_ports.insert(port);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Bind TCP ports
|
||||
for port in &tcp_ports {
|
||||
listener.add_port(*port).await?;
|
||||
}
|
||||
|
||||
self.listener_manager = Some(listener);
|
||||
|
||||
// Bind UDP ports (if any)
|
||||
if !udp_ports.is_empty() {
|
||||
let conn_tracker = self.listener_manager.as_ref().unwrap().conn_tracker().clone();
|
||||
let mut udp_mgr = UdpListenerManager::new(
|
||||
Arc::clone(&*self.route_table.load()),
|
||||
Arc::clone(&self.metrics),
|
||||
conn_tracker,
|
||||
self.cancel_token.clone(),
|
||||
);
|
||||
|
||||
for port in &udp_ports {
|
||||
udp_mgr.add_port_with_tls(*port, quic_tls_config.clone()).await?;
|
||||
}
|
||||
info!("UDP listeners started on {} ports: {:?}",
|
||||
udp_ports.len(), udp_mgr.listening_ports());
|
||||
self.udp_listener_manager = Some(udp_mgr);
|
||||
}
|
||||
|
||||
self.started = true;
|
||||
self.started_at = Some(Instant::now());
|
||||
|
||||
@@ -567,6 +615,13 @@ impl RustProxy {
|
||||
listener.graceful_stop().await;
|
||||
}
|
||||
self.listener_manager = None;
|
||||
|
||||
// Stop UDP listeners
|
||||
if let Some(ref mut udp_mgr) = self.udp_listener_manager {
|
||||
udp_mgr.stop().await;
|
||||
}
|
||||
self.udp_listener_manager = None;
|
||||
|
||||
self.started = false;
|
||||
// Reset cancel token so proxy can be restarted
|
||||
self.cancel_token = CancellationToken::new();
|
||||
@@ -681,6 +736,67 @@ impl RustProxy {
|
||||
}
|
||||
}
|
||||
|
||||
// Reconcile UDP ports
|
||||
{
|
||||
let mut new_udp_ports = HashSet::new();
|
||||
for route in &routes {
|
||||
if !route.is_enabled() { continue; }
|
||||
let transport = route.route_match.transport.as_ref();
|
||||
match transport {
|
||||
Some(rustproxy_config::TransportProtocol::Udp) |
|
||||
Some(rustproxy_config::TransportProtocol::All) => {
|
||||
for port in route.route_match.ports.to_ports() {
|
||||
new_udp_ports.insert(port);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let old_udp_ports: HashSet<u16> = self.udp_listener_manager
|
||||
.as_ref()
|
||||
.map(|u| u.listening_ports().into_iter().collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
if !new_udp_ports.is_empty() {
|
||||
// Ensure UDP manager exists
|
||||
if self.udp_listener_manager.is_none() {
|
||||
if let Some(ref listener) = self.listener_manager {
|
||||
let conn_tracker = listener.conn_tracker().clone();
|
||||
self.udp_listener_manager = Some(UdpListenerManager::new(
|
||||
Arc::clone(&new_manager),
|
||||
Arc::clone(&self.metrics),
|
||||
conn_tracker,
|
||||
self.cancel_token.clone(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref mut udp_mgr) = self.udp_listener_manager {
|
||||
udp_mgr.update_routes(Arc::clone(&new_manager));
|
||||
|
||||
// Add new UDP ports
|
||||
for port in &new_udp_ports {
|
||||
if !old_udp_ports.contains(port) {
|
||||
udp_mgr.add_port(*port).await?;
|
||||
}
|
||||
}
|
||||
// Remove old UDP ports
|
||||
for port in &old_udp_ports {
|
||||
if !new_udp_ports.contains(port) {
|
||||
udp_mgr.remove_port(*port);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if self.udp_listener_manager.is_some() {
|
||||
// All UDP routes removed — shut down UDP manager
|
||||
if let Some(ref mut udp_mgr) = self.udp_listener_manager {
|
||||
udp_mgr.stop().await;
|
||||
}
|
||||
self.udp_listener_manager = None;
|
||||
}
|
||||
}
|
||||
|
||||
// Update NFTables rules: remove old, apply new
|
||||
self.update_nftables_rules(&routes).await;
|
||||
|
||||
@@ -840,6 +956,65 @@ impl RustProxy {
|
||||
self.socket_handler_relay.read().unwrap().clone()
|
||||
}
|
||||
|
||||
/// Build a rustls ServerConfig suitable for QUIC (TLS 1.3 only, h3 ALPN).
|
||||
/// Uses the first available cert from tls_configs, or returns None if no certs available.
|
||||
fn build_quic_tls_config(
|
||||
tls_configs: &HashMap<String, TlsCertConfig>,
|
||||
) -> Option<Arc<rustls::ServerConfig>> {
|
||||
// Find the first available cert (prefer wildcard, then any)
|
||||
let cert_config = tls_configs.get("*")
|
||||
.or_else(|| tls_configs.values().next());
|
||||
|
||||
let cert_config = match cert_config {
|
||||
Some(c) => c,
|
||||
None => return None,
|
||||
};
|
||||
|
||||
// Parse cert chain from PEM
|
||||
let mut cert_reader = std::io::BufReader::new(cert_config.cert_pem.as_bytes());
|
||||
let certs: Vec<rustls::pki_types::CertificateDer<'static>> =
|
||||
rustls_pemfile::certs(&mut cert_reader)
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
|
||||
if certs.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Parse private key from PEM
|
||||
let mut key_reader = std::io::BufReader::new(cert_config.key_pem.as_bytes());
|
||||
let key = match rustls_pemfile::private_key(&mut key_reader) {
|
||||
Ok(Some(key)) => key,
|
||||
_ => return None,
|
||||
};
|
||||
|
||||
let mut tls_config = match rustls::ServerConfig::builder()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(certs, key)
|
||||
{
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
warn!("Failed to build QUIC TLS config: {}", e);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// QUIC requires h3 ALPN
|
||||
tls_config.alpn_protocols = vec![b"h3".to_vec()];
|
||||
|
||||
Some(Arc::new(tls_config))
|
||||
}
|
||||
|
||||
/// Set the Unix domain socket path for relaying UDP datagrams to TypeScript datagramHandler callbacks.
|
||||
pub async fn set_datagram_handler_relay_path(&mut self, path: Option<String>) {
|
||||
info!("Datagram handler relay path set to: {:?}", path);
|
||||
if let Some(ref udp_mgr) = self.udp_listener_manager {
|
||||
if let Some(ref p) = path {
|
||||
udp_mgr.set_datagram_handler_relay(p.clone()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load a certificate for a domain and hot-swap the TLS configuration.
|
||||
pub async fn load_certificate(
|
||||
&mut self,
|
||||
|
||||
@@ -149,6 +149,7 @@ async fn handle_request(
|
||||
"getListeningPorts" => handle_get_listening_ports(&id, proxy),
|
||||
"getNftablesStatus" => handle_get_nftables_status(&id, proxy).await,
|
||||
"setSocketHandlerRelay" => handle_set_socket_handler_relay(&id, &request.params, proxy).await,
|
||||
"setDatagramHandlerRelay" => handle_set_datagram_handler_relay(&id, &request.params, proxy).await,
|
||||
"addListeningPort" => handle_add_listening_port(&id, &request.params, proxy).await,
|
||||
"removeListeningPort" => handle_remove_listening_port(&id, &request.params, proxy).await,
|
||||
"loadCertificate" => handle_load_certificate(&id, &request.params, proxy).await,
|
||||
@@ -391,6 +392,26 @@ async fn handle_set_socket_handler_relay(
|
||||
ManagementResponse::ok(id.to_string(), serde_json::json!({}))
|
||||
}
|
||||
|
||||
async fn handle_set_datagram_handler_relay(
|
||||
id: &str,
|
||||
params: &serde_json::Value,
|
||||
proxy: &mut Option<RustProxy>,
|
||||
) -> ManagementResponse {
|
||||
let p = match proxy.as_mut() {
|
||||
Some(p) => p,
|
||||
None => return ManagementResponse::err(id.to_string(), "Proxy is not running".to_string()),
|
||||
};
|
||||
|
||||
let socket_path = params.get("socketPath")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
info!("setDatagramHandlerRelay: socket_path={:?}", socket_path);
|
||||
p.set_datagram_handler_relay_path(socket_path).await;
|
||||
|
||||
ManagementResponse::ok(id.to_string(), serde_json::json!({}))
|
||||
}
|
||||
|
||||
async fn handle_add_listening_port(
|
||||
id: &str,
|
||||
params: &serde_json::Value,
|
||||
|
||||
@@ -269,6 +269,7 @@ pub fn make_test_route(
|
||||
id: None,
|
||||
route_match: rustproxy_config::RouteMatch {
|
||||
ports: rustproxy_config::PortRange::Single(port),
|
||||
transport: None,
|
||||
domains: domain.map(|d| rustproxy_config::DomainSpec::Single(d.to_string())),
|
||||
path: None,
|
||||
client_ip: None,
|
||||
@@ -288,6 +289,7 @@ pub fn make_test_route(
|
||||
send_proxy_protocol: None,
|
||||
headers: None,
|
||||
advanced: None,
|
||||
backend_transport: None,
|
||||
priority: None,
|
||||
}]),
|
||||
tls: None,
|
||||
@@ -298,6 +300,7 @@ pub fn make_test_route(
|
||||
forwarding_engine: None,
|
||||
nftables: None,
|
||||
send_proxy_protocol: None,
|
||||
udp: None,
|
||||
},
|
||||
headers: None,
|
||||
security: None,
|
||||
|
||||
Reference in New Issue
Block a user