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

feat(admin): 完善代理节点与 OAuth 授权管理

Browse Source
This commit is contained in:
fawney19
2026-04-14 14:09:24 +08:00
parent 593640ac19
commit 861ae81ff0
44 changed files with 2757 additions and 365 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
apps/aether-gateway/src/state/core.rs
View File

@@ -4,8 +4,8 @@ use std::sync::Mutex as StdMutex;
use std::time::Duration;
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode,
StoredProxyNodeEvent,
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeTunnelStatusMutation, StoredProxyNode, StoredProxyNodeEvent,
};
use aether_http::{build_http_client, HttpClientConfig};
use aether_runtime::{
@@ -468,6 +468,26 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn create_manual_proxy_node(
&self,
mutation: &ProxyNodeManualCreateMutation,
) -> Result<Option<StoredProxyNode>, GatewayError> {
self.data
.create_manual_proxy_node(mutation)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn update_manual_proxy_node(
&self,
mutation: &ProxyNodeManualUpdateMutation,
) -> Result<Option<StoredProxyNode>, GatewayError> {
self.data
.update_manual_proxy_node(mutation)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub async fn reset_stale_proxy_node_tunnel_statuses(&self) -> std::io::Result<usize> {
self.data
.reset_stale_proxy_node_tunnel_statuses()
@@ -495,6 +515,16 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn delete_proxy_node(
&self,
node_id: &str,
) -> Result<Option<StoredProxyNode>, GatewayError> {
self.data
.delete_proxy_node(node_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn update_proxy_node_remote_config(
&self,
mutation: &aether_data::repository::proxy_nodes::ProxyNodeRemoteConfigMutation,

8
apps/aether-gateway/src/state/integrations.rs
View File

@@ -23,10 +23,7 @@ use tracing::debug;
use super::{AppState, GatewayError};
use crate::model_fetch::ModelFetchRuntimeState;
use crate::provider_transport::{
resolve_transport_proxy_snapshot_with_tunnel_affinity, GatewayProviderTransportSnapshot,
LocalResolvedOAuthRequestAuth,
};
use crate::provider_transport::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
use crate::request_candidate_runtime::{
RequestCandidateRuntimeCapabilityReader, RequestCandidateRuntimeReader,
RequestCandidateRuntimeWriter,
@@ -90,7 +87,8 @@ impl ModelFetchTransportRuntime for AppState {
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot> {
resolve_transport_proxy_snapshot_with_tunnel_affinity(self, transport).await
self.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
async fn execute_model_fetch_execution_plan(

1
apps/aether-gateway/src/state/mod.rs
View File

@@ -10,6 +10,7 @@ mod core;
mod cors;
mod integrations;
mod oauth;
mod proxy;
mod runtime;
#[cfg(test)]
mod testing;

177
apps/aether-gateway/src/state/oauth.rs
View File

@@ -3,12 +3,41 @@ use super::{
GatewayError, ProviderTransportSnapshotCacheKey, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES,
PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL,
};
use crate::provider_transport::LocalOAuthHttpExecutor;
use super::super::provider_transport;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use base64::{engine::general_purpose::STANDARD, Engine as _};
use flate2::read::{DeflateDecoder, GzDecoder};
use std::collections::BTreeMap;
use std::io::Read;
use std::time::Duration;
use aether_crypto::encrypt_python_fernet_plaintext;
const LOCAL_OAUTH_HTTP_TIMEOUT_MS: u64 = 30_000;
struct GatewayLocalOAuthHttpExecutor<'a> {
state: &'a AppState,
}
#[async_trait::async_trait]
impl<'a> provider_transport::LocalOAuthHttpExecutor for GatewayLocalOAuthHttpExecutor<'a> {
async fn execute(
&self,
provider_type: &'static str,
transport: &provider_transport::GatewayProviderTransportSnapshot,
request: &provider_transport::LocalOAuthHttpRequest,
) -> Result<
provider_transport::LocalOAuthHttpResponse,
provider_transport::LocalOAuthRefreshError,
> {
self.state
.execute_local_oauth_http_request(provider_type, transport, request)
.await
}
}
impl AppState {
pub(crate) fn clear_provider_transport_snapshot_cache(&self) {
self.provider_transport_snapshot_cache
@@ -368,12 +397,13 @@ impl AppState {
let distributed_lock = self.data.oauth_refresh_lock_runner();
let lock_owner = format!("aether-gateway-{}", std::process::id());
let mut current_transport = transport.clone();
let executor = GatewayLocalOAuthHttpExecutor { state: self };
for _ in 0..2 {
let resolution = self
.oauth_refresh
.resolve_with_result(
&self.client,
&executor,
&current_transport,
distributed_lock.as_ref(),
Some(lock_owner.as_str()),
@@ -429,12 +459,13 @@ impl AppState {
let lock_owner = format!("aether-gateway-admin-{}", std::process::id());
let mut current_transport = transport.clone();
current_transport.key.decrypted_api_key = "__placeholder__".to_string();
let executor = GatewayLocalOAuthHttpExecutor { state: self };
for _ in 0..2 {
let resolution = self
.oauth_refresh
.resolve_with_result(
&self.client,
&executor,
&current_transport,
distributed_lock.as_ref(),
Some(lock_owner.as_str()),
@@ -534,6 +565,88 @@ impl AppState {
Ok(())
}
async fn execute_local_oauth_http_request(
&self,
provider_type: &'static str,
transport: &provider_transport::GatewayProviderTransportSnapshot,
request: &provider_transport::LocalOAuthHttpRequest,
) -> Result<
provider_transport::LocalOAuthHttpResponse,
provider_transport::LocalOAuthRefreshError,
> {
if local_oauth_request_uses_direct_client(request.url.as_str()) {
let executor =
provider_transport::ReqwestLocalOAuthHttpExecutor::new(self.client.clone());
return executor.execute(provider_type, transport, request).await;
}
let body = if let Some(json_body) = request.json_body.clone() {
RequestBody::from_json(json_body)
} else {
RequestBody {
json_body: None,
body_bytes_b64: request
.body_bytes
.as_ref()
.map(|bytes| STANDARD.encode(bytes)),
body_ref: None,
}
};
let plan = ExecutionPlan {
request_id: request.request_id.to_string(),
candidate_id: None,
provider_name: Some(transport.provider.name.clone()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: request.method.as_str().to_string(),
url: request.url.clone(),
headers: request.headers.clone(),
content_type: request
.headers
.get("content-type")
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
content_encoding: None,
body,
stream: false,
client_api_format: "provider_oauth:local_refresh".to_string(),
provider_api_format: "provider_oauth:local_refresh".to_string(),
model_name: Some(provider_type.to_string()),
proxy: self
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(LOCAL_OAUTH_HTTP_TIMEOUT_MS),
read_ms: Some(LOCAL_OAUTH_HTTP_TIMEOUT_MS),
write_ms: Some(LOCAL_OAUTH_HTTP_TIMEOUT_MS),
pool_ms: Some(LOCAL_OAUTH_HTTP_TIMEOUT_MS),
total_ms: Some(LOCAL_OAUTH_HTTP_TIMEOUT_MS),
..ExecutionTimeouts::default()
}),
};
let result =
crate::execution_runtime::execute_execution_runtime_sync_plan(self, None, &plan)
.await
.map_err(
|err| provider_transport::LocalOAuthRefreshError::InvalidResponse {
provider_type,
message: match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Internal(message) => message,
},
},
)?;
Ok(provider_transport::LocalOAuthHttpResponse {
status_code: result.status_code,
body_text: local_oauth_execution_body_text(&result),
})
}
async fn wait_for_remote_oauth_refresh(
&self,
transport: &provider_transport::GatewayProviderTransportSnapshot,
@@ -564,3 +677,63 @@ impl AppState {
Ok(None)
}
}
fn local_oauth_execution_body_text(result: &aether_contracts::ExecutionResult) -> String {
result
.body
.as_ref()
.and_then(|body| local_oauth_execution_body_bytes(&result.headers, body))
.map(|bytes| String::from_utf8_lossy(&bytes).to_string())
.or_else(|| {
result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(|value| serde_json::to_string(value).ok())
})
.unwrap_or_default()
}
fn local_oauth_execution_body_bytes(
headers: &BTreeMap<String, String>,
body: &aether_contracts::ResponseBody,
) -> Option<Vec<u8>> {
let bytes = body
.body_bytes_b64
.as_deref()
.and_then(|value| STANDARD.decode(value).ok())?;
let encoding = headers
.get("content-encoding")
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase());
match encoding.as_deref() {
Some("gzip") => {
let mut decoder = GzDecoder::new(bytes.as_slice());
let mut out = Vec::new();
decoder.read_to_end(&mut out).ok()?;
Some(out)
}
Some("deflate") => {
let mut decoder = DeflateDecoder::new(bytes.as_slice());
let mut out = Vec::new();
decoder.read_to_end(&mut out).ok()?;
Some(out)
}
_ => Some(bytes),
}
}
fn local_oauth_request_uses_direct_client(url: &str) -> bool {
reqwest::Url::parse(url)
.ok()
.and_then(|parsed| parsed.host_str().map(str::to_owned))
.is_some_and(|host| {
host.eq_ignore_ascii_case("localhost")
|| host
.parse::<std::net::IpAddr>()
.map(|addr| addr.is_loopback())
.unwrap_or(false)
})
}

217
apps/aether-gateway/src/state/proxy.rs Normal file
View File

@@ -0,0 +1,217 @@
use aether_contracts::ProxySnapshot;
use aether_data::repository::proxy_nodes::proxy_node_accepts_new_tunnels;
use serde_json::{json, Map, Value};
use super::AppState;
use crate::provider_transport::{GatewayProviderTransportSnapshot, TransportTunnelAffinityLookup};
const TUNNEL_BASE_URL_EXTRA_KEY: &str = "tunnel_base_url";
const TUNNEL_OWNER_INSTANCE_ID_EXTRA_KEY: &str = "tunnel_owner_instance_id";
const TUNNEL_OWNER_OBSERVED_AT_EXTRA_KEY: &str = "tunnel_owner_observed_at_unix_secs";
impl AppState {
pub(crate) async fn read_system_proxy_node_id(&self) -> Option<String> {
self.read_system_config_json_value("system_proxy_node_id")
.await
.ok()
.flatten()
.and_then(|value| value.as_str().map(str::trim).map(ToOwned::to_owned))
.filter(|value| !value.is_empty())
}
pub(crate) async fn resolve_proxy_node_snapshot(
&self,
node_id: Option<&str>,
) -> Option<ProxySnapshot> {
let node_id = node_id.map(str::trim).filter(|value| !value.is_empty())?;
let node = self.find_proxy_node(node_id).await.ok().flatten()?;
if node.status.trim() != "online" {
return None;
}
if !proxy_node_accepts_new_tunnels(&node) {
return None;
}
if node.tunnel_mode && node.tunnel_connected {
let mut extra = Map::new();
if let Ok(Some(owner)) = self.lookup_tunnel_attachment_owner(node_id).await {
extra.insert(
TUNNEL_BASE_URL_EXTRA_KEY.to_string(),
Value::String(owner.relay_base_url),
);
extra.insert(
TUNNEL_OWNER_INSTANCE_ID_EXTRA_KEY.to_string(),
Value::String(owner.gateway_instance_id),
);
extra.insert(
TUNNEL_OWNER_OBSERVED_AT_EXTRA_KEY.to_string(),
json!(owner.observed_at_unix_secs),
);
}
return Some(ProxySnapshot {
enabled: Some(true),
mode: Some("tunnel".to_string()),
node_id: Some(node.id),
label: Some(node.name),
url: None,
extra: if extra.is_empty() {
None
} else {
Some(Value::Object(extra))
},
});
}
if !node.is_manual {
return None;
}
let proxy_url = node
.proxy_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(ProxySnapshot {
enabled: Some(true),
mode: proxy_mode_from_url(Some(proxy_url)),
node_id: Some(node.id),
label: Some(node.name),
url: proxy_url_with_node_auth(
proxy_url,
node.proxy_username.as_deref(),
node.proxy_password.as_deref(),
)
.or_else(|| Some(proxy_url.to_string())),
extra: None,
})
}
pub(crate) async fn resolve_system_proxy_snapshot(&self) -> Option<ProxySnapshot> {
let node_id = self.read_system_proxy_node_id().await;
self.resolve_proxy_node_snapshot(node_id.as_deref()).await
}
pub(crate) async fn resolve_transport_proxy_snapshot_with_tunnel_affinity(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot> {
if let Some(snapshot) = self
.resolve_proxy_snapshot_from_config(transport.key.proxy.as_ref())
.await
{
return Some(snapshot);
}
if let Some(snapshot) = self
.resolve_proxy_snapshot_from_config(transport.provider.proxy.as_ref())
.await
{
return Some(snapshot);
}
if let Some(snapshot) = self.resolve_system_proxy_snapshot().await {
return Some(snapshot);
}
self.resolve_proxy_snapshot_from_config(transport.endpoint.proxy.as_ref())
.await
}
async fn resolve_proxy_snapshot_from_config(
&self,
raw: Option<&Value>,
) -> Option<ProxySnapshot> {
let object = raw?.as_object()?;
if !proxy_enabled(object) {
return None;
}
let node_id = json_string_field(object, "node_id");
if let Some(snapshot) = self.resolve_proxy_node_snapshot(node_id.as_deref()).await {
return Some(snapshot);
}
proxy_snapshot_from_object(object)
}
}
fn proxy_enabled(object: &Map<String, Value>) -> bool {
object
.get("enabled")
.and_then(Value::as_bool)
.unwrap_or(true)
}
fn proxy_snapshot_from_object(object: &Map<String, Value>) -> Option<ProxySnapshot> {
let mode = json_string_field(object, "mode");
let node_id = json_string_field(object, "node_id");
let label = json_string_field(object, "label");
let url = json_string_field(object, "url").or_else(|| json_string_field(object, "proxy_url"));
if node_id.is_none() && url.is_none() {
return None;
}
let mut extra = Map::new();
for (key, value) in object {
if matches!(
key.as_str(),
"enabled" | "mode" | "node_id" | "label" | "url" | "proxy_url"
) {
continue;
}
extra.insert(key.clone(), value.clone());
}
Some(ProxySnapshot {
enabled: object.get("enabled").and_then(Value::as_bool),
mode,
node_id,
label,
url,
extra: if extra.is_empty() {
None
} else {
Some(Value::Object(extra))
},
})
}
fn json_string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn proxy_mode_from_url(proxy_url: Option<&str>) -> Option<String> {
let proxy_url = proxy_url?.trim();
if proxy_url.is_empty() {
return None;
}
let scheme = url::Url::parse(proxy_url)
.ok()
.map(|value| value.scheme().to_ascii_lowercase())
.unwrap_or_default();
if scheme.starts_with("socks") {
Some("socks".to_string())
} else {
Some("http".to_string())
}
}
fn proxy_url_with_node_auth(
proxy_url: &str,
username: Option<&str>,
password: Option<&str>,
) -> Option<String> {
let username = username.map(str::trim).filter(|value| !value.is_empty())?;
let mut parsed = url::Url::parse(proxy_url).ok()?;
if parsed.set_username(username).is_err() {
return None;
}
let password = password
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_default();
if parsed.set_password(Some(password)).is_err() {
return None;
}
Some(parsed.to_string())
}