mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
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:
@@ -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::duplicates::find_duplicate_provider_oauth_key;
|
||||||
use crate::handlers::admin::provider::oauth::provisioning::{
|
use crate::handlers::admin::provider::oauth::provisioning::{
|
||||||
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
|
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::oauth::state::decode_jwt_claims;
|
||||||
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
|
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
|
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
|
||||||
use crate::provider_transport::kiro::generate_machine_id;
|
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
use aether_admin::provider::oauth::{
|
use aether_admin::provider::oauth::{
|
||||||
build_kiro_batch_import_key_name, coerce_admin_provider_oauth_import_str,
|
build_kiro_batch_import_key_name, coerce_admin_provider_oauth_import_str,
|
||||||
parse_admin_provider_oauth_kiro_batch_import_entries,
|
parse_admin_provider_oauth_kiro_batch_import_entries,
|
||||||
};
|
};
|
||||||
use aether_contracts::ProxySnapshot;
|
|
||||||
use serde_json::{json, Map, Value};
|
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(
|
pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
|
||||||
state: &AdminAppState<'_>,
|
state: &AdminAppState<'_>,
|
||||||
@@ -473,7 +174,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
|
|||||||
.as_deref()
|
.as_deref()
|
||||||
.map(str::trim)
|
.map(str::trim)
|
||||||
.filter(|value| !value.is_empty())
|
.filter(|value| !value.is_empty())
|
||||||
.map(ToOwned::to_owned);
|
.map(|value| value.to_string());
|
||||||
let Some(access_token) = access_token else {
|
let Some(access_token) = access_token else {
|
||||||
failed += 1;
|
failed += 1;
|
||||||
results.push(json!({
|
results.push(json!({
|
||||||
|
|||||||
@@ -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::{
|
use super::session::{
|
||||||
attach_admin_provider_oauth_device_poll_terminal_response, AdminProviderOAuthDevicePollPayload,
|
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::errors::build_internal_control_error_response;
|
||||||
use crate::handlers::admin::provider::oauth::provisioning::{
|
use crate::handlers::admin::provider::oauth::provisioning::{
|
||||||
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
|
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
|
||||||
update_existing_provider_oauth_catalog_key,
|
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::oauth::runtime::{
|
use crate::handlers::admin::provider::oauth::runtime::{
|
||||||
provider_oauth_runtime_endpoint_for_provider,
|
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,
|
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::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 crate::GatewayError;
|
||||||
use axum::{
|
use axum::{
|
||||||
body::{Body, Bytes},
|
body::{Body, Bytes},
|
||||||
@@ -231,26 +233,85 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
|||||||
}))
|
}))
|
||||||
.into_response());
|
.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))
|
.map(|expires_in| current_unix_secs().saturating_add(expires_in))
|
||||||
.unwrap_or_else(|| current_unix_secs().saturating_add(3600));
|
.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(|claims| claims.get("email").cloned())
|
||||||
.and_then(|value| value.as_str().map(ToOwned::to_owned));
|
.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("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() {
|
if let Some(email) = email.as_ref() {
|
||||||
auth_config.insert("email".to_string(), json!(email));
|
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 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 mut replaced = false;
|
||||||
let persisted_key = if let Some(existing_key) = duplicate {
|
let persisted_key = if let Some(existing_key) = duplicate {
|
||||||
replaced = true;
|
replaced = true;
|
||||||
@@ -279,7 +341,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
|||||||
&existing_key,
|
&existing_key,
|
||||||
&access_token,
|
&access_token,
|
||||||
&auth_config,
|
&auth_config,
|
||||||
None,
|
key_proxy.clone(),
|
||||||
Some(expires_at),
|
Some(expires_at),
|
||||||
)
|
)
|
||||||
.await?
|
.await?
|
||||||
@@ -293,7 +355,10 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} 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
|
match state
|
||||||
.create_provider_oauth_catalog_key(
|
.create_provider_oauth_catalog_key(
|
||||||
&provider_id,
|
&provider_id,
|
||||||
@@ -301,7 +366,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
|||||||
&access_token,
|
&access_token,
|
||||||
&auth_config,
|
&auth_config,
|
||||||
&api_formats,
|
&api_formats,
|
||||||
None,
|
key_proxy,
|
||||||
Some(expires_at),
|
Some(expires_at),
|
||||||
)
|
)
|
||||||
.await?
|
.await?
|
||||||
|
|||||||
@@ -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"))
|
||||||
|
}
|
||||||
@@ -23,6 +23,7 @@ mod complete;
|
|||||||
mod device;
|
mod device;
|
||||||
mod helpers;
|
mod helpers;
|
||||||
mod import;
|
mod import;
|
||||||
|
mod kiro;
|
||||||
mod refresh;
|
mod refresh;
|
||||||
mod start;
|
mod start;
|
||||||
mod tasks;
|
mod tasks;
|
||||||
|
|||||||
@@ -100,6 +100,20 @@ fn sample_kiro_device_access_token(email: &str) -> String {
|
|||||||
format!("{header}.{payload}.sig")
|
format!("{header}.{payload}.sig")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn sample_kiro_device_access_token_without_email() -> String {
|
||||||
|
use base64::Engine as _;
|
||||||
|
|
||||||
|
let header =
|
||||||
|
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
|
||||||
|
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
|
||||||
|
json!({
|
||||||
|
"sub": "kiro-user-123",
|
||||||
|
})
|
||||||
|
.to_string(),
|
||||||
|
);
|
||||||
|
format!("{header}.{payload}.sig")
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_handles_admin_provider_oauth_supported_types_locally_with_trusted_admin_principal()
|
async fn gateway_handles_admin_provider_oauth_supported_types_locally_with_trusted_admin_principal()
|
||||||
{
|
{
|
||||||
@@ -348,7 +362,8 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
|
|||||||
"error_msg": null,
|
"error_msg": null,
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"));
|
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"))
|
||||||
|
.with_provider_oauth_token_url_for_tests("kiro_idc_refresh", token_url.to_string());
|
||||||
let gateway = build_router_with_state(state.clone());
|
let gateway = build_router_with_state(state.clone());
|
||||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||||
|
|
||||||
@@ -374,7 +389,7 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
|
|||||||
assert_eq!(payload["email"], "kiro@example.com");
|
assert_eq!(payload["email"], "kiro@example.com");
|
||||||
assert_eq!(payload["replaced"], false);
|
assert_eq!(payload["replaced"], false);
|
||||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
|
assert_eq!(*token_hits.lock().expect("mutex should lock"), 2);
|
||||||
|
|
||||||
let stored = state
|
let stored = state
|
||||||
.load_provider_oauth_device_session_for_tests("device_auth_session:session-123")
|
.load_provider_oauth_device_session_for_tests("device_auth_session:session-123")
|
||||||
@@ -398,7 +413,10 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
|
|||||||
.next()
|
.next()
|
||||||
.expect("persisted key should exist");
|
.expect("persisted key should exist");
|
||||||
assert_eq!(persisted.auth_type, "oauth");
|
assert_eq!(persisted.auth_type, "oauth");
|
||||||
assert_eq!(persisted.proxy, None);
|
assert_eq!(
|
||||||
|
persisted.proxy,
|
||||||
|
Some(json!({"node_id": "proxy-node-kiro", "enabled": true}))
|
||||||
|
);
|
||||||
let decrypted_api_key =
|
let decrypted_api_key =
|
||||||
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
|
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
|
||||||
.expect("api key should decrypt");
|
.expect("api key should decrypt");
|
||||||
@@ -426,6 +444,238 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
|
|||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_revalidates_kiro_device_poll_via_idc_refresh_and_backfills_email() {
|
||||||
|
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| {
|
||||||
|
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||||
|
async move {
|
||||||
|
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||||
|
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
|
let token_requests = Arc::new(Mutex::new(Vec::<String>::new()));
|
||||||
|
let token_requests_clone = Arc::clone(&token_requests);
|
||||||
|
let initial_access_token = sample_kiro_device_access_token_without_email();
|
||||||
|
let refreshed_access_token = sample_kiro_device_access_token_without_email();
|
||||||
|
let expected_refreshed_access_token = refreshed_access_token.clone();
|
||||||
|
let token_server = Router::new().route(
|
||||||
|
"/token",
|
||||||
|
post(move |request: Request| {
|
||||||
|
let token_requests_inner = Arc::clone(&token_requests_clone);
|
||||||
|
let initial_access_token_inner = initial_access_token.clone();
|
||||||
|
let refreshed_access_token_inner = refreshed_access_token.clone();
|
||||||
|
async move {
|
||||||
|
let raw_body = String::from_utf8(
|
||||||
|
to_bytes(request.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.expect("body should read")
|
||||||
|
.to_vec(),
|
||||||
|
)
|
||||||
|
.expect("body should be utf8");
|
||||||
|
token_requests_inner
|
||||||
|
.lock()
|
||||||
|
.expect("mutex should lock")
|
||||||
|
.push(raw_body.clone());
|
||||||
|
if raw_body.contains("urn:ietf:params:oauth:grant-type:device_code") {
|
||||||
|
return Json(json!({
|
||||||
|
"accessToken": initial_access_token_inner,
|
||||||
|
"refreshToken": "kiro-device-refresh-token-initial",
|
||||||
|
"expiresIn": 1800,
|
||||||
|
}))
|
||||||
|
.into_response();
|
||||||
|
}
|
||||||
|
if raw_body.contains("\"grantType\":\"refresh_token\"") {
|
||||||
|
return Json(json!({
|
||||||
|
"accessToken": refreshed_access_token_inner,
|
||||||
|
"refreshToken": "kiro-device-refresh-token-rotated",
|
||||||
|
"expiresIn": 2400,
|
||||||
|
}))
|
||||||
|
.into_response();
|
||||||
|
}
|
||||||
|
(
|
||||||
|
StatusCode::BAD_REQUEST,
|
||||||
|
Json(json!({
|
||||||
|
"error": "unexpected_request",
|
||||||
|
"body": raw_body,
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let usage_hits = Arc::new(Mutex::new(0usize));
|
||||||
|
let usage_hits_clone = Arc::clone(&usage_hits);
|
||||||
|
let usage_server = Router::new().route(
|
||||||
|
"/getUsageLimits",
|
||||||
|
get(move |_request: Request| {
|
||||||
|
let usage_hits_inner = Arc::clone(&usage_hits_clone);
|
||||||
|
async move {
|
||||||
|
*usage_hits_inner.lock().expect("mutex should lock") += 1;
|
||||||
|
Json(json!({
|
||||||
|
"subscriptionInfo": {
|
||||||
|
"subscriptionTitle": "KIRO PRO+"
|
||||||
|
},
|
||||||
|
"usageBreakdownList": [{
|
||||||
|
"currentUsageWithPrecision": 5.0,
|
||||||
|
"usageLimitWithPrecision": 20.0,
|
||||||
|
"nextDateReset": 1_900_000_000u64
|
||||||
|
}],
|
||||||
|
"desktopUserInfo": {
|
||||||
|
"email": "kiro-usage@example.com"
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut provider = sample_provider("provider-kiro", "kiro", 10);
|
||||||
|
provider.provider_type = "kiro".to_string();
|
||||||
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![provider],
|
||||||
|
vec![],
|
||||||
|
vec![],
|
||||||
|
));
|
||||||
|
|
||||||
|
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||||
|
let (token_url, token_handle) = start_server(token_server).await;
|
||||||
|
let (usage_url, usage_handle) = start_server(usage_server).await;
|
||||||
|
let state = AppState::new()
|
||||||
|
.expect("gateway should build")
|
||||||
|
.with_data_state_for_tests(
|
||||||
|
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||||
|
provider_catalog_repository.clone(),
|
||||||
|
)
|
||||||
|
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||||
|
)
|
||||||
|
.with_provider_oauth_device_session_entry_for_tests(
|
||||||
|
"session-refresh-email",
|
||||||
|
json!({
|
||||||
|
"provider_id": "provider-kiro",
|
||||||
|
"region": "us-east-1",
|
||||||
|
"client_id": "kiro-device-client",
|
||||||
|
"client_secret": "kiro-device-secret",
|
||||||
|
"device_code": "device-code-123",
|
||||||
|
"interval": 5,
|
||||||
|
"expires_at_unix_secs": 4_102_444_800u64,
|
||||||
|
"status": "pending",
|
||||||
|
"proxy_node_id": "proxy-node-kiro",
|
||||||
|
"created_at_unix_ms": 1_711_000_000u64,
|
||||||
|
"key_id": null,
|
||||||
|
"email": null,
|
||||||
|
"replaced": false,
|
||||||
|
"error_msg": null,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"))
|
||||||
|
.with_provider_oauth_token_url_for_tests("kiro_idc_refresh", token_url.to_string())
|
||||||
|
.with_provider_oauth_token_url_for_tests(
|
||||||
|
"kiro_device_email",
|
||||||
|
format!("{usage_url}/getUsageLimits"),
|
||||||
|
);
|
||||||
|
let gateway = build_router_with_state(state.clone());
|
||||||
|
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||||
|
|
||||||
|
let response = reqwest::Client::new()
|
||||||
|
.post(format!(
|
||||||
|
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-poll"
|
||||||
|
))
|
||||||
|
.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!({
|
||||||
|
"session_id": "session-refresh-email"
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("request should succeed");
|
||||||
|
|
||||||
|
let status = response.status();
|
||||||
|
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||||
|
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||||
|
assert_eq!(payload["status"], "authorized");
|
||||||
|
assert_eq!(payload["email"], "kiro-usage@example.com");
|
||||||
|
assert_eq!(payload["replaced"], false);
|
||||||
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
|
assert_eq!(*usage_hits.lock().expect("mutex should lock"), 1);
|
||||||
|
{
|
||||||
|
let requests = token_requests.lock().expect("mutex should lock");
|
||||||
|
assert_eq!(requests.len(), 2);
|
||||||
|
assert!(
|
||||||
|
requests
|
||||||
|
.iter()
|
||||||
|
.any(|body| body.contains("urn:ietf:params:oauth:grant-type:device_code")),
|
||||||
|
"requests={requests:?}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
requests
|
||||||
|
.iter()
|
||||||
|
.any(|body| body.contains("\"grantType\":\"refresh_token\"")),
|
||||||
|
"requests={requests:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let stored = state
|
||||||
|
.load_provider_oauth_device_session_for_tests("device_auth_session:session-refresh-email")
|
||||||
|
.expect("device session should persist");
|
||||||
|
let stored: serde_json::Value =
|
||||||
|
serde_json::from_str(&stored).expect("device session json should parse");
|
||||||
|
assert_eq!(stored["status"], "authorized");
|
||||||
|
assert_eq!(stored["email"], "kiro-usage@example.com");
|
||||||
|
assert_eq!(stored["replaced"], false);
|
||||||
|
let key_id = stored["key_id"]
|
||||||
|
.as_str()
|
||||||
|
.expect("key_id should be stored")
|
||||||
|
.to_string();
|
||||||
|
assert_eq!(payload["key_id"], key_id);
|
||||||
|
|
||||||
|
let persisted = provider_catalog_repository
|
||||||
|
.list_keys_by_ids(std::slice::from_ref(&key_id))
|
||||||
|
.await
|
||||||
|
.expect("keys should load")
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
.expect("persisted key should exist");
|
||||||
|
assert_eq!(persisted.auth_type, "oauth");
|
||||||
|
assert_eq!(
|
||||||
|
persisted.proxy,
|
||||||
|
Some(json!({"node_id": "proxy-node-kiro", "enabled": true}))
|
||||||
|
);
|
||||||
|
let decrypted_api_key =
|
||||||
|
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
|
||||||
|
.expect("api key should decrypt");
|
||||||
|
assert_eq!(decrypted_api_key, expected_refreshed_access_token);
|
||||||
|
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
|
||||||
|
DEVELOPMENT_ENCRYPTION_KEY,
|
||||||
|
persisted
|
||||||
|
.encrypted_auth_config
|
||||||
|
.as_deref()
|
||||||
|
.expect("auth config should exist"),
|
||||||
|
)
|
||||||
|
.expect("auth config should decrypt");
|
||||||
|
let auth_config: serde_json::Value =
|
||||||
|
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
|
||||||
|
assert_eq!(auth_config["provider_type"], "kiro");
|
||||||
|
assert_eq!(auth_config["auth_method"], "idc");
|
||||||
|
assert_eq!(
|
||||||
|
auth_config["refresh_token"],
|
||||||
|
"kiro-device-refresh-token-rotated"
|
||||||
|
);
|
||||||
|
assert_eq!(auth_config["email"], "kiro-usage@example.com");
|
||||||
|
assert_eq!(auth_config["client_id"], "kiro-device-client");
|
||||||
|
assert_eq!(auth_config["client_secret"], "kiro-device-secret");
|
||||||
|
assert_eq!(auth_config["region"], "us-east-1");
|
||||||
|
|
||||||
|
gateway_handle.abort();
|
||||||
|
usage_handle.abort();
|
||||||
|
token_handle.abort();
|
||||||
|
upstream_handle.abort();
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transition_reaches_terminal_state(
|
async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transition_reaches_terminal_state(
|
||||||
) {
|
) {
|
||||||
@@ -480,7 +730,8 @@ async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transit
|
|||||||
"error_msg": null,
|
"error_msg": null,
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"));
|
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"))
|
||||||
|
.with_provider_oauth_token_url_for_tests("kiro_idc_refresh", token_url.to_string());
|
||||||
|
|
||||||
let terminal_response = local_admin_provider_oauth_response(
|
let terminal_response = local_admin_provider_oauth_response(
|
||||||
&terminal_state,
|
&terminal_state,
|
||||||
|
|||||||
Reference in New Issue
Block a user