mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 13:10:21 +08:00
Merge branch 'fawney19:main' into main
This commit is contained in:
Generated
+5
@@ -123,10 +123,14 @@ version = "0.1.0"
|
||||
name = "aether-contracts"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"flate2",
|
||||
"hmac",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
@@ -511,6 +515,7 @@ dependencies = [
|
||||
"tower-service",
|
||||
"tracing",
|
||||
"url",
|
||||
"uuid",
|
||||
"webpki-roots 0.26.11",
|
||||
]
|
||||
|
||||
|
||||
@@ -67,6 +67,7 @@ aether-http = { path = "crates/aether-http" }
|
||||
aether-runtime = { path = "crates/aether-runtime" }
|
||||
aether-testkit = { path = "crates/aether-testkit" }
|
||||
aes = "0.8"
|
||||
aes-gcm = "0.10"
|
||||
async-stream = "0.3"
|
||||
async-trait = "0.1"
|
||||
axum = "0.8"
|
||||
|
||||
@@ -31,7 +31,7 @@ aether-task-runtime.workspace = true
|
||||
aether-usage-runtime.workspace = true
|
||||
aether-video-tasks-core.workspace = true
|
||||
aether-wallet.workspace = true
|
||||
aes-gcm = "0.10"
|
||||
aes-gcm.workspace = true
|
||||
async-stream.workspace = true
|
||||
async-trait.workspace = true
|
||||
axum = { version = "0.8", features = ["ws"] }
|
||||
|
||||
@@ -26,7 +26,6 @@ pub(crate) fn mount_public_support_routes(router: Router<AppState>) -> Router<Ap
|
||||
.route("/api/capabilities/model/{*model_path}", get(proxy_request))
|
||||
.route("/install/{*install_path}", get(proxy_request))
|
||||
.route("/install-tunnel/{*install_path}", get(proxy_request))
|
||||
.route("/install-proxy/{*install_path}", get(proxy_request))
|
||||
.route("/i/{*install_path}", get(proxy_request))
|
||||
.route("/", get(proxy_request))
|
||||
}
|
||||
|
||||
@@ -797,7 +797,6 @@ pub(super) fn classify_public_support_route(
|
||||
} else if method == http::Method::GET
|
||||
&& (has_single_segment_after_prefix(normalized_path, "/install/")
|
||||
|| has_single_segment_after_prefix(normalized_path, "/install-tunnel/")
|
||||
|| has_single_segment_after_prefix(normalized_path, "/install-proxy/")
|
||||
|| has_single_segment_after_prefix(normalized_path, "/i/"))
|
||||
{
|
||||
Some(classified(
|
||||
|
||||
@@ -38,6 +38,9 @@ use crate::{AppState, GatewayError};
|
||||
const HUB_RELAY_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
|
||||
const HUB_RELAY_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
||||
const TUNNEL_RELAY_PATH_PREFIX: &str = "/api/internal/tunnel/relay";
|
||||
const DEFAULT_TUNNEL_TIMEOUT_MS: u64 = 60_000;
|
||||
const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1;
|
||||
const MAX_TUNNEL_TIMEOUT_SECS: u64 = 300;
|
||||
pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
|
||||
let mut kinds = Vec::new();
|
||||
if err.is_connect() {
|
||||
@@ -176,6 +179,12 @@ struct RelayRequestMeta {
|
||||
method: String,
|
||||
url: String,
|
||||
headers: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "is_false")]
|
||||
stream: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
request_timeout_ms: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
stream_first_byte_timeout_ms: Option<u64>,
|
||||
timeout: u64,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
follow_redirects: Option<bool>,
|
||||
@@ -195,6 +204,13 @@ pub(crate) struct ExecutionTransportControls {
|
||||
accept_invalid_certs: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct TunnelTimeoutMetadata {
|
||||
request_timeout_ms: Option<u64>,
|
||||
stream_first_byte_timeout_ms: Option<u64>,
|
||||
legacy_timeout_secs: u64,
|
||||
}
|
||||
|
||||
pub(crate) enum DirectUpstreamResponse {
|
||||
Reqwest(reqwest::Response),
|
||||
BrowserWreq(wreq::Response),
|
||||
@@ -579,6 +595,7 @@ fn build_direct_tunnel_request_meta(
|
||||
headers: &HeaderMap,
|
||||
transport_controls: ExecutionTransportControls,
|
||||
) -> tunnel_protocol::RequestMeta {
|
||||
let timeout_metadata = resolve_tunnel_timeout_metadata(plan);
|
||||
tunnel_protocol::RequestMeta {
|
||||
provider_id: Some(plan.provider_id.clone()),
|
||||
endpoint_id: Some(plan.endpoint_id.clone()),
|
||||
@@ -586,7 +603,10 @@ fn build_direct_tunnel_request_meta(
|
||||
method: plan.method.clone(),
|
||||
url: plan.url.clone(),
|
||||
headers: header_map_to_string_map(headers).into_iter().collect(),
|
||||
timeout: resolve_relay_timeout_seconds(plan),
|
||||
stream: plan.stream,
|
||||
request_timeout_ms: timeout_metadata.request_timeout_ms,
|
||||
stream_first_byte_timeout_ms: timeout_metadata.stream_first_byte_timeout_ms,
|
||||
timeout: timeout_metadata.legacy_timeout_secs,
|
||||
follow_redirects: transport_controls.follow_redirects,
|
||||
http1_only: transport_controls.http1_only,
|
||||
transport_profile: plan.transport_profile.clone(),
|
||||
@@ -753,7 +773,8 @@ async fn send_via_tunnel_relay(
|
||||
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
|
||||
let client = build_relay_client(plan.timeouts.as_ref())?;
|
||||
let relay_url = build_relay_url(plan.proxy.as_ref(), node_id);
|
||||
let timeout_secs = resolve_relay_timeout_seconds(plan);
|
||||
let timeout_metadata = resolve_tunnel_timeout_metadata(plan);
|
||||
let timeout_secs = timeout_metadata.legacy_timeout_secs;
|
||||
let envelope = build_relay_envelope(
|
||||
RelayRequestMeta {
|
||||
provider_id: plan.provider_id.clone(),
|
||||
@@ -762,6 +783,9 @@ async fn send_via_tunnel_relay(
|
||||
method: method.as_str().to_string(),
|
||||
url: plan.url.clone(),
|
||||
headers: header_map_to_string_map(&headers),
|
||||
stream: plan.stream,
|
||||
request_timeout_ms: timeout_metadata.request_timeout_ms,
|
||||
stream_first_byte_timeout_ms: timeout_metadata.stream_first_byte_timeout_ms,
|
||||
timeout: timeout_secs,
|
||||
follow_redirects: transport_controls.follow_redirects,
|
||||
http1_only: transport_controls.http1_only,
|
||||
@@ -791,15 +815,22 @@ async fn send_via_tunnel_relay(
|
||||
.request(reqwest::Method::POST, relay_url)
|
||||
.header(reqwest::header::CONTENT_TYPE, HUB_RELAY_CONTENT_TYPE)
|
||||
.body(envelope);
|
||||
if let Some(timeout) = total_timeout {
|
||||
request = request.timeout(timeout);
|
||||
if !plan.stream {
|
||||
if let Some(timeout) = total_timeout {
|
||||
request = request.timeout(timeout);
|
||||
}
|
||||
}
|
||||
|
||||
let first_byte_timeout = if plan.stream {
|
||||
resolve_tunnel_first_byte_timeout(plan)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let started_at = Instant::now();
|
||||
let response = request
|
||||
.send()
|
||||
let response = send_relay_request(request, first_byte_timeout)
|
||||
.await
|
||||
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))?;
|
||||
.map_err(ExecutionRuntimeTransportError::RelayError)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let status_code = response.status().as_u16();
|
||||
let proxy_timing = response
|
||||
@@ -869,6 +900,21 @@ async fn send_via_tunnel_relay(
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn send_relay_request(
|
||||
request: reqwest::RequestBuilder,
|
||||
first_byte_timeout: Option<Duration>,
|
||||
) -> Result<reqwest::Response, String> {
|
||||
if let Some(timeout) = first_byte_timeout {
|
||||
return match tokio::time::timeout(timeout, request.send()).await {
|
||||
Ok(Ok(response)) => Ok(response),
|
||||
Ok(Err(error)) => Err(error.to_string()),
|
||||
Err(_) => Err("tunnel relay first byte timeout".to_string()),
|
||||
};
|
||||
}
|
||||
|
||||
request.send().await.map_err(|err| err.to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn build_request_body(
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
@@ -965,18 +1011,49 @@ fn resolve_tunnel_base_url_from_proxy(proxy: &ProxySnapshot) -> Option<String> {
|
||||
}
|
||||
|
||||
fn resolve_relay_timeout_seconds(plan: &ExecutionPlan) -> u64 {
|
||||
let ms = plan
|
||||
.timeouts
|
||||
resolve_tunnel_timeout_metadata(plan).legacy_timeout_secs
|
||||
}
|
||||
|
||||
fn resolve_tunnel_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
|
||||
plan.stream.then(|| {
|
||||
Duration::from_millis(
|
||||
resolve_selected_tunnel_timeout_ms(plan).unwrap_or(DEFAULT_TUNNEL_TIMEOUT_MS),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_tunnel_timeout_metadata(plan: &ExecutionPlan) -> TunnelTimeoutMetadata {
|
||||
TunnelTimeoutMetadata {
|
||||
request_timeout_ms: plan
|
||||
.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| timeouts.total_ms),
|
||||
stream_first_byte_timeout_ms: plan
|
||||
.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| timeouts.first_byte_ms),
|
||||
legacy_timeout_secs: timeout_ms_to_secs(
|
||||
resolve_selected_tunnel_timeout_ms(plan).unwrap_or(DEFAULT_TUNNEL_TIMEOUT_MS),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_selected_tunnel_timeout_ms(plan: &ExecutionPlan) -> Option<u64> {
|
||||
plan.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| {
|
||||
timeouts
|
||||
.read_ms
|
||||
.or(timeouts.total_ms)
|
||||
.or(timeouts.connect_ms)
|
||||
if plan.stream {
|
||||
timeouts.first_byte_ms.or(timeouts.total_ms)
|
||||
} else {
|
||||
timeouts.total_ms.or(timeouts.first_byte_ms)
|
||||
}
|
||||
})
|
||||
.unwrap_or(60_000);
|
||||
.map(|value| value.max(1))
|
||||
}
|
||||
|
||||
fn timeout_ms_to_secs(ms: u64) -> u64 {
|
||||
let secs = ms.div_ceil(1_000);
|
||||
secs.clamp(1, 300)
|
||||
secs.clamp(MIN_TUNNEL_TIMEOUT_SECS, MAX_TUNNEL_TIMEOUT_SECS)
|
||||
}
|
||||
|
||||
fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
|
||||
@@ -1522,12 +1599,12 @@ mod tests {
|
||||
use tokio::sync::watch;
|
||||
|
||||
use super::{
|
||||
build_browser_wreq_client, build_client, build_execution_response_body,
|
||||
build_request_headers, execute_sync_plan, record_manual_proxy_request_failure,
|
||||
record_manual_proxy_request_outcome, record_manual_proxy_request_success,
|
||||
record_manual_proxy_stream_error, resolve_execution_transport_controls,
|
||||
response_body_is_json, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
|
||||
ExecutionTransportControls,
|
||||
build_browser_wreq_client, build_client, build_direct_tunnel_request_meta,
|
||||
build_execution_response_body, build_request_headers, execute_sync_plan,
|
||||
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
|
||||
record_manual_proxy_request_success, record_manual_proxy_stream_error,
|
||||
resolve_execution_transport_controls, response_body_is_json, DirectSyncExecutionRuntime,
|
||||
ExecutionRuntimeTransportError, ExecutionTransportControls,
|
||||
};
|
||||
use crate::constants::{
|
||||
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
|
||||
@@ -1629,6 +1706,64 @@ mod tests {
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_meta_uses_total_timeout_for_non_stream_requests() {
|
||||
let plan = tunnel_timeout_plan(false);
|
||||
let meta = build_direct_tunnel_request_meta(
|
||||
&plan,
|
||||
&reqwest::header::HeaderMap::new(),
|
||||
ExecutionTransportControls::default(),
|
||||
);
|
||||
|
||||
assert!(!meta.stream);
|
||||
assert_eq!(meta.request_timeout_ms, Some(90_000));
|
||||
assert_eq!(meta.stream_first_byte_timeout_ms, Some(12_345));
|
||||
assert_eq!(meta.timeout, 90);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_meta_uses_first_byte_timeout_for_stream_requests() {
|
||||
let plan = tunnel_timeout_plan(true);
|
||||
let meta = build_direct_tunnel_request_meta(
|
||||
&plan,
|
||||
&reqwest::header::HeaderMap::new(),
|
||||
ExecutionTransportControls::default(),
|
||||
);
|
||||
|
||||
assert!(meta.stream);
|
||||
assert_eq!(meta.request_timeout_ms, Some(90_000));
|
||||
assert_eq!(meta.stream_first_byte_timeout_ms, Some(12_345));
|
||||
assert_eq!(meta.timeout, 13);
|
||||
}
|
||||
|
||||
fn tunnel_timeout_plan(stream: bool) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req-timeout".into(),
|
||||
candidate_id: None,
|
||||
provider_name: Some("provider".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://example.com/chat".into(),
|
||||
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
|
||||
stream,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
total_ms: Some(90_000),
|
||||
first_byte_ms: Some(12_345),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn tunnel_proxy_snapshot(base_url: String) -> ProxySnapshot {
|
||||
ProxySnapshot {
|
||||
enabled: Some(true),
|
||||
|
||||
@@ -81,6 +81,133 @@ async fn admin_monitoring_trace_request_returns_local_payload() {
|
||||
assert_eq!(payload["candidates"][0]["status_code"], json!(502));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_trace_request_resolves_usage_id_to_header_trace_id() {
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate(
|
||||
"cand-used",
|
||||
"trace-1",
|
||||
0,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(101),
|
||||
Some(33),
|
||||
Some(200),
|
||||
),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![sample_key()],
|
||||
));
|
||||
let mut usage = sample_usage(
|
||||
"usage-request-1",
|
||||
"provider-1",
|
||||
"OpenAI",
|
||||
40,
|
||||
0.02,
|
||||
"completed",
|
||||
Some(200),
|
||||
100,
|
||||
);
|
||||
usage.id = "usage-row-1".to_string();
|
||||
usage.candidate_id = Some("cand-used".to_string());
|
||||
usage.request_headers = Some(json!({
|
||||
"x-trace-id": "trace-1"
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||
let data_state =
|
||||
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
||||
request_candidates,
|
||||
usage_repository,
|
||||
)
|
||||
.with_provider_catalog_reader(provider_catalog);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
let context = request_context(
|
||||
http::Method::GET,
|
||||
"/api/admin/monitoring/trace/usage-row-1?attempted_only=true",
|
||||
);
|
||||
|
||||
let response = local_monitoring_response(&state, &context)
|
||||
.await
|
||||
.expect("handler should not error")
|
||||
.expect("route should be handled locally");
|
||||
|
||||
assert_eq!(response.status(), http::StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
|
||||
assert_eq!(payload["request_id"], json!("trace-1"));
|
||||
assert_eq!(payload["candidates"][0]["id"], json!("cand-used"));
|
||||
assert_eq!(
|
||||
payload["candidates"][0]["extra_data"]["first_byte_time_ms"],
|
||||
json!(30)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_trace_request_resolves_usage_request_id_to_metadata_trace_id() {
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate(
|
||||
"cand-used",
|
||||
"trace-2",
|
||||
0,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(101),
|
||||
Some(33),
|
||||
Some(200),
|
||||
),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![sample_key()],
|
||||
));
|
||||
let mut usage = sample_usage(
|
||||
"usage-request-2",
|
||||
"provider-1",
|
||||
"OpenAI",
|
||||
40,
|
||||
0.02,
|
||||
"completed",
|
||||
Some(200),
|
||||
100,
|
||||
);
|
||||
usage.candidate_id = Some("cand-used".to_string());
|
||||
usage.request_metadata = Some(json!({
|
||||
"trace_id": "trace-2"
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||
let data_state =
|
||||
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
||||
request_candidates,
|
||||
usage_repository,
|
||||
)
|
||||
.with_provider_catalog_reader(provider_catalog);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
let context = request_context(
|
||||
http::Method::GET,
|
||||
"/api/admin/monitoring/trace/usage-request-2",
|
||||
);
|
||||
|
||||
let response = local_monitoring_response(&state, &context)
|
||||
.await
|
||||
.expect("handler should not error")
|
||||
.expect("route should be handled locally");
|
||||
|
||||
assert_eq!(response.status(), http::StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
|
||||
assert_eq!(payload["request_id"], json!("trace-2"));
|
||||
assert_eq!(payload["candidates"][0]["id"], json!("cand-used"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_trace_request_returns_oauth_account_label_from_auth_config() {
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
|
||||
@@ -12,6 +12,7 @@ use aether_admin::observability::monitoring::{
|
||||
use aether_data_contracts::repository::{
|
||||
candidates::{DecisionTrace, RequestCandidateStatus},
|
||||
provider_catalog::StoredProviderCatalogKey,
|
||||
usage::StoredRequestUsageAudit,
|
||||
};
|
||||
use axum::{
|
||||
body::Body,
|
||||
@@ -21,12 +22,16 @@ use serde_json::{Map, Value};
|
||||
use std::collections::BTreeMap;
|
||||
use tracing::debug;
|
||||
|
||||
struct ResolvedAdminMonitoringTrace {
|
||||
trace: DecisionTrace,
|
||||
usage: Option<StoredRequestUsageAudit>,
|
||||
}
|
||||
|
||||
pub(super) async fn build_admin_monitoring_trace_request_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let admin_state = state;
|
||||
let state = state.as_ref();
|
||||
let Some(request_id) =
|
||||
admin_monitoring_trace_request_id_from_path(&request_context.request_path)
|
||||
else {
|
||||
@@ -39,11 +44,8 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
|
||||
Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
|
||||
};
|
||||
|
||||
let Some(trace) = state
|
||||
.data
|
||||
.read_decision_trace(&request_id, attempted_only)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
let Some(resolved) =
|
||||
resolve_admin_monitoring_trace(admin_state, &request_id, attempted_only).await?
|
||||
else {
|
||||
debug!(
|
||||
event_name = "admin_monitoring_request_trace_not_found",
|
||||
@@ -58,22 +60,113 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
|
||||
attempted_only,
|
||||
));
|
||||
};
|
||||
let usage = state
|
||||
.data
|
||||
.read_request_usage_audit(&request_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let key_accounts = build_admin_monitoring_key_account_display_map(admin_state, &trace).await?;
|
||||
let key_accounts =
|
||||
build_admin_monitoring_key_account_display_map(admin_state, &resolved.trace).await?;
|
||||
|
||||
Ok(
|
||||
build_admin_monitoring_trace_request_payload_response_with_key_accounts(
|
||||
&trace,
|
||||
usage.as_ref(),
|
||||
&resolved.trace,
|
||||
resolved.usage.as_ref(),
|
||||
&key_accounts,
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
async fn resolve_admin_monitoring_trace(
|
||||
state: &AdminAppState<'_>,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<ResolvedAdminMonitoringTrace>, GatewayError> {
|
||||
let app = state.as_ref();
|
||||
if let Some(trace) = app
|
||||
.data
|
||||
.read_decision_trace(request_id, attempted_only)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
{
|
||||
let usage = app
|
||||
.data
|
||||
.read_request_usage_audit(request_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
return Ok(Some(ResolvedAdminMonitoringTrace { trace, usage }));
|
||||
}
|
||||
|
||||
let mut usage_candidates = Vec::new();
|
||||
if let Some(usage) = app
|
||||
.data
|
||||
.read_request_usage_audit(request_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
{
|
||||
usage_candidates.push(usage);
|
||||
}
|
||||
if let Some(usage) = state.find_request_usage_by_id(request_id).await? {
|
||||
if !usage_candidates.iter().any(|item| item.id == usage.id) {
|
||||
usage_candidates.push(usage);
|
||||
}
|
||||
}
|
||||
|
||||
for usage in usage_candidates {
|
||||
for trace_request_id in admin_monitoring_usage_trace_request_ids(&usage) {
|
||||
if trace_request_id == request_id {
|
||||
continue;
|
||||
}
|
||||
if let Some(trace) = app
|
||||
.data
|
||||
.read_decision_trace(&trace_request_id, attempted_only)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
{
|
||||
return Ok(Some(ResolvedAdminMonitoringTrace {
|
||||
trace,
|
||||
usage: Some(usage),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn admin_monitoring_usage_trace_request_ids(usage: &StoredRequestUsageAudit) -> Vec<String> {
|
||||
let mut ids = Vec::new();
|
||||
push_non_empty_unique(&mut ids, usage.request_id.as_str());
|
||||
if let Some(trace_id) = usage.trace_id() {
|
||||
push_non_empty_unique(&mut ids, trace_id);
|
||||
}
|
||||
if let Some(trace_id) = usage_trace_id_from_headers(usage.request_headers.as_ref()) {
|
||||
push_non_empty_unique(&mut ids, trace_id.as_str());
|
||||
}
|
||||
if let Some(trace_id) = usage_trace_id_from_headers(usage.provider_request_headers.as_ref()) {
|
||||
push_non_empty_unique(&mut ids, trace_id.as_str());
|
||||
}
|
||||
ids
|
||||
}
|
||||
|
||||
fn usage_trace_id_from_headers(headers: Option<&Value>) -> Option<String> {
|
||||
let object = headers?.as_object()?;
|
||||
object.iter().find_map(|(key, value)| {
|
||||
key.eq_ignore_ascii_case(crate::constants::TRACE_ID_HEADER)
|
||||
.then(|| {
|
||||
value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
})
|
||||
.flatten()
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn push_non_empty_unique(values: &mut Vec<String>, value: &str) {
|
||||
let value = value.trim();
|
||||
if value.is_empty() || values.iter().any(|existing| existing == value) {
|
||||
return;
|
||||
}
|
||||
values.push(value.to_string());
|
||||
}
|
||||
|
||||
async fn build_admin_monitoring_key_account_display_map(
|
||||
state: &AdminAppState<'_>,
|
||||
trace: &DecisionTrace,
|
||||
|
||||
@@ -63,6 +63,10 @@ struct ProxyNodeRegisterRequest {
|
||||
proxy_version: Option<String>,
|
||||
#[serde(default)]
|
||||
tunnel_mode: Option<bool>,
|
||||
#[serde(default)]
|
||||
tunnel_security: Option<String>,
|
||||
#[serde(default)]
|
||||
tunnel_encryption_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -349,9 +353,21 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
|
||||
Ok(mutation) => mutation,
|
||||
Err(response) => return Ok(Some(response)),
|
||||
};
|
||||
let tunnel_encryption_key = mutation
|
||||
.proxy_metadata
|
||||
.as_ref()
|
||||
.and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::to_string);
|
||||
let Some(node) = state.register_proxy_node(&mutation).await? else {
|
||||
return Ok(Some(build_admin_proxy_nodes_data_unavailable_response()));
|
||||
};
|
||||
if let Some(key) = tunnel_encryption_key {
|
||||
state
|
||||
.app()
|
||||
.tunnel
|
||||
.register_secure_tunnel_key(node.id.clone(), key);
|
||||
}
|
||||
return Ok(Some(
|
||||
Json(json!({
|
||||
"node_id": node.id,
|
||||
@@ -1358,6 +1374,9 @@ fn build_tunnel_probe_relay_envelope(
|
||||
method: "GET".to_string(),
|
||||
url: probe_url.trim().to_string(),
|
||||
headers: std::collections::HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: timeout_secs,
|
||||
follow_redirects: Some(false),
|
||||
http1_only: false,
|
||||
@@ -1420,12 +1439,43 @@ fn validate_register_request(
|
||||
}
|
||||
validate_optional_object(input.hardware_info.as_ref(), "hardware_info")?;
|
||||
validate_optional_object(input.proxy_metadata.as_ref(), "proxy_metadata")?;
|
||||
let tunnel_security =
|
||||
normalize_optional_string(input.tunnel_security.as_deref(), "tunnel_security", 64)?;
|
||||
let tunnel_encryption_key = normalize_optional_string(
|
||||
input.tunnel_encryption_key.as_deref(),
|
||||
"tunnel_encryption_key",
|
||||
128,
|
||||
)?;
|
||||
|
||||
let registered_by = request_context
|
||||
.decision()
|
||||
.and_then(|decision| decision.admin_principal.as_ref())
|
||||
.map(|principal| principal.user_id.clone());
|
||||
|
||||
let mut proxy_metadata = input.proxy_metadata;
|
||||
if tunnel_security.as_deref()
|
||||
== Some(aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED)
|
||||
{
|
||||
let key = tunnel_encryption_key.as_deref().ok_or_else(|| {
|
||||
bad_request_response(
|
||||
"tunnel_encryption_key is required when tunnel_security=non_tls_required",
|
||||
)
|
||||
})?;
|
||||
aether_contracts::tunnel_security::decode_psk(key)
|
||||
.map_err(|err| bad_request_response(err.to_string()))?;
|
||||
let mut metadata = proxy_metadata
|
||||
.and_then(|value| value.as_object().cloned())
|
||||
.unwrap_or_default();
|
||||
metadata.insert(
|
||||
"tunnel_security".to_string(),
|
||||
json!({
|
||||
"mode": aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED,
|
||||
"encryption_key": key,
|
||||
}),
|
||||
);
|
||||
proxy_metadata = Some(Value::Object(metadata));
|
||||
}
|
||||
|
||||
Ok(
|
||||
aether_data::repository::proxy_nodes::ProxyNodeRegistrationMutation {
|
||||
name,
|
||||
@@ -1438,7 +1488,7 @@ fn validate_register_request(
|
||||
avg_latency_ms: input.avg_latency_ms,
|
||||
hardware_info: input.hardware_info,
|
||||
estimated_max_concurrency: input.estimated_max_concurrency,
|
||||
proxy_metadata: input.proxy_metadata,
|
||||
proxy_metadata,
|
||||
proxy_version: normalize_optional_string(
|
||||
input.proxy_version.as_deref(),
|
||||
"proxy_version",
|
||||
|
||||
@@ -16,7 +16,7 @@ const INSTALL_SESSION_TTL_SECS: u64 = 15 * 60;
|
||||
const INSTALL_SESSION_KEY_PREFIX: &str = "install:session:";
|
||||
const TUNNEL_INSTALL_SESSION_KEY_PREFIX: &str = "tunnel-install:session:";
|
||||
const TUNNEL_INSTALL_UNIX_SCRIPT_URL: &str =
|
||||
"https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.sh";
|
||||
"https://raw.githubusercontent.com/fawney19/Aether/refs/heads/main/apps/aether-tunnel/install.sh";
|
||||
const TUNNEL_INSTALL_POWERSHELL_SCRIPT_URL: &str =
|
||||
"https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.ps1";
|
||||
|
||||
@@ -59,6 +59,8 @@ struct StoredTunnelInstallSession {
|
||||
aether_url: String,
|
||||
management_token: String,
|
||||
node_name: String,
|
||||
tunnel_security: String,
|
||||
tunnel_encryption_key: String,
|
||||
expires_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
@@ -93,8 +95,7 @@ fn install_code_from_path(request_path: &str) -> Option<(String, bool)> {
|
||||
|
||||
fn tunnel_install_code_from_path(request_path: &str) -> Option<(String, bool)> {
|
||||
let raw = request_path
|
||||
.strip_prefix("/install-tunnel/")
|
||||
.or_else(|| request_path.strip_prefix("/install-proxy/"))?
|
||||
.strip_prefix("/install-tunnel/")?
|
||||
.trim()
|
||||
.trim_matches('/');
|
||||
if raw.is_empty() || raw.contains('/') {
|
||||
@@ -122,6 +123,17 @@ fn generate_install_code() -> String {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn generate_tunnel_encryption_key() -> String {
|
||||
use base64::Engine;
|
||||
|
||||
let first = uuid::Uuid::new_v4();
|
||||
let second = uuid::Uuid::new_v4();
|
||||
let mut key = [0_u8; 32];
|
||||
key[..16].copy_from_slice(first.as_bytes());
|
||||
key[16..].copy_from_slice(second.as_bytes());
|
||||
base64::engine::general_purpose::STANDARD.encode(key)
|
||||
}
|
||||
|
||||
fn unix_secs_now() -> u64 {
|
||||
chrono::Utc::now().timestamp().max(0) as u64
|
||||
}
|
||||
@@ -172,6 +184,8 @@ set -eu
|
||||
export AETHER_TUNNEL_AETHER_URL={aether_url}
|
||||
export AETHER_TUNNEL_MANAGEMENT_TOKEN={management_token}
|
||||
export AETHER_TUNNEL_NODE_NAME={node_name}
|
||||
export AETHER_TUNNEL_SECURITY={tunnel_security}
|
||||
export AETHER_TUNNEL_ENCRYPTION_KEY={tunnel_encryption_key}
|
||||
|
||||
if command -v curl >/dev/null 2>&1; then
|
||||
curl -fsSL {script_url} | sh
|
||||
@@ -185,6 +199,8 @@ fi
|
||||
aether_url = shell_single_quote(&session.aether_url),
|
||||
management_token = shell_single_quote(&session.management_token),
|
||||
node_name = shell_single_quote(&session.node_name),
|
||||
tunnel_security = shell_single_quote(&session.tunnel_security),
|
||||
tunnel_encryption_key = shell_single_quote(&session.tunnel_encryption_key),
|
||||
script_url = shell_single_quote(TUNNEL_INSTALL_UNIX_SCRIPT_URL),
|
||||
)
|
||||
}
|
||||
@@ -195,11 +211,15 @@ fn build_tunnel_powershell_script(session: &StoredTunnelInstallSession) -> Strin
|
||||
$env:AETHER_TUNNEL_AETHER_URL = {aether_url}
|
||||
$env:AETHER_TUNNEL_MANAGEMENT_TOKEN = {management_token}
|
||||
$env:AETHER_TUNNEL_NODE_NAME = {node_name}
|
||||
$env:AETHER_TUNNEL_SECURITY = {tunnel_security}
|
||||
$env:AETHER_TUNNEL_ENCRYPTION_KEY = {tunnel_encryption_key}
|
||||
irm {script_url} | iex
|
||||
"###,
|
||||
aether_url = powershell_single_quote(&session.aether_url),
|
||||
management_token = powershell_single_quote(&session.management_token),
|
||||
node_name = powershell_single_quote(&session.node_name),
|
||||
tunnel_security = powershell_single_quote(&session.tunnel_security),
|
||||
tunnel_encryption_key = powershell_single_quote(&session.tunnel_encryption_key),
|
||||
script_url = powershell_single_quote(TUNNEL_INSTALL_POWERSHELL_SCRIPT_URL),
|
||||
)
|
||||
}
|
||||
@@ -664,6 +684,8 @@ pub(crate) async fn build_proxy_node_install_session_response(
|
||||
aether_url: base_url_from_request(headers, request_context),
|
||||
management_token,
|
||||
node_name,
|
||||
tunnel_security: "non_tls_required".to_string(),
|
||||
tunnel_encryption_key: generate_tunnel_encryption_key(),
|
||||
expires_at_unix_secs,
|
||||
};
|
||||
let serialized = match serde_json::to_string(&session) {
|
||||
@@ -712,9 +734,7 @@ pub(super) async fn maybe_build_local_install_response(
|
||||
if decision.route_family.as_deref() != Some("install") {
|
||||
return None;
|
||||
}
|
||||
if request_context.request_path.starts_with("/install-tunnel/")
|
||||
|| request_context.request_path.starts_with("/install-proxy/")
|
||||
{
|
||||
if request_context.request_path.starts_with("/install-tunnel/") {
|
||||
return Some(maybe_build_local_tunnel_install_response(state, request_context).await);
|
||||
}
|
||||
let Some((code, wants_powershell)) = install_code_from_path(&request_context.request_path)
|
||||
@@ -893,6 +913,8 @@ mod tests {
|
||||
aether_url: "https://aether.example".to_string(),
|
||||
management_token: "ae-test-token".to_string(),
|
||||
node_name: "jp-proxy-01".to_string(),
|
||||
tunnel_security: "non_tls_required".to_string(),
|
||||
tunnel_encryption_key: "base64-32-bytes".to_string(),
|
||||
expires_at_unix_secs: u64::MAX,
|
||||
}
|
||||
}
|
||||
@@ -907,10 +929,6 @@ mod tests {
|
||||
tunnel_install_code_from_path("/install-tunnel/abc123.ps1"),
|
||||
Some(("abc123".to_string(), true))
|
||||
);
|
||||
assert_eq!(
|
||||
tunnel_install_code_from_path("/install-proxy/abc123"),
|
||||
Some(("abc123".to_string(), false))
|
||||
);
|
||||
assert_eq!(tunnel_install_code_from_path("/install-tunnel/a/b"), None);
|
||||
}
|
||||
|
||||
@@ -921,8 +939,10 @@ mod tests {
|
||||
assert!(script.contains("export AETHER_TUNNEL_AETHER_URL='https://aether.example'"));
|
||||
assert!(script.contains("export AETHER_TUNNEL_MANAGEMENT_TOKEN='ae-test-token'"));
|
||||
assert!(script.contains("export AETHER_TUNNEL_NODE_NAME='jp-proxy-01'"));
|
||||
assert!(script.contains("export AETHER_TUNNEL_SECURITY='non_tls_required'"));
|
||||
assert!(script.contains("export AETHER_TUNNEL_ENCRYPTION_KEY='base64-32-bytes'"));
|
||||
assert!(script.contains(
|
||||
"https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.sh"
|
||||
"https://raw.githubusercontent.com/fawney19/Aether/refs/heads/main/apps/aether-tunnel/install.sh"
|
||||
));
|
||||
assert!(!script.contains("aether-rust-pioneer"));
|
||||
assert!(!script.contains("[[servers]]"));
|
||||
@@ -935,6 +955,8 @@ mod tests {
|
||||
assert!(script.contains("$env:AETHER_TUNNEL_AETHER_URL = 'https://aether.example'"));
|
||||
assert!(script.contains("$env:AETHER_TUNNEL_MANAGEMENT_TOKEN = 'ae-test-token'"));
|
||||
assert!(script.contains("$env:AETHER_TUNNEL_NODE_NAME = 'jp-proxy-01'"));
|
||||
assert!(script.contains("$env:AETHER_TUNNEL_SECURITY = 'non_tls_required'"));
|
||||
assert!(script.contains("$env:AETHER_TUNNEL_ENCRYPTION_KEY = 'base64-32-bytes'"));
|
||||
assert!(script.contains(
|
||||
"https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.ps1"
|
||||
));
|
||||
|
||||
@@ -90,7 +90,6 @@ fn frontend_path_bypasses_static(path: &str) -> bool {
|
||||
|| path.starts_with("/.well-known/")
|
||||
|| path.starts_with("/install/")
|
||||
|| path.starts_with("/install-tunnel/")
|
||||
|| path.starts_with("/install-proxy/")
|
||||
|| path.starts_with("/i/")
|
||||
}
|
||||
|
||||
|
||||
@@ -1307,6 +1307,9 @@ mod tests {
|
||||
method: "GET".to_string(),
|
||||
url: "https://example.com".to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
|
||||
@@ -24,6 +24,8 @@ use super::AppState;
|
||||
|
||||
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
||||
const MAX_RELAY_META_LEN: usize = 256 * 1024;
|
||||
const MIN_RELAY_TIMEOUT_MS: u64 = 1;
|
||||
const MAX_RELAY_TIMEOUT_MS: u64 = 300_000;
|
||||
|
||||
struct StreamGuard {
|
||||
hub: std::sync::Arc<super::hub::HubRouter>,
|
||||
@@ -92,7 +94,7 @@ pub(crate) async fn open_direct_relay_stream(
|
||||
return Err(format!("connect: {error}"));
|
||||
}
|
||||
|
||||
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
|
||||
let wait_timeout = relay_header_timeout(&meta);
|
||||
let response_head = match stream.wait_headers(wait_timeout).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
@@ -148,6 +150,19 @@ fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
|
||||
let timeout_ms = if meta.stream {
|
||||
meta.stream_first_byte_timeout_ms
|
||||
.or(meta.request_timeout_ms)
|
||||
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
|
||||
} else {
|
||||
meta.request_timeout_ms
|
||||
.or(meta.stream_first_byte_timeout_ms)
|
||||
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
|
||||
};
|
||||
Duration::from_millis(timeout_ms.clamp(MIN_RELAY_TIMEOUT_MS, MAX_RELAY_TIMEOUT_MS))
|
||||
}
|
||||
|
||||
pub async fn relay_request(
|
||||
Path(node_id): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
@@ -322,7 +337,7 @@ pub async fn relay_request(
|
||||
finished: false,
|
||||
};
|
||||
|
||||
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
|
||||
let wait_timeout = relay_header_timeout(&meta);
|
||||
let response_head = match stream.wait_headers(wait_timeout).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
@@ -639,6 +654,9 @@ mod tests {
|
||||
method: "GET".to_string(),
|
||||
url: "https://example.com/health".to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
@@ -754,6 +772,9 @@ mod tests {
|
||||
method: "GET".to_string(),
|
||||
url: "https://example.com/headers".to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
|
||||
@@ -17,6 +17,7 @@ use axum::http::HeaderMap;
|
||||
use axum::response::{IntoResponse, Json};
|
||||
use axum::routing::{get, post};
|
||||
use axum::Router;
|
||||
use dashmap::DashMap;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::{data::GatewayDataState, middleware};
|
||||
@@ -34,6 +35,7 @@ pub struct AppState {
|
||||
data: Arc<GatewayDataState>,
|
||||
request_gate: Option<Arc<ConcurrencyGate>>,
|
||||
distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
|
||||
secure_tunnel_keys: Arc<DashMap<String, String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
@@ -55,9 +57,48 @@ impl AppState {
|
||||
data: Arc::new(GatewayDataState::disabled()),
|
||||
request_gate: None,
|
||||
distributed_request_gate: None,
|
||||
secure_tunnel_keys: Arc::new(DashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn register_secure_tunnel_key(
|
||||
&self,
|
||||
node_id: impl Into<String>,
|
||||
key: impl Into<String>,
|
||||
) {
|
||||
self.secure_tunnel_keys.insert(node_id.into(), key.into());
|
||||
}
|
||||
|
||||
pub(crate) fn secure_tunnel_key(&self, node_id: &str) -> Option<String> {
|
||||
self.secure_tunnel_keys
|
||||
.get(node_id)
|
||||
.map(|entry| entry.value().clone())
|
||||
}
|
||||
|
||||
async fn secure_tunnel_key_for_node(&self, node_id: &str) -> Option<String> {
|
||||
if let Some(key) = self.secure_tunnel_key(node_id) {
|
||||
return Some(key);
|
||||
}
|
||||
let key = self
|
||||
.data
|
||||
.find_proxy_node(node_id)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|node| {
|
||||
node.proxy_metadata.and_then(|metadata| {
|
||||
metadata
|
||||
.pointer("/tunnel_security/encryption_key")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::to_string)
|
||||
})
|
||||
});
|
||||
if let Some(key) = key.as_ref() {
|
||||
self.register_secure_tunnel_key(node_id.to_string(), key.clone());
|
||||
}
|
||||
key
|
||||
}
|
||||
|
||||
pub(crate) fn with_data(mut self, data: Arc<GatewayDataState>) -> Self {
|
||||
self.data = data;
|
||||
self
|
||||
@@ -212,11 +253,47 @@ pub async fn ws_proxy(
|
||||
|
||||
let max_streams = resolve_proxy_max_streams(&headers, state.max_streams);
|
||||
let protocol_version = resolve_proxy_protocol_version(&headers);
|
||||
let tunnel_security = headers
|
||||
.get(aether_contracts::tunnel_security::TUNNEL_SECURITY_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
let security_session = headers
|
||||
.get(aether_contracts::tunnel_security::TUNNEL_SECURITY_SESSION_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
|
||||
if node_id.is_empty() {
|
||||
warn!("proxy connection rejected: missing X-Node-ID header");
|
||||
return axum::http::StatusCode::BAD_REQUEST.into_response();
|
||||
}
|
||||
let stored_security_key = state.secure_tunnel_key_for_node(&node_id).await;
|
||||
let (security_key, security_session) = match tunnel_security.as_deref() {
|
||||
Some(aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED) => {
|
||||
match stored_security_key {
|
||||
Some(key) => {
|
||||
let Some(session) = security_session else {
|
||||
warn!(node_id = %node_id, "secure tunnel requested without a security session");
|
||||
return axum::http::StatusCode::BAD_REQUEST.into_response();
|
||||
};
|
||||
(Some(key), session)
|
||||
}
|
||||
None => {
|
||||
warn!(node_id = %node_id, "secure tunnel requested but no PSK is registered");
|
||||
return axum::http::StatusCode::UNAUTHORIZED.into_response();
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(_) => return axum::http::StatusCode::BAD_REQUEST.into_response(),
|
||||
None if stored_security_key.is_some() => {
|
||||
warn!(node_id = %node_id, "proxy connection rejected: stored secure tunnel key requires encrypted frames");
|
||||
return axum::http::StatusCode::UNAUTHORIZED.into_response();
|
||||
}
|
||||
None => (None, String::new()),
|
||||
};
|
||||
|
||||
let request_permit = match state.try_acquire_request_permit().await {
|
||||
Ok(permit) => permit,
|
||||
@@ -253,6 +330,8 @@ pub async fn ws_proxy(
|
||||
node_name,
|
||||
max_streams,
|
||||
protocol_version,
|
||||
security_key,
|
||||
security_session,
|
||||
state.proxy_conn_cfg,
|
||||
)
|
||||
.await
|
||||
|
||||
@@ -13,6 +13,8 @@ use tracing::{debug, info, warn};
|
||||
|
||||
use super::hub::{ConnConfig, HubRouter, ProxyConn, SendStatus};
|
||||
use super::protocol;
|
||||
use aether_contracts::tunnel::Frame;
|
||||
use aether_contracts::tunnel_security::{SecureFrameCodec, TunnelSecurityRole};
|
||||
|
||||
/// Maximum single frame size: 64 MB
|
||||
const MAX_FRAME_SIZE: usize = 64 * 1024 * 1024;
|
||||
@@ -24,6 +26,8 @@ pub async fn handle_proxy_connection(
|
||||
node_name: String,
|
||||
max_streams: usize,
|
||||
protocol_version: u8,
|
||||
security_key: Option<String>,
|
||||
security_session: String,
|
||||
cfg: ConnConfig,
|
||||
) {
|
||||
let conn_id = hub.alloc_conn_id();
|
||||
@@ -31,6 +35,18 @@ pub async fn handle_proxy_connection(
|
||||
|
||||
let (tx, mut rx) = bounded_queue::<Message>(cfg.outbound_queue_capacity);
|
||||
let (close_tx, mut close_rx) = watch::channel(false);
|
||||
let security = match security_key.as_deref() {
|
||||
Some(key) => {
|
||||
match SecureFrameCodec::new(key, &security_session, TunnelSecurityRole::Server) {
|
||||
Ok(codec) => Some(Arc::new(codec)),
|
||||
Err(error) => {
|
||||
warn!(conn_id, node_id = %node_id, error = %error, "secure tunnel codec initialization failed");
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
||||
let conn = Arc::new(ProxyConn::new(
|
||||
conn_id,
|
||||
@@ -46,6 +62,7 @@ pub async fn handle_proxy_connection(
|
||||
|
||||
let writer_conn_id = conn_id;
|
||||
let writer_conn = conn.clone();
|
||||
let writer_security = security.clone();
|
||||
let writer = tokio::spawn(async move {
|
||||
let mut frames_sent: u64 = 0;
|
||||
loop {
|
||||
@@ -58,6 +75,13 @@ pub async fn handle_proxy_connection(
|
||||
_ => 0,
|
||||
};
|
||||
let send_started_at = std::time::Instant::now();
|
||||
let msg = match encrypt_message(msg, writer_security.as_deref()) {
|
||||
Ok(msg) => msg,
|
||||
Err(error) => {
|
||||
warn!(conn_id = writer_conn_id, error = %error, "failed to encrypt outbound proxy frame");
|
||||
break;
|
||||
}
|
||||
};
|
||||
let send_result = tokio::time::timeout(
|
||||
Duration::from_secs(15),
|
||||
ws_tx.send(msg),
|
||||
@@ -194,7 +218,7 @@ pub async fn handle_proxy_connection(
|
||||
let reader_hub = hub.clone();
|
||||
let reader_conn = conn.clone();
|
||||
let reader = tokio::spawn(async move {
|
||||
run_proxy_reader(ws_rx, reader_hub, reader_conn, cfg.idle_timeout).await;
|
||||
run_proxy_reader(ws_rx, reader_hub, reader_conn, cfg.idle_timeout, security).await;
|
||||
});
|
||||
|
||||
let _ = reader.await;
|
||||
@@ -223,6 +247,7 @@ async fn run_proxy_reader(
|
||||
hub: Arc<HubRouter>,
|
||||
conn: Arc<ProxyConn>,
|
||||
idle_timeout: Duration,
|
||||
security: Option<Arc<SecureFrameCodec>>,
|
||||
) {
|
||||
let idle_enabled = !idle_timeout.is_zero();
|
||||
let mut oversized_count = 0u32;
|
||||
@@ -245,7 +270,14 @@ async fn run_proxy_reader(
|
||||
match msg {
|
||||
Some(Ok(Message::Binary(data))) => {
|
||||
frames_received += 1;
|
||||
let mut data = data.to_vec();
|
||||
let mut data = match decrypt_message(data, security.as_deref()) {
|
||||
Ok(data) => data,
|
||||
Err(error) => {
|
||||
warn!(conn_id = conn.id, error = %error, "failed to decrypt secure proxy frame");
|
||||
conn.request_close();
|
||||
break;
|
||||
}
|
||||
};
|
||||
if data.len() > MAX_FRAME_SIZE {
|
||||
oversized_count += 1;
|
||||
warn!(
|
||||
@@ -294,3 +326,33 @@ async fn run_proxy_reader(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn encrypt_message(
|
||||
msg: Message,
|
||||
security: Option<&SecureFrameCodec>,
|
||||
) -> Result<Message, aether_contracts::tunnel_security::TunnelSecurityError> {
|
||||
let Some(codec) = security else {
|
||||
return Ok(msg);
|
||||
};
|
||||
match msg {
|
||||
Message::Binary(data) => {
|
||||
let frame = Frame::decode(bytes::Bytes::from(data.to_vec()))
|
||||
.map_err(|_| aether_contracts::tunnel_security::TunnelSecurityError::Encrypt)?;
|
||||
Ok(Message::Binary(codec.encrypt_frame(frame)?))
|
||||
}
|
||||
other => Ok(other),
|
||||
}
|
||||
}
|
||||
|
||||
fn decrypt_message(
|
||||
data: bytes::Bytes,
|
||||
security: Option<&SecureFrameCodec>,
|
||||
) -> Result<Vec<u8>, aether_contracts::tunnel_security::TunnelSecurityError> {
|
||||
let Some(codec) = security else {
|
||||
return Ok(data.to_vec());
|
||||
};
|
||||
let frame = Frame::decode(data)
|
||||
.map_err(|_| aether_contracts::tunnel_security::TunnelSecurityError::Decrypt)?;
|
||||
let frame = codec.decrypt_frame(frame)?;
|
||||
Ok(frame.encode().to_vec())
|
||||
}
|
||||
|
||||
@@ -455,6 +455,14 @@ impl EmbeddedTunnelState {
|
||||
self.inner.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn register_secure_tunnel_key(
|
||||
&self,
|
||||
node_id: impl Into<String>,
|
||||
key: impl Into<String>,
|
||||
) {
|
||||
self.inner.register_secure_tunnel_key(node_id, key);
|
||||
}
|
||||
|
||||
pub(crate) fn has_local_proxy(&self, node_id: &str) -> bool {
|
||||
self.inner.hub.has_local_proxy(node_id)
|
||||
}
|
||||
@@ -511,6 +519,9 @@ impl EmbeddedTunnelState {
|
||||
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,
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
$ErrorActionPreference = 'Stop'
|
||||
|
||||
if ($env:AETHER_PROXY_AETHER_URL -and -not $env:AETHER_TUNNEL_AETHER_URL) {
|
||||
$env:AETHER_TUNNEL_AETHER_URL = $env:AETHER_PROXY_AETHER_URL
|
||||
}
|
||||
if ($env:AETHER_PROXY_MANAGEMENT_TOKEN -and -not $env:AETHER_TUNNEL_MANAGEMENT_TOKEN) {
|
||||
$env:AETHER_TUNNEL_MANAGEMENT_TOKEN = $env:AETHER_PROXY_MANAGEMENT_TOKEN
|
||||
}
|
||||
if ($env:AETHER_PROXY_NODE_NAME -and -not $env:AETHER_TUNNEL_NODE_NAME) {
|
||||
$env:AETHER_TUNNEL_NODE_NAME = $env:AETHER_PROXY_NODE_NAME
|
||||
}
|
||||
|
||||
irm 'https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.ps1' | iex
|
||||
@@ -1,21 +0,0 @@
|
||||
#!/bin/sh
|
||||
set -eu
|
||||
|
||||
if [ -n "${AETHER_PROXY_AETHER_URL:-}" ] && [ -z "${AETHER_TUNNEL_AETHER_URL:-}" ]; then
|
||||
export AETHER_TUNNEL_AETHER_URL="${AETHER_PROXY_AETHER_URL}"
|
||||
fi
|
||||
if [ -n "${AETHER_PROXY_MANAGEMENT_TOKEN:-}" ] && [ -z "${AETHER_TUNNEL_MANAGEMENT_TOKEN:-}" ]; then
|
||||
export AETHER_TUNNEL_MANAGEMENT_TOKEN="${AETHER_PROXY_MANAGEMENT_TOKEN}"
|
||||
fi
|
||||
if [ -n "${AETHER_PROXY_NODE_NAME:-}" ] && [ -z "${AETHER_TUNNEL_NODE_NAME:-}" ]; then
|
||||
export AETHER_TUNNEL_NODE_NAME="${AETHER_PROXY_NODE_NAME}"
|
||||
fi
|
||||
|
||||
if command -v curl >/dev/null 2>&1; then
|
||||
curl -fsSL 'https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.sh' | sh
|
||||
elif command -v wget >/dev/null 2>&1; then
|
||||
wget -qO- 'https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.sh' | sh
|
||||
else
|
||||
printf '%s\n' "[Aether Tunnel] 需要 curl 或 wget 下载安装脚本" >&2
|
||||
exit 1
|
||||
fi
|
||||
@@ -7,6 +7,10 @@ AETHER_TUNNEL_MANAGEMENT_TOKEN=ae_xxxxx
|
||||
# Node identification
|
||||
AETHER_TUNNEL_NODE_NAME=jp-proxy-01
|
||||
|
||||
# Secure non-TLS tunnel. Set non_tls_required with a key to enable secure tunnel.
|
||||
AETHER_TUNNEL_SECURITY=off
|
||||
# AETHER_TUNNEL_ENCRYPTION_KEY=base64-32-bytes
|
||||
|
||||
# Maximum request body buffered for 307/308 replay (supports K/M/G, 0 disables body replay buffering)
|
||||
AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES=5M
|
||||
|
||||
|
||||
@@ -41,6 +41,7 @@ tar = "0.4"
|
||||
socket2 = { version = "0.5", features = ["all"] }
|
||||
tower-service = "0.3"
|
||||
webpki-roots = "0.26"
|
||||
uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
aether-gateway.workspace = true
|
||||
|
||||
@@ -53,6 +53,7 @@ curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tu
|
||||
AETHER_TUNNEL_AETHER_URL="https://aether.example.com" \
|
||||
AETHER_TUNNEL_MANAGEMENT_TOKEN="ae_xxx" \
|
||||
AETHER_TUNNEL_NODE_NAME="jp-proxy-01" \
|
||||
AETHER_TUNNEL_SECURITY="off" \
|
||||
sh
|
||||
```
|
||||
|
||||
@@ -60,6 +61,7 @@ curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tu
|
||||
$env:AETHER_TUNNEL_AETHER_URL = "https://aether.example.com"
|
||||
$env:AETHER_TUNNEL_MANAGEMENT_TOKEN = "ae_xxx"
|
||||
$env:AETHER_TUNNEL_NODE_NAME = "jp-proxy-01"
|
||||
$env:AETHER_TUNNEL_SECURITY = "off"
|
||||
irm https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.ps1 | iex
|
||||
```
|
||||
|
||||
@@ -111,6 +113,8 @@ sudo aether-tunnel uninstall
|
||||
| `--aether-url` | `AETHER_TUNNEL_AETHER_URL` | **必填** | Aether 服务器地址 |
|
||||
| `--management-token` | `AETHER_TUNNEL_MANAGEMENT_TOKEN` | **必填** | 管理员 Token(`ae_xxx` 格式) |
|
||||
| `--node-name` | `AETHER_TUNNEL_NODE_NAME` | **必填** | 节点名称标识 |
|
||||
| `--tunnel-security` | `AETHER_TUNNEL_SECURITY` | `off` | Aether ↔ tunnel 通道安全模式;支持 `off` / `non_tls_required`;在 `[[servers]]` 中省略该字段且 `http://` 提供 key 时会自动按 `non_tls_required` 生效 |
|
||||
| `--tunnel-encryption-key` | `AETHER_TUNNEL_ENCRYPTION_KEY` | 空 | secure tunnel 使用的长期 PSK(base64 32-byte),每个 `[[servers]]` 节点独立配置 |
|
||||
| `--public-ip` | `AETHER_TUNNEL_PUBLIC_IP` | 自动检测 | 公网 IP |
|
||||
| `--node-region` | `AETHER_TUNNEL_NODE_REGION` | 自动检测 | 地区标识 |
|
||||
| `--heartbeat-interval` | `AETHER_TUNNEL_HEARTBEAT_INTERVAL` | `5` | 心跳间隔(秒) |
|
||||
@@ -216,13 +220,19 @@ tunnel 会在心跳兼容字段 `proxy_metadata` 中主动上报隧道稳定性
|
||||
aether_url = "https://aether-1.example.com"
|
||||
management_token = "ae_xxx"
|
||||
node_name = "jp-proxy-01"
|
||||
tunnel_security = "off"
|
||||
|
||||
[[servers]]
|
||||
aether_url = "https://aether-2.example.com"
|
||||
aether_url = "http://aether-2.example.com"
|
||||
management_token = "ae_yyy"
|
||||
node_name = "jp-proxy-02"
|
||||
tunnel_encryption_key = "base64-32-bytes"
|
||||
```
|
||||
|
||||
`tunnel_security = "non_tls_required"` 是非 TLS secure tunnel 的 MVP 配置面:它要求同时提供当前 `[[servers]]` 条目的 `tunnel_encryption_key`,后续握手使用 `node_name` / `X-Node-Id` 查找对应 PSK,不引入 `tunnel_encryption_key_id`。`wss://` 仍是推荐方案;`ws:// + secure tunnel` 只加密注册完成后的 WebSocket tunnel frame,不保护安装脚本、注册请求、`management_token` 或 PSK 的首次分发;这些 bootstrap 凭据仍必须通过 HTTPS 或其他可信通道交付。它不等价于 HTTPS 伪装,也不覆盖 tunnel ↔ origin/provider 这段链路。
|
||||
|
||||
如果 `aether_url` 使用 `http://` 且当前 `[[servers]]` 条目提供了 `tunnel_encryption_key`,省略 `tunnel_security` 时运行时会自动按 `non_tls_required` 生效;显式配置 `tunnel_security = "off"` 会关闭该自动推断。secure tunnel 会在 WebSocket tunnel 上加密所有二进制 tunnel frame;未配置 key 或显式关闭的旧节点仍按原明文协议工作。
|
||||
|
||||
## 发布新版本
|
||||
|
||||
推送 `tunnel-v*` 格式的 tag,GitHub Actions 会自动:
|
||||
|
||||
@@ -105,7 +105,7 @@ function Test-ServerExists([string]$Path, [string]$QuotedUrl, [string]$QuotedNam
|
||||
return ($FoundUrl -and $FoundName)
|
||||
}
|
||||
|
||||
function Add-ServerConfig([string]$AetherUrl, [string]$ManagementToken, [string]$NodeName) {
|
||||
function Add-ServerConfig([string]$AetherUrl, [string]$ManagementToken, [string]$NodeName, [string]$TunnelSecurity, [string]$TunnelEncryptionKey) {
|
||||
$ConfigDir = Split-Path -Parent $script:ConfigPath
|
||||
New-Item -ItemType Directory -Force -Path $ConfigDir | Out-Null
|
||||
|
||||
@@ -116,6 +116,7 @@ function Add-ServerConfig([string]$AetherUrl, [string]$ManagementToken, [string]
|
||||
$QuotedUrl = ConvertTo-TomlQuotedString $AetherUrl
|
||||
$QuotedToken = ConvertTo-TomlQuotedString $ManagementToken
|
||||
$QuotedName = ConvertTo-TomlQuotedString $NodeName
|
||||
$QuotedTunnelEncryptionKey = ConvertTo-TomlQuotedString $TunnelEncryptionKey
|
||||
|
||||
if (Test-ServerExists $script:ConfigPath $QuotedUrl $QuotedName) {
|
||||
Say "Same aether_url + node_name already exists, skipping config append: $script:ConfigPath"
|
||||
@@ -134,6 +135,13 @@ function Add-ServerConfig([string]$AetherUrl, [string]$ManagementToken, [string]
|
||||
"management_token = $QuotedToken",
|
||||
"node_name = $QuotedName"
|
||||
) -join "`n"
|
||||
if ($TunnelSecurity) {
|
||||
$QuotedTunnelSecurity = ConvertTo-TomlQuotedString $TunnelSecurity
|
||||
$Block += "`ntunnel_security = $QuotedTunnelSecurity"
|
||||
}
|
||||
if ($TunnelEncryptionKey) {
|
||||
$Block += "`ntunnel_encryption_key = $QuotedTunnelEncryptionKey"
|
||||
}
|
||||
Add-Content -Path $script:ConfigPath -Value ($Block + "`n") -Encoding UTF8
|
||||
Say "Appended [[servers]] to: $script:ConfigPath"
|
||||
}
|
||||
@@ -143,13 +151,21 @@ function Main {
|
||||
$AetherUrl = Prompt-IfEmpty 'AETHER_TUNNEL_AETHER_URL' $env:AETHER_TUNNEL_AETHER_URL 'Aether URL'
|
||||
$ManagementToken = Prompt-IfEmpty 'AETHER_TUNNEL_MANAGEMENT_TOKEN' $env:AETHER_TUNNEL_MANAGEMENT_TOKEN 'Management token (ae_xxx)'
|
||||
$NodeName = Prompt-IfEmpty 'AETHER_TUNNEL_NODE_NAME' $env:AETHER_TUNNEL_NODE_NAME 'Node name'
|
||||
$TunnelSecurity = if ($env:AETHER_TUNNEL_SECURITY) { $env:AETHER_TUNNEL_SECURITY } else { '' }
|
||||
$TunnelEncryptionKey = if ($env:AETHER_TUNNEL_ENCRYPTION_KEY) { $env:AETHER_TUNNEL_ENCRYPTION_KEY } else { '' }
|
||||
if ($TunnelSecurity -and ($TunnelSecurity -notin @('off', 'non_tls_required'))) {
|
||||
Fail 'AETHER_TUNNEL_SECURITY must be off or non_tls_required'
|
||||
}
|
||||
if (($TunnelSecurity -eq 'non_tls_required') -and -not $TunnelEncryptionKey) {
|
||||
Fail 'AETHER_TUNNEL_ENCRYPTION_KEY is required when AETHER_TUNNEL_SECURITY=non_tls_required'
|
||||
}
|
||||
|
||||
$TempDir = Join-Path ([IO.Path]::GetTempPath()) ("aether-tunnel-" + [Guid]::NewGuid().ToString('N'))
|
||||
New-Item -ItemType Directory -Force -Path $TempDir | Out-Null
|
||||
try {
|
||||
$Tag = Resolve-LatestTunnelTag
|
||||
Install-AetherTunnelBinary $Tag $TempDir
|
||||
Add-ServerConfig $AetherUrl $ManagementToken $NodeName
|
||||
Add-ServerConfig $AetherUrl $ManagementToken $NodeName $TunnelSecurity $TunnelEncryptionKey
|
||||
} finally {
|
||||
Remove-Item -Recurse -Force $TempDir -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
@@ -186,11 +186,17 @@ append_server_config() {
|
||||
aether_url="$1"
|
||||
management_token="$2"
|
||||
node_name="$3"
|
||||
tunnel_security="$4"
|
||||
tunnel_encryption_key="$5"
|
||||
|
||||
mkdir -p "$(dirname "$CONFIG_PATH")"
|
||||
quoted_url=$(toml_quote "$aether_url")
|
||||
quoted_token=$(toml_quote "$management_token")
|
||||
quoted_name=$(toml_quote "$node_name")
|
||||
quoted_encryption_key=$(toml_quote "$tunnel_encryption_key")
|
||||
if [ -n "$tunnel_security" ]; then
|
||||
quoted_security=$(toml_quote "$tunnel_security")
|
||||
fi
|
||||
|
||||
if has_legacy_single_server_keys; then
|
||||
fail "现有配置仍使用旧的顶层 aether_url/management_token,请先运行 aether-tunnel setup 迁移为 [[servers]] 后重试:$CONFIG_PATH"
|
||||
@@ -214,6 +220,12 @@ append_server_config() {
|
||||
printf 'aether_url = %s\n' "$quoted_url"
|
||||
printf 'management_token = %s\n' "$quoted_token"
|
||||
printf 'node_name = %s\n' "$quoted_name"
|
||||
if [ -n "$tunnel_security" ]; then
|
||||
printf 'tunnel_security = %s\n' "$quoted_security"
|
||||
fi
|
||||
if [ -n "$tunnel_encryption_key" ]; then
|
||||
printf 'tunnel_encryption_key = %s\n' "$quoted_encryption_key"
|
||||
fi
|
||||
} >> "$CONFIG_PATH"
|
||||
chmod 600 "$CONFIG_PATH" 2>/dev/null || true
|
||||
say "已追加 [[servers]] 到:$CONFIG_PATH"
|
||||
@@ -227,12 +239,21 @@ main() {
|
||||
aether_url=$(prompt_if_empty AETHER_TUNNEL_AETHER_URL "${AETHER_TUNNEL_AETHER_URL:-}" "Aether URL: ")
|
||||
management_token=$(prompt_if_empty AETHER_TUNNEL_MANAGEMENT_TOKEN "${AETHER_TUNNEL_MANAGEMENT_TOKEN:-}" "Management token (ae_xxx): ")
|
||||
node_name=$(prompt_if_empty AETHER_TUNNEL_NODE_NAME "${AETHER_TUNNEL_NODE_NAME:-}" "Node name: ")
|
||||
tunnel_security="${AETHER_TUNNEL_SECURITY:-}"
|
||||
tunnel_encryption_key="${AETHER_TUNNEL_ENCRYPTION_KEY:-}"
|
||||
case "$tunnel_security" in
|
||||
""|off|non_tls_required) ;;
|
||||
*) fail "AETHER_TUNNEL_SECURITY 必须是 off 或 non_tls_required" ;;
|
||||
esac
|
||||
if [ "$tunnel_security" = "non_tls_required" ] && [ -z "$tunnel_encryption_key" ]; then
|
||||
fail "AETHER_TUNNEL_SECURITY=non_tls_required 时必须设置 AETHER_TUNNEL_ENCRYPTION_KEY"
|
||||
fi
|
||||
|
||||
tag=$(resolve_latest_tunnel_tag)
|
||||
[ -n "$tag" ] || fail "没有找到可用的 tunnel-v* release"
|
||||
asset=$(detect_asset)
|
||||
install_binary "$tag" "$asset"
|
||||
append_server_config "$aether_url" "$management_token" "$node_name"
|
||||
append_server_config "$aether_url" "$management_token" "$node_name" "$tunnel_security" "$tunnel_encryption_key"
|
||||
|
||||
say "完成。运行以下命令启动/配置服务:"
|
||||
say " $INSTALL_DIR/aether-tunnel setup $CONFIG_PATH"
|
||||
|
||||
@@ -18,7 +18,10 @@ use tokio::sync::{watch, Mutex};
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::config::{Config, ServerEntry, TunnelPoolSizing};
|
||||
use crate::config::{
|
||||
effective_tunnel_security, validate_tunnel_encryption_key, Config, ServerEntry,
|
||||
TunnelPoolSizing,
|
||||
};
|
||||
use crate::net;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::{self, DynamicConfig};
|
||||
@@ -212,8 +215,31 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
||||
&entry.aether_url,
|
||||
&entry.management_token,
|
||||
));
|
||||
let tunnel_security = effective_tunnel_security(
|
||||
&entry.aether_url,
|
||||
entry.tunnel_security,
|
||||
entry.tunnel_encryption_key.as_deref(),
|
||||
);
|
||||
if tunnel_security == crate::config::TunnelSecurity::NonTlsRequired {
|
||||
if entry.aether_url.trim_start().starts_with("http://") {
|
||||
warn!(
|
||||
server = %label,
|
||||
url = %entry.aether_url,
|
||||
"secure tunnel frame encryption starts after registration; deliver install and registration credentials over HTTPS or another trusted bootstrap channel"
|
||||
);
|
||||
}
|
||||
let key = entry
|
||||
.tunnel_encryption_key
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!("tunnel_encryption_key must be set for secure non-TLS tunnel")
|
||||
})?;
|
||||
validate_tunnel_encryption_key(key)?;
|
||||
}
|
||||
match client
|
||||
.register(&config, &node_name, &public_ip, Some(&hw_info))
|
||||
.register(&config, entry, &node_name, &public_ip, Some(&hw_info))
|
||||
.await
|
||||
{
|
||||
Ok(node_id) => {
|
||||
@@ -603,7 +629,13 @@ async fn retry_failed_registration(
|
||||
}
|
||||
|
||||
match client
|
||||
.register(&state.config, &node_name, &public_ip, Some(&hw_info))
|
||||
.register(
|
||||
&state.config,
|
||||
&entry,
|
||||
&node_name,
|
||||
&public_ip,
|
||||
Some(&hw_info),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(node_id) => {
|
||||
@@ -681,6 +713,12 @@ fn build_server_context(
|
||||
server_label: label.to_string(),
|
||||
aether_url: entry.aether_url.clone(),
|
||||
management_token: entry.management_token.clone(),
|
||||
tunnel_security: effective_tunnel_security(
|
||||
&entry.aether_url,
|
||||
entry.tunnel_security,
|
||||
entry.tunnel_encryption_key.as_deref(),
|
||||
),
|
||||
tunnel_encryption_key: entry.tunnel_encryption_key.clone(),
|
||||
node_name: node_name.to_string(),
|
||||
node_id: Arc::new(RwLock::new(node_id)),
|
||||
aether_client: client,
|
||||
@@ -987,6 +1025,8 @@ mod tests {
|
||||
aether_url: gateway_base_url.clone(),
|
||||
management_token: "token".to_string(),
|
||||
node_name: Some("node-recovery".to_string()),
|
||||
tunnel_security: None,
|
||||
tunnel_encryption_key: None,
|
||||
},
|
||||
)];
|
||||
let (shutdown_tx, shutdown_rx) = watch::channel(false);
|
||||
@@ -1289,6 +1329,8 @@ mod tests {
|
||||
aether_url: state.config.aether_url.clone(),
|
||||
management_token: state.config.management_token.clone(),
|
||||
node_name: Some(state.config.node_name.clone()),
|
||||
tunnel_security: Some(state.config.tunnel_security),
|
||||
tunnel_encryption_key: state.config.tunnel_encryption_key.clone(),
|
||||
};
|
||||
let client = Arc::new(AetherClient::new(
|
||||
&state.config,
|
||||
@@ -1311,6 +1353,8 @@ mod tests {
|
||||
management_token: "token".to_string(),
|
||||
public_ip: None,
|
||||
node_name: "tunnel-test".to_string(),
|
||||
tunnel_security: crate::config::TunnelSecurity::Off,
|
||||
tunnel_encryption_key: None,
|
||||
node_region: None,
|
||||
heartbeat_interval: 1,
|
||||
allowed_ports: vec![80, 443],
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::fmt;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::Path;
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_runtime::{FileLoggingConfig, LogDestination, LogRotation, ServiceRuntimeConfig};
|
||||
@@ -247,6 +248,63 @@ impl From<TunnelLogRotationArg> for LogRotation {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TunnelSecurity {
|
||||
Off,
|
||||
NonTlsRequired,
|
||||
}
|
||||
|
||||
impl fmt::Display for TunnelSecurity {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(match self {
|
||||
TunnelSecurity::Off => "off",
|
||||
TunnelSecurity::NonTlsRequired => "non_tls_required",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for TunnelSecurity {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(value: &str) -> Result<Self, Self::Err> {
|
||||
match value.trim() {
|
||||
"off" => Ok(Self::Off),
|
||||
"non_tls_required" | "non-tls-required" => Ok(Self::NonTlsRequired),
|
||||
other => Err(format!(
|
||||
"invalid tunnel_security {other:?}; expected off or non_tls_required"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate_tunnel_encryption_key(value: &str) -> anyhow::Result<()> {
|
||||
aether_contracts::tunnel_security::decode_psk(value)
|
||||
.map(|_| ())
|
||||
.map_err(|err| anyhow::anyhow!(err))
|
||||
}
|
||||
|
||||
pub fn effective_tunnel_security(
|
||||
aether_url: &str,
|
||||
configured: Option<TunnelSecurity>,
|
||||
tunnel_encryption_key: Option<&str>,
|
||||
) -> TunnelSecurity {
|
||||
match configured {
|
||||
Some(TunnelSecurity::NonTlsRequired) => return TunnelSecurity::NonTlsRequired,
|
||||
Some(TunnelSecurity::Off) => return TunnelSecurity::Off,
|
||||
None => {}
|
||||
}
|
||||
if aether_url.trim_start().starts_with("http://")
|
||||
&& tunnel_encryption_key
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.is_some()
|
||||
{
|
||||
return TunnelSecurity::NonTlsRequired;
|
||||
}
|
||||
TunnelSecurity::Off
|
||||
}
|
||||
|
||||
/// Aether tunnel agent.
|
||||
///
|
||||
/// Deployed on overseas VPS to relay API traffic for Aether instances
|
||||
@@ -271,6 +329,18 @@ pub struct Config {
|
||||
#[arg(long, env = "AETHER_TUNNEL_NODE_NAME")]
|
||||
pub node_name: String,
|
||||
|
||||
/// Application-layer tunnel security mode.
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_TUNNEL_SECURITY",
|
||||
default_value_t = TunnelSecurity::Off
|
||||
)]
|
||||
pub tunnel_security: TunnelSecurity,
|
||||
|
||||
/// Base64-encoded 32-byte PSK used when tunnel_security=non_tls_required.
|
||||
#[arg(long, env = "AETHER_TUNNEL_ENCRYPTION_KEY")]
|
||||
pub tunnel_encryption_key: Option<String>,
|
||||
|
||||
/// Region label (e.g. ap-northeast-1)
|
||||
#[arg(long, env = "AETHER_TUNNEL_NODE_REGION")]
|
||||
pub node_region: Option<String>,
|
||||
@@ -670,6 +740,14 @@ impl Config {
|
||||
if self.node_name.trim().is_empty() {
|
||||
anyhow::bail!("node_name must not be empty");
|
||||
}
|
||||
if self.tunnel_security == TunnelSecurity::NonTlsRequired {
|
||||
let Some(key) = normalized_proxy_url(&self.tunnel_encryption_key) else {
|
||||
anyhow::bail!(
|
||||
"tunnel_encryption_key must be set when tunnel_security=non_tls_required"
|
||||
);
|
||||
};
|
||||
validate_tunnel_encryption_key(key)?;
|
||||
}
|
||||
for &port in &self.allowed_ports {
|
||||
if port == 0 {
|
||||
anyhow::bail!("allowed_ports: port 0 is not valid");
|
||||
@@ -896,6 +974,12 @@ pub struct ServerEntry {
|
||||
pub management_token: String,
|
||||
/// Per-server node name override. Falls back to the global `node_name`.
|
||||
pub node_name: Option<String>,
|
||||
/// Per-server tunnel security mode.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tunnel_security: Option<TunnelSecurity>,
|
||||
/// Per-server PSK for secure non-TLS tunnel handshakes.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tunnel_encryption_key: Option<String>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -1066,6 +1150,9 @@ impl ConfigFile {
|
||||
let first_server = self.servers.first();
|
||||
let aether_url = first_server.map(|s| s.aether_url.as_str());
|
||||
let management_token = first_server.map(|s| s.management_token.as_str());
|
||||
let tunnel_security =
|
||||
first_server.map(|s| s.tunnel_security.unwrap_or(TunnelSecurity::Off));
|
||||
let tunnel_encryption_key = first_server.and_then(|s| s.tunnel_encryption_key.as_deref());
|
||||
let node_name = self
|
||||
.node_name
|
||||
.as_deref()
|
||||
@@ -1073,6 +1160,8 @@ impl ConfigFile {
|
||||
|
||||
set!("AETHER_TUNNEL_AETHER_URL", aether_url);
|
||||
set!("AETHER_TUNNEL_MANAGEMENT_TOKEN", management_token);
|
||||
set!("AETHER_TUNNEL_SECURITY", tunnel_security);
|
||||
set!("AETHER_TUNNEL_ENCRYPTION_KEY", tunnel_encryption_key);
|
||||
set!("AETHER_TUNNEL_PUBLIC_IP", self.public_ip);
|
||||
set!("AETHER_TUNNEL_NODE_NAME", node_name);
|
||||
set!("AETHER_TUNNEL_NODE_REGION", self.node_region);
|
||||
@@ -1395,6 +1484,31 @@ tunnel_ipv6_only = false
|
||||
assert_eq!(cfg.tunnel_ipv6_only, Some(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_deserializes_server_tunnel_security_fields() {
|
||||
let cfg: ConfigFile = toml::from_str(
|
||||
r#"
|
||||
[[servers]]
|
||||
aether_url = "http://aether.example.com"
|
||||
management_token = "ae_test"
|
||||
node_name = "jp-proxy-01"
|
||||
tunnel_security = "non_tls_required"
|
||||
tunnel_encryption_key = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc="
|
||||
"#,
|
||||
)
|
||||
.expect("server tunnel security TOML");
|
||||
|
||||
assert_eq!(cfg.servers.len(), 1);
|
||||
assert_eq!(
|
||||
cfg.servers[0].tunnel_security,
|
||||
Some(TunnelSecurity::NonTlsRequired)
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.servers[0].tunnel_encryption_key.as_deref(),
|
||||
Some("BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_deserializes_upstream_proxy_url() {
|
||||
let cfg: ConfigFile = toml::from_str("upstream_proxy_url = \"http://proxy.example:8080\"")
|
||||
@@ -1608,6 +1722,117 @@ node_name = "tunnel-test"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_defaults_tunnel_security_to_off() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
]);
|
||||
|
||||
assert_eq!(config.tunnel_security, TunnelSecurity::Off);
|
||||
assert!(config.tunnel_encryption_key.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_requires_encryption_key_for_non_tls_security() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"http://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-security",
|
||||
"non_tls_required",
|
||||
]);
|
||||
|
||||
let error = config
|
||||
.validate()
|
||||
.expect_err("non_tls_required should require a PSK");
|
||||
assert!(error.to_string().contains("tunnel_encryption_key"));
|
||||
|
||||
let with_key = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"http://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-security",
|
||||
"non_tls_required",
|
||||
"--tunnel-encryption-key",
|
||||
"BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=",
|
||||
]);
|
||||
with_key
|
||||
.validate()
|
||||
.expect("non_tls_required with a PSK should validate");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_infers_non_tls_security_for_http_url_with_key() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"http://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-encryption-key",
|
||||
"BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=",
|
||||
]);
|
||||
|
||||
assert_eq!(config.tunnel_security, TunnelSecurity::Off);
|
||||
assert_eq!(
|
||||
effective_tunnel_security(
|
||||
&config.aether_url,
|
||||
None,
|
||||
config.tunnel_encryption_key.as_deref(),
|
||||
),
|
||||
TunnelSecurity::NonTlsRequired
|
||||
);
|
||||
assert_eq!(
|
||||
effective_tunnel_security(
|
||||
&config.aether_url,
|
||||
Some(TunnelSecurity::Off),
|
||||
config.tunnel_encryption_key.as_deref(),
|
||||
),
|
||||
TunnelSecurity::Off
|
||||
);
|
||||
config
|
||||
.validate()
|
||||
.expect("http URL with PSK should validate when tunnel_security is off");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_rejects_invalid_tunnel_encryption_key() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"http://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-security",
|
||||
"non_tls_required",
|
||||
"--tunnel-encryption-key",
|
||||
"not-a-valid-32-byte-key",
|
||||
]);
|
||||
|
||||
let error = config
|
||||
.validate()
|
||||
.expect_err("invalid PSK should fail validation");
|
||||
assert!(error.to_string().contains("base64-encoded 32 bytes"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_accepts_tunnel_ipv4_only() {
|
||||
let config = Config::parse_from([
|
||||
|
||||
@@ -15,9 +15,9 @@ mod upstream_client;
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use clap::{CommandFactory, FromArgMatches, Parser};
|
||||
use clap::{parser::ValueSource, CommandFactory, FromArgMatches, Parser};
|
||||
|
||||
use config::Config;
|
||||
use config::{Config, ServerEntry, TunnelSecurity};
|
||||
|
||||
/// Default config file name.
|
||||
const DEFAULT_CONFIG: &str = "aether-tunnel.toml";
|
||||
@@ -102,7 +102,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
None => {
|
||||
// No subcommand: run the tunnel with parsed config.
|
||||
let config = Config::from_arg_matches(&matches)?;
|
||||
run_tunnel(config).await
|
||||
let tunnel_security = configured_tunnel_security_from_matches(&matches, &config);
|
||||
run_tunnel(config, tunnel_security).await
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
@@ -139,7 +140,7 @@ async fn handle_setup_result(outcome: setup::SetupOutcome) -> anyhow::Result<()>
|
||||
let config = Config::try_parse_from(["aether-tunnel"])
|
||||
.map_err(|e| anyhow::anyhow!("config invalid after setup: {}", e))?;
|
||||
eprintln!(" Starting tunnel...\n");
|
||||
run_tunnel(config).await
|
||||
run_tunnel(config, None).await
|
||||
}
|
||||
setup::SetupOutcome::Cancelled => {
|
||||
eprintln!(" Setup cancelled.");
|
||||
@@ -148,8 +149,31 @@ async fn handle_setup_result(outcome: setup::SetupOutcome) -> anyhow::Result<()>
|
||||
}
|
||||
}
|
||||
|
||||
fn configured_tunnel_security_from_matches(
|
||||
matches: &clap::ArgMatches,
|
||||
config: &Config,
|
||||
) -> Option<TunnelSecurity> {
|
||||
matches
|
||||
.value_source("tunnel_security")
|
||||
.filter(|source| *source != ValueSource::DefaultValue)
|
||||
.map(|_| config.tunnel_security)
|
||||
}
|
||||
|
||||
fn single_server_entry(config: &Config, tunnel_security: Option<TunnelSecurity>) -> ServerEntry {
|
||||
ServerEntry {
|
||||
aether_url: config.aether_url.clone(),
|
||||
management_token: config.management_token.clone(),
|
||||
node_name: None,
|
||||
tunnel_security,
|
||||
tunnel_encryption_key: config.tunnel_encryption_key.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the tunnel agent, checking for managed-service conflicts first.
|
||||
async fn run_tunnel(config: Config) -> anyhow::Result<()> {
|
||||
async fn run_tunnel(
|
||||
config: Config,
|
||||
single_server_tunnel_security: Option<TunnelSecurity>,
|
||||
) -> anyhow::Result<()> {
|
||||
// Warn if a managed service is already running (would cause conflicts).
|
||||
if std::env::var_os("AETHER_TUNNEL_SERVICE_MANAGER").is_none()
|
||||
&& std::env::var_os("INVOCATION_ID").is_none()
|
||||
@@ -178,12 +202,76 @@ async fn run_tunnel(config: Config) -> anyhow::Result<()> {
|
||||
}
|
||||
file_cfg.servers.clone()
|
||||
} else {
|
||||
vec![config::ServerEntry {
|
||||
aether_url: config.aether_url.clone(),
|
||||
management_token: config.management_token.clone(),
|
||||
node_name: None,
|
||||
}]
|
||||
vec![single_server_entry(&config, single_server_tunnel_security)]
|
||||
};
|
||||
|
||||
app::run(config, servers).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn parse_config_and_security(args: &[&str]) -> (Config, Option<TunnelSecurity>) {
|
||||
let matches = build_command()
|
||||
.try_get_matches_from(args)
|
||||
.expect("arguments should parse");
|
||||
let config = Config::from_arg_matches(&matches).expect("config should parse");
|
||||
let tunnel_security = configured_tunnel_security_from_matches(&matches, &config);
|
||||
(config, tunnel_security)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_server_entry_omits_default_tunnel_security_for_auto_inference() {
|
||||
let (config, tunnel_security) = parse_config_and_security(&[
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"http://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-encryption-key",
|
||||
"BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=",
|
||||
]);
|
||||
|
||||
let entry = single_server_entry(&config, tunnel_security);
|
||||
assert_eq!(entry.tunnel_security, None);
|
||||
assert_eq!(
|
||||
config::effective_tunnel_security(
|
||||
&entry.aether_url,
|
||||
entry.tunnel_security,
|
||||
entry.tunnel_encryption_key.as_deref(),
|
||||
),
|
||||
TunnelSecurity::NonTlsRequired
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_server_entry_preserves_explicit_tunnel_security_off() {
|
||||
let (config, tunnel_security) = parse_config_and_security(&[
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"http://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-security",
|
||||
"off",
|
||||
"--tunnel-encryption-key",
|
||||
"BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=",
|
||||
]);
|
||||
|
||||
let entry = single_server_entry(&config, tunnel_security);
|
||||
assert_eq!(entry.tunnel_security, Some(TunnelSecurity::Off));
|
||||
assert_eq!(
|
||||
config::effective_tunnel_security(
|
||||
&entry.aether_url,
|
||||
entry.tunnel_security,
|
||||
entry.tunnel_encryption_key.as_deref(),
|
||||
),
|
||||
TunnelSecurity::Off
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
|
||||
use tokio::time::sleep;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::config::{effective_tunnel_security, Config, ServerEntry, TunnelSecurity};
|
||||
use crate::hardware::HardwareInfo;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -22,6 +22,10 @@ struct RegisterRequest {
|
||||
estimated_max_concurrency: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
proxy_metadata: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tunnel_security: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tunnel_encryption_key: Option<String>,
|
||||
tunnel_mode: bool,
|
||||
}
|
||||
|
||||
@@ -95,11 +99,17 @@ impl AetherClient {
|
||||
pub async fn register(
|
||||
&self,
|
||||
config: &Config,
|
||||
server: &ServerEntry,
|
||||
node_name: &str,
|
||||
public_ip: &str,
|
||||
hw: Option<&HardwareInfo>,
|
||||
) -> anyhow::Result<String> {
|
||||
let url = format!("{}/api/admin/proxy-nodes/register", self.base_url);
|
||||
let effective_security = effective_tunnel_security(
|
||||
&server.aether_url,
|
||||
server.tunnel_security,
|
||||
server.tunnel_encryption_key.as_deref(),
|
||||
);
|
||||
let body = RegisterRequest {
|
||||
name: node_name.to_string(),
|
||||
ip: public_ip.to_string(),
|
||||
@@ -111,6 +121,11 @@ impl AetherClient {
|
||||
proxy_metadata: Some(serde_json::json!({
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
})),
|
||||
tunnel_security: (effective_security == TunnelSecurity::NonTlsRequired)
|
||||
.then(|| effective_security.to_string()),
|
||||
tunnel_encryption_key: (effective_security == TunnelSecurity::NonTlsRequired)
|
||||
.then(|| server.tunnel_encryption_key.clone())
|
||||
.flatten(),
|
||||
tunnel_mode: true,
|
||||
};
|
||||
|
||||
|
||||
@@ -93,6 +93,22 @@ impl ServerTab {
|
||||
required: true,
|
||||
help: "Node name for identification in Aether dashboard",
|
||||
},
|
||||
Field {
|
||||
label: "Tunnel Security",
|
||||
key: "tunnel_security",
|
||||
value: String::new(),
|
||||
kind: FieldKind::Text,
|
||||
required: false,
|
||||
help: "off or non_tls_required; omit to auto-enable for http:// plus a key",
|
||||
},
|
||||
Field {
|
||||
label: "Tunnel Encryption Key",
|
||||
key: "tunnel_encryption_key",
|
||||
value: String::new(),
|
||||
kind: FieldKind::Secret,
|
||||
required: false,
|
||||
help: "Base64 32-byte PSK; required when Tunnel Security is non_tls_required",
|
||||
},
|
||||
],
|
||||
}
|
||||
}
|
||||
@@ -104,6 +120,12 @@ impl ServerTab {
|
||||
if let Some(ref name) = entry.node_name {
|
||||
tab.fields[2].value = name.clone();
|
||||
}
|
||||
if let Some(security) = entry.tunnel_security {
|
||||
tab.fields[3].value = security.to_string();
|
||||
}
|
||||
if let Some(ref key) = entry.tunnel_encryption_key {
|
||||
tab.fields[4].value = key.clone();
|
||||
}
|
||||
tab
|
||||
}
|
||||
}
|
||||
@@ -416,12 +438,33 @@ impl App {
|
||||
cfg.servers = self
|
||||
.server_tabs
|
||||
.iter()
|
||||
.map(|tab| ServerEntry {
|
||||
aether_url: get_tab(tab, "aether_url").unwrap_or_default(),
|
||||
management_token: get_tab(tab, "management_token").unwrap_or_default(),
|
||||
node_name: get_tab(tab, "node_name"),
|
||||
.map(|tab| {
|
||||
let tunnel_security = get_tab(tab, "tunnel_security")
|
||||
.map(|value| value.parse().map_err(anyhow::Error::msg))
|
||||
.transpose()?;
|
||||
Ok(ServerEntry {
|
||||
aether_url: get_tab(tab, "aether_url").unwrap_or_default(),
|
||||
management_token: get_tab(tab, "management_token").unwrap_or_default(),
|
||||
node_name: get_tab(tab, "node_name"),
|
||||
tunnel_security,
|
||||
tunnel_encryption_key: get_tab(tab, "tunnel_encryption_key"),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
.collect::<anyhow::Result<Vec<_>>>()?;
|
||||
for server in &cfg.servers {
|
||||
if server.tunnel_security == Some(crate::config::TunnelSecurity::NonTlsRequired)
|
||||
&& server
|
||||
.tunnel_encryption_key
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.is_none()
|
||||
{
|
||||
anyhow::bail!(
|
||||
"Tunnel Encryption Key is required when Tunnel Security is non_tls_required"
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(cfg)
|
||||
}
|
||||
|
||||
@@ -1108,11 +1151,55 @@ mod tests {
|
||||
|
||||
assert_eq!(
|
||||
server_keys,
|
||||
vec!["aether_url", "management_token", "node_name"]
|
||||
vec![
|
||||
"aether_url",
|
||||
"management_token",
|
||||
"node_name",
|
||||
"tunnel_security",
|
||||
"tunnel_encryption_key"
|
||||
]
|
||||
);
|
||||
assert_eq!(global_keys.first().copied(), Some("upstream_proxy_url"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_config_persists_tunnel_security_fields_per_server() {
|
||||
let mut app = sample_app();
|
||||
set_server_field(&mut app, "tunnel_security", "non_tls_required");
|
||||
set_server_field(&mut app, "tunnel_encryption_key", "base64-32-bytes");
|
||||
|
||||
let cfg = app.to_config().expect("config should serialize");
|
||||
assert_eq!(cfg.servers.len(), 1);
|
||||
assert_eq!(
|
||||
cfg.servers[0].tunnel_security,
|
||||
Some(crate::config::TunnelSecurity::NonTlsRequired)
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.servers[0].tunnel_encryption_key.as_deref(),
|
||||
Some("base64-32-bytes")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_config_omits_blank_tunnel_security_for_auto_inference() {
|
||||
let app = sample_app();
|
||||
|
||||
let cfg = app.to_config().expect("config should serialize");
|
||||
assert_eq!(cfg.servers.len(), 1);
|
||||
assert_eq!(cfg.servers[0].tunnel_security, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_config_rejects_non_tls_security_without_key() {
|
||||
let mut app = sample_app();
|
||||
set_server_field(&mut app, "tunnel_security", "non_tls_required");
|
||||
|
||||
let error = app
|
||||
.to_config()
|
||||
.expect_err("secure non-TLS mode should require a key");
|
||||
assert!(error.to_string().contains("Tunnel Encryption Key"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_config_enables_pretty_file_logging_with_defaults() {
|
||||
let mut app = sample_app();
|
||||
@@ -1226,6 +1313,8 @@ mod tests {
|
||||
aether_url: "https://aether-2.example.com".to_string(),
|
||||
management_token: "ae_test_2".to_string(),
|
||||
node_name: Some("jp-proxy-02".to_string()),
|
||||
tunnel_security: None,
|
||||
tunnel_encryption_key: None,
|
||||
}));
|
||||
app.active_tab = 1;
|
||||
|
||||
@@ -1245,6 +1334,8 @@ mod tests {
|
||||
aether_url: "https://aether-2.example.com".to_string(),
|
||||
management_token: "ae_test_2".to_string(),
|
||||
node_name: Some("jp-proxy-02".to_string()),
|
||||
tunnel_security: None,
|
||||
tunnel_encryption_key: None,
|
||||
}));
|
||||
app.active_tab = 0;
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ use aether_runtime::{
|
||||
use aether_runtime_state::{RuntimeSemaphore, RuntimeSemaphoreError, RuntimeSemaphoreSnapshot};
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::config::TunnelSecurity;
|
||||
use crate::hardware::RuntimeResourceMonitor;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::SharedDynamicConfig;
|
||||
@@ -44,6 +45,10 @@ pub struct ServerContext {
|
||||
pub aether_url: String,
|
||||
/// Management token for this server.
|
||||
pub management_token: String,
|
||||
/// Effective security mode for this server connection.
|
||||
pub tunnel_security: TunnelSecurity,
|
||||
/// PSK used for secure non-TLS tunnel frames.
|
||||
pub tunnel_encryption_key: Option<String>,
|
||||
/// Resolved node name at registration time (per-server override or global fallback).
|
||||
/// After startup, the active node_name is read from `dynamic` (may be updated remotely).
|
||||
#[allow(dead_code)]
|
||||
|
||||
@@ -17,6 +17,10 @@ use crate::egress_proxy::{
|
||||
};
|
||||
use crate::state::{AppState, ServerContext};
|
||||
use aether_contracts::tunnel::{CURRENT_TUNNEL_PROTOCOL_VERSION, TUNNEL_PROTOCOL_VERSION_HEADER};
|
||||
use aether_contracts::tunnel_security::{
|
||||
SecureFrameCodec, TunnelSecurityRole, TUNNEL_SECURITY_HEADER, TUNNEL_SECURITY_NON_TLS_REQUIRED,
|
||||
TUNNEL_SECURITY_SESSION_HEADER,
|
||||
};
|
||||
|
||||
use super::{dispatcher, heartbeat, writer};
|
||||
|
||||
@@ -45,16 +49,29 @@ pub async fn connect_and_run(
|
||||
// Build WebSocket request with auth headers
|
||||
let mut request = ws_url.clone().into_client_request()?;
|
||||
let headers = request.headers_mut();
|
||||
headers.insert(
|
||||
"Authorization",
|
||||
http::HeaderValue::from_str(&format!("Bearer {}", server.management_token))?,
|
||||
);
|
||||
if server.tunnel_security != crate::config::TunnelSecurity::NonTlsRequired {
|
||||
headers.insert(
|
||||
"Authorization",
|
||||
http::HeaderValue::from_str(&format!("Bearer {}", server.management_token))?,
|
||||
);
|
||||
}
|
||||
headers.insert(
|
||||
TUNNEL_PROTOCOL_VERSION_HEADER,
|
||||
http::HeaderValue::from_str(&CURRENT_TUNNEL_PROTOCOL_VERSION.to_string())?,
|
||||
);
|
||||
let node_id = server.node_id.read().unwrap().clone();
|
||||
headers.insert("X-Node-Id", http::HeaderValue::from_str(&node_id)?);
|
||||
let security_session = uuid::Uuid::new_v4().simple().to_string();
|
||||
if server.tunnel_security == crate::config::TunnelSecurity::NonTlsRequired {
|
||||
headers.insert(
|
||||
TUNNEL_SECURITY_HEADER,
|
||||
http::HeaderValue::from_static(TUNNEL_SECURITY_NON_TLS_REQUIRED),
|
||||
);
|
||||
headers.insert(
|
||||
TUNNEL_SECURITY_SESSION_HEADER,
|
||||
http::HeaderValue::from_str(&security_session)?,
|
||||
);
|
||||
}
|
||||
// Use dynamic node_name (may be updated by remote config) instead of
|
||||
// the static server.node_name, so that remote name changes take effect
|
||||
// on the next reconnect.
|
||||
@@ -119,6 +136,19 @@ pub async fn connect_and_run(
|
||||
handshake_timeout.as_millis()
|
||||
)
|
||||
})??;
|
||||
let security = if server.tunnel_security == crate::config::TunnelSecurity::NonTlsRequired {
|
||||
let key = server
|
||||
.tunnel_encryption_key
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("secure tunnel requires tunnel_encryption_key"))?;
|
||||
Some(Arc::new(SecureFrameCodec::new(
|
||||
key,
|
||||
&security_session,
|
||||
TunnelSecurityRole::Client,
|
||||
)?))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let stale_timeout = state
|
||||
.config
|
||||
.tunnel_stale_timeout()
|
||||
@@ -146,10 +176,11 @@ pub async fn connect_and_run(
|
||||
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
|
||||
|
||||
// Spawn writer task (with WebSocket ping keepalive)
|
||||
let (frame_tx, mut writer_handle) = writer::spawn_writer_with_metrics(
|
||||
let (frame_tx, mut writer_handle) = writer::spawn_writer_with_metrics_and_security(
|
||||
ws_sink,
|
||||
ping_interval,
|
||||
Some(Arc::clone(&server.tunnel_metrics)),
|
||||
security.clone(),
|
||||
);
|
||||
let drain_signal = spawn_drain_signal(conn_idx, frame_tx.clone(), drain.clone());
|
||||
|
||||
@@ -174,13 +205,14 @@ pub async fn connect_and_run(
|
||||
let state_clone = Arc::clone(state);
|
||||
let server_clone = Arc::clone(server);
|
||||
let outcome = tokio::select! {
|
||||
result = dispatcher::run(
|
||||
result = dispatcher::run_with_security(
|
||||
state_clone,
|
||||
server_clone,
|
||||
ws_read,
|
||||
frame_tx.clone(),
|
||||
hb_handle,
|
||||
drain.clone(),
|
||||
security.clone(),
|
||||
) => {
|
||||
match result {
|
||||
Ok(()) => Ok(TunnelOutcome::Disconnected),
|
||||
|
||||
@@ -18,6 +18,7 @@ use super::heartbeat::HeartbeatHandle;
|
||||
use super::protocol::{decompress_if_gzip, Frame, MsgType, RequestMeta};
|
||||
use super::stream_handler;
|
||||
use super::writer::FrameSender;
|
||||
use aether_contracts::tunnel_security::SecureFrameCodec;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum StreamDispatchStatus {
|
||||
@@ -27,13 +28,32 @@ enum StreamDispatchStatus {
|
||||
}
|
||||
|
||||
/// Run the dispatcher loop, reading from the WebSocket stream.
|
||||
#[allow(dead_code)]
|
||||
pub async fn run<S>(
|
||||
state: Arc<AppState>,
|
||||
server: Arc<ServerContext>,
|
||||
ws_stream: S,
|
||||
frame_tx: FrameSender,
|
||||
heartbeat: HeartbeatHandle,
|
||||
drain: watch::Receiver<bool>,
|
||||
) -> Result<(), anyhow::Error>
|
||||
where
|
||||
S: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
|
||||
+ Unpin
|
||||
+ Send
|
||||
+ 'static,
|
||||
{
|
||||
run_with_security(state, server, ws_stream, frame_tx, heartbeat, drain, None).await
|
||||
}
|
||||
|
||||
pub async fn run_with_security<S>(
|
||||
state: Arc<AppState>,
|
||||
server: Arc<ServerContext>,
|
||||
mut ws_stream: S,
|
||||
frame_tx: FrameSender,
|
||||
heartbeat: HeartbeatHandle,
|
||||
mut drain: watch::Receiver<bool>,
|
||||
security: Option<Arc<SecureFrameCodec>>,
|
||||
) -> Result<(), anyhow::Error>
|
||||
where
|
||||
S: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
|
||||
@@ -130,6 +150,19 @@ where
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let frame = match security.as_deref() {
|
||||
Some(codec) => match codec.decrypt_frame(frame) {
|
||||
Ok(frame) => frame,
|
||||
Err(e) => {
|
||||
warn!(error = %e, "failed to decrypt secure tunnel frame");
|
||||
server
|
||||
.tunnel_metrics
|
||||
.record_error("secure_frame_decrypt_error", &e.to_string());
|
||||
break None;
|
||||
}
|
||||
},
|
||||
None => frame,
|
||||
};
|
||||
|
||||
match frame.msg_type {
|
||||
MsgType::RequestHeaders => {
|
||||
|
||||
@@ -425,6 +425,8 @@ mod tests {
|
||||
server_label: "heartbeat-test".to_string(),
|
||||
aether_url: config.aether_url.clone(),
|
||||
management_token: config.management_token.clone(),
|
||||
tunnel_security: config.tunnel_security,
|
||||
tunnel_encryption_key: config.tunnel_encryption_key.clone(),
|
||||
node_name: config.node_name.clone(),
|
||||
node_id: Arc::new(RwLock::new("node-123".to_string())),
|
||||
aether_client: Arc::new(AetherClient::new(
|
||||
|
||||
@@ -392,6 +392,9 @@ mod tests {
|
||||
method: "GET".to_string(),
|
||||
url: "http://127.0.0.1:80/blocked".to_string(),
|
||||
headers: std::collections::HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 5,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
@@ -476,6 +479,8 @@ mod tests {
|
||||
server_label: "gateway-owned-tunnel".to_string(),
|
||||
aether_url: config.aether_url.clone(),
|
||||
management_token: config.management_token.clone(),
|
||||
tunnel_security: config.tunnel_security,
|
||||
tunnel_encryption_key: config.tunnel_encryption_key.clone(),
|
||||
node_name: config.node_name.clone(),
|
||||
node_id: Arc::new(std::sync::RwLock::new(node_id.to_string())),
|
||||
aether_client: Arc::new(AetherClient::new(
|
||||
@@ -496,6 +501,8 @@ mod tests {
|
||||
management_token: "token".to_string(),
|
||||
public_ip: None,
|
||||
node_name: "tunnel-test".to_string(),
|
||||
tunnel_security: crate::config::TunnelSecurity::Off,
|
||||
tunnel_encryption_key: None,
|
||||
node_region: None,
|
||||
heartbeat_interval: 1,
|
||||
allowed_ports: vec![80, 443],
|
||||
|
||||
@@ -39,10 +39,10 @@ const SLOW_STREAM_LOG_THRESHOLD: Duration = Duration::from_secs(2);
|
||||
const SUCCESS_LOG_SAMPLE_MODULO: u32 = 256;
|
||||
const REQUEST_BODY_SPOOL_QUEUE_CAPACITY: usize = 64;
|
||||
|
||||
/// Minimum allowed upstream request timeout (seconds).
|
||||
const MIN_TIMEOUT_SECS: u64 = 5;
|
||||
/// Maximum allowed upstream request timeout (seconds).
|
||||
const MAX_TIMEOUT_SECS: u64 = 300;
|
||||
/// Minimum allowed upstream request timeout (milliseconds).
|
||||
const MIN_TIMEOUT_MS: u64 = 1;
|
||||
/// Maximum allowed upstream request timeout (milliseconds).
|
||||
const MAX_TIMEOUT_MS: u64 = 300_000;
|
||||
/// Match reqwest's default redirect budget so direct execution and tunnel relay
|
||||
/// fail at the same point instead of diverging after a different number of hops.
|
||||
const MAX_REDIRECTS: usize = 10;
|
||||
@@ -95,6 +95,12 @@ struct PreparedRequestBody {
|
||||
replay_body: ReplayableRequestBody,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct RequestTimeouts {
|
||||
first_byte_timeout: Duration,
|
||||
response_body_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct RequestBodyReplayState {
|
||||
budget_bytes: usize,
|
||||
@@ -588,6 +594,45 @@ fn remaining_timeout(deadline: Instant) -> Option<Duration> {
|
||||
deadline.checked_duration_since(Instant::now())
|
||||
}
|
||||
|
||||
fn resolve_request_timeouts(meta: &RequestMeta) -> RequestTimeouts {
|
||||
let first_byte_timeout = if meta.stream {
|
||||
meta.stream_first_byte_timeout_ms
|
||||
.or(meta.request_timeout_ms)
|
||||
.map(timeout_duration_from_ms)
|
||||
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
|
||||
} else {
|
||||
meta.request_timeout_ms
|
||||
.or(meta.stream_first_byte_timeout_ms)
|
||||
.map(timeout_duration_from_ms)
|
||||
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
|
||||
};
|
||||
|
||||
let response_body_timeout = if meta.stream {
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
meta.request_timeout_ms
|
||||
.or(meta.stream_first_byte_timeout_ms)
|
||||
.map(timeout_duration_from_ms)
|
||||
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)),
|
||||
)
|
||||
};
|
||||
|
||||
RequestTimeouts {
|
||||
first_byte_timeout,
|
||||
response_body_timeout,
|
||||
}
|
||||
}
|
||||
|
||||
fn timeout_duration_from_ms(ms: u64) -> Duration {
|
||||
Duration::from_millis(ms.clamp(MIN_TIMEOUT_MS, MAX_TIMEOUT_MS))
|
||||
}
|
||||
|
||||
fn timeout_duration_from_legacy_secs(secs: u64) -> Duration {
|
||||
let ms = secs.saturating_mul(1_000);
|
||||
timeout_duration_from_ms(ms)
|
||||
}
|
||||
|
||||
async fn spool_request_body(
|
||||
mut body_rx: mpsc::Receiver<TunnelFrame>,
|
||||
mut spool_tx: mpsc::Sender<SpoolBodyEvent>,
|
||||
@@ -888,7 +933,7 @@ async fn relay_upstream_response<B>(
|
||||
redirect_count: usize,
|
||||
request_body_mode: &'static str,
|
||||
emit_proxy_timing_header: bool,
|
||||
deadline: Instant,
|
||||
response_body_deadline: Option<Instant>,
|
||||
) -> Option<Duration>
|
||||
where
|
||||
B: hyper::body::Body<Data = Bytes> + Send + Unpin + 'static,
|
||||
@@ -956,28 +1001,8 @@ where
|
||||
|
||||
let mut stream = response.into_body().into_data_stream();
|
||||
loop {
|
||||
let Some(remaining) = remaining_timeout(deadline) else {
|
||||
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
|
||||
let error_message = "upstream response body timeout".to_string();
|
||||
log_stream_failure(
|
||||
stream_log_context(
|
||||
server,
|
||||
stream_id,
|
||||
method,
|
||||
Some(request_url),
|
||||
redirect_count,
|
||||
request_body_size.load(Ordering::Relaxed),
|
||||
),
|
||||
&error_message,
|
||||
total_elapsed,
|
||||
);
|
||||
send_error(frame_tx, stream_id, &error_message).await;
|
||||
return Some(total_elapsed);
|
||||
};
|
||||
|
||||
let chunk_result = match tokio::time::timeout(remaining, stream.next()).await {
|
||||
Ok(chunk_result) => chunk_result,
|
||||
Err(_) => {
|
||||
let chunk_result = if let Some(deadline) = response_body_deadline {
|
||||
let Some(remaining) = remaining_timeout(deadline) else {
|
||||
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
|
||||
let error_message = "upstream response body timeout".to_string();
|
||||
log_stream_failure(
|
||||
@@ -994,7 +1019,31 @@ where
|
||||
);
|
||||
send_error(frame_tx, stream_id, &error_message).await;
|
||||
return Some(total_elapsed);
|
||||
};
|
||||
|
||||
match tokio::time::timeout(remaining, stream.next()).await {
|
||||
Ok(chunk_result) => chunk_result,
|
||||
Err(_) => {
|
||||
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
|
||||
let error_message = "upstream response body timeout".to_string();
|
||||
log_stream_failure(
|
||||
stream_log_context(
|
||||
server,
|
||||
stream_id,
|
||||
method,
|
||||
Some(request_url),
|
||||
redirect_count,
|
||||
request_body_size.load(Ordering::Relaxed),
|
||||
),
|
||||
&error_message,
|
||||
total_elapsed,
|
||||
);
|
||||
send_error(frame_tx, stream_id, &error_message).await;
|
||||
return Some(total_elapsed);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
stream.next().await
|
||||
};
|
||||
|
||||
let Some(chunk_result) = chunk_result else {
|
||||
@@ -1272,16 +1321,19 @@ async fn handle_stream_inner(
|
||||
}
|
||||
}
|
||||
|
||||
let deadline = Instant::now()
|
||||
+ Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
|
||||
let overall_start = Instant::now();
|
||||
let request_timeouts = resolve_request_timeouts(&meta);
|
||||
let first_byte_deadline = overall_start + request_timeouts.first_byte_timeout;
|
||||
let response_body_deadline = request_timeouts
|
||||
.response_body_timeout
|
||||
.map(|timeout| overall_start + timeout);
|
||||
let follow_redirects = follow_redirects_enabled(&meta);
|
||||
let mut current_headers = sanitize_upstream_headers(&meta.headers);
|
||||
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
|
||||
let first_byte_timeout = request_timeouts.first_byte_timeout;
|
||||
let request_body_size = Arc::new(AtomicUsize::new(0));
|
||||
let request_has_body = request_likely_has_body(¤t_method, &meta.headers);
|
||||
let replay_budget_bytes = state.config.redirect_replay_budget_bytes;
|
||||
let can_buffer_redirect_body = request_has_body && follow_redirects && replay_budget_bytes > 0;
|
||||
let overall_start = Instant::now();
|
||||
let request_body_mode = if can_buffer_redirect_body {
|
||||
"buffered_fixed"
|
||||
} else if request_has_body {
|
||||
@@ -1293,7 +1345,7 @@ async fn handle_stream_inner(
|
||||
let buffered_body = match collect_request_body_for_replay(
|
||||
body_rx,
|
||||
Arc::clone(&request_body_size),
|
||||
deadline,
|
||||
first_byte_deadline,
|
||||
replay_budget_bytes,
|
||||
)
|
||||
.await
|
||||
@@ -1321,7 +1373,12 @@ async fn handle_stream_inner(
|
||||
replay_body: replay_body_from_buffered(buffered_body, replay_budget_bytes),
|
||||
}
|
||||
} else if request_has_body {
|
||||
prepare_request_body(body_rx, Arc::clone(&request_body_size), deadline, 0)
|
||||
prepare_request_body(
|
||||
body_rx,
|
||||
Arc::clone(&request_body_size),
|
||||
first_byte_deadline,
|
||||
0,
|
||||
)
|
||||
} else {
|
||||
PreparedRequestBody {
|
||||
first_request_body: Some(build_streaming_request_body(
|
||||
@@ -1341,7 +1398,7 @@ async fn handle_stream_inner(
|
||||
let mut next_request_body = None::<upstream_client::UpstreamRequestBody>;
|
||||
|
||||
loop {
|
||||
let Some(remaining) = remaining_timeout(deadline) else {
|
||||
let Some(remaining) = remaining_timeout(first_byte_deadline) else {
|
||||
log_stream_failure(
|
||||
stream_log_context(
|
||||
server,
|
||||
@@ -1369,7 +1426,7 @@ async fn handle_stream_inner(
|
||||
current_method.clone(),
|
||||
¤t_headers,
|
||||
request_body,
|
||||
remaining.min(timeout),
|
||||
remaining.min(first_byte_timeout),
|
||||
meta.http1_only,
|
||||
)
|
||||
.await
|
||||
@@ -1419,7 +1476,7 @@ async fn handle_stream_inner(
|
||||
redirects_followed,
|
||||
request_body_mode,
|
||||
state.config.emit_proxy_timing_header,
|
||||
deadline,
|
||||
response_body_deadline,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -1431,7 +1488,7 @@ async fn handle_stream_inner(
|
||||
} => match prepare_redirect_request_body(
|
||||
prepared_body.replay_body.clone(),
|
||||
body_mode,
|
||||
deadline,
|
||||
first_byte_deadline,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -1459,7 +1516,7 @@ async fn handle_stream_inner(
|
||||
redirects_followed,
|
||||
request_body_mode,
|
||||
state.config.emit_proxy_timing_header,
|
||||
deadline,
|
||||
response_body_deadline,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -1515,7 +1572,7 @@ async fn handle_stream_inner(
|
||||
redirects_followed,
|
||||
request_body_mode,
|
||||
state.config.emit_proxy_timing_header,
|
||||
deadline,
|
||||
response_body_deadline,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -1850,6 +1907,34 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_request_timeouts_use_first_byte_without_response_body_deadline() {
|
||||
let mut meta = sample_request_meta();
|
||||
meta.stream = true;
|
||||
meta.request_timeout_ms = Some(90_000);
|
||||
meta.stream_first_byte_timeout_ms = Some(12_345);
|
||||
|
||||
let timeouts = resolve_request_timeouts(&meta);
|
||||
|
||||
assert_eq!(timeouts.first_byte_timeout, Duration::from_millis(12_345));
|
||||
assert!(timeouts.response_body_timeout.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_stream_request_timeouts_use_total_for_response_body_deadline() {
|
||||
let mut meta = sample_request_meta();
|
||||
meta.request_timeout_ms = Some(90_000);
|
||||
meta.stream_first_byte_timeout_ms = Some(12_345);
|
||||
|
||||
let timeouts = resolve_request_timeouts(&meta);
|
||||
|
||||
assert_eq!(timeouts.first_byte_timeout, Duration::from_millis(90_000));
|
||||
assert_eq!(
|
||||
timeouts.response_body_timeout,
|
||||
Some(Duration::from_millis(90_000))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_redirect_changes_post_to_get_for_302() {
|
||||
let current_url = url::Url::parse("https://redirect.test/start").expect("url");
|
||||
@@ -2073,7 +2158,7 @@ mod tests {
|
||||
0,
|
||||
"empty",
|
||||
true,
|
||||
Instant::now(),
|
||||
Some(Instant::now()),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -2086,6 +2171,51 @@ mod tests {
|
||||
assert_eq!(server.metrics.stream_errors.load(Ordering::Acquire), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_response_body_without_total_deadline_allows_late_chunk() {
|
||||
let state = sample_state(None, None);
|
||||
let server = sample_server(&state);
|
||||
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||
let request_url = url::Url::parse("https://example.com/stream").expect("url");
|
||||
let request_body_size = AtomicUsize::new(0);
|
||||
let body = Body::from_stream(stream::once(async {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
Ok::<Bytes, std::convert::Infallible>(Bytes::from_static(b"late"))
|
||||
}));
|
||||
let response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(body)
|
||||
.expect("response");
|
||||
|
||||
relay_upstream_response(
|
||||
&server,
|
||||
14,
|
||||
&hyper::Method::GET,
|
||||
&request_url,
|
||||
&frame_tx,
|
||||
response,
|
||||
0,
|
||||
Duration::ZERO,
|
||||
upstream_client::RequestTiming::default(),
|
||||
&request_body_size,
|
||||
0,
|
||||
"empty",
|
||||
true,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
|
||||
assert!(
|
||||
result.error.is_none(),
|
||||
"unexpected error: {:?}",
|
||||
result.error
|
||||
);
|
||||
assert_eq!(result.response.expect("response metadata").status, 200);
|
||||
assert_eq!(result.body, Bytes::from_static(b"late"));
|
||||
assert_eq!(server.metrics.stream_errors.load(Ordering::Acquire), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn follows_redirects_when_explicitly_enabled_for_replayable_post_requests() {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
@@ -2427,6 +2557,9 @@ mod tests {
|
||||
method: "GET".to_string(),
|
||||
url: "https://example.com/ok".to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
@@ -2491,6 +2624,8 @@ mod tests {
|
||||
server_label: "server".to_string(),
|
||||
aether_url: config.aether_url.clone(),
|
||||
management_token: config.management_token.clone(),
|
||||
tunnel_security: config.tunnel_security,
|
||||
tunnel_encryption_key: config.tunnel_encryption_key.clone(),
|
||||
node_name: config.node_name.clone(),
|
||||
node_id: Arc::new(std::sync::RwLock::new("node-1".to_string())),
|
||||
aether_client: Arc::new(AetherClient::new(
|
||||
@@ -2511,6 +2646,8 @@ mod tests {
|
||||
management_token: "token".to_string(),
|
||||
public_ip: None,
|
||||
node_name: "tunnel-test".to_string(),
|
||||
tunnel_security: crate::config::TunnelSecurity::Off,
|
||||
tunnel_encryption_key: None,
|
||||
node_region: None,
|
||||
heartbeat_interval: 30,
|
||||
allowed_ports: vec![80, 443],
|
||||
|
||||
@@ -20,6 +20,7 @@ use tracing::{debug, error, trace};
|
||||
use crate::state::TunnelMetrics;
|
||||
|
||||
use super::protocol::Frame;
|
||||
use aether_contracts::tunnel_security::SecureFrameCodec;
|
||||
|
||||
const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 64;
|
||||
const NORMAL_PRIORITY_QUEUE_CAPACITY: usize = 256;
|
||||
@@ -89,10 +90,23 @@ where
|
||||
}
|
||||
|
||||
/// Spawn the writer task with optional tunnel metrics instrumentation.
|
||||
#[allow(dead_code)]
|
||||
pub fn spawn_writer_with_metrics<S>(
|
||||
sink: S,
|
||||
ping_interval: Duration,
|
||||
tunnel_metrics: Option<Arc<TunnelMetrics>>,
|
||||
) -> (FrameSender, JoinHandle<()>)
|
||||
where
|
||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||
{
|
||||
spawn_writer_with_metrics_and_security(sink, ping_interval, tunnel_metrics, None)
|
||||
}
|
||||
|
||||
pub fn spawn_writer_with_metrics_and_security<S>(
|
||||
mut sink: S,
|
||||
ping_interval: Duration,
|
||||
tunnel_metrics: Option<Arc<TunnelMetrics>>,
|
||||
security: Option<Arc<SecureFrameCodec>>,
|
||||
) -> (FrameSender, JoinHandle<()>)
|
||||
where
|
||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||
@@ -109,7 +123,14 @@ where
|
||||
|
||||
loop {
|
||||
if let Ok(frame) = high_rx.try_recv() {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
|
||||
if !write_frame(
|
||||
&mut sink,
|
||||
frame,
|
||||
tunnel_metrics.as_deref(),
|
||||
security.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
@@ -123,7 +144,7 @@ where
|
||||
frame = high_rx.recv(), if high_open => {
|
||||
match frame {
|
||||
Some(frame) => {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref(), security.as_deref()).await {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -143,7 +164,7 @@ where
|
||||
frame = normal_rx.recv(), if normal_open => {
|
||||
match frame {
|
||||
Some(frame) => {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
|
||||
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref(), security.as_deref()).await {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -175,14 +196,31 @@ fn classify_frame_priority(frame: &Frame) -> FramePriority {
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_frame<S>(sink: &mut S, frame: Frame, tunnel_metrics: Option<&TunnelMetrics>) -> bool
|
||||
async fn write_frame<S>(
|
||||
sink: &mut S,
|
||||
frame: Frame,
|
||||
tunnel_metrics: Option<&TunnelMetrics>,
|
||||
security: Option<&SecureFrameCodec>,
|
||||
) -> bool
|
||||
where
|
||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||
{
|
||||
let stream_id = frame.stream_id;
|
||||
let msg_type = frame.msg_type;
|
||||
let flags = frame.flags;
|
||||
let data = frame.encode();
|
||||
let data = match security {
|
||||
Some(codec) => match codec.encrypt_frame(frame) {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
error!(error = %e, "failed to encrypt tunnel frame");
|
||||
if let Some(metrics) = tunnel_metrics {
|
||||
metrics.record_error("secure_frame_encrypt_error", &e.to_string());
|
||||
}
|
||||
return false;
|
||||
}
|
||||
},
|
||||
None => frame.encode(),
|
||||
};
|
||||
let wire_len = data.len().max(HEADER_SIZE);
|
||||
if let Err(e) = sink.send(Message::Binary(data.into())).await {
|
||||
error!(
|
||||
|
||||
@@ -385,7 +385,7 @@ fn resolve_admin_monitoring_usage_candidate_id(
|
||||
trace: &DecisionTrace,
|
||||
usage: &StoredRequestUsageAudit,
|
||||
) -> Option<String> {
|
||||
if usage.request_id.trim() != trace.request_id {
|
||||
if !admin_monitoring_usage_matches_trace(usage, trace.request_id.as_str()) {
|
||||
return None;
|
||||
}
|
||||
|
||||
@@ -413,6 +413,43 @@ fn resolve_admin_monitoring_usage_candidate_id(
|
||||
.map(|item| item.candidate.id.clone())
|
||||
}
|
||||
|
||||
fn admin_monitoring_usage_matches_trace(
|
||||
usage: &StoredRequestUsageAudit,
|
||||
trace_request_id: &str,
|
||||
) -> bool {
|
||||
let trace_request_id = trace_request_id.trim();
|
||||
if trace_request_id.is_empty() {
|
||||
return false;
|
||||
}
|
||||
usage.request_id.trim() == trace_request_id
|
||||
|| usage
|
||||
.trace_id()
|
||||
.is_some_and(|value| value.trim() == trace_request_id)
|
||||
|| admin_monitoring_headers_contain_trace_id(
|
||||
usage.request_headers.as_ref(),
|
||||
trace_request_id,
|
||||
)
|
||||
|| admin_monitoring_headers_contain_trace_id(
|
||||
usage.provider_request_headers.as_ref(),
|
||||
trace_request_id,
|
||||
)
|
||||
}
|
||||
|
||||
fn admin_monitoring_headers_contain_trace_id(
|
||||
headers: Option<&Value>,
|
||||
trace_request_id: &str,
|
||||
) -> bool {
|
||||
headers.and_then(Value::as_object).is_some_and(|object| {
|
||||
object.iter().any(|(key, value)| {
|
||||
key.eq_ignore_ascii_case("x-trace-id")
|
||||
&& value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| value == trace_request_id)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn admin_monitoring_candidate_status_rank(status: RequestCandidateStatus) -> u8 {
|
||||
match status {
|
||||
RequestCandidateStatus::Success => 7,
|
||||
|
||||
@@ -7,8 +7,12 @@ repository.workspace = true
|
||||
description = "Shared contracts for Python and Rust Aether components"
|
||||
|
||||
[dependencies]
|
||||
aes-gcm.workspace = true
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
flate2.workspace = true
|
||||
hmac.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
@@ -3,6 +3,7 @@ mod frame;
|
||||
mod plan;
|
||||
mod result;
|
||||
pub mod tunnel;
|
||||
pub mod tunnel_security;
|
||||
mod usage;
|
||||
|
||||
pub use error::{ExecutionError, ExecutionErrorKind, ExecutionPhase};
|
||||
|
||||
@@ -15,6 +15,7 @@ pub const CURRENT_TUNNEL_PROTOCOL_VERSION_STR: &str = "2";
|
||||
pub mod flags {
|
||||
pub const END_STREAM: u8 = 0x01;
|
||||
pub const GZIP_COMPRESSED: u8 = 0x02;
|
||||
pub const ENCRYPTED: u8 = crate::tunnel_security::FLAG_ENCRYPTED;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -65,6 +66,7 @@ pub const HEARTBEAT_DATA: u8 = MsgType::HeartbeatData as u8;
|
||||
pub const HEARTBEAT_ACK: u8 = MsgType::HeartbeatAck as u8;
|
||||
pub const FLAG_END_STREAM: u8 = flags::END_STREAM;
|
||||
pub const FLAG_GZIP_COMPRESSED: u8 = flags::GZIP_COMPRESSED;
|
||||
pub const FLAG_ENCRYPTED: u8 = flags::ENCRYPTED;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct FrameHeader {
|
||||
@@ -183,6 +185,12 @@ pub struct RequestMeta {
|
||||
pub method: String,
|
||||
pub url: String,
|
||||
pub headers: std::collections::HashMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "is_false")]
|
||||
pub stream: bool,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_timeout_ms: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stream_first_byte_timeout_ms: Option<u64>,
|
||||
#[serde(default = "default_timeout", deserialize_with = "deserialize_timeout")]
|
||||
pub timeout: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
|
||||
@@ -0,0 +1,338 @@
|
||||
use aes_gcm::aead::{Aead, Payload};
|
||||
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
|
||||
use base64::Engine;
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use crate::tunnel::{Frame, MsgType, HEADER_SIZE};
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
pub const TUNNEL_SECURITY_HEADER: &str = "x-aether-tunnel-security";
|
||||
pub const TUNNEL_SECURITY_SESSION_HEADER: &str = "x-aether-tunnel-security-session";
|
||||
pub const TUNNEL_SECURITY_NON_TLS_REQUIRED: &str = "non_tls_required";
|
||||
pub const FLAG_ENCRYPTED: u8 = 0x04;
|
||||
|
||||
const CONTEXT: &[u8] = b"aether-tunnel-secure-v1";
|
||||
const CLIENT_TO_SERVER_LABEL: &[u8] = b"client-to-server";
|
||||
const SERVER_TO_CLIENT_LABEL: &[u8] = b"server-to-client";
|
||||
const CLIENT_TO_SERVER_NONCE_PREFIX: [u8; 4] = *b"c2s1";
|
||||
const SERVER_TO_CLIENT_NONCE_PREFIX: [u8; 4] = *b"s2c1";
|
||||
const SEQUENCE_LEN: usize = 8;
|
||||
const NONCE_LEN: usize = 12;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TunnelSecurityRole {
|
||||
Client,
|
||||
Server,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum TunnelSecurityError {
|
||||
#[error("tunnel_encryption_key must be base64-encoded 32 bytes")]
|
||||
InvalidKey,
|
||||
#[error("tunnel security session id must not be empty")]
|
||||
InvalidSession,
|
||||
#[error("secure tunnel frame is missing encrypted flag")]
|
||||
MissingEncryptedFlag,
|
||||
#[error("secure tunnel frame payload is too short")]
|
||||
PayloadTooShort,
|
||||
#[error("secure tunnel frame sequence is not the expected next value")]
|
||||
UnexpectedSequence,
|
||||
#[error("secure tunnel frame encryption failed")]
|
||||
Encrypt,
|
||||
#[error("secure tunnel frame decryption failed")]
|
||||
Decrypt,
|
||||
}
|
||||
|
||||
pub struct SecureFrameCodec {
|
||||
seal: Aes256Gcm,
|
||||
open: Aes256Gcm,
|
||||
seal_prefix: [u8; 4],
|
||||
open_prefix: [u8; 4],
|
||||
next_sequence: AtomicU64,
|
||||
next_open_sequence: AtomicU64,
|
||||
}
|
||||
|
||||
impl SecureFrameCodec {
|
||||
pub fn new(
|
||||
key: &str,
|
||||
session_id: &str,
|
||||
role: TunnelSecurityRole,
|
||||
) -> Result<Self, TunnelSecurityError> {
|
||||
let psk = decode_psk(key)?;
|
||||
let session_id = session_id.trim();
|
||||
if session_id.is_empty() {
|
||||
return Err(TunnelSecurityError::InvalidSession);
|
||||
}
|
||||
|
||||
let client_to_server = derive_key(&psk, session_id.as_bytes(), CLIENT_TO_SERVER_LABEL);
|
||||
let server_to_client = derive_key(&psk, session_id.as_bytes(), SERVER_TO_CLIENT_LABEL);
|
||||
let (seal_key, open_key, seal_prefix, open_prefix) = match role {
|
||||
TunnelSecurityRole::Client => (
|
||||
client_to_server,
|
||||
server_to_client,
|
||||
CLIENT_TO_SERVER_NONCE_PREFIX,
|
||||
SERVER_TO_CLIENT_NONCE_PREFIX,
|
||||
),
|
||||
TunnelSecurityRole::Server => (
|
||||
server_to_client,
|
||||
client_to_server,
|
||||
SERVER_TO_CLIENT_NONCE_PREFIX,
|
||||
CLIENT_TO_SERVER_NONCE_PREFIX,
|
||||
),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
seal: Aes256Gcm::new_from_slice(&seal_key)
|
||||
.map_err(|_| TunnelSecurityError::InvalidKey)?,
|
||||
open: Aes256Gcm::new_from_slice(&open_key)
|
||||
.map_err(|_| TunnelSecurityError::InvalidKey)?,
|
||||
seal_prefix,
|
||||
open_prefix,
|
||||
next_sequence: AtomicU64::new(0),
|
||||
next_open_sequence: AtomicU64::new(0),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn encrypt_frame(&self, frame: Frame) -> Result<Bytes, TunnelSecurityError> {
|
||||
let sequence = self.next_sequence.fetch_add(1, Ordering::Relaxed);
|
||||
let nonce_bytes = nonce_bytes(self.seal_prefix, sequence);
|
||||
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||
let clear_flags = frame.flags & !FLAG_ENCRYPTED;
|
||||
let aad = frame_aad(frame.stream_id, frame.msg_type, clear_flags);
|
||||
let ciphertext = self
|
||||
.seal
|
||||
.encrypt(
|
||||
nonce,
|
||||
Payload {
|
||||
msg: &frame.payload,
|
||||
aad: &aad,
|
||||
},
|
||||
)
|
||||
.map_err(|_| TunnelSecurityError::Encrypt)?;
|
||||
|
||||
let mut payload = BytesMut::with_capacity(SEQUENCE_LEN + ciphertext.len());
|
||||
payload.put_u64(sequence);
|
||||
payload.extend_from_slice(&ciphertext);
|
||||
Ok(Frame::new(
|
||||
frame.stream_id,
|
||||
frame.msg_type,
|
||||
clear_flags | FLAG_ENCRYPTED,
|
||||
payload.freeze(),
|
||||
)
|
||||
.encode())
|
||||
}
|
||||
|
||||
pub fn decrypt_frame(&self, frame: Frame) -> Result<Frame, TunnelSecurityError> {
|
||||
if frame.flags & FLAG_ENCRYPTED == 0 {
|
||||
return Err(TunnelSecurityError::MissingEncryptedFlag);
|
||||
}
|
||||
if frame.payload.len() < SEQUENCE_LEN {
|
||||
return Err(TunnelSecurityError::PayloadTooShort);
|
||||
}
|
||||
|
||||
let mut payload = frame.payload.clone();
|
||||
let sequence = payload.get_u64();
|
||||
let expected_sequence = self.next_open_sequence.load(Ordering::Relaxed);
|
||||
if sequence != expected_sequence {
|
||||
return Err(TunnelSecurityError::UnexpectedSequence);
|
||||
}
|
||||
let nonce_bytes = nonce_bytes(self.open_prefix, sequence);
|
||||
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||
let clear_flags = frame.flags & !FLAG_ENCRYPTED;
|
||||
let aad = frame_aad(frame.stream_id, frame.msg_type, clear_flags);
|
||||
let plaintext = self
|
||||
.open
|
||||
.decrypt(
|
||||
nonce,
|
||||
Payload {
|
||||
msg: &payload,
|
||||
aad: &aad,
|
||||
},
|
||||
)
|
||||
.map_err(|_| TunnelSecurityError::Decrypt)?;
|
||||
self.next_open_sequence
|
||||
.store(expected_sequence.wrapping_add(1), Ordering::Relaxed);
|
||||
|
||||
Ok(Frame::new(
|
||||
frame.stream_id,
|
||||
frame.msg_type,
|
||||
clear_flags,
|
||||
Bytes::from(plaintext),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode_psk(key: &str) -> Result<[u8; 32], TunnelSecurityError> {
|
||||
let decoded = base64::engine::general_purpose::STANDARD
|
||||
.decode(key.trim())
|
||||
.map_err(|_| TunnelSecurityError::InvalidKey)?;
|
||||
decoded
|
||||
.try_into()
|
||||
.map_err(|_| TunnelSecurityError::InvalidKey)
|
||||
}
|
||||
|
||||
fn derive_key(psk: &[u8; 32], session_id: &[u8], label: &[u8]) -> [u8; 32] {
|
||||
let mut mac = <HmacSha256 as Mac>::new_from_slice(psk).expect("HMAC accepts 32-byte PSK");
|
||||
mac.update(CONTEXT);
|
||||
mac.update(&[0]);
|
||||
mac.update(session_id);
|
||||
mac.update(&[0]);
|
||||
mac.update(label);
|
||||
mac.finalize().into_bytes().into()
|
||||
}
|
||||
|
||||
fn nonce_bytes(prefix: [u8; 4], sequence: u64) -> [u8; NONCE_LEN] {
|
||||
let mut nonce = [0_u8; NONCE_LEN];
|
||||
nonce[..4].copy_from_slice(&prefix);
|
||||
nonce[4..].copy_from_slice(&sequence.to_be_bytes());
|
||||
nonce
|
||||
}
|
||||
|
||||
fn frame_aad(stream_id: u32, msg_type: MsgType, clear_flags: u8) -> [u8; HEADER_SIZE - 4] {
|
||||
let mut aad = [0_u8; HEADER_SIZE - 4];
|
||||
aad[..4].copy_from_slice(&stream_id.to_be_bytes());
|
||||
aad[4] = msg_type as u8;
|
||||
aad[5] = clear_flags;
|
||||
aad
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::tunnel::{Frame, MsgType};
|
||||
|
||||
fn test_key() -> String {
|
||||
base64::engine::general_purpose::STANDARD.encode([7_u8; 32])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secure_frame_round_trips_between_roles() {
|
||||
let client = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Client)
|
||||
.expect("client codec");
|
||||
let server = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Server)
|
||||
.expect("server codec");
|
||||
let frame = Frame::new(3, MsgType::RequestBody, 0, Bytes::from_static(b"secret"));
|
||||
|
||||
let encrypted = client.encrypt_frame(frame).expect("encrypt");
|
||||
assert!(!encrypted.windows(b"secret".len()).any(|w| w == b"secret"));
|
||||
|
||||
let wire = Frame::decode(encrypted).expect("wire frame");
|
||||
assert_ne!(wire.payload, Bytes::from_static(b"secret"));
|
||||
assert_ne!(wire.flags & FLAG_ENCRYPTED, 0);
|
||||
let decrypted = server.decrypt_frame(wire).expect("decrypt");
|
||||
|
||||
assert_eq!(decrypted.stream_id, 3);
|
||||
assert_eq!(decrypted.msg_type, MsgType::RequestBody);
|
||||
assert_eq!(decrypted.flags, 0);
|
||||
assert_eq!(decrypted.payload, Bytes::from_static(b"secret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secure_frame_rejects_wrong_session() {
|
||||
let client = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Client)
|
||||
.expect("client codec");
|
||||
let server = SecureFrameCodec::new(&test_key(), "session-2", TunnelSecurityRole::Server)
|
||||
.expect("server codec");
|
||||
let encrypted = client
|
||||
.encrypt_frame(Frame::new(
|
||||
1,
|
||||
MsgType::RequestBody,
|
||||
0,
|
||||
Bytes::from_static(b"secret"),
|
||||
))
|
||||
.expect("encrypt");
|
||||
let wire = Frame::decode(encrypted).expect("wire frame");
|
||||
|
||||
assert!(matches!(
|
||||
server.decrypt_frame(wire),
|
||||
Err(TunnelSecurityError::Decrypt)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secure_frame_rejects_replayed_sequence() {
|
||||
let client = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Client)
|
||||
.expect("client codec");
|
||||
let server = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Server)
|
||||
.expect("server codec");
|
||||
let encrypted = client
|
||||
.encrypt_frame(Frame::new(
|
||||
1,
|
||||
MsgType::RequestBody,
|
||||
0,
|
||||
Bytes::from_static(b"secret"),
|
||||
))
|
||||
.expect("encrypt");
|
||||
let wire = Frame::decode(encrypted).expect("wire frame");
|
||||
|
||||
server.decrypt_frame(wire.clone()).expect("first decrypt");
|
||||
assert!(matches!(
|
||||
server.decrypt_frame(wire),
|
||||
Err(TunnelSecurityError::UnexpectedSequence)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secure_frame_rejects_out_of_order_sequence_without_advancing() {
|
||||
let client = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Client)
|
||||
.expect("client codec");
|
||||
let server = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Server)
|
||||
.expect("server codec");
|
||||
let first = Frame::decode(
|
||||
client
|
||||
.encrypt_frame(Frame::new(
|
||||
1,
|
||||
MsgType::RequestBody,
|
||||
0,
|
||||
Bytes::from_static(b"first"),
|
||||
))
|
||||
.expect("encrypt first"),
|
||||
)
|
||||
.expect("first wire frame");
|
||||
let second = Frame::decode(
|
||||
client
|
||||
.encrypt_frame(Frame::new(
|
||||
1,
|
||||
MsgType::RequestBody,
|
||||
0,
|
||||
Bytes::from_static(b"second"),
|
||||
))
|
||||
.expect("encrypt second"),
|
||||
)
|
||||
.expect("second wire frame");
|
||||
|
||||
assert!(matches!(
|
||||
server.decrypt_frame(second),
|
||||
Err(TunnelSecurityError::UnexpectedSequence)
|
||||
));
|
||||
assert_eq!(
|
||||
server.decrypt_frame(first).expect("first decrypt").payload,
|
||||
Bytes::from_static(b"first")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secure_frame_uses_session_in_key_derivation() {
|
||||
let session_a = "node-1:connection-a";
|
||||
let session_b = "node-1:connection-b";
|
||||
let client_a = SecureFrameCodec::new(&test_key(), session_a, TunnelSecurityRole::Client)
|
||||
.expect("client codec a");
|
||||
let client_b = SecureFrameCodec::new(&test_key(), session_b, TunnelSecurityRole::Client)
|
||||
.expect("client codec b");
|
||||
let frame = Frame::new(
|
||||
7,
|
||||
MsgType::RequestBody,
|
||||
0,
|
||||
Bytes::from_static(b"same payload"),
|
||||
);
|
||||
|
||||
let encrypted_a = client_a.encrypt_frame(frame.clone()).expect("encrypt a");
|
||||
let encrypted_b = client_b.encrypt_frame(frame).expect("encrypt b");
|
||||
|
||||
assert_ne!(encrypted_a, encrypted_b);
|
||||
}
|
||||
}
|
||||
@@ -9,11 +9,12 @@ use uuid::Uuid;
|
||||
use super::types::{
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
log_reported_tunnel_error_event, normalize_proxy_metadata,
|
||||
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
|
||||
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
|
||||
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
|
||||
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
|
||||
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
StoredProxyNodeMetricsBucket, TunnelMetricsSample, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
@@ -590,6 +591,10 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
|
||||
previous_proxy_metadata.as_ref(),
|
||||
normalized_proxy_metadata,
|
||||
);
|
||||
if let Some(value) = normalized_proxy_metadata {
|
||||
node.proxy_metadata = Some(value);
|
||||
}
|
||||
|
||||
@@ -4,11 +4,12 @@ use sqlx::{mysql::MySqlRow, Row};
|
||||
use super::types::{
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
log_reported_tunnel_error_event, normalize_proxy_metadata,
|
||||
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
|
||||
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
|
||||
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
|
||||
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
|
||||
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
StoredProxyNodeMetricsBucket, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
@@ -779,10 +780,15 @@ WHERE is_manual = 0
|
||||
if let Some(value) = mutation.avg_latency_ms {
|
||||
node.avg_latency_ms = Some(value);
|
||||
}
|
||||
if let Some(value) = normalize_proxy_metadata(
|
||||
let normalized_proxy_metadata = normalize_proxy_metadata(
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
) {
|
||||
);
|
||||
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
|
||||
previous_proxy_metadata.as_ref(),
|
||||
normalized_proxy_metadata,
|
||||
);
|
||||
if let Some(value) = normalized_proxy_metadata {
|
||||
node.proxy_metadata = Some(value);
|
||||
}
|
||||
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {
|
||||
|
||||
@@ -6,11 +6,12 @@ use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||
use super::types::{
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
log_reported_tunnel_error_event, normalize_proxy_metadata,
|
||||
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
|
||||
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
|
||||
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
|
||||
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
|
||||
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
StoredProxyNodeMetricsBucket, TunnelMetricsSample, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::{
|
||||
@@ -1174,6 +1175,10 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
);
|
||||
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
|
||||
existing.proxy_metadata.as_ref(),
|
||||
normalized_proxy_metadata,
|
||||
);
|
||||
|
||||
sqlx::query(APPLY_HEARTBEAT_SQL)
|
||||
.bind(&mutation.node_id)
|
||||
|
||||
@@ -4,11 +4,12 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
|
||||
use super::types::{
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
log_reported_tunnel_error_event, normalize_proxy_metadata,
|
||||
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
|
||||
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
|
||||
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
|
||||
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
|
||||
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
|
||||
StoredProxyNodeMetricsBucket, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
@@ -770,10 +771,15 @@ WHERE is_manual = 0
|
||||
if let Some(value) = mutation.avg_latency_ms {
|
||||
node.avg_latency_ms = Some(value);
|
||||
}
|
||||
if let Some(value) = normalize_proxy_metadata(
|
||||
let normalized_proxy_metadata = normalize_proxy_metadata(
|
||||
mutation.proxy_metadata.as_ref(),
|
||||
mutation.proxy_version.as_deref(),
|
||||
) {
|
||||
);
|
||||
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
|
||||
previous_proxy_metadata.as_ref(),
|
||||
normalized_proxy_metadata,
|
||||
);
|
||||
if let Some(value) = normalized_proxy_metadata {
|
||||
node.proxy_metadata = Some(value);
|
||||
}
|
||||
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {
|
||||
|
||||
@@ -485,6 +485,34 @@ pub fn normalize_proxy_metadata(
|
||||
}
|
||||
}
|
||||
|
||||
pub fn preserve_proxy_metadata_tunnel_security(
|
||||
previous_proxy_metadata: Option<&Value>,
|
||||
next_proxy_metadata: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
let Some(tunnel_security) = previous_proxy_metadata
|
||||
.and_then(|value| value.get("tunnel_security"))
|
||||
.filter(|value| value.is_object())
|
||||
.cloned()
|
||||
else {
|
||||
return next_proxy_metadata;
|
||||
};
|
||||
|
||||
match next_proxy_metadata {
|
||||
Some(Value::Object(mut metadata)) => {
|
||||
metadata
|
||||
.entry("tunnel_security".to_string())
|
||||
.or_insert(tunnel_security);
|
||||
Some(Value::Object(metadata))
|
||||
}
|
||||
Some(value) => Some(value),
|
||||
None => {
|
||||
let mut metadata = serde_json::Map::new();
|
||||
metadata.insert("tunnel_security".to_string(), tunnel_security);
|
||||
Some(Value::Object(metadata))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_tunnel_metrics_counters(
|
||||
proxy_metadata: Option<&Value>,
|
||||
) -> Option<TunnelMetricsCounters> {
|
||||
@@ -772,10 +800,10 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
||||
normalize_proxy_node_scheduling_state, proxy_node_accepts_new_tunnels,
|
||||
proxy_reported_version, reconcile_remote_config_after_heartbeat,
|
||||
remote_config_scheduling_state, remote_config_upgrade_target, ProxyNodeMetricsStep,
|
||||
StoredProxyNode,
|
||||
normalize_proxy_node_scheduling_state, preserve_proxy_metadata_tunnel_security,
|
||||
proxy_node_accepts_new_tunnels, proxy_reported_version,
|
||||
reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
|
||||
remote_config_upgrade_target, ProxyNodeMetricsStep, StoredProxyNode,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -934,6 +962,40 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_secure_tunnel_metadata_across_heartbeat_metadata_refresh() {
|
||||
let previous = json!({
|
||||
"version": "1.0.0",
|
||||
"tunnel_security": {
|
||||
"mode": "non_tls_required",
|
||||
"encryption_key": "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc="
|
||||
}
|
||||
});
|
||||
let next = json!({
|
||||
"version": "1.0.1",
|
||||
"tunnel_metrics": {"connect_successes": 1}
|
||||
});
|
||||
|
||||
let merged = preserve_proxy_metadata_tunnel_security(Some(&previous), Some(next))
|
||||
.expect("metadata should remain present");
|
||||
assert_eq!(
|
||||
merged
|
||||
.pointer("/tunnel_security/mode")
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("non_tls_required")
|
||||
);
|
||||
assert_eq!(
|
||||
merged
|
||||
.pointer("/tunnel_security/encryption_key")
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=")
|
||||
);
|
||||
assert_eq!(
|
||||
merged.pointer("/tunnel_metrics/connect_successes"),
|
||||
Some(&json!(1))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_timestamps_to_metric_buckets() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -609,6 +609,9 @@ fn relay_envelope() -> Vec<u8> {
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
|
||||
@@ -117,6 +117,9 @@ fn relay_envelope() -> Vec<u8> {
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
stream: true,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
|
||||
@@ -253,6 +253,9 @@ fn relay_envelope() -> Vec<u8> {
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
stream: false,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 30,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
|
||||
@@ -537,9 +537,9 @@
|
||||
<!-- 请求链路追踪卡片 -->
|
||||
<div>
|
||||
<HorizontalRequestTimeline
|
||||
v-if="showTimeline && (detail.request_id || detail.id)"
|
||||
v-if="showTimeline && traceTimelineRequestId"
|
||||
ref="timelineRef"
|
||||
:request-id="detail.request_id || detail.id"
|
||||
:request-id="traceTimelineRequestId"
|
||||
:override-status-code="detail.status_code"
|
||||
:request-status="detail.status"
|
||||
:request-api-format="detail.api_format || null"
|
||||
@@ -996,6 +996,17 @@ function getNestedString(record: JsonRecord | null, ...path: string[]): string |
|
||||
return typeof value === 'string' && value.trim() ? value.trim() : null
|
||||
}
|
||||
|
||||
function getCaseInsensitiveString(record: JsonRecord | null, key: string): string | null {
|
||||
if (!record) return null
|
||||
const normalizedKey = key.toLowerCase()
|
||||
for (const [name, value] of Object.entries(record)) {
|
||||
if (name.toLowerCase() === normalizedKey && typeof value === 'string' && value.trim()) {
|
||||
return value.trim()
|
||||
}
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
function normalizeCacheTtlPricing(value: unknown): CacheTTLPriceEntry[] {
|
||||
if (!Array.isArray(value)) return []
|
||||
return value
|
||||
@@ -1095,6 +1106,19 @@ const traceRequestMetadata = computed<Record<string, unknown> | null>(() => {
|
||||
return meta as Record<string, unknown>
|
||||
})
|
||||
|
||||
const traceRecord = computed<Record<string, unknown> | null>(() =>
|
||||
asRecord(detail.value?.trace ?? null),
|
||||
)
|
||||
|
||||
const traceTimelineRequestId = computed(() =>
|
||||
getNestedString(traceRecord.value, 'trace_id')
|
||||
?? getNestedString(traceRequestMetadata.value, 'trace_id')
|
||||
?? getCaseInsensitiveString(asRecord(detail.value?.request_headers ?? null), 'x-trace-id')
|
||||
?? detail.value?.request_id
|
||||
?? detail.value?.id
|
||||
?? null,
|
||||
)
|
||||
|
||||
const metadataPanelData = computed<Record<string, unknown> | null>(() => {
|
||||
if (!detail.value) return null
|
||||
|
||||
|
||||
Reference in New Issue
Block a user