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

Complete secure tunnel encryption support

Browse Source
This commit is contained in:
RWDai
2026-05-22 09:47:28 +08:00
parent 4f49dd5943
commit 2e701a90c9
24 changed files with 720 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/Cargo.toml
View File

@@ -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"] }

49
apps/aether-gateway/src/handlers/admin/system/proxy_nodes.rs
View File

@@ -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,
@@ -1420,12 +1436,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 +1485,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",

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

@@ -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,42 @@ 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())
.unwrap_or(node_id.as_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 = match tunnel_security.as_deref() {
Some(aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED) => {
match stored_security_key {
Some(key) => Some(key),
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,
};
let request_permit = match state.try_acquire_request_permit().await {
Ok(permit) => permit,
@@ -253,6 +325,8 @@ pub async fn ws_proxy(
node_name,
max_streams,
protocol_version,
security_key,
security_session,
state.proxy_conn_cfg,
)
.await

66
apps/aether-gateway/src/tunnel/embedded/proxy_conn.rs
View File

@@ -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)?.into()))
}
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())
}

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

@@ -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)
}