refactor(oauth): 抽出 Kiro 刷新逻辑为共享模块并在 device-poll 复用

- 将 batch/kiro_import 中的 Kiro refresh/IDC helpers 抽到 dispatch/kiro.rs 共享
- device-poll 新增 IDC refresh 复核、JWT 缺失时通过 usage 接口回填 email、代理节点写入 key.proxy
- 补充对应集成测试覆盖 refresh 复核与 email 回填路径
This commit is contained in:
fawney19
2026-04-20 01:06:51 +08:00
parent c2b7b6c231
commit 13bf222b8d
5 changed files with 768 additions and 330 deletions

View File

@@ -1,4 +1,8 @@
use super::parse::{AdminProviderOAuthBatchImportEntry, AdminProviderOAuthBatchImportOutcome};
use super::super::kiro::{
admin_provider_oauth_kiro_refresh_base_url_override,
refresh_admin_provider_oauth_kiro_auth_config,
};
use super::parse::AdminProviderOAuthBatchImportOutcome;
use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider_oauth_key;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
@@ -11,315 +15,12 @@ use crate::handlers::admin::provider::oauth::runtime::{
use crate::handlers::admin::provider::oauth::state::decode_jwt_claims;
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
use crate::provider_transport::kiro::generate_machine_id;
use crate::GatewayError;
use aether_admin::provider::oauth::{
build_kiro_batch_import_key_name, coerce_admin_provider_oauth_import_str,
parse_admin_provider_oauth_kiro_batch_import_entries,
};
use aether_contracts::ProxySnapshot;
use serde_json::{json, Map, Value};
use std::collections::BTreeSet;
use std::time::{SystemTime, UNIX_EPOCH};
const KIRO_IDC_AMZ_USER_AGENT: &str =
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
fn admin_provider_oauth_kiro_refresh_base_url_override(
state: &AdminAppState<'_>,
override_key: &str,
) -> Option<String> {
let override_url = state.provider_oauth_token_url(override_key, "");
let normalized = override_url.trim();
(!normalized.is_empty()).then(|| normalized.to_string())
}
fn admin_provider_oauth_kiro_build_refresh_url(
auth_config: &AdminKiroAuthConfig,
override_base_url: Option<&str>,
path: &str,
default_host: impl FnOnce(&str) -> String,
) -> String {
if let Some(base_url) = override_base_url
.map(str::trim)
.filter(|value| !value.is_empty())
{
return format!("{}/{}", base_url.trim_end_matches('/'), path);
}
let region = auth_config.effective_auth_region();
default_host(region)
}
fn admin_provider_oauth_kiro_effective_host(url: &str, fallback_host: String) -> String {
reqwest::Url::parse(url)
.ok()
.and_then(|value| value.host_str().map(ToOwned::to_owned))
.unwrap_or(fallback_host)
}
fn admin_provider_oauth_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
if machine_id.trim().is_empty() {
format!("KiroIDE-{kiro_version}")
} else {
format!("KiroIDE-{kiro_version}-{machine_id}")
}
}
fn admin_provider_oauth_kiro_refresh_expires_at(payload: &Value) -> u64 {
let expires_in = payload
.get("expiresIn")
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_str()?.parse::<u64>().ok())
})
.unwrap_or(3600);
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default()
.saturating_add(expires_in)
}
fn admin_provider_oauth_kiro_refresh_response_json(
body_text: &str,
json_body: Option<Value>,
) -> Result<Value, String> {
json_body
.or_else(|| serde_json::from_str::<Value>(body_text).ok())
.ok_or_else(|| "refresh 接口返回了非 JSON 响应".to_string())
}
fn admin_provider_oauth_kiro_refresh_error_detail(
status: http::StatusCode,
body_text: &str,
) -> String {
let detail = body_text.trim();
if detail.is_empty() {
format!("HTTP {}", status.as_u16())
} else {
detail.to_string()
}
}
async fn refresh_admin_provider_oauth_kiro_auth_config(
state: &AdminAppState<'_>,
auth_config: &AdminKiroAuthConfig,
proxy: Option<ProxySnapshot>,
social_refresh_base_url: Option<&str>,
idc_refresh_base_url: Option<&str>,
) -> Result<AdminKiroAuthConfig, String> {
if auth_config.is_idc_auth() {
let fallback_host = format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region());
let url = admin_provider_oauth_kiro_build_refresh_url(
auth_config,
idc_refresh_base_url,
"token",
|region| format!("https://oidc.{region}.amazonaws.com/token"),
);
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
let headers = reqwest::header::HeaderMap::from_iter([
(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(&host)
.map_err(|_| "IDC host 无效".to_string())?,
),
(
reqwest::header::HeaderName::from_static("x-amz-user-agent"),
reqwest::header::HeaderValue::from_static(KIRO_IDC_AMZ_USER_AGENT),
),
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_static("node"),
),
(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("*/*"),
),
]);
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_batch_refresh:idc",
reqwest::Method::POST,
&url,
&headers,
Some("application/json"),
Some(json!({
"clientId": auth_config
.client_id
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"clientSecret": auth_config
.client_secret
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"grantType": "refresh_token",
})),
None,
proxy.clone(),
)
.await
.map_err(|err| format!("IDC refresh 请求失败: {err}"))?;
if !response.status.is_success() {
return Err(format!(
"IDC refresh 失败: {}",
admin_provider_oauth_kiro_refresh_error_detail(
response.status,
&response.body_text
)
));
}
let payload = admin_provider_oauth_kiro_refresh_response_json(
&response.body_text,
response.json_body,
)?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| "IDC refresh 返回了空 accessToken".to_string())?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id = generate_machine_id(auth_config, None);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
return Ok(refreshed);
}
let machine_id = generate_machine_id(auth_config, None)
.ok_or_else(|| "缺少 machine_id 种子,无法刷新 social token".to_string())?;
let fallback_host = format!(
"prod.{}.auth.desktop.kiro.dev",
auth_config.effective_auth_region()
);
let url = admin_provider_oauth_kiro_build_refresh_url(
auth_config,
social_refresh_base_url,
"refreshToken",
|region| format!("https://prod.{region}.auth.desktop.kiro.dev/refreshToken"),
);
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
let user_agent =
admin_provider_oauth_kiro_ide_tag(auth_config.effective_kiro_version(), &machine_id);
let headers = reqwest::header::HeaderMap::from_iter([
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_str(&user_agent)
.map_err(|_| "Kiro User-Agent 无效".to_string())?,
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(&host)
.map_err(|_| "Kiro host 无效".to_string())?,
),
(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("application/json, text/plain, */*"),
),
(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
),
(
reqwest::header::CONNECTION,
reqwest::header::HeaderValue::from_static("close"),
),
(
reqwest::header::ACCEPT_ENCODING,
reqwest::header::HeaderValue::from_static("gzip, compress, deflate, br"),
),
]);
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_batch_refresh:social",
reqwest::Method::POST,
&url,
&headers,
Some("application/json"),
Some(json!({
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
})),
None,
proxy,
)
.await
.map_err(|err| format!("social refresh 请求失败: {err}"))?;
if !response.status.is_success() {
return Err(format!(
"social refresh 失败: {}",
admin_provider_oauth_kiro_refresh_error_detail(response.status, &response.body_text)
));
}
let payload =
admin_provider_oauth_kiro_refresh_response_json(&response.body_text, response.json_body)?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| "social refresh 返回了空 accessToken".to_string())?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id = Some(machine_id);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
if let Some(profile_arn) = payload
.get("profileArn")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.profile_arn = Some(profile_arn.to_string());
}
Ok(refreshed)
}
pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
state: &AdminAppState<'_>,
@@ -473,7 +174,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
.map(|value| value.to_string());
let Some(access_token) = access_token else {
failed += 1;
results.push(json!({

View File

@@ -1,11 +1,13 @@
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,
};
use super::session::{
attach_admin_provider_oauth_device_poll_terminal_response, AdminProviderOAuthDevicePollPayload,
};
use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider_oauth_key;
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider,
@@ -16,7 +18,7 @@ use crate::handlers::admin::provider::oauth::state::{
current_unix_secs, decode_jwt_claims, json_non_empty_string, json_u64_value,
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_device_poll_provider_id;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig, AdminRequestContext};
use crate::GatewayError;
use axum::{
body::{Body, Bytes},
@@ -231,26 +233,85 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
}))
.into_response());
};
let expires_at = json_u64_value(token_result.get("expiresIn"))
let initial_expires_at = json_u64_value(token_result.get("expiresIn"))
.map(|expires_in| current_unix_secs().saturating_add(expires_in))
.unwrap_or_else(|| current_unix_secs().saturating_add(3600));
let email = decode_jwt_claims(&access_token)
let social_refresh_base_url =
admin_provider_oauth_kiro_refresh_base_url_override(state, "kiro_social_refresh");
let idc_refresh_base_url =
admin_provider_oauth_kiro_refresh_base_url_override(state, "kiro_idc_refresh");
let mut refreshed_auth_config = match refresh_admin_provider_oauth_kiro_auth_config(
state,
&AdminKiroAuthConfig {
auth_method: Some("idc".to_string()),
refresh_token: Some(refresh_token.clone()),
expires_at: Some(initial_expires_at),
profile_arn: None,
region: Some(session.region.clone()),
auth_region: Some(session.region.clone()),
api_region: None,
client_id: Some(session.client_id.clone()),
client_secret: Some(session.client_secret.clone()),
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: Some(access_token.clone()),
},
request_proxy.clone(),
social_refresh_base_url.as_deref(),
idc_refresh_base_url.as_deref(),
)
.await
{
Ok(config) => config,
Err(detail) => {
return Ok(Json(json!({
"status": "error",
"error": format!("token 验证失败: {detail}"),
"replaced": false,
}))
.into_response());
}
};
if refreshed_auth_config.auth_method.is_none() {
refreshed_auth_config.auth_method = Some("idc".to_string());
}
let Some(access_token) = refreshed_auth_config
.access_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
else {
return Ok(Json(json!({
"status": "error",
"error": "token 验证失败: accessToken 为空",
"replaced": false,
}))
.into_response());
};
let expires_at = refreshed_auth_config
.expires_at
.unwrap_or_else(|| current_unix_secs().saturating_add(3600));
let mut email = decode_jwt_claims(&access_token)
.and_then(|claims| claims.get("email").cloned())
.and_then(|value| value.as_str().map(ToOwned::to_owned));
if email.is_none() {
email = fetch_admin_provider_oauth_kiro_email(
state,
&refreshed_auth_config,
request_proxy.clone(),
)
.await;
}
let mut auth_config = serde_json::Map::new();
let mut auth_config = refreshed_auth_config
.to_json_value()
.as_object()
.cloned()
.unwrap_or_default();
auth_config.insert("provider_type".to_string(), json!("kiro"));
auth_config.insert("auth_method".to_string(), json!("idc"));
auth_config.insert("refresh_token".to_string(), json!(refresh_token.clone()));
auth_config.insert("client_id".to_string(), json!(session.client_id.clone()));
auth_config.insert(
"client_secret".to_string(),
json!(session.client_secret.clone()),
);
auth_config.insert("region".to_string(), json!(session.region.clone()));
auth_config.insert("auth_region".to_string(), json!(session.region.clone()));
auth_config.insert("access_token".to_string(), json!(access_token.clone()));
auth_config.insert("expires_at".to_string(), json!(expires_at));
if let Some(email) = email.as_ref() {
auth_config.insert("email".to_string(), json!(email));
}
@@ -271,6 +332,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
};
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(session.proxy_node_id.as_deref());
let mut replaced = false;
let persisted_key = if let Some(existing_key) = duplicate {
replaced = true;
@@ -279,7 +341,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&existing_key,
&access_token,
&auth_config,
None,
key_proxy.clone(),
Some(expires_at),
)
.await?
@@ -293,7 +355,10 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
}
}
} else {
let key_name = build_kiro_device_key_name(email.as_deref(), Some(&refresh_token));
let key_name = build_kiro_device_key_name(
email.as_deref(),
refreshed_auth_config.refresh_token.as_deref(),
);
match state
.create_provider_oauth_catalog_key(
&provider_id,
@@ -301,7 +366,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy,
Some(expires_at),
)
.await?

View File

@@ -0,0 +1,420 @@
use crate::handlers::admin::provider::oauth::state::{current_unix_secs, json_non_empty_string};
use crate::handlers::admin::provider::shared::payloads::{
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
};
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
use crate::provider_transport::kiro::{build_kiro_request_auth_from_config, KiroRequestAuth};
use aether_contracts::ProxySnapshot;
use serde_json::{json, Value};
use std::time::{SystemTime, UNIX_EPOCH};
use url::form_urlencoded;
const KIRO_IDC_AMZ_USER_AGENT: &str =
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
pub(super) fn admin_provider_oauth_kiro_refresh_base_url_override(
state: &AdminAppState<'_>,
override_key: &str,
) -> Option<String> {
let override_url = state.app().provider_oauth_token_url(override_key, "");
let normalized = override_url.trim();
(!normalized.is_empty()).then(|| normalized.to_string())
}
fn admin_provider_oauth_kiro_build_refresh_url(
auth_config: &AdminKiroAuthConfig,
override_base_url: Option<&str>,
path: &str,
default_host: impl FnOnce(&str) -> String,
) -> String {
if let Some(base_url) = override_base_url
.map(str::trim)
.filter(|value| !value.is_empty())
{
return format!("{}/{}", base_url.trim_end_matches('/'), path);
}
let region = auth_config.effective_auth_region();
default_host(region)
}
fn admin_provider_oauth_kiro_effective_host(url: &str, fallback_host: String) -> String {
reqwest::Url::parse(url)
.ok()
.and_then(|value| value.host_str().map(ToOwned::to_owned))
.unwrap_or(fallback_host)
}
fn admin_provider_oauth_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
if machine_id.trim().is_empty() {
format!("KiroIDE-{kiro_version}")
} else {
format!("KiroIDE-{kiro_version}-{machine_id}")
}
}
fn admin_provider_oauth_kiro_refresh_expires_at(payload: &Value) -> u64 {
let expires_in = payload
.get("expiresIn")
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_str()?.parse::<u64>().ok())
})
.unwrap_or(3600);
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default()
.saturating_add(expires_in)
}
fn admin_provider_oauth_kiro_refresh_response_json(
body_text: &str,
json_body: Option<Value>,
) -> Result<Value, String> {
json_body
.or_else(|| serde_json::from_str::<Value>(body_text).ok())
.ok_or_else(|| "refresh 接口返回了非 JSON 响应".to_string())
}
fn admin_provider_oauth_kiro_refresh_error_detail(
status: http::StatusCode,
body_text: &str,
) -> String {
let detail = body_text.trim();
if detail.is_empty() {
format!("HTTP {}", status.as_u16())
} else {
detail.to_string()
}
}
pub(super) async fn refresh_admin_provider_oauth_kiro_auth_config(
state: &AdminAppState<'_>,
auth_config: &AdminKiroAuthConfig,
proxy: Option<ProxySnapshot>,
social_refresh_base_url: Option<&str>,
idc_refresh_base_url: Option<&str>,
) -> Result<AdminKiroAuthConfig, String> {
if auth_config.is_idc_auth() {
let fallback_host = format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region());
let url = admin_provider_oauth_kiro_build_refresh_url(
auth_config,
idc_refresh_base_url,
"token",
|region| format!("https://oidc.{region}.amazonaws.com/token"),
);
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
let headers = reqwest::header::HeaderMap::from_iter([
(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(&host)
.map_err(|_| "IDC host 无效".to_string())?,
),
(
reqwest::header::HeaderName::from_static("x-amz-user-agent"),
reqwest::header::HeaderValue::from_static(KIRO_IDC_AMZ_USER_AGENT),
),
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_static("node"),
),
(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("*/*"),
),
]);
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_batch_refresh:idc",
reqwest::Method::POST,
&url,
&headers,
Some("application/json"),
Some(json!({
"clientId": auth_config
.client_id
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"clientSecret": auth_config
.client_secret
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"grantType": "refresh_token",
})),
None,
proxy.clone(),
)
.await
.map_err(|err| format!("IDC refresh 请求失败: {err}"))?;
if !response.status.is_success() {
return Err(format!(
"IDC refresh 失败: {}",
admin_provider_oauth_kiro_refresh_error_detail(
response.status,
&response.body_text
)
));
}
let payload = admin_provider_oauth_kiro_refresh_response_json(
&response.body_text,
response.json_body,
)?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| "IDC refresh 返回了空 accessToken".to_string())?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id =
crate::provider_transport::kiro::generate_machine_id(auth_config, None);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
return Ok(refreshed);
}
let machine_id = crate::provider_transport::kiro::generate_machine_id(auth_config, None)
.ok_or_else(|| "缺少 machine_id 种子,无法刷新 social token".to_string())?;
let fallback_host = format!(
"prod.{}.auth.desktop.kiro.dev",
auth_config.effective_auth_region()
);
let url = admin_provider_oauth_kiro_build_refresh_url(
auth_config,
social_refresh_base_url,
"refreshToken",
|region| format!("https://prod.{region}.auth.desktop.kiro.dev/refreshToken"),
);
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
let user_agent =
admin_provider_oauth_kiro_ide_tag(auth_config.effective_kiro_version(), &machine_id);
let headers = reqwest::header::HeaderMap::from_iter([
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_str(&user_agent)
.map_err(|_| "Kiro User-Agent 无效".to_string())?,
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(&host)
.map_err(|_| "Kiro host 无效".to_string())?,
),
(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("application/json, text/plain, */*"),
),
(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
),
(
reqwest::header::CONNECTION,
reqwest::header::HeaderValue::from_static("close"),
),
(
reqwest::header::ACCEPT_ENCODING,
reqwest::header::HeaderValue::from_static("gzip, compress, deflate, br"),
),
]);
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_batch_refresh:social",
reqwest::Method::POST,
&url,
&headers,
Some("application/json"),
Some(json!({
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
})),
None,
proxy,
)
.await
.map_err(|err| format!("social refresh 请求失败: {err}"))?;
if !response.status.is_success() {
return Err(format!(
"social refresh 失败: {}",
admin_provider_oauth_kiro_refresh_error_detail(response.status, &response.body_text)
));
}
let payload =
admin_provider_oauth_kiro_refresh_response_json(&response.body_text, response.json_body)?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| "social refresh 返回了空 accessToken".to_string())?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id = Some(machine_id);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
if let Some(profile_arn) = payload
.get("profileArn")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.profile_arn = Some(profile_arn.to_string());
}
Ok(refreshed)
}
fn build_kiro_usage_url(auth: &KiroRequestAuth) -> String {
let host = format!(
"q.{}.amazonaws.com",
auth.auth_config.effective_api_region()
);
let mut serializer = form_urlencoded::Serializer::new(String::new());
serializer.append_pair("origin", "AI_EDITOR");
serializer.append_pair("resourceType", "AGENTIC_REQUEST");
serializer.append_pair("isEmailRequired", "true");
if let Some(profile_arn) = auth.auth_config.profile_arn_for_payload() {
serializer.append_pair("profileArn", profile_arn);
}
format!(
"https://{host}{KIRO_USAGE_LIMITS_PATH}?{}",
serializer.finish()
)
}
fn build_kiro_usage_headers(
auth: &KiroRequestAuth,
host: &str,
) -> Result<reqwest::header::HeaderMap, String> {
let kiro_version = auth.auth_config.effective_kiro_version();
let machine_id = auth.machine_id.trim();
let ide_tag = admin_provider_oauth_kiro_ide_tag(kiro_version, machine_id);
Ok(reqwest::header::HeaderMap::from_iter([
(
reqwest::header::HeaderName::from_static("x-amz-user-agent"),
reqwest::header::HeaderValue::from_str(&format!(
"aws-sdk-js/{KIRO_USAGE_SDK_VERSION} {ide_tag}"
))
.map_err(|_| "Kiro usage x-amz-user-agent 无效".to_string())?,
),
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_str(&format!(
"aws-sdk-js/{KIRO_USAGE_SDK_VERSION} ua/2.1 os/other#unknown lang/js md/nodejs#22.21.1 api/codewhispererruntime#1.0.0 m/N,E {ide_tag}"
))
.map_err(|_| "Kiro usage User-Agent 无效".to_string())?,
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(host)
.map_err(|_| "Kiro usage host 无效".to_string())?,
),
(
reqwest::header::HeaderName::from_static("amz-sdk-invocation-id"),
reqwest::header::HeaderValue::from_str(&uuid::Uuid::new_v4().to_string())
.map_err(|_| "Kiro usage invocation id 无效".to_string())?,
),
(
reqwest::header::HeaderName::from_static("amz-sdk-request"),
reqwest::header::HeaderValue::from_static("attempt=1; max=1"),
),
(
reqwest::header::AUTHORIZATION,
reqwest::header::HeaderValue::from_str(auth.value.as_str())
.map_err(|_| "Kiro usage authorization 无效".to_string())?,
),
(
reqwest::header::CONNECTION,
reqwest::header::HeaderValue::from_static("close"),
),
]))
}
pub(super) async fn fetch_admin_provider_oauth_kiro_email(
state: &AdminAppState<'_>,
auth_config: &AdminKiroAuthConfig,
proxy: Option<ProxySnapshot>,
) -> Option<String> {
let request_auth = build_kiro_request_auth_from_config(auth_config.clone(), None)?;
let default_url = build_kiro_usage_url(&request_auth);
let url = state
.app()
.provider_oauth_token_url("kiro_device_email", &default_url);
let host = reqwest::Url::parse(&url)
.ok()
.and_then(|value| value.host_str().map(ToOwned::to_owned))
.unwrap_or_else(|| {
format!(
"q.{}.amazonaws.com",
request_auth.auth_config.effective_api_region()
)
});
let headers = build_kiro_usage_headers(&request_auth, &host).ok()?;
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_device_email",
reqwest::Method::GET,
&url,
&headers,
None,
None,
None,
proxy,
)
.await
.ok()?;
if !response.status.is_success() {
return None;
}
let payload = response
.json_body
.or_else(|| serde_json::from_str::<Value>(&response.body_text).ok())?;
let metadata =
aether_admin::provider::quota::parse_kiro_usage_response(&payload, current_unix_secs())?;
json_non_empty_string(metadata.get("email"))
}

View File

@@ -23,6 +23,7 @@ mod complete;
mod device;
mod helpers;
mod import;
mod kiro;
mod refresh;
mod start;
mod tasks;