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

feat(proxy): 实现代理节点批量升级回滚、隧道重定向跟随及远程配置管理

Browse Source 
核心功能：
- 新增代理节点批量升级回滚工作流，支持分批升级、健康探针、跳过/重试/取消等操作
- proxy 隧道流处理器支持 HTTP 重定向跟随（最多 10 跳），区分 307/308 可重播与不可重播请求体
- proxy 协议新增 follow_redirects / http1_only 字段，网关侧同步支持
- 新增代理节点远端配置变更接口（名称、允许端口、调度状态、升级目标等）
- 新增代理节点注册/反注册/心跳的 Admin API，及节点过期清理维护任务
- gateway 隧道 owner-relay 支持流式代理大请求体，新增 5 MiB 默认限制
- 新增 ProxyNodeRegistrationMutation / ProxyNodeRemoteConfigMutation 数据类型
- proxy 配置新增重定向重播预算、心跳间隔等参数，TUI 安装向导同步更新
- 前端 ProxyNodes 页面新增批量升级操作面板及滚动进度展示
This commit is contained in:
fawney19
2026-04-12 16:02:38 +08:00
parent 7c5bb7f383
commit 9703840a36
83 changed files with 11832 additions and 1520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
apps/aether-gateway/src/tunnel/embedded/control_plane.rs
View File

@@ -6,7 +6,7 @@ use std::sync::Arc;
type HeartbeatAckCallback =
dyn Fn(Vec<u8>) -> BoxFuture<'static, Result<Vec<u8>, String>> + Send + Sync;
type NodeStatusCallback =
dyn Fn(String, bool, usize) -> BoxFuture<'static, Result<(), String>> + Send + Sync;
dyn Fn(String, bool, usize, u64) -> BoxFuture<'static, Result<(), String>> + Send + Sync;
enum ControlPlaneMode {
Disabled,
@@ -51,7 +51,7 @@ impl ControlPlaneClient {
where
HeartbeatAck:
Fn(Vec<u8>) -> BoxFuture<'static, Result<Vec<u8>, String>> + Send + Sync + 'static,
PushNodeStatus: Fn(String, bool, usize) -> BoxFuture<'static, Result<(), String>>
PushNodeStatus: Fn(String, bool, usize, u64) -> BoxFuture<'static, Result<(), String>>
+ Send
+ Sync
+ 'static,
@@ -103,6 +103,7 @@ impl ControlPlaneClient {
node_id: &str,
connected: bool,
conn_count: usize,
observed_at_unix_secs: u64,
) -> Result<(), String> {
match self.inner.as_ref() {
ControlPlaneMode::Disabled => Ok(()),
@@ -120,6 +121,7 @@ impl ControlPlaneClient {
"node_id": node_id,
"connected": connected,
"conn_count": conn_count,
"observed_at_unix_secs": observed_at_unix_secs,
}))
.send()
.await
@@ -135,7 +137,15 @@ impl ControlPlaneClient {
}
ControlPlaneMode::Local {
push_node_status, ..
} => push_node_status(node_id.to_string(), connected, conn_count).await,
} => {
push_node_status(
node_id.to_string(),
connected,
conn_count,
observed_at_unix_secs,
)
.await
}
}
}
}

116
apps/aether-gateway/src/tunnel/embedded/hub.rs
View File

@@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_runtime::{BoundedQueueSender, MetricKind, MetricSample, QueueSendError};
use axum::extract::ws::Message;
@@ -324,10 +324,46 @@ pub struct HubRouter {
next_conn_id: AtomicU64,
next_local_stream_id: AtomicU64,
control_plane: ControlPlaneClient,
node_status_tx: mpsc::UnboundedSender<NodeStatusEvent>,
}
#[derive(Debug)]
struct NodeStatusEvent {
node_id: String,
connected: bool,
conn_count: usize,
observed_at_unix_secs: u64,
}
impl HubRouter {
pub fn new(control_plane: ControlPlaneClient) -> Arc<Self> {
let (node_status_tx, mut node_status_rx) = mpsc::unbounded_channel::<NodeStatusEvent>();
let worker_control_plane = control_plane.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
while let Some(event) = node_status_rx.recv().await {
if let Err(error) = worker_control_plane
.push_node_status(
&event.node_id,
event.connected,
event.conn_count,
event.observed_at_unix_secs,
)
.await
{
warn!(
node_id = %event.node_id,
connected = event.connected,
conn_count = event.conn_count,
observed_at_unix_secs = event.observed_at_unix_secs,
error = %error,
"failed to push node status to app control plane"
);
}
}
});
}
Arc::new(Self {
proxy_conns: RwLock::new(HashMap::new()),
proxy_conns_by_id: DashMap::new(),
@@ -336,6 +372,7 @@ impl HubRouter {
next_conn_id: AtomicU64::new(1),
next_local_stream_id: AtomicU64::new(1),
control_plane,
node_status_tx,
})
}
@@ -405,21 +442,21 @@ impl HubRouter {
}
fn notify_node_status(&self, node_id: String, connected: bool, conn_count: usize) {
let control_plane = self.control_plane.clone();
tokio::spawn(async move {
if let Err(error) = control_plane
.push_node_status(&node_id, connected, conn_count)
.await
{
warn!(
node_id = %node_id,
connected = connected,
conn_count = conn_count,
error = %error,
"failed to push node status to app control plane"
);
}
});
let event = NodeStatusEvent {
node_id,
connected,
conn_count,
observed_at_unix_secs: current_unix_secs(),
};
if let Err(error) = self.node_status_tx.send(event) {
warn!(
node_id = %error.0.node_id,
connected = error.0.connected,
conn_count = error.0.conn_count,
observed_at_unix_secs = error.0.observed_at_unix_secs,
"node status worker unavailable"
);
}
}
fn get_proxy_conn(&self, node_id: &str) -> Option<Arc<ProxyConn>> {
@@ -750,8 +787,12 @@ impl HubRouter {
let ack_payload = match self.control_plane.heartbeat_ack(&payload).await {
Ok(payload) => payload,
Err(error) => {
warn!(proxy_conn_id = proxy_conn_id, error = %error, "control-plane heartbeat callback failed");
b"{}".to_vec()
warn!(
proxy_conn_id = proxy_conn_id,
error = %error,
"control-plane heartbeat callback failed; keeping heartbeat pending"
);
return;
}
};
if let Some(pc) = self.proxy_conns_by_id.get(&proxy_conn_id) {
@@ -796,6 +837,13 @@ impl HubRouter {
}
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[derive(serde::Serialize)]
pub struct HubStats {
pub proxy_connections: usize,
@@ -845,6 +893,8 @@ mod tests {
url: "https://example.com".to_string(),
headers: HashMap::new(),
timeout: 30,
follow_redirects: None,
http1_only: false,
}
}
@@ -924,4 +974,34 @@ mod tests {
assert_eq!(second_header.msg_type, protocol::REQUEST_BODY);
assert_ne!(second_header.flags & protocol::FLAG_END_STREAM, 0);
}
#[tokio::test]
async fn heartbeat_callback_failure_does_not_send_fake_ack() {
let hub = HubRouter::new(ControlPlaneClient::local(
|_payload| Box::pin(async { Err("db unavailable".to_string()) }),
|_node_id, _connected, _conn_count, _observed_at_unix_secs| Box::pin(async { Ok(()) }),
));
let (proxy_tx, mut proxy_rx) = bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
let proxy = Arc::new(ProxyConn::new(
300,
"node-3".to_string(),
"Node 3".to_string(),
proxy_tx,
proxy_close_tx,
16,
));
hub.register_proxy(proxy);
let payload = serde_json::to_vec(&serde_json::json!({
"node_id": "node-3",
"heartbeat_id": 99u64,
}))
.expect("payload should serialize");
let mut frame = protocol::encode_frame(1, protocol::HEARTBEAT_DATA, 0, &payload);
hub.handle_proxy_frame(300, &mut frame).await;
assert!(proxy_rx.try_recv().is_err());
}
}

211
apps/aether-gateway/src/tunnel/embedded/local_relay.rs
View File

@@ -14,6 +14,7 @@ use futures_util::StreamExt;
use tracing::warn;
use crate::api::response::apply_streaming_response_headers;
use crate::maintenance::record_proxy_upgrade_traffic_success;
use super::hub::{LocalBodyEvent, LocalStream};
use super::protocol;
@@ -222,6 +223,13 @@ pub async fn relay_request(
);
}
};
if let Err(error) = record_proxy_upgrade_traffic_success(state.data.as_ref(), &node_id).await {
warn!(
node_id = %node_id,
error = %error,
"failed to record proxy upgrade traffic confirmation"
);
}
let Some(mut body_rx) = stream.take_body_receiver() else {
state
@@ -340,12 +348,25 @@ fn tunnel_error_response(status: StatusCode, kind: &str, message: &str) -> Respo
#[cfg(test)]
mod tests {
use super::super::{AppState, ConnConfig, ControlPlaneClient};
use super::super::hub::ProxyConn;
use super::super::{protocol, AppState, ConnConfig, ControlPlaneClient};
use super::{relay_request, Body, 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;
fn test_app_state() -> AppState {
AppState::new(
@@ -359,6 +380,49 @@ mod tests {
)
}
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()
@@ -400,4 +464,149 @@ mod tests {
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,
)));
let meta = protocol::RequestMeta {
method: "GET".to_string(),
url: "https://example.com/health".to_string(),
headers: HashMap::new(),
timeout: 30,
follow_redirects: None,
http1_only: false,
};
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());
}
}

11
apps/aether-gateway/src/tunnel/embedded/mod.rs
View File

@@ -19,8 +19,10 @@ use axum::routing::{get, post};
use axum::Router;
use tracing::warn;
use crate::data::GatewayDataState;
pub use control_plane::ControlPlaneClient;
pub use hub::{ConnConfig, HubRouter};
pub use hub::{ConnConfig, HubRouter, ProxyConn};
pub use local_relay::relay_request;
#[derive(Clone)]
@@ -28,6 +30,7 @@ pub struct AppState {
pub hub: Arc<HubRouter>,
pub proxy_conn_cfg: ConnConfig,
pub max_streams: usize,
data: Arc<GatewayDataState>,
request_gate: Option<Arc<ConcurrencyGate>>,
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
}
@@ -48,11 +51,17 @@ impl AppState {
hub: HubRouter::new(control_plane),
proxy_conn_cfg,
max_streams,
data: Arc::new(GatewayDataState::disabled()),
request_gate: None,
distributed_request_gate: None,
}
}
pub(crate) fn with_data(mut self, data: Arc<GatewayDataState>) -> Self {
self.data = data;
self
}
pub fn with_request_concurrency_limit(mut self, limit: Option<usize>) -> Self {
self.request_gate = limit
.filter(|limit| *limit > 0)

211
apps/aether-gateway/src/tunnel/mod.rs
View File

@@ -1,6 +1,9 @@
mod embedded;
use std::collections::HashMap;
use std::fmt;
use std::io;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
@@ -11,11 +14,13 @@ use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode,
};
use aether_runtime::MetricSample;
use axum::body::{to_bytes, Body};
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::extract::ws::WebSocketUpgrade;
use axum::extract::{ConnectInfo, Path, Request, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::IntoResponse;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::json;
use tracing::warn;
@@ -28,6 +33,7 @@ use super::error::GatewayError;
use super::headers::{extract_or_generate_trace_id, should_skip_request_header};
use super::AppState;
pub(crate) use embedded::ProxyConn as TunnelProxyConn;
pub use embedded::{
build_router_with_state as build_tunnel_runtime_router_with_state, protocol as tunnel_protocol,
AppState as TunnelRuntimeState, ConnConfig as TunnelConnConfig,
@@ -45,6 +51,7 @@ const DEFAULT_PING_INTERVAL_SECS: u64 = 15;
const DEFAULT_MAX_STREAMS: usize = 2048;
const DEFAULT_OUTBOUND_QUEUE_CAPACITY: usize = 128;
const DEFAULT_ATTACHMENT_TTL_SECS: u64 = 90;
const DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES: usize = 5_242_880;
const TUNNEL_ATTACHMENT_KEY_PREFIX: &str = "tunnel.attachments.";
const TUNNEL_ATTACHMENT_REDIS_KEY_PREFIX: &str = "tunnel:attachments:";
const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID";
@@ -55,6 +62,8 @@ const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS";
struct InternalTunnelHeartbeatRequest {
node_id: String,
#[serde(default)]
heartbeat_id: Option<u64>,
#[serde(default)]
heartbeat_interval: Option<i32>,
#[serde(default)]
active_connections: Option<i32>,
@@ -183,6 +192,7 @@ impl TunnelAttachmentDirectory {
node_id: &str,
connected: bool,
conn_count: usize,
observed_at_unix_secs: u64,
) -> Result<(), String> {
let node_id = node_id.trim();
if node_id.is_empty() {
@@ -203,7 +213,7 @@ impl TunnelAttachmentDirectory {
gateway_instance_id: self.identity.instance_id.clone(),
relay_base_url: relay_base_url.clone(),
conn_count,
observed_at_unix_secs: current_unix_secs(),
observed_at_unix_secs,
},
)
.await
@@ -404,7 +414,8 @@ impl EmbeddedTunnelState {
outbound_queue_capacity: DEFAULT_OUTBOUND_QUEUE_CAPACITY,
},
DEFAULT_MAX_STREAMS,
),
)
.with_data(data),
attachment_directory,
}
}
@@ -434,6 +445,39 @@ impl EmbeddedTunnelState {
self.inner.hub.stats().to_metric_samples()
}
pub(crate) async fn probe_node_url(
&self,
node_id: &str,
url: &str,
timeout_secs: u64,
) -> Result<u16, String> {
let timeout_secs = timeout_secs.clamp(5, 60);
let meta = tunnel_protocol::RequestMeta {
method: "GET".to_string(),
url: url.trim().to_string(),
headers: HashMap::new(),
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,
};
let stream = self.inner.hub.open_local_stream(node_id, &meta)?;
let stream_id = stream.id;
let result = async {
self.inner
.hub
.push_local_request_body(stream_id, Bytes::new(), true)?;
let response = stream
.wait_headers(Duration::from_secs(timeout_secs))
.await?;
Ok(response.status)
}
.await;
self.inner
.hub
.cancel_local_stream(stream_id, "tunnel health probe completed");
result
}
pub(crate) fn local_instance_id(&self) -> &str {
self.attachment_directory.local_instance_id()
}
@@ -579,14 +623,26 @@ fn build_embedded_control_plane(
Ok(ack)
})
},
move |node_id, connected, conn_count| {
move |node_id, connected, conn_count, observed_at_unix_secs| {
let data = Arc::clone(&node_status_data);
let directory = node_status_directory.clone();
Box::pin(async move {
apply_embedded_tunnel_node_status(data.as_ref(), &node_id, connected, conn_count)
.await?;
apply_embedded_tunnel_node_status(
data.as_ref(),
&node_id,
connected,
conn_count,
Some(observed_at_unix_secs),
)
.await?;
if let Err(error) = directory
.sync_node_status(data.as_ref(), &node_id, connected, conn_count)
.sync_node_status(
data.as_ref(),
&node_id,
connected,
conn_count,
observed_at_unix_secs,
)
.await
{
warn!(error = %error, node_id = %node_id, "failed to sync tunnel attachment");
@@ -606,9 +662,16 @@ async fn forward_relay_request_to_owner(
) -> Result<axum::http::Response<Body>, GatewayError> {
let owner_url = build_owner_relay_url(&owner.relay_base_url, node_id)?;
let (parts, body) = request.into_parts();
let body = to_bytes(body, usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_limit = owner_relay_body_limit_bytes(state.data.as_ref()).await;
if request_content_length_exceeds_limit(&parts.headers, body_limit) {
return build_local_http_error_response(
trace_id,
None,
StatusCode::PAYLOAD_TOO_LARGE,
&format!("tunnel relay body exceeds {body_limit} bytes"),
);
}
let limit_exceeded = Arc::new(AtomicBool::new(false));
let mut upstream_request = state.client.post(owner_url);
for (name, value) in &parts.headers {
@@ -630,11 +693,30 @@ async fn forward_relay_request_to_owner(
upstream_request = upstream_request.header(TRACE_ID_HEADER, trace_id);
}
let upstream_response = upstream_request
.body(body)
let upstream_response = match upstream_request
.body(build_owner_relay_request_body(
body,
body_limit,
Arc::clone(&limit_exceeded),
))
.send()
.await
.map_err(|err| GatewayError::Internal(format!("owner tunnel relay failed: {err}")))?;
{
Ok(response) => response,
Err(err) if limit_exceeded.load(Ordering::SeqCst) => {
return build_local_http_error_response(
trace_id,
None,
StatusCode::PAYLOAD_TOO_LARGE,
&format!("tunnel relay body exceeds {body_limit} bytes"),
);
}
Err(err) => {
return Err(GatewayError::Internal(format!(
"owner tunnel relay failed: {err}"
)));
}
};
build_client_response(upstream_response, trace_id, None)
}
@@ -656,6 +738,57 @@ fn build_owner_relay_url(relay_base_url: &str, node_id: &str) -> Result<String,
Ok(url.to_string())
}
async fn owner_relay_body_limit_bytes(data: &GatewayDataState) -> usize {
data.find_system_config_value("max_request_body_size")
.await
.ok()
.flatten()
.and_then(|value| value.as_u64())
.and_then(|value| usize::try_from(value).ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES)
}
fn request_content_length_exceeds_limit(headers: &HeaderMap, body_limit: usize) -> bool {
headers
.get(http::header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.and_then(|value| usize::try_from(value).ok())
.is_some_and(|value| value > body_limit)
}
fn build_owner_relay_request_body(
body: Body,
body_limit: usize,
limit_exceeded: Arc<AtomicBool>,
) -> reqwest::Body {
let mut body_stream = body.into_data_stream();
reqwest::Body::wrap_stream(stream! {
let mut forwarded = 0usize;
while let Some(next_chunk) = body_stream.next().await {
match next_chunk {
Ok(chunk) => {
forwarded = forwarded.saturating_add(chunk.len());
if forwarded > body_limit {
limit_exceeded.store(true, Ordering::SeqCst);
yield Err::<Bytes, io::Error>(io::Error::new(
io::ErrorKind::InvalidData,
format!("tunnel relay body exceeds {body_limit} bytes"),
));
break;
}
yield Ok::<Bytes, io::Error>(chunk);
}
Err(err) => {
yield Err::<Bytes, io::Error>(io::Error::other(err));
break;
}
}
}
})
}
fn tunnel_attachment_key(node_id: &str) -> String {
format!("{TUNNEL_ATTACHMENT_KEY_PREFIX}{}", node_id.trim())
}
@@ -719,7 +852,10 @@ async fn apply_embedded_tunnel_heartbeat(
.map_err(|err| format!("heartbeat sync failed: {err}"))?
.ok_or_else(|| format!("heartbeat sync failed: ProxyNode {node_id} 不存在"))?;
Ok(build_embedded_tunnel_heartbeat_ack(&node))
Ok(build_embedded_tunnel_heartbeat_ack(
&node,
payload.heartbeat_id,
))
}
async fn apply_embedded_tunnel_node_status(
@@ -727,13 +863,14 @@ async fn apply_embedded_tunnel_node_status(
node_id: &str,
connected: bool,
conn_count: usize,
observed_at_unix_secs: Option<u64>,
) -> Result<(), String> {
let mutation = ProxyNodeTunnelStatusMutation {
node_id: node_id.trim().to_string(),
connected,
conn_count: conn_count.min(i32::MAX as usize) as i32,
detail: None,
observed_at_unix_secs: None,
observed_at_unix_secs,
};
data.update_proxy_node_tunnel_status(&mutation)
@@ -742,22 +879,26 @@ async fn apply_embedded_tunnel_node_status(
.map_err(|err| format!("node status sync failed: {err}"))
}
fn build_embedded_tunnel_heartbeat_ack(node: &StoredProxyNode) -> Vec<u8> {
let Some(remote_config) = node.remote_config.as_ref() else {
return b"{}".to_vec();
};
fn build_embedded_tunnel_heartbeat_ack(
node: &StoredProxyNode,
heartbeat_id: Option<u64>,
) -> Vec<u8> {
let mut payload = serde_json::Map::new();
payload.insert("remote_config".to_string(), remote_config.clone());
payload.insert("config_version".to_string(), json!(node.config_version));
if let Some(upgrade_to) = remote_config
.as_object()
.and_then(|value| value.get("upgrade_to"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
payload.insert("upgrade_to".to_string(), json!(upgrade_to));
if let Some(heartbeat_id) = heartbeat_id {
payload.insert("heartbeat_id".to_string(), json!(heartbeat_id));
}
if let Some(remote_config) = node.remote_config.as_ref() {
payload.insert("remote_config".to_string(), remote_config.clone());
payload.insert("config_version".to_string(), json!(node.config_version));
if let Some(upgrade_to) = remote_config
.as_object()
.and_then(|value| value.get("upgrade_to"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
payload.insert("upgrade_to".to_string(), json!(upgrade_to));
}
}
serde_json::to_vec(&serde_json::Value::Object(payload)).unwrap_or_else(|_| b"{}".to_vec())
@@ -857,6 +998,7 @@ mod tests {
&data,
br#"{
"node_id": "node-123",
"heartbeat_id": 42,
"heartbeat_interval": 45,
"active_connections": 5,
"total_requests": 9,
@@ -873,6 +1015,7 @@ mod tests {
let payload: serde_json::Value =
serde_json::from_slice(&ack).expect("ack payload should parse");
assert_eq!(payload["heartbeat_id"], 42);
assert_eq!(payload["config_version"], 7);
assert_eq!(payload["upgrade_to"], "1.2.3");
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
@@ -906,7 +1049,7 @@ mod tests {
)]));
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository));
apply_embedded_tunnel_node_status(&data, "node-123", true, 4)
apply_embedded_tunnel_node_status(&data, "node-123", true, 4, Some(1_800_000_123))
.await
.expect("node status should succeed");
@@ -917,6 +1060,7 @@ mod tests {
.expect("node should exist");
assert_eq!(node.status, "online");
assert_eq!(node.tunnel_connected, true);
assert_eq!(node.tunnel_connected_at_unix_secs, Some(1_800_000_123));
}
#[tokio::test]
@@ -929,7 +1073,7 @@ mod tests {
);
directory
.sync_node_status(&data, "node-123", true, 2)
.sync_node_status(&data, "node-123", true, 2, 1_800_000_010)
.await
.expect("attachment should sync");
let record = directory
@@ -940,9 +1084,10 @@ mod tests {
assert_eq!(record.gateway_instance_id, "gateway-a");
assert_eq!(record.relay_base_url, "http://gateway-a.internal");
assert_eq!(record.conn_count, 2);
assert_eq!(record.observed_at_unix_secs, 1_800_000_010);
directory
.sync_node_status(&data, "node-123", false, 0)
.sync_node_status(&data, "node-123", false, 0, 1_800_000_011)
.await
.expect("attachment should clear");
assert!(directory