i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

Implement transport profile routing

Browse Source
This commit is contained in:
fawney19
2026-05-05 22:21:23 +08:00
parent aacab1a90c
commit f959f02d40
86 changed files with 860 additions and 318 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
apps/aether-proxy/src/app.rs
View File

@@ -134,12 +134,11 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
config.dns_cache_capacity,
));
// Build Hyper client for tunnel upstream requests (shared).
// DNS still flows through validated addresses from DnsCache, while the
// custom connector exposes per-request connect/TLS timing when available.
let upstream_client = upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_http1_client =
upstream_client::build_http1_only_upstream_client(&config, Arc::clone(&dns_cache));
let config = Arc::new(config);
// Build a profile-keyed Hyper client pool for tunnel upstream requests.
let upstream_client_pool =
upstream_client::UpstreamClientPool::new(Arc::clone(&config), Arc::clone(&dns_cache));
// Register with each Aether server and build per-server contexts.
// Wrapped in Arc<Mutex> so retry_failed_registrations can append later.
@@ -196,10 +195,9 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
// Build shared application state
let tunnel_tls_config = Arc::new(crate::tunnel::client::build_tls_config());
let mut state = AppState {
config: Arc::new(config),
config,
dns_cache,
upstream_client,
upstream_http1_client,
upstream_client_pool,
tunnel_tls_config,
stream_gate: None,
distributed_stream_gate: None,
@@ -916,15 +914,12 @@ mod tests {
fn sample_state(config: Config) -> Arc<ProxyAppState> {
let config = Arc::new(config);
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_http1_client =
upstream_client::build_http1_only_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_client_pool =
upstream_client::UpstreamClientPool::new(Arc::clone(&config), Arc::clone(&dns_cache));
Arc::new(ProxyAppState {
config,
dns_cache,
upstream_client,
upstream_http1_client,
upstream_client_pool,
tunnel_tls_config: Arc::new(crate::tunnel::client::build_tls_config()),
stream_gate: None,
distributed_stream_gate: None,

8
apps/aether-proxy/src/state.rs
View File

@@ -13,17 +13,15 @@ use crate::config::Config;
use crate::registration::client::AetherClient;
use crate::runtime::SharedDynamicConfig;
use crate::target_filter::DnsCache;
use crate::upstream_client::UpstreamClient;
use crate::upstream_client::UpstreamClientPool;
/// Central application state shared across all servers/tunnels.
pub struct AppState {
pub config: Arc<Config>,
/// DNS cache for upstream target resolution (shared).
pub dns_cache: Arc<DnsCache>,
/// Hyper client for tunnel upstream requests with validated DNS and connection timing.
pub upstream_client: UpstreamClient,
/// Dedicated Hyper client that forces HTTP/1.1 for upstreams that break on H2/ALPN.
pub upstream_http1_client: UpstreamClient,
/// Profile-keyed upstream client pool used by tunnel requests.
pub upstream_client_pool: UpstreamClientPool,
/// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect).
pub tunnel_tls_config: Arc<rustls::ClientConfig>,
/// Optional per-process stream admission gate.

13
apps/aether-proxy/src/tunnel/mod.rs
View File

@@ -379,12 +379,16 @@ mod tests {
fn relay_probe_envelope() -> Vec<u8> {
let meta = protocol::RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "GET".to_string(),
url: "http://127.0.0.1:80/blocked".to_string(),
headers: std::collections::HashMap::new(),
timeout: 5,
follow_redirects: None,
http1_only: false,
transport_profile: None,
};
let meta_json =
serde_json::to_vec(&meta).expect("tunnel relay probe metadata should serialize");
@@ -446,15 +450,12 @@ mod tests {
fn sample_state(config: Config) -> Arc<ProxyAppState> {
let config = Arc::new(config);
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_http1_client =
upstream_client::build_http1_only_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_client_pool =
upstream_client::UpstreamClientPool::new(Arc::clone(&config), Arc::clone(&dns_cache));
Arc::new(ProxyAppState {
config,
dns_cache,
upstream_client,
upstream_http1_client,
upstream_client_pool,
tunnel_tls_config: Arc::new(crate::tunnel::client::build_tls_config()),
stream_gate: None,
distributed_stream_gate: None,

93
apps/aether-proxy/src/tunnel/stream_handler.rs
View File

@@ -726,6 +726,7 @@ fn resolve_redirect<B>(
async fn execute_upstream_request(
state: &AppState,
server: &ServerContext,
meta: &RequestMeta,
current_url: &url::Url,
method: hyper::Method,
headers: &[(String, String)],
@@ -756,11 +757,14 @@ async fn execute_upstream_request(
}
let dns_ms = dns_start.elapsed().as_millis() as u64;
let client = if http1_only {
&state.upstream_http1_client
} else {
&state.upstream_client
};
let client_key = upstream_client::upstream_client_pool_key(
meta.provider_id.as_deref(),
meta.endpoint_id.as_deref(),
meta.key_id.as_deref(),
meta.transport_profile.as_ref(),
http1_only,
);
let client = state.upstream_client_pool.get_or_build(client_key)?;
let mut request = hyper::Request::builder()
.method(method)
@@ -1026,15 +1030,16 @@ async fn relay_upstream_response(
}
#[cfg(test)]
fn upstream_client_for_request<'a>(
state: &'a AppState,
fn upstream_client_pool_key_for_request(
meta: &RequestMeta,
) -> &'a upstream_client::UpstreamClient {
if meta.http1_only {
&state.upstream_http1_client
} else {
&state.upstream_client
}
) -> upstream_client::UpstreamClientPoolKey {
upstream_client::upstream_client_pool_key(
meta.provider_id.as_deref(),
meta.endpoint_id.as_deref(),
meta.key_id.as_deref(),
meta.transport_profile.as_ref(),
meta.http1_only,
)
}
/// Handle a single stream: receive body, execute upstream, send response.
@@ -1235,6 +1240,7 @@ async fn handle_stream_inner(
let response_ctx = match execute_upstream_request(
state,
server,
&meta,
&current_url,
current_method.clone(),
&current_headers,
@@ -1671,19 +1677,40 @@ mod tests {
#[test]
fn selects_http1_only_client_when_request_metadata_requires_it() {
let state = sample_state(None, None);
let default_meta = sample_request_meta();
assert!(std::ptr::eq(
upstream_client_for_request(state.as_ref(), &default_meta),
&state.upstream_client
));
assert_eq!(
upstream_client_pool_key_for_request(&default_meta).http_mode,
"auto"
);
let mut http1_meta = sample_request_meta();
http1_meta.http1_only = true;
assert!(std::ptr::eq(
upstream_client_for_request(state.as_ref(), &http1_meta),
&state.upstream_http1_client
));
assert_eq!(
upstream_client_pool_key_for_request(&http1_meta).http_mode,
"http1_only"
);
}
#[test]
fn upstream_client_pool_key_isolates_accounts() {
let mut first = sample_request_meta();
first.provider_id = Some("provider-1".to_string());
first.endpoint_id = Some("endpoint-1".to_string());
first.key_id = Some("key-1".to_string());
first.transport_profile = Some(aether_contracts::ResolvedTransportProfile {
profile_id: "profile-a".to_string(),
backend: "reqwest_rustls".to_string(),
http_mode: "auto".to_string(),
pool_scope: "key".to_string(),
extra: None,
});
let mut second = first.clone();
second.key_id = Some("key-2".to_string());
assert_ne!(
upstream_client_pool_key_for_request(&first),
upstream_client_pool_key_for_request(&second)
);
}
#[test]
@@ -2210,12 +2237,16 @@ mod tests {
fn sample_request_meta() -> RequestMeta {
RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "GET".to_string(),
url: "https://example.com/ok".to_string(),
headers: HashMap::new(),
timeout: 30,
follow_redirects: None,
http1_only: false,
transport_profile: None,
}
}
@@ -2226,15 +2257,12 @@ mod tests {
ensure_rustls_provider();
let config = Arc::new(sample_config());
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_http1_client =
upstream_client::build_http1_only_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_client_pool =
upstream_client::UpstreamClientPool::new(Arc::clone(&config), Arc::clone(&dns_cache));
Arc::new(AppState {
config,
dns_cache,
upstream_client,
upstream_http1_client,
upstream_client_pool,
tunnel_tls_config: Arc::new(build_tls_config()),
stream_gate,
distributed_stream_gate,
@@ -2259,15 +2287,12 @@ mod tests {
fn sample_state_with_config(config: Config) -> Arc<AppState> {
let config = Arc::new(config);
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_http1_client =
upstream_client::build_http1_only_upstream_client(&config, Arc::clone(&dns_cache));
let upstream_client_pool =
upstream_client::UpstreamClientPool::new(Arc::clone(&config), Arc::clone(&dns_cache));
Arc::new(AppState {
config,
dns_cache,
upstream_client,
upstream_http1_client,
upstream_client_pool,
tunnel_tls_config: Arc::new(build_tls_config()),
stream_gate: None,
distributed_stream_gate: None,

159
apps/aether-proxy/src/upstream_client.rs
View File

@@ -1,12 +1,18 @@
use std::collections::HashMap;
use std::convert::Infallible;
use std::future::Future;
use std::io;
use std::net::IpAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::task::{Context, Poll};
use std::time::Duration;
use aether_contracts::{
ResolvedTransportProfile, TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS,
TRANSPORT_HTTP_MODE_HTTP1_ONLY,
};
use bytes::Bytes;
use futures_util::Stream;
use http_body_util::combinators::UnsyncBoxBody;
@@ -37,6 +43,115 @@ type TlsStream = TokioIo<tokio_rustls::client::TlsStream<TcpStream>>;
pub type UpstreamRequestBody = UnsyncBoxBody<Bytes, io::Error>;
pub type UpstreamClient = Client<InstrumentedConnector, UpstreamRequestBody>;
const DEFAULT_PROFILE_ID: &str = "default";
const DEFAULT_BACKEND: &str = TRANSPORT_BACKEND_HYPER_RUSTLS;
const DEFAULT_HTTP_MODE: &str = "auto";
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct UpstreamClientPoolKey {
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub profile_id: String,
pub backend: String,
pub http_mode: String,
}
#[derive(Clone)]
pub struct UpstreamClientPool {
config: Arc<Config>,
dns_cache: Arc<DnsCache>,
clients: Arc<Mutex<HashMap<UpstreamClientPoolKey, UpstreamClient>>>,
}
impl UpstreamClientPool {
pub fn new(config: Arc<Config>, dns_cache: Arc<DnsCache>) -> Self {
Self {
config,
dns_cache,
clients: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn get_or_build(&self, key: UpstreamClientPoolKey) -> Result<UpstreamClient, String> {
if let Some(client) = self
.clients
.lock()
.expect("client pool lock")
.get(&key)
.cloned()
{
return Ok(client);
}
validate_proxy_transport_backend(&key.backend)?;
let http1_only = key
.http_mode
.eq_ignore_ascii_case(TRANSPORT_HTTP_MODE_HTTP1_ONLY);
let client = build_upstream_client_with_protocol(
&self.config,
Arc::clone(&self.dns_cache),
http1_only,
);
self.clients
.lock()
.expect("client pool lock")
.insert(key, client.clone());
Ok(client)
}
}
pub fn upstream_client_pool_key(
provider_id: Option<&str>,
endpoint_id: Option<&str>,
key_id: Option<&str>,
profile: Option<&ResolvedTransportProfile>,
http1_only: bool,
) -> UpstreamClientPoolKey {
let profile_http_mode = profile
.map(|profile| profile.http_mode.trim())
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_HTTP_MODE);
let http_mode = if http1_only {
TRANSPORT_HTTP_MODE_HTTP1_ONLY
} else {
profile_http_mode
};
UpstreamClientPoolKey {
provider_id: normalized_pool_key_part(provider_id),
endpoint_id: normalized_pool_key_part(endpoint_id),
key_id: normalized_pool_key_part(key_id),
profile_id: profile
.map(|profile| profile.profile_id.trim())
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_PROFILE_ID)
.to_string(),
backend: profile
.map(|profile| profile.backend.trim())
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_BACKEND)
.to_string(),
http_mode: http_mode.to_string(),
}
}
fn normalized_pool_key_part(value: Option<&str>) -> String {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("-")
.to_string()
}
fn validate_proxy_transport_backend(backend: &str) -> Result<(), String> {
if backend.eq_ignore_ascii_case(TRANSPORT_BACKEND_HYPER_RUSTLS)
|| backend.eq_ignore_ascii_case(TRANSPORT_BACKEND_REQWEST_RUSTLS)
{
return Ok(());
}
Err(format!("unsupported transport profile backend: {backend}"))
}
pub fn stream_request_body<S>(stream: S) -> UpstreamRequestBody
where
S: Stream<Item = Result<Frame<Bytes>, io::Error>> + Send + 'static,
@@ -181,17 +296,6 @@ impl Service<Uri> for InstrumentedConnector {
}
}
pub fn build_upstream_client(config: &Config, dns_cache: Arc<DnsCache>) -> UpstreamClient {
build_upstream_client_with_protocol(config, dns_cache, false)
}
pub fn build_http1_only_upstream_client(
config: &Config,
dns_cache: Arc<DnsCache>,
) -> UpstreamClient {
build_upstream_client_with_protocol(config, dns_cache, true)
}
fn build_upstream_client_with_protocol(
config: &Config,
dns_cache: Arc<DnsCache>,
@@ -424,6 +528,7 @@ impl rt::Write for MaybeHttpsStream {
#[cfg(test)]
mod tests {
use super::*;
use aether_contracts::ResolvedTransportProfile;
use hyper::Response;
#[test]
@@ -476,4 +581,36 @@ mod tests {
assert_eq!(timing.response_wait_ms, 320);
assert!(!timing.connection_reused);
}
#[test]
fn upstream_client_pool_key_includes_profile_identity() {
let profile = ResolvedTransportProfile {
profile_id: "profile-a".to_string(),
backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.to_string(),
http_mode: "auto".to_string(),
pool_scope: "key".to_string(),
extra: None,
};
let pool_key = upstream_client_pool_key(
Some("provider-1"),
Some("endpoint-1"),
Some("key-1"),
Some(&profile),
false,
);
assert_eq!(pool_key.provider_id, "provider-1");
assert_eq!(pool_key.endpoint_id, "endpoint-1");
assert_eq!(pool_key.key_id, "key-1");
assert_eq!(pool_key.profile_id, "profile-a");
assert_eq!(pool_key.backend, TRANSPORT_BACKEND_REQWEST_RUSTLS);
assert_eq!(pool_key.http_mode, "auto");
}
#[test]
fn upstream_client_pool_rejects_unsupported_backend() {
let error = validate_proxy_transport_backend("utls").unwrap_err();
assert!(error.contains("unsupported transport profile backend"));
}
}