mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 18:59:50 +08:00
924 lines
32 KiB
Rust
924 lines
32 KiB
Rust
use std::io;
|
|
use std::net::SocketAddr;
|
|
use std::time::Duration;
|
|
|
|
use aether_contracts::tunnel::{
|
|
resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta,
|
|
TUNNEL_RELAY_FORWARDED_BY_HEADER,
|
|
};
|
|
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
|
|
use async_stream::stream;
|
|
use axum::body::{Body, Bytes};
|
|
use axum::extract::{ConnectInfo, Path, Request, State};
|
|
use axum::http::{HeaderMap, HeaderName, HeaderValue, Response, StatusCode};
|
|
use axum::response::IntoResponse;
|
|
use bytes::BytesMut;
|
|
use futures_util::StreamExt;
|
|
use tokio::sync::mpsc;
|
|
use tracing::warn;
|
|
|
|
use crate::api::response::apply_streaming_response_headers;
|
|
use crate::headers::should_skip_response_header;
|
|
use crate::maintenance::record_proxy_upgrade_traffic_success;
|
|
|
|
use super::hub::{LocalBodyEvent, LocalStream};
|
|
use super::protocol;
|
|
use super::AppState;
|
|
|
|
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
|
|
|
struct StreamGuard {
|
|
hub: std::sync::Arc<super::hub::HubRouter>,
|
|
stream_id: u64,
|
|
finished: bool,
|
|
}
|
|
|
|
impl Drop for StreamGuard {
|
|
fn drop(&mut self) {
|
|
if !self.finished {
|
|
self.hub
|
|
.cancel_local_stream(self.stream_id, "local relay client dropped");
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) struct DirectRelayResponse {
|
|
status: u16,
|
|
headers: Vec<(String, String)>,
|
|
body_rx: mpsc::Receiver<LocalBodyEvent>,
|
|
request_guard: StreamGuard,
|
|
_request_permit: Option<AdmissionPermit>,
|
|
}
|
|
|
|
impl DirectRelayResponse {
|
|
pub(crate) fn status(&self) -> u16 {
|
|
self.status
|
|
}
|
|
|
|
pub(crate) fn headers(&self) -> &[(String, String)] {
|
|
&self.headers
|
|
}
|
|
|
|
pub(crate) async fn next_chunk(&mut self) -> Result<Option<Bytes>, String> {
|
|
let event = self.body_rx.recv().await;
|
|
match event {
|
|
Some(LocalBodyEvent::Chunk(chunk)) => Ok(Some(chunk)),
|
|
Some(LocalBodyEvent::End) | None => {
|
|
self.request_guard.finished = true;
|
|
Ok(None)
|
|
}
|
|
Some(LocalBodyEvent::Error(error)) => {
|
|
self.request_guard.finished = true;
|
|
Err(error)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn open_direct_relay_stream(
|
|
state: &AppState,
|
|
node_id: &str,
|
|
meta: protocol::RequestMeta,
|
|
body: Bytes,
|
|
) -> Result<DirectRelayResponse, String> {
|
|
let request_permit = state
|
|
.try_acquire_request_permit()
|
|
.await
|
|
.map_err(map_request_admission_error)?;
|
|
let stream = state
|
|
.hub
|
|
.open_local_stream(node_id, &meta)
|
|
.await
|
|
.map_err(|error| format!("connect: {error}"))?;
|
|
if let Err(error) = state
|
|
.hub
|
|
.push_local_request_body(stream.id, body, true)
|
|
.await
|
|
{
|
|
state.hub.cancel_local_stream(stream.id, &error);
|
|
return Err(format!("connect: {error}"));
|
|
}
|
|
|
|
let wait_timeout = relay_header_timeout(&meta);
|
|
let response_head = match stream.wait_headers(wait_timeout).await {
|
|
Ok(response) => response,
|
|
Err(error) => {
|
|
state.hub.cancel_local_stream(stream.id, &error);
|
|
return Err(format!("timeout: {error}"));
|
|
}
|
|
};
|
|
if let Err(error) = record_proxy_upgrade_traffic_success(state.data.as_ref(), node_id).await {
|
|
warn!(
|
|
node_id = %node_id,
|
|
error = %error,
|
|
"failed to record proxy upgrade traffic confirmation"
|
|
);
|
|
}
|
|
|
|
let Some(body_rx) = stream.take_body_receiver() else {
|
|
state
|
|
.hub
|
|
.cancel_local_stream(stream.id, "missing relay response body receiver");
|
|
return Err("relay: missing relay response body receiver".to_string());
|
|
};
|
|
|
|
Ok(DirectRelayResponse {
|
|
status: response_head.status,
|
|
headers: response_head.headers,
|
|
body_rx,
|
|
request_guard: StreamGuard {
|
|
hub: state.hub.clone(),
|
|
stream_id: stream.id,
|
|
finished: false,
|
|
},
|
|
_request_permit: request_permit,
|
|
})
|
|
}
|
|
|
|
fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
|
|
match error {
|
|
super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Saturated {
|
|
..
|
|
})
|
|
| super::RequestAdmissionError::Distributed(
|
|
aether_runtime_state::RuntimeSemaphoreError::Saturated { .. },
|
|
)
|
|
| super::RequestAdmissionError::Distributed(
|
|
aether_runtime_state::RuntimeSemaphoreError::Unavailable { .. },
|
|
) => "overloaded: hub relay overloaded".to_string(),
|
|
super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Closed {
|
|
..
|
|
}) => "overloaded: hub relay gate closed".to_string(),
|
|
super::RequestAdmissionError::Distributed(
|
|
aether_runtime_state::RuntimeSemaphoreError::InvalidConfiguration(_),
|
|
) => "overloaded: hub relay distributed gate invalid".to_string(),
|
|
}
|
|
}
|
|
|
|
fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
|
|
Duration::from_millis(resolve_tunnel_request_timeouts(meta).first_byte_ms)
|
|
}
|
|
|
|
fn is_rollout_probe_request(headers: &HeaderMap, forwarded_by_gateway: bool) -> bool {
|
|
forwarded_by_gateway
|
|
&& headers
|
|
.get(crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER)
|
|
.and_then(|value| value.to_str().ok())
|
|
.map(str::trim)
|
|
.is_some_and(|value| value == crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE)
|
|
}
|
|
|
|
pub async fn relay_request(
|
|
Path(node_id): Path<String>,
|
|
State(state): State<AppState>,
|
|
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
|
request: Request,
|
|
) -> impl IntoResponse {
|
|
let forwarded_by_gateway = request
|
|
.headers()
|
|
.get(TUNNEL_RELAY_FORWARDED_BY_HEADER)
|
|
.and_then(|value| value.to_str().ok())
|
|
.map(str::trim)
|
|
.is_some_and(|value| !value.is_empty());
|
|
let rollout_probe = is_rollout_probe_request(request.headers(), forwarded_by_gateway);
|
|
if !addr.ip().is_loopback() && !forwarded_by_gateway {
|
|
return tunnel_error_response(
|
|
StatusCode::FORBIDDEN,
|
|
"forbidden",
|
|
"local relay only accepts loopback requests",
|
|
);
|
|
}
|
|
|
|
let request_permit = match state.try_acquire_request_permit().await {
|
|
Ok(permit) => permit,
|
|
Err(super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Saturated {
|
|
..
|
|
}))
|
|
| Err(super::RequestAdmissionError::Distributed(
|
|
aether_runtime_state::RuntimeSemaphoreError::Saturated { .. },
|
|
))
|
|
| Err(super::RequestAdmissionError::Distributed(
|
|
aether_runtime_state::RuntimeSemaphoreError::Unavailable { .. },
|
|
)) => {
|
|
return tunnel_error_response(
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
"overloaded",
|
|
"hub relay overloaded",
|
|
);
|
|
}
|
|
Err(super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Closed {
|
|
..
|
|
})) => {
|
|
return tunnel_error_response(
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
"overloaded",
|
|
"hub relay gate closed",
|
|
);
|
|
}
|
|
Err(super::RequestAdmissionError::Distributed(
|
|
aether_runtime_state::RuntimeSemaphoreError::InvalidConfiguration(_),
|
|
)) => {
|
|
return tunnel_error_response(
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
"overloaded",
|
|
"hub relay distributed gate invalid",
|
|
);
|
|
}
|
|
};
|
|
|
|
let mut body_stream = request.into_body().into_data_stream();
|
|
let mut envelope_buf = BytesMut::new();
|
|
let mut meta: Option<protocol::RequestMeta> = None;
|
|
let mut stream: Option<std::sync::Arc<LocalStream>> = None;
|
|
|
|
while let Some(chunk_result) = body_stream.next().await {
|
|
let chunk = match chunk_result {
|
|
Ok(chunk) => chunk,
|
|
Err(error) => {
|
|
if let Some(active_stream) = &stream {
|
|
state
|
|
.hub
|
|
.cancel_local_stream(active_stream.id, "failed to read relay request body");
|
|
}
|
|
warn!(error = %error, "failed to read local relay request body");
|
|
return release_permit_response(
|
|
tunnel_error_response(
|
|
StatusCode::BAD_GATEWAY,
|
|
"relay",
|
|
"failed to read relay request body",
|
|
),
|
|
request_permit,
|
|
);
|
|
}
|
|
};
|
|
|
|
if stream.is_none() {
|
|
envelope_buf.extend_from_slice(&chunk);
|
|
let Some((parsed_meta, body_offset)) =
|
|
(match try_decode_tunnel_relay_request_meta(&envelope_buf) {
|
|
Ok(result) => result,
|
|
Err(error) => {
|
|
return release_permit_response(
|
|
tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error),
|
|
request_permit,
|
|
);
|
|
}
|
|
})
|
|
else {
|
|
continue;
|
|
};
|
|
|
|
let opened_stream = match state.hub.open_local_stream(&node_id, &parsed_meta).await {
|
|
Ok(stream) => stream,
|
|
Err(error) => {
|
|
return release_permit_response(
|
|
tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
|
|
request_permit,
|
|
);
|
|
}
|
|
};
|
|
|
|
if envelope_buf.len() > body_offset {
|
|
let first_body_chunk = Bytes::copy_from_slice(&envelope_buf[body_offset..]);
|
|
if let Err(error) = state
|
|
.hub
|
|
.push_local_request_body(opened_stream.id, first_body_chunk, false)
|
|
.await
|
|
{
|
|
state.hub.cancel_local_stream(opened_stream.id, &error);
|
|
return release_permit_response(
|
|
tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
|
|
request_permit,
|
|
);
|
|
}
|
|
}
|
|
|
|
envelope_buf.clear();
|
|
meta = Some(parsed_meta);
|
|
stream = Some(opened_stream);
|
|
continue;
|
|
}
|
|
|
|
let Some(active_stream) = &stream else {
|
|
continue;
|
|
};
|
|
if let Err(error) = state
|
|
.hub
|
|
.push_local_request_body(active_stream.id, chunk, false)
|
|
.await
|
|
{
|
|
state.hub.cancel_local_stream(active_stream.id, &error);
|
|
return release_permit_response(
|
|
tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
|
|
request_permit,
|
|
);
|
|
}
|
|
}
|
|
|
|
let (meta, stream) = match (meta, stream) {
|
|
(Some(meta), Some(stream)) => (meta, stream),
|
|
_ => {
|
|
return release_permit_response(
|
|
tunnel_error_response(
|
|
StatusCode::BAD_REQUEST,
|
|
"bad_request",
|
|
"relay envelope metadata truncated",
|
|
),
|
|
request_permit,
|
|
);
|
|
}
|
|
};
|
|
|
|
if let Err(error) = state
|
|
.hub
|
|
.push_local_request_body(stream.id, Bytes::new(), true)
|
|
.await
|
|
{
|
|
state.hub.cancel_local_stream(stream.id, &error);
|
|
return release_permit_response(
|
|
tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
|
|
request_permit,
|
|
);
|
|
}
|
|
|
|
let request_guard = StreamGuard {
|
|
hub: state.hub.clone(),
|
|
stream_id: stream.id,
|
|
finished: false,
|
|
};
|
|
|
|
let wait_timeout = relay_header_timeout(&meta);
|
|
let response_head = match stream.wait_headers(wait_timeout).await {
|
|
Ok(response) => response,
|
|
Err(error) => {
|
|
state.hub.cancel_local_stream(stream.id, &error);
|
|
return release_permit_response(
|
|
tunnel_error_response(StatusCode::GATEWAY_TIMEOUT, "timeout", &error),
|
|
request_permit,
|
|
);
|
|
}
|
|
};
|
|
if !rollout_probe {
|
|
if let Err(error) =
|
|
record_proxy_upgrade_traffic_success(state.data.as_ref(), &node_id).await
|
|
{
|
|
warn!(
|
|
node_id = %node_id,
|
|
error = %error,
|
|
"failed to record proxy upgrade traffic confirmation"
|
|
);
|
|
}
|
|
}
|
|
|
|
let Some(mut body_rx) = stream.take_body_receiver() else {
|
|
state
|
|
.hub
|
|
.cancel_local_stream(stream.id, "missing relay response body receiver");
|
|
return release_permit_response(
|
|
tunnel_error_response(
|
|
StatusCode::BAD_GATEWAY,
|
|
"relay",
|
|
"missing relay response body receiver",
|
|
),
|
|
request_permit,
|
|
);
|
|
};
|
|
|
|
let hub = state.hub.clone();
|
|
let stream_id = stream.id;
|
|
let body_stream = stream! {
|
|
let mut guard = request_guard;
|
|
guard.hub = hub;
|
|
guard.stream_id = stream_id;
|
|
while let Some(event) = body_rx.recv().await {
|
|
match event {
|
|
LocalBodyEvent::Chunk(chunk) => yield Ok::<Bytes, io::Error>(chunk),
|
|
LocalBodyEvent::End => {
|
|
guard.finished = true;
|
|
break;
|
|
}
|
|
LocalBodyEvent::Error(error) => {
|
|
guard.finished = true;
|
|
yield Err(io::Error::other(error));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
guard.finished = true;
|
|
};
|
|
|
|
let mut builder = Response::builder().status(response_head.status);
|
|
if let Some(headers) = builder.headers_mut() {
|
|
append_headers(headers, &response_head.headers);
|
|
apply_streaming_response_headers(headers);
|
|
}
|
|
match builder.body(Body::from_stream(body_stream)) {
|
|
Ok(response) => maybe_hold_axum_response_permit(response, request_permit),
|
|
Err(error) => {
|
|
warn!(error = %error, "failed to build relay response");
|
|
release_permit_response(
|
|
tunnel_error_response(
|
|
StatusCode::BAD_GATEWAY,
|
|
"relay",
|
|
"failed to build relay response",
|
|
),
|
|
request_permit,
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn release_permit_response(
|
|
response: Response<Body>,
|
|
_request_permit: Option<AdmissionPermit>,
|
|
) -> Response<Body> {
|
|
response
|
|
}
|
|
|
|
fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
|
|
for (name, value) in headers {
|
|
if should_skip_local_relay_response_header(name) {
|
|
continue;
|
|
}
|
|
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
|
|
continue;
|
|
};
|
|
let Ok(value) = HeaderValue::from_str(value) else {
|
|
continue;
|
|
};
|
|
target.append(name, value);
|
|
}
|
|
}
|
|
|
|
fn should_skip_local_relay_response_header(name: &str) -> bool {
|
|
should_skip_response_header(name) || name.eq_ignore_ascii_case("content-length")
|
|
}
|
|
|
|
fn tunnel_error_response(status: StatusCode, kind: &str, message: &str) -> Response<Body> {
|
|
let mut builder = Response::builder().status(status);
|
|
if let Some(headers) = builder.headers_mut() {
|
|
headers.insert(
|
|
HeaderName::from_static(TUNNEL_ERROR_HEADER),
|
|
HeaderValue::from_str(kind).unwrap_or_else(|_| HeaderValue::from_static("relay")),
|
|
);
|
|
headers.insert(
|
|
axum::http::header::CONTENT_TYPE,
|
|
HeaderValue::from_static("text/plain; charset=utf-8"),
|
|
);
|
|
}
|
|
builder
|
|
.body(Body::from(message.to_string()))
|
|
.unwrap_or_else(|_| Response::new(Body::from("relay error")))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::super::hub::ProxyConn;
|
|
use super::super::{protocol, AppState, ConnConfig, ControlPlaneClient};
|
|
use super::{
|
|
is_rollout_probe_request, relay_header_timeout, relay_request, Body, HeaderMap, Request,
|
|
SocketAddr, StatusCode, TUNNEL_ERROR_HEADER,
|
|
};
|
|
use crate::data::GatewayDataState;
|
|
use crate::maintenance::start_proxy_upgrade_rollout;
|
|
use aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER;
|
|
use aether_data::repository::proxy_nodes::{
|
|
InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, ProxyNodeWriteRepository,
|
|
StoredProxyNode,
|
|
};
|
|
use axum::extract::ws::Message;
|
|
use axum::extract::{ConnectInfo, Path, State};
|
|
use axum::response::IntoResponse;
|
|
use bytes::Bytes;
|
|
use serde_json::json;
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use tokio::sync::watch;
|
|
|
|
#[test]
|
|
fn rollout_probe_marker_is_only_trusted_from_a_forwarding_gateway() {
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert(
|
|
crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER,
|
|
crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE
|
|
.parse()
|
|
.expect("probe marker should be a valid header"),
|
|
);
|
|
|
|
assert!(!is_rollout_probe_request(&headers, false));
|
|
assert!(is_rollout_probe_request(&headers, true));
|
|
}
|
|
|
|
fn test_app_state() -> AppState {
|
|
AppState::new(
|
|
ControlPlaneClient::disabled(),
|
|
ConnConfig {
|
|
ping_interval: Duration::from_secs(15),
|
|
idle_timeout: Duration::from_secs(0),
|
|
outbound_queue_capacity: 128,
|
|
},
|
|
128,
|
|
)
|
|
}
|
|
|
|
#[test]
|
|
fn relay_header_timeout_ignores_request_timeout_for_stream_requests() {
|
|
let meta = protocol::RequestMeta {
|
|
provider_id: None,
|
|
endpoint_id: None,
|
|
key_id: None,
|
|
method: "GET".to_string(),
|
|
url: "https://example.com/stream".to_string(),
|
|
headers: HashMap::new(),
|
|
stream: true,
|
|
request_timeout_ms: Some(90_000),
|
|
stream_first_byte_timeout_ms: None,
|
|
timeout: 7,
|
|
follow_redirects: None,
|
|
http1_only: false,
|
|
transport_profile: None,
|
|
};
|
|
|
|
assert_eq!(relay_header_timeout(&meta), Duration::from_secs(7));
|
|
}
|
|
|
|
#[test]
|
|
fn relay_header_timeout_keeps_the_protocol_maximum_for_non_stream_requests() {
|
|
let meta = protocol::RequestMeta {
|
|
provider_id: None,
|
|
endpoint_id: None,
|
|
key_id: None,
|
|
method: "POST".to_string(),
|
|
url: "https://example.com/responses".to_string(),
|
|
headers: HashMap::new(),
|
|
stream: false,
|
|
request_timeout_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS),
|
|
stream_first_byte_timeout_ms: None,
|
|
timeout: 60,
|
|
follow_redirects: None,
|
|
http1_only: false,
|
|
transport_profile: None,
|
|
};
|
|
|
|
assert_eq!(
|
|
relay_header_timeout(&meta),
|
|
Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
|
|
);
|
|
}
|
|
|
|
fn sample_connected_proxy_node(node_id: &str) -> StoredProxyNode {
|
|
StoredProxyNode::new(
|
|
node_id.to_string(),
|
|
format!("proxy-{node_id}"),
|
|
"127.0.0.1".to_string(),
|
|
0,
|
|
false,
|
|
"online".to_string(),
|
|
30,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
true,
|
|
true,
|
|
0,
|
|
)
|
|
.expect("node should build")
|
|
.with_runtime_fields(
|
|
Some("test".to_string()),
|
|
None,
|
|
Some(1_800_000_000),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
Some(1_800_000_000),
|
|
None,
|
|
Some(1_800_000_000),
|
|
Some(1_800_000_000),
|
|
)
|
|
}
|
|
|
|
fn encode_relay_envelope(meta: &protocol::RequestMeta, body: &[u8]) -> Vec<u8> {
|
|
let meta_bytes = serde_json::to_vec(meta).expect("meta should serialize");
|
|
let mut payload = Vec::with_capacity(4 + meta_bytes.len() + body.len());
|
|
payload.extend_from_slice(&(meta_bytes.len() as u32).to_be_bytes());
|
|
payload.extend_from_slice(&meta_bytes);
|
|
payload.extend_from_slice(body);
|
|
payload
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn relay_rejects_non_loopback_without_forwarded_header() {
|
|
let request = Request::builder()
|
|
.body(Body::empty())
|
|
.expect("request should build");
|
|
let response = relay_request(
|
|
Path("node-123".to_string()),
|
|
State(test_app_state()),
|
|
ConnectInfo(SocketAddr::from(([10, 0, 0, 1], 4242))),
|
|
request,
|
|
)
|
|
.await
|
|
.into_response();
|
|
|
|
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn relay_accepts_forwarded_gateway_request_from_non_loopback() {
|
|
let request = Request::builder()
|
|
.header(TUNNEL_RELAY_FORWARDED_BY_HEADER, "gateway-a")
|
|
.body(Body::empty())
|
|
.expect("request should build");
|
|
let response = relay_request(
|
|
Path("node-123".to_string()),
|
|
State(test_app_state()),
|
|
ConnectInfo(SocketAddr::from(([10, 0, 0, 1], 4242))),
|
|
request,
|
|
)
|
|
.await
|
|
.into_response();
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
assert_eq!(
|
|
response
|
|
.headers()
|
|
.get(TUNNEL_ERROR_HEADER)
|
|
.and_then(|value| value.to_str().ok()),
|
|
Some("bad_request")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn relay_records_real_traffic_confirmation_for_upgrade_rollout() {
|
|
let mut node = sample_connected_proxy_node("node-123");
|
|
node.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![node]));
|
|
let data = Arc::new(
|
|
GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new()),
|
|
);
|
|
|
|
let started = start_proxy_upgrade_rollout(data.as_ref(), "2.0.0".to_string(), 1, 0, None)
|
|
.await
|
|
.expect("rollout should start");
|
|
assert_eq!(started.node_ids, vec!["node-123".to_string()]);
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-123".to_string(),
|
|
heartbeat_interval: None,
|
|
active_connections: Some(1),
|
|
total_requests_delta: Some(1),
|
|
avg_latency_ms: Some(2.0),
|
|
failed_requests_delta: Some(0),
|
|
dns_failures_delta: Some(0),
|
|
stream_errors_delta: Some(0),
|
|
proxy_metadata: Some(json!({"version": "2.0.0"})),
|
|
proxy_version: Some("2.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should succeed");
|
|
|
|
let observed = start_proxy_upgrade_rollout(data.as_ref(), "2.0.0".to_string(), 1, 0, None)
|
|
.await
|
|
.expect("rollout should observe version confirmation");
|
|
assert!(observed.blocked);
|
|
assert_eq!(observed.pending_node_ids, vec!["node-123".to_string()]);
|
|
|
|
let state = test_app_state().with_data(Arc::clone(&data));
|
|
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
|
|
let (proxy_close_tx, _) = watch::channel(false);
|
|
state.hub.register_proxy(Arc::new(ProxyConn::new(
|
|
500,
|
|
"node-123".to_string(),
|
|
"Node 123".to_string(),
|
|
proxy_tx,
|
|
proxy_close_tx,
|
|
16,
|
|
2,
|
|
)));
|
|
|
|
let meta = protocol::RequestMeta {
|
|
provider_id: None,
|
|
endpoint_id: None,
|
|
key_id: None,
|
|
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,
|
|
transport_profile: None,
|
|
};
|
|
let request = Request::builder()
|
|
.body(Body::from(encode_relay_envelope(&meta, &[])))
|
|
.expect("request should build");
|
|
|
|
let relay_state = state.clone();
|
|
let relay_task = tokio::spawn(async move {
|
|
relay_request(
|
|
Path("node-123".to_string()),
|
|
State(relay_state),
|
|
ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 4242))),
|
|
request,
|
|
)
|
|
.await
|
|
.into_response()
|
|
});
|
|
|
|
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
|
|
Message::Binary(data) => data,
|
|
other => panic!("unexpected message: {other:?}"),
|
|
};
|
|
let request_header = protocol::FrameHeader::parse(&request_headers)
|
|
.expect("request header frame should parse");
|
|
assert_eq!(request_header.msg_type, protocol::REQUEST_HEADERS);
|
|
|
|
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
|
|
Message::Binary(data) => data,
|
|
other => panic!("unexpected message: {other:?}"),
|
|
};
|
|
let request_body_header =
|
|
protocol::FrameHeader::parse(&request_body).expect("request body frame should parse");
|
|
assert_eq!(request_body_header.msg_type, protocol::REQUEST_BODY);
|
|
|
|
let response_meta = protocol::ResponseMeta {
|
|
status: 200,
|
|
headers: vec![("content-type".to_string(), "text/plain".to_string())],
|
|
};
|
|
let response_payload =
|
|
serde_json::to_vec(&response_meta).expect("response meta should serialize");
|
|
let mut response_headers_frame = protocol::encode_frame(
|
|
request_header.stream_id,
|
|
protocol::RESPONSE_HEADERS,
|
|
0,
|
|
&response_payload,
|
|
);
|
|
state
|
|
.hub
|
|
.handle_proxy_frame(500, &mut response_headers_frame)
|
|
.await;
|
|
|
|
let mut response_body_frame = protocol::encode_frame(
|
|
request_header.stream_id,
|
|
protocol::RESPONSE_BODY,
|
|
0,
|
|
Bytes::new().as_ref(),
|
|
);
|
|
state
|
|
.hub
|
|
.handle_proxy_frame(500, &mut response_body_frame)
|
|
.await;
|
|
let mut response_end_frame =
|
|
protocol::encode_frame(request_header.stream_id, protocol::STREAM_END, 0, &[]);
|
|
state
|
|
.hub
|
|
.handle_proxy_frame(500, &mut response_end_frame)
|
|
.await;
|
|
|
|
let response = relay_task.await.expect("relay task should complete");
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
|
|
.await
|
|
.expect("response body should read");
|
|
assert!(body.is_empty());
|
|
|
|
let rollout_entry = data
|
|
.list_system_config_entries()
|
|
.await
|
|
.expect("system config list should succeed")
|
|
.into_iter()
|
|
.find(|entry| entry.key == "proxy_node_upgrade_rollout")
|
|
.expect("rollout entry should exist");
|
|
let tracked_nodes = rollout_entry.value["tracked_nodes"]
|
|
.as_array()
|
|
.expect("tracked nodes should be an array");
|
|
assert_eq!(tracked_nodes.len(), 1);
|
|
assert!(tracked_nodes[0]["version_confirmed_at_unix_secs"].is_u64());
|
|
assert!(tracked_nodes[0]["traffic_confirmed_at_unix_secs"].is_u64());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn relay_strips_hop_by_hop_and_stale_length_headers_from_proxy_response() {
|
|
let state = test_app_state();
|
|
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
|
|
let (proxy_close_tx, _) = watch::channel(false);
|
|
state.hub.register_proxy(Arc::new(ProxyConn::new(
|
|
501,
|
|
"node-123".to_string(),
|
|
"Node 123".to_string(),
|
|
proxy_tx,
|
|
proxy_close_tx,
|
|
16,
|
|
2,
|
|
)));
|
|
|
|
let meta = protocol::RequestMeta {
|
|
provider_id: None,
|
|
endpoint_id: None,
|
|
key_id: None,
|
|
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,
|
|
transport_profile: None,
|
|
};
|
|
let request = Request::builder()
|
|
.body(Body::from(encode_relay_envelope(&meta, &[])))
|
|
.expect("request should build");
|
|
|
|
let relay_state = state.clone();
|
|
let relay_task = tokio::spawn(async move {
|
|
relay_request(
|
|
Path("node-123".to_string()),
|
|
State(relay_state),
|
|
ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 4242))),
|
|
request,
|
|
)
|
|
.await
|
|
.into_response()
|
|
});
|
|
|
|
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
|
|
Message::Binary(data) => data,
|
|
other => panic!("unexpected message: {other:?}"),
|
|
};
|
|
let request_header = protocol::FrameHeader::parse(&request_headers)
|
|
.expect("request header frame should parse");
|
|
assert_eq!(request_header.msg_type, protocol::REQUEST_HEADERS);
|
|
|
|
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
|
|
Message::Binary(data) => data,
|
|
other => panic!("unexpected message: {other:?}"),
|
|
};
|
|
let request_body_header =
|
|
protocol::FrameHeader::parse(&request_body).expect("request body frame should parse");
|
|
assert_eq!(request_body_header.msg_type, protocol::REQUEST_BODY);
|
|
|
|
let response_meta = protocol::ResponseMeta {
|
|
status: 200,
|
|
headers: vec![
|
|
("content-length".to_string(), "999".to_string()),
|
|
("transfer-encoding".to_string(), "chunked".to_string()),
|
|
("connection".to_string(), "keep-alive".to_string()),
|
|
("content-type".to_string(), "text/plain".to_string()),
|
|
(
|
|
"x-proxy-timing".to_string(),
|
|
"{\"mode\":\"tunnel\"}".to_string(),
|
|
),
|
|
],
|
|
};
|
|
let response_payload =
|
|
serde_json::to_vec(&response_meta).expect("response meta should serialize");
|
|
let mut response_headers_frame = protocol::encode_frame(
|
|
request_header.stream_id,
|
|
protocol::RESPONSE_HEADERS,
|
|
0,
|
|
&response_payload,
|
|
);
|
|
state
|
|
.hub
|
|
.handle_proxy_frame(501, &mut response_headers_frame)
|
|
.await;
|
|
|
|
let mut response_end_frame =
|
|
protocol::encode_frame(request_header.stream_id, protocol::STREAM_END, 0, &[]);
|
|
state
|
|
.hub
|
|
.handle_proxy_frame(501, &mut response_end_frame)
|
|
.await;
|
|
|
|
let response = relay_task.await.expect("relay task should complete");
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert!(response.headers().get("content-length").is_none());
|
|
assert!(response.headers().get("transfer-encoding").is_none());
|
|
assert!(response.headers().get("connection").is_none());
|
|
assert_eq!(
|
|
response
|
|
.headers()
|
|
.get("content-type")
|
|
.and_then(|value| value.to_str().ok()),
|
|
Some("text/plain")
|
|
);
|
|
assert_eq!(
|
|
response
|
|
.headers()
|
|
.get("x-proxy-timing")
|
|
.and_then(|value| value.to_str().ok()),
|
|
Some("{\"mode\":\"tunnel\"}")
|
|
);
|
|
}
|
|
}
|