i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 18:07:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(vscodex): add remote Codex collaboration module

Browse Source
This commit is contained in:
fawney
2026-09-01 20:25:35 +08:00
parent 5a69cfe40d
commit 30a75832f8
102 changed files with 38569 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/api/backend/public.rs
+5 -1
View File
@@ -1,7 +1,10 @@
use axum::routing::get;
use axum::Router;
use crate::{handlers::proxy::proxy_request, state::AppState};
use crate::{
handlers::{proxy::proxy_request, public::vscodex_ws_proxy},
state::AppState,
};
pub(crate) fn mount_public_support_routes(router: Router<AppState>) -> Router<AppState> {
router
@@ -26,6 +29,7 @@ pub(crate) fn mount_public_support_routes(router: Router<AppState>) -> Router<Ap
.route("/api/capabilities", get(proxy_request))
.route("/api/capabilities/user-configurable", get(proxy_request))
.route("/api/capabilities/model/{*model_path}", get(proxy_request))
.route("/api/vscodex/ws", get(vscodex_ws_proxy))
.route("/install/{*install_path}", get(proxy_request))
.route("/install-tunnel/{*install_path}", get(proxy_request))
.route("/i/{*install_path}", get(proxy_request))
apps/aether-gateway/src/control/route/public_support.rs
+59
View File
@@ -520,6 +520,65 @@ pub(super) fn classify_public_support_route(
"aether:ccswitch_usage",
false,
))
} else if method == http::Method::POST
&& matches!(normalized_path, "/api/vscodex/pair" | "/api/vscodex/pair/")
{
Some(classified(
"public_support",
"vscodex",
"pairing_exchange",
"public:vscodex",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
"/api/users/me/vscodex/devices" | "/api/users/me/vscodex/devices/"
)
{
Some(classified(
"public_support",
"users_me",
"vscodex_devices_list",
"user:self",
false,
))
} else if method == http::Method::POST
&& matches!(
normalized_path,
"/api/users/me/vscodex/pairings" | "/api/users/me/vscodex/pairings/"
)
{
Some(classified(
"public_support",
"users_me",
"vscodex_pairing_create",
"user:self",
false,
))
} else if method == http::Method::POST
&& matches!(
normalized_path,
"/api/users/me/vscodex/ws-tickets" | "/api/users/me/vscodex/ws-tickets/"
)
{
Some(classified(
"public_support",
"users_me",
"vscodex_ws_ticket_create",
"user:self",
false,
))
} else if method == http::Method::DELETE
&& has_single_segment_after_prefix(normalized_path, "/api/users/me/vscodex/devices/")
{
Some(classified(
"public_support",
"users_me",
"vscodex_device_delete",
"user:self",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
apps/aether-gateway/src/control/tests/public_support.rs
+63
View File
@@ -440,6 +440,26 @@ fn classifies_users_me_routes_as_public_support_route() {
"/api/users/me/available-models",
"available_models",
),
(
http::Method::GET,
"/api/users/me/vscodex/devices",
"vscodex_devices_list",
),
(
http::Method::POST,
"/api/users/me/vscodex/pairings",
"vscodex_pairing_create",
),
(
http::Method::DELETE,
"/api/users/me/vscodex/devices/device-1",
"vscodex_device_delete",
),
(
http::Method::POST,
"/api/users/me/vscodex/ws-tickets",
"vscodex_ws_ticket_create",
),
(
http::Method::PUT,
"/api/users/me/model-capabilities",
@@ -496,6 +516,49 @@ fn classifies_users_me_routes_as_public_support_route() {
}
}
#[test]
fn vscodex_post_routes_buffer_request_body() {
let headers = headers(&[]);
for path in [
"/api/vscodex/pair",
"/api/users/me/vscodex/pairings",
"/api/users/me/vscodex/ws-tickets",
] {
let uri: Uri = path.parse().expect("uri should parse");
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
.expect("route should classify");
let context = GatewayPublicRequestContext::from_request_parts(
"trace-vscodex",
&http::Method::POST,
&uri,
&headers,
Some(decision),
);
assert!(
local_proxy_route_requires_buffered_body(&context),
"{path} should buffer its JSON body"
);
}
}
#[test]
fn classifies_public_vscodex_pairing_exchange() {
let headers = headers(&[]);
let uri: Uri = "/api/vscodex/pair".parse().expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_class.as_deref(), Some("public_support"));
assert_eq!(decision.route_family.as_deref(), Some("vscodex"));
assert_eq!(decision.route_kind.as_deref(), Some("pairing_exchange"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("public:vscodex")
);
assert!(!decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_ccswitch_usage_as_api_key_public_support_route() {
let headers = headers(&[]);
apps/aether-gateway/src/handlers/proxy/mod.rs
+1
View File
@@ -1352,6 +1352,7 @@ async fn proxy_request_inner(
.extensions
.get::<crate::middleware::CfConnectingIp>()
.map(|value| value.0.as_str()),
client_ip,
local_proxy_body.as_ref(),
)
.await
apps/aether-gateway/src/handlers/public/mod.rs
+1 -1
View File
@@ -28,5 +28,5 @@ pub(crate) use self::support::{
build_api_key_install_session_response, build_proxy_node_install_session_response,
build_unhandled_public_support_response, matches_model_mapping_for_models,
maybe_build_local_admin_announcements_response, maybe_build_local_public_support_response,
CreateApiKeyInstallSessionRequest,
vscodex_ws_proxy, CreateApiKeyInstallSessionRequest,
};
apps/aether-gateway/src/handlers/public/support.rs
+10
View File
@@ -48,6 +48,8 @@ mod support_payment;
mod support_test_connection;
#[path = "support/user_me.rs"]
mod support_user_me;
#[path = "support/user_me_vscodex.rs"]
mod support_vscodex;
#[path = "support/wallet.rs"]
mod support_wallet;
@@ -89,6 +91,8 @@ use self::support_oauth::maybe_build_local_oauth_response;
use self::support_payment::maybe_build_local_payment_callback_response;
use self::support_test_connection::maybe_build_local_test_connection_response;
use self::support_user_me::maybe_build_local_users_me_response;
pub(crate) use self::support_vscodex::vscodex_ws_proxy;
use self::support_vscodex::{handle_users_me_vscodex_request, maybe_build_local_vscodex_response};
use self::support_wallet::{
build_wallet_balance_payload_for_auth_scope, build_wallet_balance_payload_for_user,
build_wallet_live_today_usage_payload_for_api_key,
@@ -121,6 +125,7 @@ pub(crate) async fn maybe_build_local_public_support_response(
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
cf_connecting_ip: Option<&str>,
client_ip: std::net::IpAddr,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
@@ -192,6 +197,11 @@ pub(crate) async fn maybe_build_local_public_support_response(
.await;
}
if decision.route_family.as_deref() == Some("vscodex") {
return maybe_build_local_vscodex_response(state, request_context, client_ip, request_body)
.await;
}
if decision.route_family.as_deref() == Some("install") {
return maybe_build_local_install_response(state, request_context).await;
}
apps/aether-gateway/src/handlers/public/support/user_me.rs
+3 -2
View File
@@ -2,8 +2,9 @@ use super::{
auth_password_policy_level, base_url_from_request, build_auth_error_response,
build_auth_wallet_summary_payload, decrypt_catalog_secret_with_fallbacks,
encrypt_catalog_secret_with_fallbacks, handle_auth_me,
handle_users_me_api_key_install_session_create, query_param_optional_bool, query_param_value,
resolve_authenticated_local_user, sanitize_public_model_config_for_user, unix_secs_to_rfc3339,
handle_users_me_api_key_install_session_create, handle_users_me_vscodex_request,
query_param_optional_bool, query_param_value, resolve_authenticated_local_user,
sanitize_public_model_config_for_user, unix_secs_to_rfc3339,
users_me_api_key_install_sessions_path_matches, validate_auth_register_password, AppState,
AuthenticatedLocalUserContext, GatewayPublicRequestContext, PUBLIC_CAPABILITY_DEFINITIONS,
};
apps/aether-gateway/src/handlers/public/support/user_me_routes.rs
+12 -3
View File
@@ -18,9 +18,10 @@ use super::{
handle_users_me_preferences_put, handle_users_me_providers_get, handle_users_me_referral_get,
handle_users_me_sessions_get, handle_users_me_update_session, handle_users_me_usage_active_get,
handle_users_me_usage_get, handle_users_me_usage_heatmap_get,
handle_users_me_usage_interval_timeline_get, users_me_api_key_capabilities_path_matches,
users_me_api_key_detail_path_matches, users_me_api_key_install_sessions_path_matches,
users_me_api_key_providers_path_matches, users_me_management_token_detail_path_matches,
handle_users_me_usage_interval_timeline_get, handle_users_me_vscodex_request,
users_me_api_key_capabilities_path_matches, users_me_api_key_detail_path_matches,
users_me_api_key_install_sessions_path_matches, users_me_api_key_providers_path_matches,
users_me_management_token_detail_path_matches,
users_me_management_token_regenerate_path_matches,
users_me_management_token_toggle_path_matches, users_me_management_tokens_root,
users_me_session_detail_path_matches, AppState, GatewayPublicRequestContext,
@@ -55,6 +56,14 @@ pub(crate) async fn maybe_build_local_users_me_response(
{
Some(handle_users_me_delete_other_sessions(state, request_context, headers).await)
}
Some(
"vscodex_devices_list"
| "vscodex_pairing_create"
| "vscodex_device_delete"
| "vscodex_ws_ticket_create",
) => Some(
handle_users_me_vscodex_request(state, request_context, headers, request_body).await,
),
Some("session_delete")
if users_me_session_detail_path_matches(&request_context.request_path) =>
{
apps/aether-gateway/src/handlers/public/support/user_me_vscodex.rs
+931
View File
@@ -0,0 +1,931 @@
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::{Arc, LazyLock, Mutex};
use std::time::Duration;
use axum::body::{Body, Bytes};
use axum::extract::{
ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage, WebSocket, WebSocketUpgrade},
ConnectInfo, State,
};
use axum::http::{self, header};
use axum::response::{IntoResponse, Response};
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::{json, Map, Value};
use tokio::sync::Semaphore;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::protocol::{
CloseFrame as TungsteniteCloseFrame, WebSocketConfig,
};
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
use tracing::warn;
use super::{
build_auth_error_response, build_auth_json_response, module_available_from_env,
resolve_authenticated_local_user, AppState, GatewayPublicRequestContext,
};
const VSCODEX_ENABLED_ENV: &str = "AETHER_VSCODEX_ENABLED";
const VSCODEX_INTERNAL_URL_ENV: &str = "AETHER_VSCODEX_INTERNAL_URL";
const VSCODEX_INTERNAL_TOKEN_ENV: &str = "AETHER_VSCODEX_INTERNAL_TOKEN";
const VSCODEX_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
const VSCODEX_MAX_RESPONSE_BYTES: usize = 1024 * 1024;
const VSCODEX_WS_MAX_MESSAGE_BYTES: usize = 16 * 1024 * 1024;
const VSCODEX_WS_MAX_CONNECTIONS: usize = 256;
const VSCODEX_WS_MAX_CONNECTIONS_PER_IP: usize = 16;
const VSCODEX_DEVICE_PATH_PREFIX: &str = "/api/users/me/vscodex/devices/";
const VSCODEX_CLIENT_IP_HEADER: &str = "x-aether-client-ip";
static VSCODEX_HTTP_CLIENT: LazyLock<Result<reqwest::Client, reqwest::Error>> =
LazyLock::new(|| {
reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
});
static VSCODEX_WS_CONNECTIONS: LazyLock<Arc<Semaphore>> =
LazyLock::new(|| Arc::new(Semaphore::new(VSCODEX_WS_MAX_CONNECTIONS)));
static VSCODEX_WS_CONNECTIONS_BY_IP: LazyLock<Arc<VscodexWsIpConnectionLimiter>> =
LazyLock::new(|| {
Arc::new(VscodexWsIpConnectionLimiter::new(
VSCODEX_WS_MAX_CONNECTIONS_PER_IP,
))
});
#[derive(Debug)]
struct VscodexWsIpConnectionLimiter {
max_connections: usize,
active: Mutex<HashMap<IpAddr, usize>>,
}
impl VscodexWsIpConnectionLimiter {
fn new(max_connections: usize) -> Self {
Self {
max_connections: max_connections.max(1),
active: Mutex::new(HashMap::new()),
}
}
fn try_acquire(self: &Arc<Self>, client_ip: IpAddr) -> Option<VscodexWsIpConnectionPermit> {
let mut active = self
.active
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let current = active.get(&client_ip).copied().unwrap_or_default();
if current >= self.max_connections {
return None;
}
active.insert(client_ip, current.saturating_add(1));
Some(VscodexWsIpConnectionPermit {
limiter: Arc::clone(self),
client_ip,
})
}
fn release(&self, client_ip: IpAddr) {
let mut active = self
.active
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let Some(current) = active.get_mut(&client_ip) else {
return;
};
if *current <= 1 {
active.remove(&client_ip);
} else {
*current -= 1;
}
}
#[cfg(test)]
fn active_ip_count(&self) -> usize {
self.active
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.len()
}
}
#[derive(Debug)]
struct VscodexWsIpConnectionPermit {
limiter: Arc<VscodexWsIpConnectionLimiter>,
client_ip: IpAddr,
}
impl Drop for VscodexWsIpConnectionPermit {
fn drop(&mut self) {
self.limiter.release(self.client_ip);
}
}
#[derive(Debug)]
struct VscodexSidecarConfig {
base_url: reqwest::Url,
authorization: reqwest::header::HeaderValue,
http_client: reqwest::Client,
}
#[derive(Debug, Default, Deserialize)]
struct CreatePairingRequest {
name: Option<String>,
}
#[derive(Debug, Deserialize)]
struct CreateWsTicketRequest {
device_id: String,
}
#[derive(Debug, Deserialize)]
struct ExchangePairingRequest {
code: String,
name: Option<String>,
}
pub(crate) async fn vscodex_ws_proxy(
State(state): State<AppState>,
ConnectInfo(remote_addr): ConnectInfo<std::net::SocketAddr>,
ws: WebSocketUpgrade,
headers: http::HeaderMap,
) -> Response<Body> {
let request_permit = match state.try_acquire_request_permit().await {
Ok(value) => value,
Err(err) => {
warn!(error = ?err, "VS Codex WebSocket request admission rejected");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"服务繁忙，请稍后重试",
false,
);
}
};
let client_ip = crate::headers::effective_client_ip(&headers, &remote_addr);
match state.admin_security_ip_blacklisted(client_ip).await {
Ok(true) => {
return build_auth_error_response(
http::StatusCode::FORBIDDEN,
"当前 IP 已被禁止访问",
false,
)
}
Ok(false) => {}
Err(err) => warn!(
client_ip = %client_ip,
error = ?err,
"VS Codex WebSocket IP blacklist check failed open"
),
}
let connection_permit = match Arc::clone(&VSCODEX_WS_CONNECTIONS).try_acquire_owned() {
Ok(value) => value,
Err(_) => {
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 连接数已达上限",
false,
)
}
};
// Only active connections have entries, and each already owns one of the 256 global slots.
let ip_connection_permit = match VSCODEX_WS_CONNECTIONS_BY_IP.try_acquire(client_ip) {
Some(value) => value,
None => {
warn!(
client_ip = %client_ip,
limit = VSCODEX_WS_MAX_CONNECTIONS_PER_IP,
"VS Codex per-IP WebSocket connection limit reached"
);
let mut response = build_auth_error_response(
http::StatusCode::TOO_MANY_REQUESTS,
"当前 IP 的 VS Codex 连接数已达上限",
false,
);
response
.headers_mut()
.insert(header::RETRY_AFTER, http::HeaderValue::from_static("1"));
return response;
}
};
let config = match load_vscodex_sidecar_config() {
Ok(Some(value)) => value,
Ok(None) => {
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务未启用",
false,
)
}
Err(detail) => {
warn!(error = %detail, "VS Codex WebSocket sidecar configuration is invalid");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
let sidecar_url = match build_vscodex_websocket_url(&config.base_url) {
Ok(value) => value,
Err(detail) => {
warn!(error = %detail, "could not build VS Codex sidecar WebSocket URL");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
let mut sidecar_request = match sidecar_url.as_str().into_client_request() {
Ok(value) => value,
Err(err) => {
warn!(error = %err, "could not build VS Codex sidecar WebSocket request");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
for header_name in [header::ORIGIN, header::SEC_WEBSOCKET_PROTOCOL] {
if let Some(value) = headers.get(&header_name) {
sidecar_request
.headers_mut()
.insert(header_name, value.clone());
}
}
let mut sidecar_config = WebSocketConfig::default();
sidecar_config.max_message_size = Some(VSCODEX_WS_MAX_MESSAGE_BYTES);
sidecar_config.max_frame_size = Some(VSCODEX_WS_MAX_MESSAGE_BYTES);
let (sidecar_socket, sidecar_response) = match tokio::time::timeout(
VSCODEX_REQUEST_TIMEOUT,
tokio_tungstenite::connect_async_with_config(sidecar_request, Some(sidecar_config), true),
)
.await
{
Ok(Ok(value)) => value,
Ok(Err(err)) => {
warn!(error = %err, "VS Codex sidecar WebSocket handshake failed");
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务暂时不可用",
false,
);
}
Err(_) => {
warn!("VS Codex sidecar WebSocket handshake timed out");
return build_auth_error_response(
http::StatusCode::GATEWAY_TIMEOUT,
"VS Codex 服务请求超时",
false,
);
}
};
let selected_protocol = sidecar_response
.headers()
.get(header::SEC_WEBSOCKET_PROTOCOL)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let ws = ws
.max_message_size(VSCODEX_WS_MAX_MESSAGE_BYTES)
.max_frame_size(VSCODEX_WS_MAX_MESSAGE_BYTES);
let ws = match selected_protocol {
Some(protocol) => ws.protocols([protocol]),
None => ws,
};
drop(request_permit);
ws.on_upgrade(move |browser_socket| async move {
let _connection_permit = connection_permit;
bridge_vscodex_websockets(browser_socket, sidecar_socket, ip_connection_permit).await;
})
}
pub(super) async fn maybe_build_local_vscodex_response(
_state: &AppState,
request_context: &GatewayPublicRequestContext,
client_ip: std::net::IpAddr,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
if decision.route_family.as_deref() != Some("vscodex") {
return None;
}
if decision.route_kind.as_deref() != Some("pairing_exchange")
|| !matches!(
request_context.request_path.as_str(),
"/api/vscodex/pair" | "/api/vscodex/pair/"
)
{
return Some(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"VS Codex 接口不存在",
false,
));
}
let config = match load_vscodex_sidecar_config() {
Ok(Some(value)) => value,
Ok(None) => {
return Some(build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务未启用",
false,
))
}
Err(detail) => {
warn!(
error = %detail,
"VS Codex sidecar configuration is invalid"
);
return Some(build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
));
}
};
let payload = match parse_pairing_exchange_request(request_body) {
Ok(value) => value,
Err(response) => return Some(response),
};
let url = match append_vscodex_sidecar_path(&config.base_url, &["v1", "pairings", "exchange"]) {
Ok(value) => value,
Err(detail) => {
warn!(error = %detail, "could not build VS Codex pairing exchange URL");
return Some(build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
));
}
};
let request =
build_authenticated_sidecar_request(&config, reqwest::Method::POST, url, Some(payload))
.header(VSCODEX_CLIENT_IP_HEADER, client_ip.to_string());
Some(send_vscodex_sidecar_request(request, "public", "pairing_exchange").await)
}
pub(super) async fn handle_users_me_vscodex_request(
state: &AppState,
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
request_body: Option<&Bytes>,
) -> Response<Body> {
let auth = match resolve_authenticated_local_user(state, request_context, headers).await {
Ok(value) => value,
Err(response) => return response,
};
let config = match load_vscodex_sidecar_config() {
Ok(Some(value)) => value,
Ok(None) => {
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务未启用",
false,
)
}
Err(detail) => {
warn!(
user_id = %auth.user.id,
error = %detail,
"VS Codex sidecar configuration is invalid"
);
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
let Some(route_kind) = request_context
.control_decision
.as_ref()
.and_then(|decision| decision.route_kind.as_deref())
else {
return build_auth_error_response(
http::StatusCode::NOT_FOUND,
"VS Codex 接口不存在",
false,
);
};
let request = match build_vscodex_sidecar_request(
&config,
&auth.user.id,
route_kind,
&request_context.request_path,
request_body,
) {
Ok(value) => value,
Err(response) => return response,
};
send_vscodex_sidecar_request(request, &auth.user.id, route_kind).await
}
fn load_vscodex_sidecar_config() -> Result<Option<VscodexSidecarConfig>, String> {
if !module_available_from_env(VSCODEX_ENABLED_ENV, false) {
return Ok(None);
}
let raw_url = required_env(VSCODEX_INTERNAL_URL_ENV)?;
let base_url = reqwest::Url::parse(&raw_url)
.map_err(|err| format!("{VSCODEX_INTERNAL_URL_ENV} is invalid: {err}"))?;
if !matches!(base_url.scheme(), "http" | "https")
|| !base_url.has_host()
|| !base_url.username().is_empty()
|| base_url.password().is_some()
|| base_url.query().is_some()
|| base_url.fragment().is_some()
|| base_url.cannot_be_a_base()
{
return Err(format!(
"{VSCODEX_INTERNAL_URL_ENV} must be an HTTP(S) base URL without credentials, query, or fragment"
));
}
let token = required_env(VSCODEX_INTERNAL_TOKEN_ENV)?;
let authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}"))
.map_err(|_| format!("{VSCODEX_INTERNAL_TOKEN_ENV} is not a valid HTTP credential"))?;
let http_client = VSCODEX_HTTP_CLIENT
.as_ref()
.map_err(|err| format!("could not initialize VS Codex HTTP client: {err}"))?
.clone();
Ok(Some(VscodexSidecarConfig {
base_url,
authorization,
http_client,
}))
}
fn required_env(key: &str) -> Result<String, String> {
std::env::var(key)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.ok_or_else(|| format!("{key} is required"))
}
fn build_vscodex_sidecar_request(
config: &VscodexSidecarConfig,
user_id: &str,
route_kind: &str,
request_path: &str,
request_body: Option<&Bytes>,
) -> Result<reqwest::RequestBuilder, Response<Body>> {
let (method, suffix, payload) = match route_kind {
"vscodex_devices_list" => (reqwest::Method::GET, vec!["devices"], None),
"vscodex_pairing_create" => (
reqwest::Method::POST,
vec!["pairings"],
Some(parse_pairing_request(request_body)?),
),
"vscodex_device_delete" => {
let Some(device_id) = vscodex_device_id_from_path(request_path) else {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"设备标识无效",
false,
));
};
(reqwest::Method::DELETE, vec!["devices", device_id], None)
}
"vscodex_ws_ticket_create" => (
reqwest::Method::POST,
vec!["ws-tickets"],
Some(parse_ws_ticket_request(request_body)?),
),
_ => {
return Err(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"VS Codex 接口不存在",
false,
))
}
};
let url = build_vscodex_sidecar_url(&config.base_url, user_id, &suffix).map_err(|detail| {
warn!(user_id = %user_id, error = %detail, "could not build VS Codex sidecar URL");
build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
)
})?;
Ok(build_authenticated_sidecar_request(
config, method, url, payload,
))
}
fn build_authenticated_sidecar_request(
config: &VscodexSidecarConfig,
method: reqwest::Method,
url: reqwest::Url,
payload: Option<Value>,
) -> reqwest::RequestBuilder {
let mut request = config
.http_client
.request(method, url)
.header(header::AUTHORIZATION, config.authorization.clone())
.header(header::ACCEPT, "application/json")
.timeout(VSCODEX_REQUEST_TIMEOUT);
if let Some(payload) = payload {
request = request.json(&payload);
}
request
}
fn build_vscodex_sidecar_url(
base_url: &reqwest::Url,
user_id: &str,
suffix: &[&str],
) -> Result<reqwest::Url, String> {
let mut segments = vec!["internal", "v1", "users", user_id];
segments.extend(suffix.iter().copied());
append_vscodex_sidecar_path(base_url, &segments)
}
fn append_vscodex_sidecar_path(
base_url: &reqwest::Url,
suffix: &[&str],
) -> Result<reqwest::Url, String> {
let mut url = base_url.clone();
let mut path_segments = url
.path_segments_mut()
.map_err(|_| "VS Codex sidecar URL cannot contain path segments".to_string())?;
path_segments.pop_if_empty();
path_segments.extend(suffix.iter().copied());
drop(path_segments);
Ok(url)
}
fn build_vscodex_websocket_url(base_url: &reqwest::Url) -> Result<reqwest::Url, String> {
let mut url = append_vscodex_sidecar_path(base_url, &["api", "vscodex", "ws"])?;
let scheme = match url.scheme() {
"http" => "ws",
"https" => "wss",
_ => return Err("VS Codex sidecar URL must use HTTP(S)".to_string()),
};
url.set_scheme(scheme)
.map_err(|_| "could not convert VS Codex sidecar URL to WebSocket".to_string())?;
Ok(url)
}
fn parse_pairing_request(request_body: Option<&Bytes>) -> Result<Value, Response<Body>> {
let payload = parse_json_request::<CreatePairingRequest>(request_body, true)?;
let mut object = Map::new();
if let Some(name) = payload.name {
object.insert("name".to_string(), Value::String(name));
}
Ok(Value::Object(object))
}
fn parse_ws_ticket_request(request_body: Option<&Bytes>) -> Result<Value, Response<Body>> {
let payload = parse_json_request::<CreateWsTicketRequest>(request_body, false)?;
let device_id = payload.device_id.trim();
if !valid_vscodex_device_id(device_id) {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"设备标识无效",
false,
));
}
Ok(json!({ "device_id": device_id }))
}
fn parse_pairing_exchange_request(request_body: Option<&Bytes>) -> Result<Value, Response<Body>> {
let payload = parse_json_request::<ExchangePairingRequest>(request_body, false)?;
let code = payload.code.trim();
if code.is_empty() || code.len() > 256 || code.chars().any(char::is_control) {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"配对码无效",
false,
));
}
let mut object = Map::from_iter([("code".to_string(), Value::String(code.to_string()))]);
if let Some(name) = payload.name {
object.insert("name".to_string(), Value::String(name));
}
Ok(Value::Object(object))
}
fn parse_json_request<T>(
request_body: Option<&Bytes>,
empty_object_allowed: bool,
) -> Result<T, Response<Body>>
where
T: serde::de::DeserializeOwned,
{
let body = request_body.filter(|body| !body.is_empty());
let result = match body {
Some(body) => serde_json::from_slice(body),
None if empty_object_allowed => serde_json::from_slice(b"{}"),
None => {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"缺少请求体",
false,
))
}
};
result.map_err(|_| {
build_auth_error_response(http::StatusCode::BAD_REQUEST, "请求数据验证失败", false)
})
}
fn vscodex_device_id_from_path(path: &str) -> Option<&str> {
let trimmed = path.trim_end_matches('/');
let device_id = trimmed.strip_prefix(VSCODEX_DEVICE_PATH_PREFIX)?;
if device_id.contains('/') || !valid_vscodex_device_id(device_id) {
return None;
}
Some(device_id)
}
fn valid_vscodex_device_id(value: &str) -> bool {
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
}
async fn send_vscodex_sidecar_request(
request: reqwest::RequestBuilder,
request_scope: &str,
operation: &str,
) -> Response<Body> {
let mut upstream = match request.send().await {
Ok(value) => value,
Err(err) => {
warn!(
request_scope = %request_scope,
operation = %operation,
error = %err,
"VS Codex sidecar request failed"
);
let (status, detail) = if err.is_timeout() {
(http::StatusCode::GATEWAY_TIMEOUT, "VS Codex 服务请求超时")
} else {
(http::StatusCode::BAD_GATEWAY, "VS Codex 服务暂时不可用")
};
return build_auth_error_response(status, detail, false);
}
};
let status = http::StatusCode::from_u16(upstream.status().as_u16())
.unwrap_or(http::StatusCode::BAD_GATEWAY);
if matches!(
status,
http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN
) {
warn!(
request_scope = %request_scope,
operation = %operation,
upstream_status = status.as_u16(),
"VS Codex sidecar rejected gateway credentials"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务鉴权失败",
false,
);
}
if status.is_redirection() {
warn!(
request_scope = %request_scope,
operation = %operation,
upstream_status = status.as_u16(),
"VS Codex sidecar returned an unexpected redirect"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
if status == http::StatusCode::NO_CONTENT {
return vscodex_no_store_response(status.into_response(), None);
}
let mut response_body = Vec::new();
while let Some(chunk) = match upstream.chunk().await {
Ok(value) => value,
Err(err) => {
warn!(
request_scope = %request_scope,
operation = %operation,
error = %err,
"could not read VS Codex sidecar response"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
} {
if response_body.len().saturating_add(chunk.len()) > VSCODEX_MAX_RESPONSE_BYTES {
warn!(
request_scope = %request_scope,
operation = %operation,
"VS Codex sidecar response exceeded the size limit"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
response_body.extend_from_slice(&chunk);
}
if response_body.is_empty() {
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
let payload = match serde_json::from_slice(&response_body) {
Ok(value) => value,
Err(err) => {
warn!(
request_scope = %request_scope,
operation = %operation,
upstream_status = status.as_u16(),
error = %err,
"VS Codex sidecar returned non-JSON data"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
};
let retry_after = upstream.headers().get(header::RETRY_AFTER).cloned();
vscodex_no_store_response(build_auth_json_response(status, payload, None), retry_after)
}
fn vscodex_no_store_response(
mut response: Response<Body>,
retry_after: Option<http::HeaderValue>,
) -> Response<Body> {
response.headers_mut().insert(
header::CACHE_CONTROL,
http::HeaderValue::from_static("no-store"),
);
if let Some(retry_after) = retry_after {
response
.headers_mut()
.insert(header::RETRY_AFTER, retry_after);
}
response
}
async fn bridge_vscodex_websockets<S>(
browser_socket: WebSocket,
sidecar_socket: S,
ip_connection_permit: VscodexWsIpConnectionPermit,
) where
S: futures_util::Stream<
Item = Result<TungsteniteMessage, tokio_tungstenite::tungstenite::Error>,
> + futures_util::Sink<TungsteniteMessage, Error = tokio_tungstenite::tungstenite::Error>
+ Unpin
+ Send
+ 'static,
{
let (mut browser_tx, mut browser_rx) = browser_socket.split();
let (mut sidecar_tx, mut sidecar_rx) = sidecar_socket.split();
let mut ip_connection_permit = Some(ip_connection_permit);
loop {
tokio::select! {
browser_message = browser_rx.next() => {
match browser_message {
Some(Ok(message)) => {
let close = matches!(message, AxumMessage::Close(_));
if let Err(err) = sidecar_tx.send(axum_to_tungstenite_message(message)).await {
warn!(error = %err, "could not forward VS Codex browser WebSocket frame");
break;
}
if close {
break;
}
}
Some(Err(err)) => {
warn!(error = %err, "VS Codex browser WebSocket read failed");
break;
}
None => break,
}
}
sidecar_message = sidecar_rx.next() => {
match sidecar_message {
Some(Ok(TungsteniteMessage::Frame(_))) => continue,
Some(Ok(message)) => {
if ip_connection_permit.is_some() && vscodex_ws_authentication_succeeded(&message) {
ip_connection_permit.take();
}
let close = matches!(message, TungsteniteMessage::Close(_));
if let Err(err) = browser_tx.send(tungstenite_to_axum_message(message)).await {
warn!(error = %err, "could not forward VS Codex sidecar WebSocket frame");
break;
}
if close {
break;
}
}
Some(Err(err)) => {
warn!(error = %err, "VS Codex sidecar WebSocket read failed");
break;
}
None => break,
}
}
}
}
let _ = sidecar_tx.close().await;
let _ = browser_tx.close().await;
}
fn vscodex_ws_authentication_succeeded(message: &TungsteniteMessage) -> bool {
let TungsteniteMessage::Text(text) = message else {
return false;
};
serde_json::from_str::<Value>(text.as_ref())
.ok()
.and_then(|payload| {
payload
.get("type")
.and_then(Value::as_str)
.map(str::to_string)
})
.as_deref()
== Some("auth.ok")
}
fn axum_to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
match message {
AxumMessage::Text(text) => TungsteniteMessage::Text(text.to_string().into()),
AxumMessage::Binary(bytes) => TungsteniteMessage::Binary(bytes),
AxumMessage::Ping(bytes) => TungsteniteMessage::Ping(bytes),
AxumMessage::Pong(bytes) => TungsteniteMessage::Pong(bytes),
AxumMessage::Close(frame) => {
TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
code: frame.code.into(),
reason: frame.reason.to_string().into(),
}))
}
}
}
fn tungstenite_to_axum_message(message: TungsteniteMessage) -> AxumMessage {
match message {
TungsteniteMessage::Text(text) => AxumMessage::Text(text.to_string().into()),
TungsteniteMessage::Binary(bytes) => AxumMessage::Binary(bytes),
TungsteniteMessage::Ping(bytes) => AxumMessage::Ping(bytes),
TungsteniteMessage::Pong(bytes) => AxumMessage::Pong(bytes),
TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
code: frame.code.into(),
reason: frame.reason.to_string().into(),
})),
TungsteniteMessage::Frame(_) => AxumMessage::Close(None),
}
}
#[cfg(test)]
mod tests {
use super::{vscodex_ws_authentication_succeeded, VscodexWsIpConnectionLimiter};
use std::sync::Arc;
use tokio_tungstenite::tungstenite::Message;
#[test]
fn vscodex_ws_ip_limiter_releases_and_removes_inactive_ips() {
let limiter = Arc::new(VscodexWsIpConnectionLimiter::new(1));
let client_ip = "198.51.100.10".parse().expect("IP should parse");
let permit = limiter
.try_acquire(client_ip)
.expect("first connection should acquire");
assert_eq!(limiter.active_ip_count(), 1);
assert!(limiter.try_acquire(client_ip).is_none());
drop(permit);
assert_eq!(limiter.active_ip_count(), 0);
assert!(limiter.try_acquire(client_ip).is_some());
}
#[test]
fn vscodex_ws_ip_limiter_releases_only_after_sidecar_auth_success() {
assert!(vscodex_ws_authentication_succeeded(&Message::Text(
r#"{"type":"auth.ok","role":"operator"}"#.into()
)));
assert!(!vscodex_ws_authentication_succeeded(&Message::Text(
r#"{"type":"auth","token":"client-controlled"}"#.into()
)));
assert!(!vscodex_ws_authentication_succeeded(&Message::Binary(
br#"{"type":"auth.ok"}"#.to_vec().into()
)));
}
}
apps/aether-gateway/src/handlers/shared/request_utils.rs
+7 -1
View File
@@ -495,8 +495,14 @@ pub(crate) fn public_support_local_requires_buffered_body(
Some(
"api_keys_create"
| "api_key_install_session_create"
| "management_tokens_create",
| "management_tokens_create"
| "vscodex_pairing_create"
| "vscodex_ws_ticket_create",
),
) | (
Some("vscodex"),
http::Method::POST,
Some("pairing_exchange"),
) | (
Some("wallet"),
http::Method::POST,
apps/aether-gateway/src/tests/frontdoor/public_support.rs
+2
View File
@@ -49,6 +49,8 @@ use chrono::{TimeZone, Utc};
#[path = "public_support/dashboard.rs"]
mod dashboard;
#[path = "public_support/vscodex.rs"]
mod vscodex;
#[tokio::test]
async fn gateway_handles_public_announcements_list_without_proxying_upstream() {
apps/aether-gateway/src/tests/frontdoor/public_support/vscodex.rs
+578
View File
@@ -0,0 +1,578 @@
use super::{
any, build_router_with_state, build_test_auth_token, json, sample_auth_session,
sample_auth_user, sample_auth_wallet, set_test_env_var, start_auth_gateway_with_state,
start_server, AppState, Arc, Json, Mutex, Request, Router, StatusCode, Utc,
};
use axum::extract::ws::{Message as AxumWsMessage, WebSocketUpgrade};
use axum::response::IntoResponse;
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
#[derive(Debug, Clone, PartialEq)]
struct CapturedSidecarRequest {
method: http::Method,
path: String,
authorization: Option<String>,
client_ip: Option<String>,
body: Option<serde_json::Value>,
}
#[test]
fn gateway_authenticates_and_proxies_vscodex_bff_routes() {
std::thread::Builder::new()
.name("vscodex-gateway-test".to_string())
.stack_size(32 * 1024 * 1024)
.spawn(|| {
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.thread_stack_size(32 * 1024 * 1024)
.build()
.expect("test runtime should build")
.block_on(run_vscodex_gateway_integration());
})
.expect("test thread should spawn")
.join()
.expect("test thread should complete");
}
async fn run_vscodex_gateway_integration() {
let captured_requests = Arc::new(Mutex::new(Vec::<CapturedSidecarRequest>::new()));
let captured_requests_for_handler = Arc::clone(&captured_requests);
let captured_ws_handshake = Arc::new(Mutex::new(None::<(Option<String>, Option<String>)>));
let captured_ws_handshake_for_handler = Arc::clone(&captured_ws_handshake);
let sidecar = Router::new()
.route(
"/api/vscodex/ws",
any(move |ws: WebSocketUpgrade, headers: http::HeaderMap| {
let captured_ws_handshake = Arc::clone(&captured_ws_handshake_for_handler);
async move {
let origin = headers
.get(http::header::ORIGIN)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let authorization = headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
*captured_ws_handshake
.lock()
.expect("WebSocket handshake store should lock") =
Some((origin, authorization));
ws.protocols(["vscodex.v1"])
.on_upgrade(|mut socket| async move {
while let Some(Ok(message)) = socket.next().await {
match message {
AxumWsMessage::Text(text) => {
let response = if text
== r#"{"type":"auth","token":"test-auth-ok"}"# {
r#"{"type":"auth.ok","role":"operator"}"#.to_string()
} else {
format!("echo:{text}")
};
if socket
.send(AxumWsMessage::Text(response.into()))
.await
.is_err()
{
break;
}
}
AxumWsMessage::Binary(bytes) => {
if socket.send(AxumWsMessage::Binary(bytes)).await.is_err()
{
break;
}
}
AxumWsMessage::Close(frame) => {
let _ = socket.send(AxumWsMessage::Close(frame)).await;
break;
}
AxumWsMessage::Ping(_) | AxumWsMessage::Pong(_) => {}
}
}
})
}
}),
)
.route(
"/{*path}",
any(move |request: Request| {
let captured_requests = Arc::clone(&captured_requests_for_handler);
async move {
let method = request.method().clone();
let path = request.uri().path().to_string();
let authorization = request
.headers()
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let client_ip = request
.headers()
.get("x-aether-client-ip")
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let body = axum::body::to_bytes(request.into_body(), 1024 * 1024)
.await
.expect("sidecar request body should be readable");
let body = (!body.is_empty()).then(|| {
serde_json::from_slice(&body).expect("sidecar request body should be JSON")
});
captured_requests
.lock()
.expect("captured request store should lock")
.push(CapturedSidecarRequest {
method: method.clone(),
path: path.clone(),
authorization,
client_ip,
body,
});
let (status, payload) = match (method, path.as_str()) {
(http::Method::GET, "/internal/v1/users/user-auth-1/devices") => (
StatusCode::OK,
json!({ "devices": [{ "id": "host-1", "name": "My Mac" }] }),
),
(http::Method::POST, "/internal/v1/users/user-auth-1/pairings") => {
(StatusCode::CREATED, json!({ "code": "PAIR-123" }))
}
(http::Method::POST, "/internal/v1/users/user-auth-1/ws-tickets") => (
StatusCode::CREATED,
json!({
"ticket": "ticket-123",
"ws_url": "wss://aether.example/vscodex/ws"
}),
),
(http::Method::DELETE, "/internal/v1/users/user-auth-1/devices/host-1") => {
(StatusCode::NO_CONTENT, json!({}))
}
(
http::Method::DELETE,
"/internal/v1/users/user-auth-1/devices/missing",
) => (StatusCode::NOT_FOUND, json!({ "detail": "设备不存在" })),
(
http::Method::DELETE,
"/internal/v1/users/user-auth-1/devices/internal-denied",
) => (
StatusCode::UNAUTHORIZED,
json!({ "detail": "internal token invalid" }),
),
(
http::Method::DELETE,
"/internal/v1/users/user-auth-1/devices/redirect",
) => (StatusCode::TEMPORARY_REDIRECT, json!({ "redirect": true })),
(
http::Method::DELETE,
"/internal/v1/users/user-auth-1/devices/empty-ok",
) => return StatusCode::OK.into_response(),
(http::Method::POST, "/v1/pairings/exchange") => (
StatusCode::CREATED,
json!({ "device_id": "host-2", "device_token": "host-secret" }),
),
_ => (
StatusCode::NOT_FOUND,
json!({ "detail": "unexpected path" }),
),
};
let mut response = (status, Json(payload)).into_response();
if status == StatusCode::TEMPORARY_REDIRECT {
response.headers_mut().insert(
http::header::LOCATION,
"/redirect-must-not-be-followed".parse().unwrap(),
);
}
response
}
}),
);
let (sidecar_url, sidecar_handle) = start_server(sidecar).await;
let _enabled = set_test_env_var("AETHER_VSCODEX_ENABLED", "true");
let _internal_url = set_test_env_var("AETHER_VSCODEX_INTERNAL_URL", &sidecar_url);
let _internal_token = set_test_env_var("AETHER_VSCODEX_INTERNAL_TOKEN", "sidecar-secret");
let now = Utc::now();
let user = sample_auth_user(now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!(user.role)),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".to_string(), json!("session-vscodex")),
]),
now + chrono::Duration::hours(1),
);
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_state(
user,
sample_auth_wallet("user-auth-1", now),
[sample_auth_session(
"user-auth-1",
"session-vscodex",
"browser-device-vscodex",
"refresh-vscodex",
now,
)],
)
.await;
let client = reqwest::Client::new();
let unauthenticated = client
.get(format!("{gateway_url}/api/users/me/vscodex/devices"))
.send()
.await
.expect("unauthenticated request should complete");
assert_eq!(unauthenticated.status(), StatusCode::UNAUTHORIZED);
assert!(captured_requests
.lock()
.expect("captured request store should lock")
.is_empty());
let devices = client
.get(format!("{gateway_url}/api/users/me/vscodex/devices"))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("devices request should complete");
assert_eq!(devices.status(), StatusCode::OK);
let devices_payload: serde_json::Value =
devices.json().await.expect("devices body should be JSON");
assert_eq!(devices_payload["devices"][0]["id"], "host-1");
let pairing = client
.post(format!("{gateway_url}/api/users/me/vscodex/pairings"))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.json(&json!({ "name": "My Mac", "user_id": "attacker" }))
.send()
.await
.expect("pairing request should complete");
assert_eq!(pairing.status(), StatusCode::CREATED);
assert_eq!(
pairing
.headers()
.get(http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok()),
Some("no-store")
);
let pairing_payload: serde_json::Value =
pairing.json().await.expect("pairing body should be JSON");
assert_eq!(pairing_payload["code"], "PAIR-123");
let ticket = client
.post(format!("{gateway_url}/api/users/me/vscodex/ws-tickets"))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.json(&json!({ "device_id": "host-1", "user_id": "attacker" }))
.send()
.await
.expect("ticket request should complete");
assert_eq!(ticket.status(), StatusCode::CREATED);
assert_eq!(
ticket
.headers()
.get(http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok()),
Some("no-store")
);
let ticket_payload: serde_json::Value =
ticket.json().await.expect("ticket body should be JSON");
assert_eq!(ticket_payload["ticket"], "ticket-123");
assert_eq!(ticket_payload["ws_url"], "wss://aether.example/vscodex/ws");
let deleted = client
.delete(format!("{gateway_url}/api/users/me/vscodex/devices/host-1"))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("delete request should complete");
assert_eq!(deleted.status(), StatusCode::NO_CONTENT);
assert_eq!(
deleted
.headers()
.get(http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok()),
Some("no-store")
);
let missing = client
.delete(format!(
"{gateway_url}/api/users/me/vscodex/devices/missing"
))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("missing device request should complete");
assert_eq!(missing.status(), StatusCode::NOT_FOUND);
let missing_payload: serde_json::Value =
missing.json().await.expect("missing body should be JSON");
assert_eq!(missing_payload["detail"], "设备不存在");
let internal_denied = client
.delete(format!(
"{gateway_url}/api/users/me/vscodex/devices/internal-denied"
))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("internal auth failure request should complete");
assert_eq!(internal_denied.status(), StatusCode::BAD_GATEWAY);
let internal_denied_payload: serde_json::Value = internal_denied
.json()
.await
.expect("internal auth failure body should be JSON");
assert_eq!(internal_denied_payload["detail"], "VS Codex 服务鉴权失败");
let redirected = client
.delete(format!(
"{gateway_url}/api/users/me/vscodex/devices/redirect"
))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("redirecting sidecar request should complete");
assert_eq!(redirected.status(), StatusCode::BAD_GATEWAY);
let empty_success = client
.delete(format!(
"{gateway_url}/api/users/me/vscodex/devices/empty-ok"
))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("empty sidecar success should complete");
assert_eq!(empty_success.status(), StatusCode::BAD_GATEWAY);
let pairing_exchange = client
.post(format!("{gateway_url}/api/vscodex/pair"))
.header("x-aether-client-ip", "203.0.113.99")
.json(&json!({
"code": "PAIR-123",
"name": "Office Mac",
"user_id": "attacker",
"device_token": "stolen"
}))
.send()
.await
.expect("public pairing exchange should complete");
assert_eq!(pairing_exchange.status(), StatusCode::CREATED);
assert_eq!(
pairing_exchange
.headers()
.get(http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok()),
Some("no-store")
);
let pairing_exchange_payload: serde_json::Value = pairing_exchange
.json()
.await
.expect("pairing exchange body should be JSON");
assert_eq!(pairing_exchange_payload["device_id"], "host-2");
assert_eq!(pairing_exchange_payload["device_token"], "host-secret");
let mut websocket_request = format!("{gateway_url}/api/vscodex/ws")
.replace("http://", "ws://")
.into_client_request()
.expect("WebSocket request should build");
websocket_request.headers_mut().insert(
http::header::ORIGIN,
"https://aether.example".parse().unwrap(),
);
websocket_request.headers_mut().insert(
http::header::AUTHORIZATION,
"Bearer browser-aether-jwt".parse().unwrap(),
);
websocket_request.headers_mut().insert(
http::header::SEC_WEBSOCKET_PROTOCOL,
"vscodex.v1".parse().unwrap(),
);
let (mut websocket, websocket_response) = tokio_tungstenite::connect_async(websocket_request)
.await
.expect("gateway WebSocket should connect");
assert_eq!(
websocket_response
.headers()
.get(http::header::SEC_WEBSOCKET_PROTOCOL)
.and_then(|value| value.to_str().ok()),
Some("vscodex.v1")
);
websocket
.send(TungsteniteMessage::Text(
"{\"type\":\"auth\",\"ticket\":\"one-time-ticket\"}".into(),
))
.await
.expect("ticket frame should send");
let echoed = websocket
.next()
.await
.expect("echoed frame should arrive")
.expect("echoed frame should be valid");
assert_eq!(
echoed,
TungsteniteMessage::Text("echo:{\"type\":\"auth\",\"ticket\":\"one-time-ticket\"}".into())
);
websocket.close(None).await.expect("WebSocket should close");
assert_eq!(
captured_ws_handshake
.lock()
.expect("WebSocket handshake store should lock")
.clone(),
Some((Some("https://aether.example".to_string()), None))
);
let limited_gateway = build_router_with_state(
AppState::new()
.expect("limited gateway state should build")
.with_request_concurrency_limit(1),
);
let (limited_gateway_url, limited_gateway_handle) = start_server(limited_gateway).await;
let limited_ws_url =
format!("{limited_gateway_url}/api/vscodex/ws").replace("http://", "ws://");
let limited_ws_request = || {
let mut request = limited_ws_url
.as_str()
.into_client_request()
.expect("limited WebSocket request should build");
request
.headers_mut()
.insert("x-real-ip", "198.51.100.50".parse().unwrap());
request
};
let mut held_websockets = Vec::new();
for index in 0..16 {
let (websocket, _) = tokio_tungstenite::connect_async(limited_ws_request())
.await
.unwrap_or_else(|err| panic!("limited WebSocket {index} should connect: {err}"));
held_websockets.push(websocket);
}
let per_ip_limit_error = tokio_tungstenite::connect_async(limited_ws_request())
.await
.expect_err("seventeenth WebSocket from one IP should be rejected");
match per_ip_limit_error {
tokio_tungstenite::tungstenite::Error::Http(response) => {
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(
response
.headers()
.get(http::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok()),
Some("1")
);
}
other => panic!("expected HTTP per-IP limit rejection, got {other:?}"),
}
held_websockets[0]
.send(TungsteniteMessage::Text(
r#"{"type":"auth","token":"test-auth-ok"}"#.into(),
))
.await
.expect("test authentication frame should send");
let auth_ok =
tokio::time::timeout(std::time::Duration::from_secs(1), held_websockets[0].next())
.await
.expect("test authentication response should arrive in time")
.expect("test authentication response should contain a frame")
.expect("test authentication response should be valid");
assert_eq!(
auth_ok,
TungsteniteMessage::Text(r#"{"type":"auth.ok","role":"operator"}"#.into())
);
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
let (mut replacement_websocket, _) = tokio_tungstenite::connect_async(limited_ws_request())
.await
.expect("sidecar auth success should release one pending per-IP slot");
replacement_websocket
.close(None)
.await
.expect("replacement WebSocket should close");
for mut websocket in held_websockets {
websocket
.close(None)
.await
.expect("held WebSocket should close");
}
limited_gateway_handle.abort();
let blacklisted_gateway = build_router_with_state(
AppState::new()
.expect("blacklisted gateway state should build")
.with_admin_security_blacklist_for_tests([(
"127.0.0.1".to_string(),
"blocked".to_string(),
)]),
);
let (blacklisted_gateway_url, blacklisted_gateway_handle) =
start_server(blacklisted_gateway).await;
let blacklisted_ws_url =
format!("{blacklisted_gateway_url}/api/vscodex/ws").replace("http://", "ws://");
let blacklisted_error = tokio_tungstenite::connect_async(&blacklisted_ws_url)
.await
.expect_err("blacklisted WebSocket should be rejected");
match blacklisted_error {
tokio_tungstenite::tungstenite::Error::Http(response) => {
assert_eq!(response.status(), StatusCode::FORBIDDEN)
}
other => panic!("expected HTTP blacklist rejection, got {other:?}"),
}
blacklisted_gateway_handle.abort();
let requests = captured_requests
.lock()
.expect("captured request store should lock")
.clone();
assert_eq!(requests.len(), 9);
assert!(requests
.iter()
.all(|request| request.authorization.as_deref() == Some("Bearer sidecar-secret")));
assert!(requests[..8]
.iter()
.all(|request| request.path.starts_with("/internal/v1/users/user-auth-1/")));
assert_eq!(requests[8].path, "/v1/pairings/exchange");
assert_eq!(requests[1].body, Some(json!({ "name": "My Mac" })));
assert_eq!(requests[2].body, Some(json!({ "device_id": "host-1" })));
assert_eq!(
requests[8].body,
Some(json!({ "code": "PAIR-123", "name": "Office Mac" }))
);
assert_eq!(requests[8].client_ip.as_deref(), Some("127.0.0.1"));
assert!(requests[..8]
.iter()
.all(|request| request.client_ip.is_none()));
let _disabled = set_test_env_var("AETHER_VSCODEX_ENABLED", "false");
let disabled = client
.get(format!("{gateway_url}/api/users/me/vscodex/devices"))
.bearer_auth(&access_token)
.header("x-client-device-id", "browser-device-vscodex")
.send()
.await
.expect("disabled feature request should complete");
assert_eq!(disabled.status(), StatusCode::SERVICE_UNAVAILABLE);
let disabled_payload: serde_json::Value =
disabled.json().await.expect("disabled body should be JSON");
assert_eq!(disabled_payload["detail"], "VS Codex 服务未启用");
assert_eq!(
captured_requests
.lock()
.expect("captured request store should lock")
.len(),
9
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
sidecar_handle.abort();
}