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

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
@@ -1463,7 +1463,7 @@ mod tests {
&auth_config,
Some(0),
),
"antigravity_[email protected]"
"[email protected]"
);
}
@@ -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(),
@@ -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]
@@ -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 {
@@ -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") {
@@ -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
);
}
}
@@ -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())