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

fix: harden OAuth identity and cookies and correct quota and JSON display

Browse Source
This commit is contained in:
elky
2026-09-08 10:51:25 +08:00
parent 7113d04f8a
commit cf8ea19856
24 changed files with 1327 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs
+1 -1
View File
@@ -1463,7 +1463,7 @@ mod tests {
&auth_config,
Some(0),
),
"antigravity_[email protected]"
"[email protected]"
);
}
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/poll.rs
+33 -5
View File
@@ -1,3 +1,4 @@
use super::super::helpers::admin_provider_oauth_key_name_from_auth_config;
use super::super::kiro::{
admin_provider_oauth_kiro_refresh_base_url_override, fetch_admin_provider_oauth_kiro_email,
refresh_admin_provider_oauth_kiro_auth_config,
@@ -79,7 +80,7 @@ fn kiro_social_key_name(
.collect::<String>()
})
.unwrap_or_else(|| "unknown".to_string());
format!("kiro_{fallback} ({provider})")
format!("账号_{fallback} ({provider})")
}
fn kiro_social_poll_error_response(error: impl Into<String>) -> Response<Body> {
@@ -1004,10 +1005,11 @@ async fn handle_admin_provider_oauth_windsurf_browser_device_poll(
}
}
} else {
let key_name = email
.as_deref()
.map(|email| format!("windsurf_{email}"))
.unwrap_or_else(|| format!("windsurf_{}", current_unix_secs()));
let key_name = admin_provider_oauth_key_name_from_auth_config(
&provider.provider_type,
&auth_config,
None,
);
match state
.create_provider_oauth_catalog_key(
&provider.id,
@@ -1356,6 +1358,32 @@ mod tests {
use crate::control::GatewayAdminPrincipalContext;
use aether_data::repository::provider_oauth::StoredAdminProviderOAuthDeviceSession;
#[test]
fn kiro_social_key_name_preserves_email_and_auth_method() {
assert_eq!(
super::kiro_social_key_name(
Some(" [email protected] "),
Some("Github"),
Some("refresh-token-1"),
),
"[email protected] (Github)"
);
}
#[test]
fn kiro_social_key_name_without_email_uses_generic_account_prefix() {
for email in [None, Some(""), Some(" ")] {
assert_eq!(
super::kiro_social_key_name(email, Some("Google"), Some("refresh-token-1")),
"账号_154f43 (Google)"
);
assert_eq!(
super::kiro_social_key_name(email, None, None),
"账号_unknown (social)"
);
}
}
fn device_session() -> StoredAdminProviderOAuthDeviceSession {
StoredAdminProviderOAuthDeviceSession {
session_id: "device-session-1".to_string(),
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/helpers.rs
+125 -13
View File
@@ -52,13 +52,12 @@ pub(super) fn admin_provider_oauth_key_name_from_auth_config(
auth_config: &Map<String, Value>,
batch_index: Option<usize>,
) -> String {
let provider_type = provider_type.trim();
if let Some(email) = trimmed_auth_config_string(auth_config, "email") {
return format!("{provider_type}_{email}");
return email;
}
if provider_type.eq_ignore_ascii_case("grok") {
if provider_type.trim().eq_ignore_ascii_case("grok") {
if let Some(user_id) = trimmed_auth_config_string(auth_config, "user_id") {
return format!("grok_{user_id}");
return user_id;
}
}
@@ -68,7 +67,7 @@ pub(super) fn admin_provider_oauth_key_name_from_auth_config(
.map(|duration| duration.as_secs())
.unwrap_or(0);
match batch_index {
Some(index) => format!("{provider_type}_{timestamp}_{index}"),
Some(index) => format!("账号_{timestamp}_{index}"),
None => format!("账号_{timestamp}"),
}
}
@@ -87,6 +86,106 @@ mod tests {
use super::*;
use serde_json::{json, Map};
const PROVIDER_TYPES: &[&str] = &[
"codex",
" Codex ",
"claude_code",
"chatgpt_web",
"gemini_cli",
"antigravity",
"grok",
" Grok ",
"kiro",
"windsurf",
];
#[test]
fn default_key_name_uses_email_without_provider_prefix() {
let mut auth_config = Map::new();
auth_config.insert("email".to_string(), json!(" [email protected] "));
for provider_type in PROVIDER_TYPES {
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
"[email protected]"
);
}
}
}
#[test]
fn antigravity_default_key_name_uses_email_without_provider_prefix() {
for email in [" [email protected] ", "[email protected]"] {
let mut auth_config = Map::new();
auth_config.insert("email".to_string(), json!(email));
for provider_type in ["antigravity", " Antigravity "] {
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
email.trim()
);
}
}
}
}
#[test]
fn default_key_name_preserves_email_with_provider_prefix() {
for provider_type in PROVIDER_TYPES {
let email = format!("{}[email protected]", provider_type.trim());
let mut auth_config = Map::new();
auth_config.insert("email".to_string(), json!(email));
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
email
);
}
}
}
#[test]
fn default_key_name_without_email_uses_generic_account_name() {
for email in [None, Some(""), Some(" ")] {
let mut auth_config = Map::new();
if let Some(email) = email {
auth_config.insert("email".to_string(), json!(email));
}
for provider_type in PROVIDER_TYPES {
for batch_index in [None, Some(3)] {
let name = admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
);
let suffix = name.strip_prefix("账号_").expect("generic account prefix");
let timestamp = if batch_index.is_some() {
suffix.strip_suffix("_3").expect("batch index suffix")
} else {
suffix
};
assert!(timestamp.parse::<u64>().is_ok());
}
}
}
}
#[test]
fn grok_default_key_name_uses_full_user_id() {
let mut auth_config = Map::new();
@@ -95,10 +194,18 @@ mod tests {
json!("1619039a-0191-4e0a-a490-8f4ad21262c9"),
);
assert_eq!(
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, None),
"grok_1619039a-0191-4e0a-a490-8f4ad21262c9"
);
for provider_type in ["grok", " Grok "] {
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
"1619039a-0191-4e0a-a490-8f4ad21262c9"
);
}
}
}
#[test]
@@ -109,17 +216,22 @@ mod tests {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, None),
"grok_grok@example.com"
"[email protected]"
);
}
#[test]
fn batch_default_key_name_keeps_existing_timestamp_shape() {
fn batch_default_key_name_keeps_distinct_indexes_without_provider_prefix() {
let auth_config = Map::new();
let name = admin_provider_oauth_key_name_from_auth_config("codex", &auth_config, Some(3));
let name = admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, Some(3));
let other_name =
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, Some(4));
assert!(name.starts_with("codex_"));
assert!(name.starts_with("账号_"));
assert!(name.ends_with("_3"));
assert!(other_name.starts_with("账号_"));
assert!(other_name.ends_with("_4"));
assert_ne!(name, other_name);
}
#[test]
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
-62
View File
@@ -5,7 +5,6 @@ use super::shared::{
quota_key_auto_removed, quota_refresh_success_invalid_state,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::AdminImportProviderModelsRequest;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::{
@@ -24,63 +23,6 @@ use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
fn antigravity_discovered_model_ids(metadata_update: Option<&serde_json::Value>) -> Vec<String> {
metadata_update
.and_then(|value| value.pointer("/antigravity/quota_by_model"))
.and_then(serde_json::Value::as_object)
.into_iter()
.flat_map(|models| models.keys())
.map(String::as_str)
.filter(|model_id| aether_model_fetch::antigravity_model_id_is_routable(model_id))
.map(ToOwned::to_owned)
.collect()
}
async fn sync_antigravity_discovered_models(
state: &AdminAppState<'_>,
provider_id: &str,
metadata_update: Option<&serde_json::Value>,
) {
if !state.has_global_model_data_reader() || !state.has_global_model_data_writer() {
return;
}
let model_ids = antigravity_discovered_model_ids(metadata_update);
if model_ids.is_empty() {
return;
}
let result = state
.build_admin_import_provider_models_payload(
provider_id,
AdminImportProviderModelsRequest {
model_ids,
tiered_pricing: None,
price_per_request: None,
},
)
.await;
match result {
Ok(payload) => {
let errors = payload
.get("errors")
.and_then(serde_json::Value::as_array)
.map(Vec::len)
.unwrap_or(0);
if errors > 0 {
warn!(
provider_id,
errors, "Antigravity discovered-model catalog sync completed with item errors"
);
}
}
Err(error) => warn!(
provider_id,
error = %error,
"Antigravity discovered-model catalog sync failed"
),
}
}
async fn execute_antigravity_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
@@ -380,10 +322,6 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
continue;
}
if status == "success" {
sync_antigravity_discovered_models(state, &provider.id, metadata_update.as_ref()).await;
}
if status == "success" {
success_count += 1;
} else {
apps/aether-gateway/src/handlers/public/support.rs
+27
View File
@@ -30,6 +30,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
mod support_announcements;
#[path = "support/auth.rs"]
mod support_auth;
#[path = "support/auth_cookie_policy.rs"]
mod support_auth_cookie_policy;
#[path = "support/billing.rs"]
mod support_billing;
#[path = "support/ccswitch.rs"]
@@ -133,6 +135,31 @@ pub(crate) async fn maybe_build_local_public_support_response(
remote_addr: &std::net::SocketAddr,
client_ip: std::net::IpAddr,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let response = build_local_public_support_response(
state,
request_context,
headers,
remote_addr,
client_ip,
request_body,
)
.await?;
Some(support_auth_cookie_policy::finalize_refresh_cookie(
response,
headers,
request_context.host_header.as_deref(),
remote_addr,
))
}
async fn build_local_public_support_response(
state: &AppState,
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
remote_addr: &std::net::SocketAddr,
client_ip: std::net::IpAddr,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
if decision.route_class.as_deref() != Some("public_support") {
apps/aether-gateway/src/handlers/public/support/auth_cookie_policy.rs
+424
View File
@@ -0,0 +1,424 @@
use super::support_auth::{auth_refresh_cookie_name, auth_refresh_cookie_secure};
use axum::body::Body;
use axum::http::{header, HeaderMap, HeaderValue, Response};
use std::net::SocketAddr;
use url::Url;
pub(super) fn finalize_refresh_cookie(
mut response: Response<Body>,
headers: &HeaderMap,
host_header: Option<&str>,
remote_addr: &SocketAddr,
) -> Response<Body> {
if !response.headers().contains_key(header::SET_COOKIE) {
return response;
}
let cookie_name = auth_refresh_cookie_name();
let explicit_secure = std::env::var("AUTH_REFRESH_COOKIE_SECURE").ok();
let public_base_url = std::env::var("AETHER_PUBLIC_BASE_URL")
.ok()
.or_else(|| std::env::var("PUBLIC_BASE_URL").ok());
let secure = refresh_cookie_secure_for_request(
headers,
host_header,
crate::headers::trusted_proxy_ip(remote_addr.ip()),
explicit_secure.as_deref(),
public_base_url.as_deref(),
auth_refresh_cookie_secure(),
);
let cookies = response
.headers()
.get_all(header::SET_COOKIE)
.iter()
.map(|cookie| rewrite_refresh_cookie(cookie, &cookie_name, secure))
.collect::<Vec<_>>();
response.headers_mut().remove(header::SET_COOKIE);
for cookie in cookies {
response.headers_mut().append(header::SET_COOKIE, cookie);
}
response
}
fn refresh_cookie_secure_for_request(
headers: &HeaderMap,
host_header: Option<&str>,
trusted_proxy: bool,
explicit_secure: Option<&str>,
public_base_url: Option<&str>,
fallback_secure: bool,
) -> bool {
if let Some(value) = explicit_secure {
return !value.trim().eq_ignore_ascii_case("false");
}
let origin = single_header(headers, header::ORIGIN.as_str()).and_then(parse_origin);
let public_url = public_base_url.and_then(parse_http_url);
let forwarded_proto = trusted_proxy.then(|| forwarded_proto(headers)).flatten();
if origin.as_ref().is_some_and(|url| url.scheme() == "https")
|| public_url
.as_ref()
.is_some_and(|url| url.scheme() == "https")
|| forwarded_proto == Some("https")
{
return true;
}
if trusted_proxy && headers.contains_key("x-forwarded-proto") {
return forwarded_proto != Some("http");
}
if public_url
.as_ref()
.is_some_and(|url| url.scheme() == "http")
{
return false;
}
if let (Some(origin), Some(host)) = (origin, host_header) {
let request_origin = parse_origin(&format!("{}://{host}", origin.scheme()));
if request_origin.is_some_and(|url| url.origin() == origin.origin()) {
return false;
}
}
fallback_secure
}
fn single_header<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
let mut values = headers.get_all(name).iter();
let value = values.next()?.to_str().ok()?.trim();
(values.next().is_none() && !value.is_empty()).then_some(value)
}
fn parse_origin(value: &str) -> Option<Url> {
let url = parse_http_url(value)?;
(url.path() == "/").then_some(url)
}
fn parse_http_url(value: &str) -> Option<Url> {
let url = Url::parse(value.trim()).ok()?;
(matches!(url.scheme(), "http" | "https")
&& url.host_str().is_some()
&& url.username().is_empty()
&& url.password().is_none()
&& url.query().is_none()
&& url.fragment().is_none())
.then_some(url)
}
fn forwarded_proto(headers: &HeaderMap) -> Option<&'static str> {
let value = headers
.get_all("x-forwarded-proto")
.iter()
.last()?
.to_str()
.ok()?
.rsplit(',')
.next()?
.trim();
if value.eq_ignore_ascii_case("https") {
Some("https")
} else if value.eq_ignore_ascii_case("http") {
Some("http")
} else {
None
}
}
fn rewrite_refresh_cookie(cookie: &HeaderValue, cookie_name: &str, secure: bool) -> HeaderValue {
let Ok(value) = cookie.to_str() else {
return cookie.clone();
};
let mut attributes = value.split(';').map(str::trim);
let Some(pair) = attributes.next() else {
return cookie.clone();
};
if pair.split_once('=').map(|(name, _)| name) != Some(cookie_name) {
return cookie.clone();
}
let secure =
secure || cookie_name.starts_with("__Secure-") || cookie_name.starts_with("__Host-");
let mut parts = vec![pair.to_string()];
for attribute in attributes {
if attribute.eq_ignore_ascii_case("Secure") {
continue;
}
if !secure
&& attribute.split_once('=').is_some_and(|(name, value)| {
name.trim().eq_ignore_ascii_case("SameSite")
&& value.trim().eq_ignore_ascii_case("None")
})
{
parts.push("SameSite=Lax".to_string());
} else {
parts.push(attribute.to_string());
}
}
if secure {
parts.push("Secure".to_string());
}
let Ok(mut rewritten) = HeaderValue::from_str(&parts.join("; ")) else {
return cookie.clone();
};
rewritten.set_sensitive(cookie.is_sensitive());
rewritten
}
#[cfg(test)]
mod tests {
use super::{refresh_cookie_secure_for_request, rewrite_refresh_cookie};
use axum::http::{header, HeaderMap, HeaderValue};
fn headers(origin: Option<&str>, forwarded_proto: Option<&str>) -> HeaderMap {
let mut headers = HeaderMap::new();
if let Some(origin) = origin {
headers.insert(header::ORIGIN, HeaderValue::from_str(origin).unwrap());
}
if let Some(proto) = forwarded_proto {
headers.insert("x-forwarded-proto", HeaderValue::from_str(proto).unwrap());
}
headers
}
#[test]
fn refresh_cookie_auto_detects_same_origin_http_and_https() {
for (origin, host, secure) in [
("http://aether.test:8084", "aether.test:8084", false),
("http://aether.test", "aether.test:80", false),
("http://[2001:db8::1]:8084", "[2001:db8::1]:8084", false),
("https://aether.test", "aether.test", true),
] {
assert_eq!(
refresh_cookie_secure_for_request(
&headers(Some(origin), None),
Some(host),
false,
None,
None,
true,
),
secure,
"{origin}",
);
}
assert!(refresh_cookie_secure_for_request(
&headers(Some("https://aether.test"), None),
Some("aether.test"),
false,
None,
None,
false,
));
}
#[test]
fn refresh_cookie_does_not_infer_http_from_other_or_invalid_origins() {
for origin in [
"http://other.test",
"http://aether.test:8085",
"null",
"http://[email protected]:8084",
"http://aether.test:8084/path",
"http://aether.test:8084?query",
"http://aether.test:8084#fragment",
"http://aether.test:8084, https://aether.test:8084",
"file:///tmp/test",
] {
assert!(
refresh_cookie_secure_for_request(
&headers(Some(origin), None),
Some("aether.test:8084"),
false,
None,
None,
true,
),
"{origin}"
);
}
let mut duplicate = headers(Some("http://aether.test:8084"), None);
duplicate.append(
header::ORIGIN,
HeaderValue::from_static("https://aether.test:8084"),
);
assert!(refresh_cookie_secure_for_request(
&duplicate,
Some("aether.test:8084"),
false,
None,
None,
true,
));
}
#[test]
fn refresh_cookie_only_trusts_forwarded_protocol_from_trusted_peers() {
for (proto, trusted, secure) in [
("http", true, false),
("https", true, true),
("http", false, true),
("https", false, true),
("https, http", true, false),
("http, https", true, true),
("ftp", true, true),
("http,", true, true),
] {
assert_eq!(
refresh_cookie_secure_for_request(
&headers(None, Some(proto)),
Some("aether.test"),
trusted,
None,
None,
true,
),
secure,
"{proto}, trusted={trusted}"
);
}
let mut chained = headers(None, Some("http, http"));
chained.append("x-forwarded-proto", HeaderValue::from_static("https"));
assert!(refresh_cookie_secure_for_request(
&chained,
Some("aether.test"),
true,
None,
None,
true,
));
}
#[test]
fn refresh_cookie_https_evidence_prevents_automatic_downgrade() {
for (origin, proto, public_url) in [
("https://aether.test", "http", None),
("http://aether.test", "https", None),
("http://aether.test", "http", Some("https://aether.test")),
] {
assert!(refresh_cookie_secure_for_request(
&headers(Some(origin), Some(proto)),
Some("aether.test"),
true,
None,
public_url,
true,
));
}
}
#[test]
fn refresh_cookie_preserves_explicit_overrides_and_unknown_defaults() {
for (explicit, secure) in [
("true", true),
("FALSE", false),
("invalid", true),
("", true),
] {
assert_eq!(
refresh_cookie_secure_for_request(
&headers(Some("http://aether.test"), None),
Some("aether.test"),
false,
Some(explicit),
None,
true,
),
secure
);
}
assert!(!refresh_cookie_secure_for_request(
&headers(Some("https://aether.test"), None),
Some("aether.test"),
false,
Some("false"),
None,
true,
));
for fallback in [false, true] {
assert_eq!(
refresh_cookie_secure_for_request(
&HeaderMap::new(),
Some("aether.test"),
false,
None,
None,
fallback,
),
fallback
);
}
}
#[test]
fn refresh_cookie_accepts_an_explicit_public_http_origin() {
assert!(!refresh_cookie_secure_for_request(
&HeaderMap::new(),
Some("internal:8084"),
false,
None,
Some("http://aether.test"),
true,
));
assert!(refresh_cookie_secure_for_request(
&headers(Some("https://aether.test"), None),
Some("internal:8084"),
false,
None,
Some("http://aether.test"),
true,
));
}
#[test]
fn refresh_cookie_rewrite_preserves_secret_path_expiry_and_httponly() {
let mut cookie = HeaderValue::from_static(
"aether_refresh_token=secret; Path=/api/auth; HttpOnly; SameSite=None; Max-Age=604800; Secure",
);
cookie.set_sensitive(true);
let rewritten = rewrite_refresh_cookie(&cookie, "aether_refresh_token", false);
assert_eq!(
rewritten.to_str().unwrap(),
"aether_refresh_token=secret; Path=/api/auth; HttpOnly; SameSite=Lax; Max-Age=604800"
);
assert!(rewritten.is_sensitive());
assert_eq!(
rewrite_refresh_cookie(&cookie, "aether_refresh_token", true),
cookie
);
}
#[test]
fn refresh_cookie_rewrite_also_clears_http_cookies() {
let cookie = HeaderValue::from_static(
"aether_refresh_token=; Path=/api/auth; HttpOnly; SameSite=None; Max-Age=0; Secure",
);
assert_eq!(
rewrite_refresh_cookie(&cookie, "aether_refresh_token", false)
.to_str()
.unwrap(),
"aether_refresh_token=; Path=/api/auth; HttpOnly; SameSite=Lax; Max-Age=0"
);
}
#[test]
fn refresh_cookie_rewrite_preserves_other_cookies_and_strict_policy() {
let unrelated = HeaderValue::from_static("oauth_binding=secret; Path=/; Secure; HttpOnly");
assert_eq!(
rewrite_refresh_cookie(&unrelated, "aether_refresh_token", false),
unrelated
);
let strict = HeaderValue::from_static(
"custom_refresh=secret; Path=/api/auth; HttpOnly; SameSite=Strict",
);
assert_eq!(
rewrite_refresh_cookie(&strict, "custom_refresh", false),
strict
);
assert!(rewrite_refresh_cookie(&strict, "custom_refresh", true)
.to_str()
.unwrap()
.ends_with("; Secure"));
let prefixed =
HeaderValue::from_static("__Secure-refresh=secret; HttpOnly; SameSite=None; Secure");
assert_eq!(
rewrite_refresh_cookie(&prefixed, "__Secure-refresh", false),
prefixed
);
}
}
apps/aether-gateway/src/handlers/public/support/auth_helpers.rs
+1 -1
View File
@@ -248,7 +248,7 @@ pub(super) fn auth_verification_send_cooldown_seconds() -> i64 {
.unwrap_or(60)
}
pub(super) fn auth_refresh_cookie_name() -> String {
pub(crate) fn auth_refresh_cookie_name() -> String {
std::env::var("AUTH_REFRESH_COOKIE_NAME")
.ok()
.map(|value| value.trim().to_string())
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+25 -14
View File
@@ -8,7 +8,7 @@ use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::global_models::{
AdminProviderModelListQuery, GlobalModelReadRepository,
AdminGlobalModelListQuery, AdminProviderModelListQuery, GlobalModelReadRepository,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
@@ -20,8 +20,9 @@ use http::StatusCode;
use serde_json::json;
use super::super::super::{
build_router_with_state, build_state_with_execution_runtime_override, sample_bound_auth_config,
sample_bound_key, sample_endpoint, sample_key, sample_proxy_node, start_server, AppState,
build_router_with_state, build_state_with_execution_runtime_override, sample_admin_global_model,
sample_bound_auth_config, sample_bound_key, sample_endpoint, sample_key, sample_proxy_node,
start_server, AppState,
};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
@@ -2421,7 +2422,15 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
)],
vec![key],
));
let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::default());
let existing_global_model = sample_admin_global_model(
"global-claude-sonnet-4",
"claude-sonnet-4",
"Claude Sonnet 4",
);
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::default()
.with_admin_global_models(vec![existing_global_model.clone()]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
@@ -2533,7 +2542,16 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
.and_then(|value| value.get("remaining_fraction")),
Some(&json!(0.25))
);
let imported_provider_models = global_model_repository
let global_models = global_model_repository
.list_admin_global_models(&AdminGlobalModelListQuery {
limit: 100,
..Default::default()
})
.await
.expect("global models should read after quota refresh");
assert_eq!(global_models.total, 1);
assert_eq!(global_models.items, vec![existing_global_model]);
let provider_models = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: "provider-antigravity".to_string(),
is_active: None,
@@ -2541,15 +2559,8 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
limit: 100,
})
.await
.expect("imported Antigravity provider models should read");
let imported_model_names = imported_provider_models
.iter()
.map(|model| model.provider_model_name.as_str())
.collect::<std::collections::BTreeSet<_>>();
assert!(imported_model_names.contains("claude-sonnet-4"));
assert!(imported_model_names.contains("gemini-2.5-pro"));
assert!(imported_model_names.contains("gemini-3.7-flash-tiered"));
assert!(!imported_model_names.contains("chat_23310"));
.expect("Antigravity provider models should read after quota refresh");
assert!(provider_models.is_empty());
assert_eq!(
reloaded[0]
.upstream_metadata
apps/aether-gateway/src/tests/control/admin/oauth.rs
+57 -10
View File
@@ -3916,11 +3916,47 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad
fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email() {
run_admin_oauth_test(
"gateway_names_new_antigravity_oauth_account_from_google_userinfo_email",
gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl,
|| {
assert_antigravity_oauth_account_uses_google_userinfo_email(
"complete",
json!({
"callback_url": "http://localhost:51121/oauth2callback?code=antigravity-code-123&state=cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"
}),
)
},
);
}
async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl() {
#[test]
fn gateway_names_imported_antigravity_oauth_account_from_google_userinfo_email() {
run_admin_oauth_test(
"gateway_names_imported_antigravity_oauth_account_from_google_userinfo_email",
|| {
assert_antigravity_oauth_account_uses_google_userinfo_email(
"import-refresh-token",
json!({"refresh_token": "antigravity-import-refresh-token"}),
)
},
);
}
#[test]
fn gateway_names_batch_imported_antigravity_oauth_account_from_google_userinfo_email() {
run_admin_oauth_test(
"gateway_names_batch_imported_antigravity_oauth_account_from_google_userinfo_email",
|| {
assert_antigravity_oauth_account_uses_google_userinfo_email(
"batch-import",
json!({"credentials": "antigravity-import-refresh-token"}),
)
},
);
}
async fn assert_antigravity_oauth_account_uses_google_userinfo_email(
operation: &str,
request_body: Value,
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().fallback(any(move |_request: Request| {
@@ -4025,15 +4061,13 @@ async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/complete"
"{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/{operation}"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"callback_url": "http://localhost:51121/oauth2callback?code=antigravity-code-123&state=cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"
}))
.json(&request_body)
.send()
.await
.expect("request should succeed");
@@ -4041,9 +4075,22 @@ async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_
let status = response.status();
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "antigravity");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
let account_result = if operation == "batch-import" {
assert_eq!(payload["total"], 1);
assert_eq!(payload["success"], 1, "payload={payload}");
assert_eq!(payload["failed"], 0);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(
payload["results"][0]["key_name"],
"[email protected]"
);
&payload["results"][0]
} else {
assert_eq!(payload["provider_type"], "antigravity");
assert_eq!(payload["email"], "[email protected]");
&payload
};
assert_eq!(account_result["replaced"], false);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
assert_eq!(*user_info_hits.lock().expect("mutex should lock"), 1);
assert_eq!(
@@ -4055,7 +4102,7 @@ async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let key_id = payload["key_id"]
let key_id = account_result["key_id"]
.as_str()
.expect("created key id should be returned")
.to_string();
apps/aether-gateway/src/tests/frontdoor/public_support.rs
+2
View File
@@ -50,6 +50,8 @@ use chrono::{TimeZone, Utc};
const TEST_EMAIL_VERIFICATION_TOKEN: &str =
"test-email-verification-token-00000000000000000000000000000000";
#[path = "public_support/auth_cookie.rs"]
mod auth_cookie;
#[path = "public_support/dashboard.rs"]
mod dashboard;
#[path = "public_support/vscodex.rs"]
apps/aether-gateway/src/tests/frontdoor/public_support/auth_cookie.rs
+124
View File
@@ -0,0 +1,124 @@
use super::{sample_auth_user, sample_auth_wallet, start_auth_gateway_with_state};
use axum::http::{header, StatusCode};
use chrono::Utc;
use serde_json::json;
fn refresh_cookie(response: &reqwest::Response, secure: bool) -> String {
let cookie = response
.headers()
.get(header::SET_COOKIE)
.unwrap()
.to_str()
.unwrap();
assert!(cookie.starts_with("aether_refresh_token="));
assert!(cookie.contains("HttpOnly"));
assert!(cookie.contains("Path=/api/auth"));
assert_eq!(
cookie
.split(';')
.any(|attribute| attribute.trim() == "Secure"),
secure
);
if !secure {
assert!(!cookie.contains("SameSite=None"));
assert!(cookie.contains("SameSite=Lax"));
}
assert_eq!(
response.headers().get(header::CACHE_CONTROL).unwrap(),
"no-store"
);
cookie.to_string()
}
#[tokio::test]
async fn gateway_auth_refresh_cookie_roundtrip_adapts_to_http_and_https() {
for (origin_scheme, forwarded_proto, secure) in [
("http", None, false),
("https", None, true),
("https", Some("https"), true),
("http", Some("https"), true),
] {
let now = Utc::now();
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_state(
sample_auth_user(now),
sample_auth_wallet("user-auth-1", now),
[],
)
.await;
let mut headers = reqwest::header::HeaderMap::new();
headers.insert(
header::ORIGIN,
gateway_url
.replacen("http:", &format!("{origin_scheme}:"), 1)
.parse()
.unwrap(),
);
headers.insert(
"x-client-device-id",
"cookie-roundtrip-device".parse().unwrap(),
);
if let Some(proto) = forwarded_proto {
headers.insert("x-forwarded-proto", proto.parse().unwrap());
}
let client = reqwest::Client::builder()
.default_headers(headers)
.build()
.unwrap();
let login = client.post(format!("{gateway_url}/api/auth/login"))
.json(&json!({ "email": "[email protected]", "password": "secret123", "auth_type": "local" }))
.send().await.unwrap();
assert_eq!(login.status(), StatusCode::OK);
let mut cookie = refresh_cookie(&login, secure);
for _ in 0..3 {
let refreshed = client
.post(format!("{gateway_url}/api/auth/refresh"))
.header(header::COOKIE, cookie.split(';').next().unwrap())
.send()
.await
.unwrap();
assert_eq!(refreshed.status(), StatusCode::OK);
let rotated = refresh_cookie(&refreshed, secure);
assert_ne!(rotated, cookie);
cookie = rotated;
let payload: serde_json::Value = refreshed.json().await.unwrap();
let current_user = client
.get(format!("{gateway_url}/api/auth/me"))
.bearer_auth(payload["access_token"].as_str().unwrap())
.send()
.await
.unwrap();
assert_eq!(current_user.status(), StatusCode::OK);
}
let logout = client
.post(format!("{gateway_url}/api/auth/logout"))
.header(header::COOKIE, cookie.split(';').next().unwrap())
.send()
.await
.unwrap();
assert_eq!(logout.status(), StatusCode::OK);
assert!(refresh_cookie(&logout, secure).contains("Max-Age=0"));
let revoked = client
.post(format!("{gateway_url}/api/auth/refresh"))
.header(header::COOKIE, cookie.split(';').next().unwrap())
.send()
.await
.unwrap();
assert_eq!(revoked.status(), StatusCode::UNAUTHORIZED);
assert!(refresh_cookie(&revoked, secure).contains("Max-Age=0"));
let missing = client
.post(format!("{gateway_url}/api/auth/refresh"))
.send()
.await
.unwrap();
assert_eq!(missing.status(), StatusCode::UNAUTHORIZED);
assert!(refresh_cookie(&missing, secure).contains("Max-Age=0"));
assert_eq!(*upstream_hits.lock().unwrap(), 0);
gateway_handle.abort();
upstream_handle.abort();
}
}