mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 10:57:03 +08:00
refactor(workspace): enforce layered crate boundaries
This commit is contained in:
@@ -159,6 +159,15 @@ fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
|
||||
Duration::from_millis(resolve_tunnel_request_timeouts(meta).first_byte_ms)
|
||||
}
|
||||
|
||||
fn is_rollout_probe_request(headers: &HeaderMap, forwarded_by_gateway: bool) -> bool {
|
||||
forwarded_by_gateway
|
||||
&& headers
|
||||
.get(crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| value == crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE)
|
||||
}
|
||||
|
||||
pub async fn relay_request(
|
||||
Path(node_id): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
@@ -171,6 +180,7 @@ pub async fn relay_request(
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
let rollout_probe = is_rollout_probe_request(request.headers(), forwarded_by_gateway);
|
||||
if !addr.ip().is_loopback() && !forwarded_by_gateway {
|
||||
return tunnel_error_response(
|
||||
StatusCode::FORBIDDEN,
|
||||
@@ -348,12 +358,16 @@ pub async fn relay_request(
|
||||
);
|
||||
}
|
||||
};
|
||||
if let Err(error) = record_proxy_upgrade_traffic_success(state.data.as_ref(), &node_id).await {
|
||||
warn!(
|
||||
node_id = %node_id,
|
||||
error = %error,
|
||||
"failed to record proxy upgrade traffic confirmation"
|
||||
);
|
||||
if !rollout_probe {
|
||||
if let Err(error) =
|
||||
record_proxy_upgrade_traffic_success(state.data.as_ref(), &node_id).await
|
||||
{
|
||||
warn!(
|
||||
node_id = %node_id,
|
||||
error = %error,
|
||||
"failed to record proxy upgrade traffic confirmation"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let Some(mut body_rx) = stream.take_body_receiver() else {
|
||||
@@ -462,8 +476,8 @@ mod tests {
|
||||
use super::super::hub::ProxyConn;
|
||||
use super::super::{protocol, AppState, ConnConfig, ControlPlaneClient};
|
||||
use super::{
|
||||
relay_header_timeout, relay_request, Body, Request, SocketAddr, StatusCode,
|
||||
TUNNEL_ERROR_HEADER,
|
||||
is_rollout_probe_request, relay_header_timeout, relay_request, Body, HeaderMap, Request,
|
||||
SocketAddr, StatusCode, TUNNEL_ERROR_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::maintenance::start_proxy_upgrade_rollout;
|
||||
@@ -482,6 +496,20 @@ mod tests {
|
||||
use std::time::Duration;
|
||||
use tokio::sync::watch;
|
||||
|
||||
#[test]
|
||||
fn rollout_probe_marker_is_only_trusted_from_a_forwarding_gateway() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER,
|
||||
crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE
|
||||
.parse()
|
||||
.expect("probe marker should be a valid header"),
|
||||
);
|
||||
|
||||
assert!(!is_rollout_probe_request(&headers, false));
|
||||
assert!(is_rollout_probe_request(&headers, true));
|
||||
}
|
||||
|
||||
fn test_app_state() -> AppState {
|
||||
AppState::new(
|
||||
ControlPlaneClient::disabled(),
|
||||
|
||||
@@ -6,6 +6,9 @@ mod proxy_conn;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_gateway_tunnel::{
|
||||
resolve_proxy_max_streams, resolve_proxy_node_name, resolve_proxy_protocol_version,
|
||||
};
|
||||
use aether_runtime::{
|
||||
hold_admission_permit_until, prometheus_response, service_up_sample, AdmissionPermit,
|
||||
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, MetricKind, MetricLabel, MetricSample,
|
||||
@@ -17,7 +20,6 @@ use axum::http::HeaderMap;
|
||||
use axum::response::{IntoResponse, Json};
|
||||
use axum::routing::{get, post};
|
||||
use axum::Router;
|
||||
use base64::Engine as _;
|
||||
use dashmap::DashMap;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -335,121 +337,3 @@ pub async fn ws_proxy(
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn resolve_proxy_max_streams(headers: &HeaderMap, fallback: usize) -> usize {
|
||||
headers
|
||||
.get("x-tunnel-max-streams")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(fallback)
|
||||
.clamp(1, 2048)
|
||||
}
|
||||
|
||||
fn resolve_proxy_node_name(headers: &HeaderMap, node_id: &str) -> String {
|
||||
if let Some(decoded) = headers
|
||||
.get(aether_contracts::tunnel::TUNNEL_NODE_NAME_B64_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| {
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(value.trim())
|
||||
.ok()
|
||||
})
|
||||
.and_then(|bytes| String::from_utf8(bytes).ok())
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty() && value.chars().count() <= 100)
|
||||
{
|
||||
return decoded;
|
||||
}
|
||||
|
||||
headers
|
||||
.get("x-node-name")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(node_id)
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn resolve_proxy_protocol_version(headers: &HeaderMap) -> u8 {
|
||||
headers
|
||||
.get(aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u8>().ok())
|
||||
.filter(|value| *value >= 1)
|
||||
.map(|value| value.min(aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION))
|
||||
.unwrap_or(1)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::http::{HeaderMap, HeaderValue};
|
||||
use base64::Engine as _;
|
||||
|
||||
use super::{
|
||||
resolve_proxy_max_streams, resolve_proxy_node_name, resolve_proxy_protocol_version,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn proxy_max_streams_honors_small_advertised_capacity() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-tunnel-max-streams", HeaderValue::from_static("8"));
|
||||
|
||||
assert_eq!(resolve_proxy_max_streams(&headers, 128), 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_max_streams_caps_unreasonably_large_capacity() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-tunnel-max-streams", HeaderValue::from_static("9999"));
|
||||
|
||||
assert_eq!(resolve_proxy_max_streams(&headers, 128), 2048);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_protocol_version_defaults_to_v1_when_header_missing() {
|
||||
let headers = HeaderMap::new();
|
||||
assert_eq!(resolve_proxy_protocol_version(&headers), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_protocol_version_reads_advertised_version() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
|
||||
HeaderValue::from_static("2"),
|
||||
);
|
||||
|
||||
assert_eq!(resolve_proxy_protocol_version(&headers), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_node_name_reads_legacy_ascii_header() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-node-name", HeaderValue::from_static("edge-1"));
|
||||
|
||||
assert_eq!(resolve_proxy_node_name(&headers, "node-1"), "edge-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_node_name_decodes_base64_header() {
|
||||
let mut headers = HeaderMap::new();
|
||||
let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode("日本节点");
|
||||
headers.insert(
|
||||
aether_contracts::tunnel::TUNNEL_NODE_NAME_B64_HEADER,
|
||||
HeaderValue::from_str(&encoded).expect("encoded header value should parse"),
|
||||
);
|
||||
|
||||
assert_eq!(resolve_proxy_node_name(&headers, "node-1"), "日本节点");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_node_name_falls_back_to_node_id_for_invalid_base64() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
aether_contracts::tunnel::TUNNEL_NODE_NAME_B64_HEADER,
|
||||
HeaderValue::from_static("not valid"),
|
||||
);
|
||||
|
||||
assert_eq!(resolve_proxy_node_name(&headers, "node-1"), "node-1");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,23 +1 @@
|
||||
use bytes::Bytes;
|
||||
|
||||
pub use aether_contracts::tunnel::{
|
||||
decode_payload, encode_connection_close, encode_frame, encode_goaway, encode_goaway_v3,
|
||||
encode_hello, encode_load_report, encode_ping, encode_pong, encode_reset_stream,
|
||||
encode_settings, encode_stream_error, encode_window_update, frame_payload_by_header,
|
||||
ConnectionClosePayload, FrameHeader, GoAwayPayload, HelloPayload, LoadReportPayload,
|
||||
RequestMeta, ResetStreamPayload, ResponseMeta, SettingsPayload, WindowUpdatePayload,
|
||||
CONNECTION_CLOSE, FLAG_END_STREAM, FLAG_GZIP_COMPRESSED, GOAWAY, HEADER_SIZE, HEARTBEAT_ACK,
|
||||
HEARTBEAT_DATA, HELLO, LOAD_REPORT, PING, PONG, REQUEST_BODY, REQUEST_HEADERS, RESET_STREAM,
|
||||
RESPONSE_BODY, RESPONSE_HEADERS, SETTINGS, STREAM_END, STREAM_ERROR, WINDOW_UPDATE,
|
||||
};
|
||||
|
||||
pub fn compress_payload(payload: &[u8]) -> Result<(Vec<u8>, u8), std::io::Error> {
|
||||
let (compressed, flags) =
|
||||
aether_contracts::tunnel::compress_payload(Bytes::copy_from_slice(payload));
|
||||
Ok((compressed.to_vec(), flags))
|
||||
}
|
||||
|
||||
pub fn raw_payload(payload: &[u8]) -> (Vec<u8>, u8) {
|
||||
let (payload, flags) = aether_contracts::tunnel::raw_payload(Bytes::copy_from_slice(payload));
|
||||
(payload.to_vec(), flags)
|
||||
}
|
||||
pub use aether_gateway_tunnel::embedded::protocol::*;
|
||||
|
||||
@@ -14,6 +14,7 @@ use aether_contracts::tunnel::{
|
||||
use aether_data::repository::proxy_nodes::{
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode,
|
||||
};
|
||||
use aether_gateway_tunnel::EmbeddedTunnelDefaults;
|
||||
use aether_runtime::MetricSample;
|
||||
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
|
||||
use async_stream::stream;
|
||||
@@ -36,6 +37,11 @@ use super::error::GatewayError;
|
||||
use super::headers::{extract_or_generate_trace_id, should_skip_request_header};
|
||||
use super::AppState;
|
||||
|
||||
pub(crate) use aether_gateway_tunnel::{
|
||||
is_tunnel_heartbeat_path, is_tunnel_node_status_path, TunnelAttachmentRecord,
|
||||
DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, PROXY_TUNNEL_PATH,
|
||||
TUNNEL_HEARTBEAT_PATH, TUNNEL_NODE_STATUS_PATH, TUNNEL_RELAY_PATH_PATTERN, TUNNEL_ROUTE_FAMILY,
|
||||
};
|
||||
pub(crate) use embedded::DirectRelayResponse;
|
||||
pub(crate) use embedded::ProxyConn as TunnelProxyConn;
|
||||
pub use embedded::{
|
||||
@@ -44,24 +50,14 @@ pub use embedded::{
|
||||
ControlPlaneClient as TunnelControlPlaneClient,
|
||||
};
|
||||
|
||||
pub(crate) const PROXY_TUNNEL_PATH: &str = "/api/internal/proxy-tunnel";
|
||||
pub(crate) const TUNNEL_HEARTBEAT_PATH: &str = "/api/internal/tunnel/heartbeat";
|
||||
pub(crate) const TUNNEL_NODE_STATUS_PATH: &str = "/api/internal/tunnel/node-status";
|
||||
pub(crate) const TUNNEL_RELAY_PATH_PATTERN: &str = "/api/internal/tunnel/relay/{node_id}";
|
||||
pub(crate) const TUNNEL_ROUTE_FAMILY: &str = "tunnel_manage";
|
||||
|
||||
const DEFAULT_PROXY_IDLE_TIMEOUT_MS: u64 = 0;
|
||||
const DEFAULT_PING_INTERVAL_MS: u64 = 15_000;
|
||||
const DEFAULT_MAX_STREAMS: usize = 2048;
|
||||
const DEFAULT_OUTBOUND_QUEUE_CAPACITY: usize = 512;
|
||||
const DEFAULT_ATTACHMENT_TTL_SECS: u64 = 90;
|
||||
const DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES: usize = 5_242_880;
|
||||
const DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES: usize = 64 * 1024;
|
||||
const TUNNEL_ATTACHMENT_KEY_PREFIX: &str = "tunnel.attachments.";
|
||||
const TUNNEL_ATTACHMENT_REDIS_KEY_PREFIX: &str = "tunnel:attachments:";
|
||||
const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID";
|
||||
const TUNNEL_RELAY_BASE_URL_ENV: &str = "AETHER_TUNNEL_RELAY_BASE_URL";
|
||||
const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS";
|
||||
pub(crate) const TUNNEL_RELAY_ROLLOUT_PROBE_HEADER: &str = "x-aether-tunnel-rollout-probe";
|
||||
pub(crate) const TUNNEL_RELAY_ROLLOUT_PROBE_VALUE: &str = "1";
|
||||
|
||||
pub(crate) async fn send_owner_forward_request(
|
||||
request: reqwest::RequestBuilder,
|
||||
@@ -118,14 +114,6 @@ pub(crate) struct TunnelInstanceIdentity {
|
||||
attachment_ttl_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub(crate) struct TunnelAttachmentRecord {
|
||||
pub(crate) gateway_instance_id: String,
|
||||
pub(crate) relay_base_url: String,
|
||||
pub(crate) conn_count: usize,
|
||||
pub(crate) observed_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct TunnelAttachmentDirectory {
|
||||
identity: Arc<TunnelInstanceIdentity>,
|
||||
@@ -263,11 +251,7 @@ impl TunnelAttachmentDirectory {
|
||||
let Some(record) = self.read_attachment_record(data, node_id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let is_expired = record
|
||||
.observed_at_unix_secs
|
||||
.saturating_add(self.identity.attachment_ttl_secs)
|
||||
< current_unix_secs();
|
||||
if is_expired || record.relay_base_url.trim().is_empty() {
|
||||
if !record.is_routable(current_unix_secs(), self.identity.attachment_ttl_secs) {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(record))
|
||||
@@ -281,7 +265,7 @@ impl TunnelAttachmentDirectory {
|
||||
let Some(record) = self.read_attachment_record(data, node_id).await? else {
|
||||
return Ok(());
|
||||
};
|
||||
if record.gateway_instance_id == self.identity.instance_id {
|
||||
if record.is_owned_by(&self.identity.instance_id) {
|
||||
self.delete_attachment_record(data, node_id).await?;
|
||||
}
|
||||
Ok(())
|
||||
@@ -462,15 +446,16 @@ impl EmbeddedTunnelState {
|
||||
data: Arc<GatewayDataState>,
|
||||
attachment_directory: TunnelAttachmentDirectory,
|
||||
) -> Self {
|
||||
let defaults = EmbeddedTunnelDefaults::default();
|
||||
Self {
|
||||
inner: TunnelAppState::new(
|
||||
build_embedded_control_plane(Arc::clone(&data), attachment_directory.clone()),
|
||||
ConnConfig {
|
||||
ping_interval: Duration::from_millis(DEFAULT_PING_INTERVAL_MS),
|
||||
idle_timeout: Duration::from_millis(DEFAULT_PROXY_IDLE_TIMEOUT_MS),
|
||||
outbound_queue_capacity: DEFAULT_OUTBOUND_QUEUE_CAPACITY,
|
||||
ping_interval: defaults.ping_interval,
|
||||
idle_timeout: defaults.proxy_idle_timeout,
|
||||
outbound_queue_capacity: defaults.outbound_queue_capacity,
|
||||
},
|
||||
DEFAULT_MAX_STREAMS,
|
||||
defaults.max_streams,
|
||||
)
|
||||
.with_data(data),
|
||||
attachment_directory,
|
||||
@@ -531,6 +516,53 @@ impl EmbeddedTunnelState {
|
||||
.status)
|
||||
}
|
||||
|
||||
pub(crate) async fn probe_node_url_routed(
|
||||
&self,
|
||||
state: &AppState,
|
||||
node_id: &str,
|
||||
url: &str,
|
||||
timeout_secs: u64,
|
||||
) -> Result<u16, String> {
|
||||
if self.has_local_proxy(node_id) {
|
||||
return self.probe_node_url(node_id, url, timeout_secs).await;
|
||||
}
|
||||
|
||||
let Some(owner) = self
|
||||
.lookup_attachment_owner(state.data.as_ref(), node_id)
|
||||
.await?
|
||||
else {
|
||||
return self.probe_node_url(node_id, url, timeout_secs).await;
|
||||
};
|
||||
if owner.gateway_instance_id == self.local_instance_id() {
|
||||
self.clear_local_attachment_if_stale(state.data.as_ref(), node_id)
|
||||
.await?;
|
||||
return self.probe_node_url(node_id, url, timeout_secs).await;
|
||||
}
|
||||
|
||||
let timeout_secs = timeout_secs.clamp(5, 60);
|
||||
let owner_url = build_owner_relay_url(&owner.relay_base_url, node_id)
|
||||
.map_err(|error| format!("invalid owner tunnel probe URL: {error:?}"))?;
|
||||
let payload = encode_tunnel_relay_envelope(&build_tunnel_probe_meta(url, timeout_secs))?;
|
||||
let response = state
|
||||
.owner_forward_client
|
||||
.post(owner_url)
|
||||
.header(TUNNEL_RELAY_FORWARDED_BY_HEADER, self.local_instance_id())
|
||||
.header(
|
||||
TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
|
||||
owner.gateway_instance_id.as_str(),
|
||||
)
|
||||
.header(
|
||||
TUNNEL_RELAY_ROLLOUT_PROBE_HEADER,
|
||||
TUNNEL_RELAY_ROLLOUT_PROBE_VALUE,
|
||||
)
|
||||
.timeout(Duration::from_secs(timeout_secs))
|
||||
.body(payload)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| format!("owner tunnel probe failed: {error}"))?;
|
||||
Ok(response.status().as_u16())
|
||||
}
|
||||
|
||||
pub(crate) async fn probe_node_url_with_response(
|
||||
&self,
|
||||
node_id: &str,
|
||||
@@ -538,21 +570,7 @@ impl EmbeddedTunnelState {
|
||||
timeout_secs: u64,
|
||||
) -> Result<TunnelProbeResponse, String> {
|
||||
let timeout_secs = timeout_secs.clamp(5, 60);
|
||||
let meta = tunnel_protocol::RequestMeta {
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
method: "GET".to_string(),
|
||||
url: url.trim().to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: timeout_secs,
|
||||
follow_redirects: Some(false),
|
||||
http1_only: false,
|
||||
transport_profile: None,
|
||||
};
|
||||
let meta = build_tunnel_probe_meta(url, timeout_secs);
|
||||
let stream = self.inner.hub.open_local_stream(node_id, &meta).await?;
|
||||
let stream_id = stream.id;
|
||||
let result = async {
|
||||
@@ -623,6 +641,35 @@ impl EmbeddedTunnelState {
|
||||
}
|
||||
}
|
||||
|
||||
fn build_tunnel_probe_meta(url: &str, timeout_secs: u64) -> tunnel_protocol::RequestMeta {
|
||||
tunnel_protocol::RequestMeta {
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
method: "GET".to_string(),
|
||||
url: url.trim().to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: timeout_secs,
|
||||
follow_redirects: Some(false),
|
||||
http1_only: false,
|
||||
transport_profile: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_tunnel_relay_envelope(meta: &tunnel_protocol::RequestMeta) -> Result<Vec<u8>, String> {
|
||||
let meta = serde_json::to_vec(meta)
|
||||
.map_err(|error| format!("failed to encode tunnel probe metadata: {error}"))?;
|
||||
let meta_len = u32::try_from(meta.len())
|
||||
.map_err(|_| "tunnel probe metadata exceeds relay envelope limit".to_string())?;
|
||||
let mut payload = Vec::with_capacity(4usize.saturating_add(meta.len()));
|
||||
payload.extend_from_slice(&meta_len.to_be_bytes());
|
||||
payload.extend_from_slice(&meta);
|
||||
Ok(payload)
|
||||
}
|
||||
|
||||
impl Default for EmbeddedTunnelState {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
@@ -631,11 +678,15 @@ impl Default for EmbeddedTunnelState {
|
||||
|
||||
impl fmt::Debug for EmbeddedTunnelState {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let defaults = EmbeddedTunnelDefaults::default();
|
||||
f.debug_struct("EmbeddedTunnelState")
|
||||
.field("proxy_idle_timeout_ms", &DEFAULT_PROXY_IDLE_TIMEOUT_MS)
|
||||
.field("ping_interval_ms", &DEFAULT_PING_INTERVAL_MS)
|
||||
.field("max_streams", &DEFAULT_MAX_STREAMS)
|
||||
.field("outbound_queue_capacity", &DEFAULT_OUTBOUND_QUEUE_CAPACITY)
|
||||
.field(
|
||||
"proxy_idle_timeout_ms",
|
||||
&defaults.proxy_idle_timeout.as_millis(),
|
||||
)
|
||||
.field("ping_interval_ms", &defaults.ping_interval.as_millis())
|
||||
.field("max_streams", &defaults.max_streams)
|
||||
.field("outbound_queue_capacity", &defaults.outbound_queue_capacity)
|
||||
.field(
|
||||
"instance_id",
|
||||
&self.attachment_directory.local_instance_id(),
|
||||
@@ -714,14 +765,6 @@ pub(crate) async fn relay_request(
|
||||
.into_response())
|
||||
}
|
||||
|
||||
pub(crate) fn is_tunnel_heartbeat_path(path: &str) -> bool {
|
||||
path == TUNNEL_HEARTBEAT_PATH
|
||||
}
|
||||
|
||||
pub(crate) fn is_tunnel_node_status_path(path: &str) -> bool {
|
||||
path == TUNNEL_NODE_STATUS_PATH
|
||||
}
|
||||
|
||||
fn build_embedded_control_plane(
|
||||
data: Arc<GatewayDataState>,
|
||||
attachment_directory: TunnelAttachmentDirectory,
|
||||
@@ -1134,17 +1177,25 @@ fn parse_embedded_tunnel_heartbeat_request(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status, current_unix_secs,
|
||||
prepare_owner_relay_request_body, tunnel_attachment_key, GatewayDataState,
|
||||
apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status,
|
||||
build_tunnel_probe_meta, current_unix_secs, encode_tunnel_relay_envelope,
|
||||
prepare_owner_relay_request_body, tunnel_attachment_key, AppState, GatewayDataState,
|
||||
TunnelAttachmentDirectory, TunnelAttachmentRecord,
|
||||
};
|
||||
use aether_contracts::tunnel::{
|
||||
try_decode_tunnel_relay_request_meta, TUNNEL_RELAY_FORWARDED_BY_HEADER,
|
||||
TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
|
||||
};
|
||||
use aether_data::repository::proxy_nodes::{
|
||||
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::routing::post;
|
||||
use axum::Router;
|
||||
use serde_json::json;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
fn sample_proxy_node(node_id: &str) -> StoredProxyNode {
|
||||
StoredProxyNode::new(
|
||||
@@ -1183,6 +1234,99 @@ mod tests {
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn routed_tunnel_probe_builds_a_valid_owner_relay_envelope() {
|
||||
let meta = build_tunnel_probe_meta("https://probe.example/health", 7);
|
||||
let envelope = encode_tunnel_relay_envelope(&meta).expect("probe should encode");
|
||||
let (decoded, body_offset) = try_decode_tunnel_relay_request_meta(&envelope)
|
||||
.expect("probe envelope should decode")
|
||||
.expect("probe envelope should contain complete metadata");
|
||||
|
||||
assert_eq!(decoded.method, "GET");
|
||||
assert_eq!(decoded.url, "https://probe.example/health");
|
||||
assert_eq!(decoded.timeout, 7);
|
||||
assert_eq!(body_offset, envelope.len());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn routed_tunnel_probe_forwards_to_the_attachment_owner() {
|
||||
let captured = Arc::new(Mutex::new(None::<(HeaderMap, Bytes)>));
|
||||
let captured_for_route = Arc::clone(&captured);
|
||||
let app = Router::new().route(
|
||||
"/api/internal/tunnel/relay/{node_id}",
|
||||
post(move |headers: HeaderMap, body: Bytes| {
|
||||
let captured = Arc::clone(&captured_for_route);
|
||||
async move {
|
||||
*captured.lock().expect("capture lock") = Some((headers, body));
|
||||
StatusCode::NO_CONTENT
|
||||
}
|
||||
}),
|
||||
);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("owner listener should bind");
|
||||
let owner_base_url = format!("http://{}", listener.local_addr().expect("owner address"));
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("owner server should run");
|
||||
});
|
||||
|
||||
let owner = TunnelAttachmentRecord {
|
||||
gateway_instance_id: "gateway-b".to_string(),
|
||||
relay_base_url: owner_base_url,
|
||||
conn_count: 1,
|
||||
observed_at_unix_secs: current_unix_secs(),
|
||||
};
|
||||
let data = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
|
||||
tunnel_attachment_key("node-remote"),
|
||||
serde_json::to_value(owner).expect("owner record should serialize"),
|
||||
)]);
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(data)
|
||||
.with_tunnel_identity("gateway-a", Some("http://gateway-a.internal"));
|
||||
|
||||
let status = state
|
||||
.tunnel
|
||||
.probe_node_url_routed(&state, "node-remote", "https://probe.example/health", 5)
|
||||
.await
|
||||
.expect("remote owner probe should succeed");
|
||||
assert_eq!(status, StatusCode::NO_CONTENT.as_u16());
|
||||
|
||||
let (headers, body) = captured
|
||||
.lock()
|
||||
.expect("capture lock")
|
||||
.take()
|
||||
.expect("owner should receive the probe");
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(TUNNEL_RELAY_FORWARDED_BY_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-a")
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(TUNNEL_RELAY_OWNER_INSTANCE_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-b")
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(super::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(super::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE)
|
||||
);
|
||||
let (meta, body_offset) = try_decode_tunnel_relay_request_meta(&body)
|
||||
.expect("owner probe envelope should decode")
|
||||
.expect("owner probe metadata should be complete");
|
||||
assert_eq!(meta.url, "https://probe.example/health");
|
||||
assert_eq!(body_offset, body.len());
|
||||
|
||||
server.abort();
|
||||
let _ = server.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn owner_relay_body_preparation_rejects_invalid_metadata() {
|
||||
let mut envelope = Vec::new();
|
||||
|
||||
Reference in New Issue
Block a user