mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
Improve provider OAuth device flow
This commit is contained in:
@@ -172,6 +172,9 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
auth_config.insert("provider_type".to_string(), json!("kiro"));
|
||||
if let Some(provider) = coerce_admin_provider_oauth_import_str(entry.get("provider")) {
|
||||
auth_config.insert("provider".to_string(), json!(provider));
|
||||
}
|
||||
let email = decode_jwt_claims(
|
||||
refreshed_auth_config
|
||||
.access_token
|
||||
|
||||
@@ -4,7 +4,7 @@ use crate::handlers::admin::provider::oauth::runtime::provider_oauth_runtime_end
|
||||
use crate::handlers::admin::provider::oauth::state::{
|
||||
build_admin_provider_oauth_backend_unavailable_response, current_unix_secs,
|
||||
default_kiro_device_start_url, generate_provider_oauth_nonce, json_non_empty_string,
|
||||
json_u64_value, normalize_kiro_device_region,
|
||||
json_u64_value, normalize_kiro_device_region, provider_oauth_pkce_s256,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_device_authorize_provider_id;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
@@ -19,6 +19,105 @@ use axum::{
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use url::{form_urlencoded, Url};
|
||||
|
||||
const KIRO_SOCIAL_PORTAL_SIGNIN_URL: &str = "https://app.kiro.dev/signin";
|
||||
const KIRO_SOCIAL_AUTH_EXPIRES_IN_SECS: u64 = 600;
|
||||
const KIRO_SOCIAL_AUTH_POLL_INTERVAL_SECS: u64 = 5;
|
||||
const KIRO_SOCIAL_MANUAL_CALLBACK_PORT: u16 = 49153;
|
||||
const KIRO_SOCIAL_ALLOWED_CALLBACK_PORTS: &[u16] = &[
|
||||
3128, 4649, 6588, 8008, 9091, 49153, 50153, 51153, 52153, 53153,
|
||||
];
|
||||
|
||||
fn normalize_kiro_device_auth_type(raw: Option<&str>) -> String {
|
||||
match raw
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("identity_center")
|
||||
.to_ascii_lowercase()
|
||||
.as_str()
|
||||
{
|
||||
"google" => "google".to_string(),
|
||||
"github" | "git_hub" | "git-hub" => "github".to_string(),
|
||||
"builderid" | "builder_id" | "builder-id" | "builder" => "builder_id".to_string(),
|
||||
"identitycenter" | "identity_center" | "identity-center" | "idc" | "enterprise" => {
|
||||
"identity_center".to_string()
|
||||
}
|
||||
_ => "identity_center".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn kiro_social_provider_id(auth_type: &str) -> Option<&'static str> {
|
||||
match auth_type {
|
||||
"google" => Some("Google"),
|
||||
"github" => Some("Github"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn default_kiro_social_redirect_uri() -> String {
|
||||
format!("http://localhost:{KIRO_SOCIAL_MANUAL_CALLBACK_PORT}")
|
||||
}
|
||||
|
||||
fn normalize_kiro_social_redirect_uri(raw: Option<&str>) -> Result<String, &'static str> {
|
||||
let Some(raw) = raw.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return Ok(default_kiro_social_redirect_uri());
|
||||
};
|
||||
|
||||
let url = Url::parse(raw).map_err(|_| "Kiro social redirect_uri 必须是合法 URL")?;
|
||||
if url.scheme() != "http" {
|
||||
return Err("Kiro social redirect_uri 必须使用 http://localhost");
|
||||
}
|
||||
if !url
|
||||
.host_str()
|
||||
.is_some_and(|host| host.eq_ignore_ascii_case("localhost"))
|
||||
{
|
||||
return Err("Kiro social redirect_uri 必须使用 localhost");
|
||||
}
|
||||
let Some(port) = url.port() else {
|
||||
return Err("Kiro social redirect_uri 必须包含端口");
|
||||
};
|
||||
if !KIRO_SOCIAL_ALLOWED_CALLBACK_PORTS.contains(&port) {
|
||||
return Err("Kiro social redirect_uri 端口不是 Kiro 允许的回调端口");
|
||||
}
|
||||
if !matches!(url.path(), "" | "/") || url.query().is_some() || url.fragment().is_some() {
|
||||
return Err("Kiro social redirect_uri 只能是 http://localhost:{port}");
|
||||
}
|
||||
|
||||
Ok(format!("http://localhost:{port}"))
|
||||
}
|
||||
|
||||
fn build_kiro_social_authorization_url(
|
||||
portal_url: &str,
|
||||
auth_type: &str,
|
||||
redirect_uri: &str,
|
||||
code_challenge: &str,
|
||||
state: &str,
|
||||
) -> String {
|
||||
if let Ok(mut url) = Url::parse(portal_url) {
|
||||
url.query_pairs_mut()
|
||||
.append_pair("state", state)
|
||||
.append_pair("code_challenge", code_challenge)
|
||||
.append_pair("code_challenge_method", "S256")
|
||||
.append_pair("redirect_uri", redirect_uri)
|
||||
.append_pair("redirect_from", "KiroIDE")
|
||||
.append_pair("login_option", auth_type);
|
||||
return url.to_string();
|
||||
}
|
||||
|
||||
let mut serializer = form_urlencoded::Serializer::new(String::new());
|
||||
serializer.append_pair("state", state);
|
||||
serializer.append_pair("code_challenge", code_challenge);
|
||||
serializer.append_pair("code_challenge_method", "S256");
|
||||
serializer.append_pair("redirect_uri", redirect_uri);
|
||||
serializer.append_pair("redirect_from", "KiroIDE");
|
||||
serializer.append_pair("login_option", auth_type);
|
||||
format!(
|
||||
"{}?{}",
|
||||
portal_url.trim_end_matches('?'),
|
||||
serializer.finish()
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) async fn handle_admin_provider_oauth_device_authorize(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -86,6 +185,83 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_type = normalize_kiro_device_auth_type(payload.auth_type.as_deref());
|
||||
if let Some(social_provider) = kiro_social_provider_id(&auth_type) {
|
||||
let redirect_uri = match normalize_kiro_social_redirect_uri(payload.redirect_uri.as_deref())
|
||||
{
|
||||
Ok(redirect_uri) => redirect_uri,
|
||||
Err(message) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
message,
|
||||
));
|
||||
}
|
||||
};
|
||||
let code_verifier = generate_provider_oauth_nonce();
|
||||
let code_challenge = provider_oauth_pkce_s256(&code_verifier);
|
||||
let session_id = generate_provider_oauth_nonce();
|
||||
let portal_url =
|
||||
state.provider_oauth_token_url("kiro_social_portal", KIRO_SOCIAL_PORTAL_SIGNIN_URL);
|
||||
let authorization_url = build_kiro_social_authorization_url(
|
||||
&portal_url,
|
||||
&auth_type,
|
||||
&redirect_uri,
|
||||
&code_challenge,
|
||||
&session_id,
|
||||
);
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let session = StoredAdminProviderOAuthDeviceSession {
|
||||
provider_id: provider_id.clone(),
|
||||
region: "us-east-1".to_string(),
|
||||
client_id: String::new(),
|
||||
client_secret: String::new(),
|
||||
device_code: String::new(),
|
||||
auth_type: Some("social".to_string()),
|
||||
social_provider: Some(social_provider.to_string()),
|
||||
code_verifier: Some(code_verifier),
|
||||
redirect_uri: Some(redirect_uri.clone()),
|
||||
machine_id: Some(uuid::Uuid::new_v4().to_string().to_ascii_lowercase()),
|
||||
interval: KIRO_SOCIAL_AUTH_POLL_INTERVAL_SECS,
|
||||
expires_at_unix_secs: now_unix_secs.saturating_add(KIRO_SOCIAL_AUTH_EXPIRES_IN_SECS),
|
||||
status: "pending".to_string(),
|
||||
proxy_node_id: payload
|
||||
.proxy_node_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
created_at_unix_ms: now_unix_secs,
|
||||
key_id: None,
|
||||
email: None,
|
||||
replaced: false,
|
||||
error_msg: None,
|
||||
};
|
||||
if let Err(response) = state
|
||||
.save_provider_oauth_device_session(
|
||||
&session_id,
|
||||
&session,
|
||||
KIRO_SOCIAL_AUTH_EXPIRES_IN_SECS
|
||||
.saturating_add(KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
return Ok(Json(json!({
|
||||
"session_id": session_id,
|
||||
"user_code": "",
|
||||
"verification_uri": portal_url,
|
||||
"verification_uri_complete": authorization_url,
|
||||
"expires_in": KIRO_SOCIAL_AUTH_EXPIRES_IN_SECS,
|
||||
"interval": KIRO_SOCIAL_AUTH_POLL_INTERVAL_SECS,
|
||||
"auth_type": auth_type,
|
||||
"redirect_uri": redirect_uri,
|
||||
"callback_required": true,
|
||||
}))
|
||||
.into_response());
|
||||
}
|
||||
|
||||
let region = normalize_kiro_device_region(Some(payload.region.as_str())).ok_or_else(|| {
|
||||
build_internal_control_error_response(http::StatusCode::BAD_REQUEST, "region 格式无效")
|
||||
});
|
||||
@@ -178,6 +354,11 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
|
||||
client_id,
|
||||
client_secret,
|
||||
device_code,
|
||||
auth_type: Some("idc".to_string()),
|
||||
social_provider: None,
|
||||
code_verifier: None,
|
||||
redirect_uri: None,
|
||||
machine_id: None,
|
||||
interval,
|
||||
expires_at_unix_secs: now_unix_secs.saturating_add(expires_in),
|
||||
status: "pending".to_string(),
|
||||
|
||||
@@ -16,17 +16,209 @@ use crate::handlers::admin::provider::oauth::runtime::{
|
||||
use crate::handlers::admin::provider::oauth::state::{
|
||||
build_admin_provider_oauth_backend_unavailable_response, build_kiro_device_key_name,
|
||||
current_unix_secs, decode_jwt_claims, json_non_empty_string, json_u64_value,
|
||||
parse_provider_oauth_callback_params,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_device_poll_provider_id;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig, AdminRequestContext};
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data::repository::provider_oauth::StoredAdminProviderOAuthDeviceSession;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
|
||||
};
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
http,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Value};
|
||||
use url::{form_urlencoded, Url};
|
||||
|
||||
const KIRO_SOCIAL_TOKEN_URL: &str = "https://prod.us-east-1.auth.desktop.kiro.dev/oauth/token";
|
||||
const KIRO_SOCIAL_AUTH_KIRO_VERSION: &str = "0.6.18";
|
||||
|
||||
fn kiro_device_session_is_social(session: &StoredAdminProviderOAuthDeviceSession) -> bool {
|
||||
session
|
||||
.auth_type
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("social"))
|
||||
|| session
|
||||
.social_provider
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn kiro_social_key_name(
|
||||
email: Option<&str>,
|
||||
social_provider: Option<&str>,
|
||||
refresh_token: Option<&str>,
|
||||
) -> String {
|
||||
let provider = social_provider
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("social");
|
||||
if let Some(email) = email.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
return format!("{email} ({provider})");
|
||||
}
|
||||
let fallback = refresh_token
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| {
|
||||
use sha2::{Digest, Sha256};
|
||||
let digest = Sha256::digest(value.as_bytes());
|
||||
digest[..3]
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect::<String>()
|
||||
})
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
format!("kiro_{fallback} ({provider})")
|
||||
}
|
||||
|
||||
fn kiro_social_poll_error_response(error: impl Into<String>) -> Response<Body> {
|
||||
Json(json!({
|
||||
"status": "error",
|
||||
"error": error.into(),
|
||||
"replaced": false,
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn kiro_social_provider_from_login_option(login_option: Option<&str>) -> Option<&'static str> {
|
||||
match login_option
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?
|
||||
.to_ascii_lowercase()
|
||||
.as_str()
|
||||
{
|
||||
"google" => Some("Google"),
|
||||
"github" | "git_hub" | "git-hub" => Some("Github"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn kiro_social_token_redirect_uri(
|
||||
session_redirect_uri: &str,
|
||||
callback_url: &str,
|
||||
login_option: Option<&str>,
|
||||
) -> String {
|
||||
let base = session_redirect_uri
|
||||
.trim()
|
||||
.trim_end_matches('/')
|
||||
.to_string();
|
||||
let Some(login_option) = login_option
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return base;
|
||||
};
|
||||
|
||||
let Ok(base_url) = Url::parse(&base) else {
|
||||
return base;
|
||||
};
|
||||
let Ok(callback_url) = Url::parse(callback_url.trim()) else {
|
||||
return base;
|
||||
};
|
||||
|
||||
let base_path = base_url.path().trim_end_matches('/');
|
||||
let callback_path = callback_url.path();
|
||||
let suffix = if base_path.is_empty() || base_path == "/" {
|
||||
callback_path.to_string()
|
||||
} else if callback_path == base_path {
|
||||
String::new()
|
||||
} else {
|
||||
callback_path
|
||||
.strip_prefix(&format!("{base_path}/"))
|
||||
.map(|value| {
|
||||
if value.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("/{value}")
|
||||
}
|
||||
})
|
||||
.unwrap_or_default()
|
||||
};
|
||||
|
||||
let mut redirect_uri = if suffix.is_empty() {
|
||||
base
|
||||
} else {
|
||||
format!("{base}{suffix}")
|
||||
};
|
||||
redirect_uri.push('?');
|
||||
redirect_uri.push_str(
|
||||
&form_urlencoded::Serializer::new(String::new())
|
||||
.append_pair("login_option", login_option)
|
||||
.finish(),
|
||||
);
|
||||
redirect_uri
|
||||
}
|
||||
|
||||
async fn exchange_admin_provider_oauth_kiro_social_code(
|
||||
state: &AdminAppState<'_>,
|
||||
code: &str,
|
||||
code_verifier: &str,
|
||||
redirect_uri: &str,
|
||||
machine_id: &str,
|
||||
proxy: Option<ProxySnapshot>,
|
||||
) -> Result<Value, String> {
|
||||
let url = state.provider_oauth_token_url("kiro_social_token", KIRO_SOCIAL_TOKEN_URL);
|
||||
let host = reqwest::Url::parse(&url)
|
||||
.ok()
|
||||
.and_then(|value| value.host_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| "prod.us-east-1.auth.desktop.kiro.dev".to_string());
|
||||
let user_agent = format!("KiroIDE-{KIRO_SOCIAL_AUTH_KIRO_VERSION}-{machine_id}");
|
||||
let headers = reqwest::header::HeaderMap::from_iter([
|
||||
(
|
||||
reqwest::header::CONTENT_TYPE,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
(
|
||||
reqwest::header::USER_AGENT,
|
||||
reqwest::header::HeaderValue::from_str(&user_agent)
|
||||
.map_err(|_| "Kiro social User-Agent 无效".to_string())?,
|
||||
),
|
||||
(
|
||||
reqwest::header::HOST,
|
||||
reqwest::header::HeaderValue::from_str(&host)
|
||||
.map_err(|_| "Kiro social host 无效".to_string())?,
|
||||
),
|
||||
]);
|
||||
let response = state
|
||||
.execute_admin_provider_oauth_http_request(
|
||||
"kiro_social_token",
|
||||
reqwest::Method::POST,
|
||||
&url,
|
||||
&headers,
|
||||
Some("application/json"),
|
||||
Some(json!({
|
||||
"code": code,
|
||||
"code_verifier": code_verifier,
|
||||
"redirect_uri": redirect_uri,
|
||||
})),
|
||||
None,
|
||||
proxy,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| format!("Kiro social token 请求失败: {err}"))?;
|
||||
if !response.status.is_success() {
|
||||
let detail = response.body_text.trim();
|
||||
return Err(if detail.is_empty() {
|
||||
format!("HTTP {}", response.status.as_u16())
|
||||
} else {
|
||||
detail.to_string()
|
||||
});
|
||||
}
|
||||
response
|
||||
.json_body
|
||||
.or_else(|| serde_json::from_str::<Value>(&response.body_text).ok())
|
||||
.ok_or_else(|| "Kiro social token 返回了非 JSON 响应".to_string())
|
||||
}
|
||||
|
||||
pub(super) async fn handle_admin_provider_oauth_device_poll(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -146,6 +338,19 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
||||
)
|
||||
.await;
|
||||
|
||||
if kiro_device_session_is_social(&session) {
|
||||
return handle_admin_provider_oauth_kiro_social_device_poll(
|
||||
state,
|
||||
&provider,
|
||||
&endpoints,
|
||||
request_proxy,
|
||||
session_id,
|
||||
session,
|
||||
payload.callback_url.as_deref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let token_result = match state
|
||||
.poll_admin_kiro_device_token(
|
||||
&session.region,
|
||||
@@ -424,3 +629,295 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
||||
.into_response(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn handle_admin_provider_oauth_kiro_social_device_poll(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
request_proxy: Option<ProxySnapshot>,
|
||||
session_id: &str,
|
||||
mut session: StoredAdminProviderOAuthDeviceSession,
|
||||
callback_url: Option<&str>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let callback_url = callback_url
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let Some(callback_url) = callback_url else {
|
||||
return Ok(Json(json!({"status": "pending", "replaced": false})).into_response());
|
||||
};
|
||||
|
||||
let callback_params = parse_provider_oauth_callback_params(callback_url);
|
||||
if let Some(error) = callback_params.get("error").map(String::as_str) {
|
||||
let error_description = callback_params
|
||||
.get("error_description")
|
||||
.map(String::as_str)
|
||||
.unwrap_or("用户拒绝授权");
|
||||
session.status = "error".to_string();
|
||||
session.error_msg = Some(format!("{error}: {error_description}"));
|
||||
let _ = state
|
||||
.save_provider_oauth_device_session(session_id, &session, 30)
|
||||
.await;
|
||||
return Ok(attach_admin_provider_oauth_device_poll_terminal_response(
|
||||
session_id,
|
||||
"error",
|
||||
kiro_social_poll_error_response(format!("{error}: {error_description}")),
|
||||
));
|
||||
}
|
||||
|
||||
let Some(code) = callback_params
|
||||
.get("code")
|
||||
.map(String::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(kiro_social_poll_error_response("回调 URL 缺少 code"));
|
||||
};
|
||||
let Some(callback_state) = callback_params
|
||||
.get("state")
|
||||
.map(String::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(kiro_social_poll_error_response("回调 URL 缺少 state"));
|
||||
};
|
||||
if callback_state != session_id {
|
||||
return Ok(kiro_social_poll_error_response("回调 state 与会话不匹配"));
|
||||
}
|
||||
|
||||
let Some(code_verifier) = session
|
||||
.code_verifier
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(kiro_social_poll_error_response("会话缺少 code_verifier"));
|
||||
};
|
||||
let redirect_uri = session
|
||||
.redirect_uri
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("kiro://kiro.kiroAgent/authenticate-success");
|
||||
let login_option = callback_params.get("login_option").map(String::as_str);
|
||||
let token_redirect_uri =
|
||||
kiro_social_token_redirect_uri(redirect_uri, callback_url, login_option);
|
||||
let machine_id = session
|
||||
.machine_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("unknown");
|
||||
|
||||
let token_result = match exchange_admin_provider_oauth_kiro_social_code(
|
||||
state,
|
||||
code,
|
||||
code_verifier,
|
||||
&token_redirect_uri,
|
||||
machine_id,
|
||||
request_proxy.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(payload) => payload,
|
||||
Err(detail) => {
|
||||
session.status = "error".to_string();
|
||||
session.error_msg = Some(format!("token exchange 失败: {detail}"));
|
||||
let _ = state
|
||||
.save_provider_oauth_device_session(session_id, &session, 30)
|
||||
.await;
|
||||
return Ok(attach_admin_provider_oauth_device_poll_terminal_response(
|
||||
session_id,
|
||||
"error",
|
||||
kiro_social_poll_error_response(format!("token exchange 失败: {detail}")),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let Some(access_token) = json_non_empty_string(
|
||||
token_result
|
||||
.get("accessToken")
|
||||
.or_else(|| token_result.get("access_token")),
|
||||
) else {
|
||||
return Ok(kiro_social_poll_error_response(
|
||||
"token 响应缺少 accessToken 或 refreshToken",
|
||||
));
|
||||
};
|
||||
let Some(refresh_token) = json_non_empty_string(
|
||||
token_result
|
||||
.get("refreshToken")
|
||||
.or_else(|| token_result.get("refresh_token")),
|
||||
) else {
|
||||
return Ok(kiro_social_poll_error_response(
|
||||
"token 响应缺少 accessToken 或 refreshToken",
|
||||
));
|
||||
};
|
||||
let expires_at = json_u64_value(
|
||||
token_result
|
||||
.get("expiresIn")
|
||||
.or_else(|| token_result.get("expires_in")),
|
||||
)
|
||||
.map(|expires_in| current_unix_secs().saturating_add(expires_in))
|
||||
.unwrap_or_else(|| current_unix_secs().saturating_add(3600));
|
||||
let social_provider = kiro_social_provider_from_login_option(login_option)
|
||||
.or_else(|| {
|
||||
session
|
||||
.social_provider
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
})
|
||||
.unwrap_or("Google")
|
||||
.to_string();
|
||||
session.social_provider = Some(social_provider.clone());
|
||||
let profile_arn = json_non_empty_string(
|
||||
token_result
|
||||
.get("profileArn")
|
||||
.or_else(|| token_result.get("profile_arn")),
|
||||
);
|
||||
let auth_config = AdminKiroAuthConfig {
|
||||
auth_method: Some("social".to_string()),
|
||||
refresh_token: Some(refresh_token.clone()),
|
||||
expires_at: Some(expires_at),
|
||||
profile_arn,
|
||||
region: Some("us-east-1".to_string()),
|
||||
auth_region: Some("us-east-1".to_string()),
|
||||
api_region: None,
|
||||
client_id: None,
|
||||
client_secret: None,
|
||||
machine_id: session.machine_id.clone(),
|
||||
kiro_version: Some(KIRO_SOCIAL_AUTH_KIRO_VERSION.to_string()),
|
||||
system_version: None,
|
||||
node_version: None,
|
||||
access_token: Some(access_token.clone()),
|
||||
};
|
||||
|
||||
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, &auth_config, request_proxy.clone()).await;
|
||||
}
|
||||
|
||||
let mut auth_config_object = auth_config
|
||||
.to_json_value()
|
||||
.as_object()
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
auth_config_object.insert("provider_type".to_string(), json!("kiro"));
|
||||
auth_config_object.insert("provider".to_string(), json!(social_provider));
|
||||
if let Some(email) = email.as_ref() {
|
||||
auth_config_object.insert("email".to_string(), json!(email));
|
||||
}
|
||||
if let Some(id_token) = json_non_empty_string(
|
||||
token_result
|
||||
.get("idToken")
|
||||
.or_else(|| token_result.get("id_token")),
|
||||
) {
|
||||
auth_config_object.insert("id_token".to_string(), json!(id_token));
|
||||
}
|
||||
if let Some(token_type) = json_non_empty_string(
|
||||
token_result
|
||||
.get("tokenType")
|
||||
.or_else(|| token_result.get("token_type")),
|
||||
) {
|
||||
auth_config_object.insert("token_type".to_string(), json!(token_type));
|
||||
}
|
||||
|
||||
let duplicate = match state
|
||||
.find_duplicate_provider_oauth_key(&provider.id, &auth_config_object, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
return Ok(Json(json!({
|
||||
"status": "error",
|
||||
"error": detail,
|
||||
"replaced": false,
|
||||
}))
|
||||
.into_response());
|
||||
}
|
||||
};
|
||||
|
||||
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;
|
||||
match state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider.provider_type,
|
||||
&access_token,
|
||||
&auth_config_object,
|
||||
&api_formats,
|
||||
key_proxy.clone(),
|
||||
Some(expires_at),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let key_name = kiro_social_key_name(
|
||||
email.as_deref(),
|
||||
session.social_provider.as_deref(),
|
||||
Some(&refresh_token),
|
||||
);
|
||||
match state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider.id,
|
||||
&provider.provider_type,
|
||||
&key_name,
|
||||
&access_token,
|
||||
&auth_config_object,
|
||||
&api_formats,
|
||||
key_proxy,
|
||||
Some(expires_at),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
spawn_provider_oauth_account_state_refresh_after_update(
|
||||
state.cloned_app(),
|
||||
provider.clone(),
|
||||
persisted_key.id.clone(),
|
||||
request_proxy.clone(),
|
||||
);
|
||||
|
||||
session.status = "authorized".to_string();
|
||||
session.key_id = Some(persisted_key.id.clone());
|
||||
session.email = email.clone();
|
||||
session.replaced = replaced;
|
||||
session.error_msg = None;
|
||||
let _ = state
|
||||
.save_provider_oauth_device_session(session_id, &session, 60)
|
||||
.await;
|
||||
|
||||
Ok(attach_admin_provider_oauth_device_poll_terminal_response(
|
||||
session_id,
|
||||
"authorized",
|
||||
Json(json!({
|
||||
"status": "authorized",
|
||||
"key_id": persisted_key.id,
|
||||
"email": email,
|
||||
"replaced": replaced,
|
||||
}))
|
||||
.into_response(),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -11,12 +11,15 @@ pub(super) struct AdminProviderOAuthDeviceAuthorizePayload {
|
||||
pub(super) start_url: String,
|
||||
#[serde(default = "default_kiro_device_region")]
|
||||
pub(super) region: String,
|
||||
pub(super) auth_type: Option<String>,
|
||||
pub(super) redirect_uri: Option<String>,
|
||||
pub(super) proxy_node_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(super) struct AdminProviderOAuthDevicePollPayload {
|
||||
pub(super) session_id: String,
|
||||
pub(super) callback_url: Option<String>,
|
||||
}
|
||||
|
||||
pub(super) fn attach_admin_provider_oauth_device_poll_terminal_response(
|
||||
|
||||
@@ -170,6 +170,7 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
let new_email = normalize_provider_oauth_identity_value(auth_config.get("email"));
|
||||
let new_user_id = normalize_provider_oauth_identity_value(auth_config.get("user_id"));
|
||||
let new_auth_method = normalize_provider_oauth_identity_value(auth_config.get("auth_method"));
|
||||
let new_kiro_provider = normalize_provider_oauth_identity_value(auth_config.get("provider"));
|
||||
|
||||
if new_email.is_none() && new_user_id.is_none() {
|
||||
return Ok(None);
|
||||
@@ -198,6 +199,8 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
normalize_provider_oauth_identity_value(existing_auth_config.get("user_id"));
|
||||
let existing_auth_method =
|
||||
normalize_provider_oauth_identity_value(existing_auth_config.get("auth_method"));
|
||||
let existing_kiro_provider =
|
||||
normalize_provider_oauth_identity_value(existing_auth_config.get("provider"));
|
||||
|
||||
let mut is_duplicate = false;
|
||||
let codex_identity_match =
|
||||
@@ -237,6 +240,10 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
.as_deref()
|
||||
.zip(existing_auth_method.as_deref())
|
||||
.is_some_and(|(left, right)| left.eq_ignore_ascii_case(right))
|
||||
&& new_kiro_provider
|
||||
.as_deref()
|
||||
.zip(existing_kiro_provider.as_deref())
|
||||
.is_none_or(|(left, right)| left.eq_ignore_ascii_case(right))
|
||||
{
|
||||
is_duplicate = true;
|
||||
}
|
||||
|
||||
@@ -321,6 +321,97 @@ async fn gateway_handles_admin_provider_oauth_device_authorize_locally_with_trus
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_oauth_device_authorize_for_kiro_google_social() {
|
||||
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 state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
))
|
||||
.with_provider_oauth_token_url_for_tests(
|
||||
"kiro_social_portal",
|
||||
"https://portal.example.com/signin",
|
||||
);
|
||||
|
||||
let response = local_admin_provider_oauth_response(
|
||||
&state,
|
||||
http::Method::POST,
|
||||
"/api/admin/provider-oauth/providers/provider-kiro/device-authorize",
|
||||
Some(json!({
|
||||
"auth_type": "google"
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
|
||||
let session_id = payload["session_id"]
|
||||
.as_str()
|
||||
.expect("session_id should exist");
|
||||
assert_eq!(payload["auth_type"], "google");
|
||||
assert_eq!(payload["callback_required"], true);
|
||||
assert_eq!(payload["redirect_uri"], "http://localhost:49153");
|
||||
|
||||
let authorization_url = payload["verification_uri_complete"]
|
||||
.as_str()
|
||||
.expect("authorization url should exist");
|
||||
let parsed = url::Url::parse(authorization_url).expect("authorization url should parse");
|
||||
let params = parsed
|
||||
.query_pairs()
|
||||
.into_owned()
|
||||
.collect::<std::collections::BTreeMap<_, _>>();
|
||||
assert_eq!(
|
||||
parsed.as_str().split('?').next(),
|
||||
Some("https://portal.example.com/signin")
|
||||
);
|
||||
assert_eq!(
|
||||
params.get("redirect_uri").map(String::as_str),
|
||||
Some("http://localhost:49153")
|
||||
);
|
||||
assert_eq!(params.get("state").map(String::as_str), Some(session_id));
|
||||
assert_eq!(
|
||||
params.get("code_challenge_method").map(String::as_str),
|
||||
Some("S256")
|
||||
);
|
||||
assert_eq!(
|
||||
params.get("redirect_from").map(String::as_str),
|
||||
Some("KiroIDE")
|
||||
);
|
||||
assert_eq!(
|
||||
params.get("login_option").map(String::as_str),
|
||||
Some("google")
|
||||
);
|
||||
assert!(params
|
||||
.get("code_challenge")
|
||||
.is_some_and(|value| !value.is_empty()));
|
||||
|
||||
let stored = state
|
||||
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
|
||||
.expect("device session should be stored");
|
||||
let stored: serde_json::Value =
|
||||
serde_json::from_str(&stored).expect("device session json should parse");
|
||||
assert_eq!(stored["provider_id"], "provider-kiro");
|
||||
assert_eq!(stored["auth_type"], "social");
|
||||
assert_eq!(stored["social_provider"], "Google");
|
||||
assert_eq!(stored["redirect_uri"], "http://localhost:49153");
|
||||
assert!(stored["code_verifier"]
|
||||
.as_str()
|
||||
.is_some_and(|value| !value.is_empty()));
|
||||
assert!(stored["machine_id"]
|
||||
.as_str()
|
||||
.is_some_and(|value| !value.is_empty()));
|
||||
assert_eq!(stored["status"], "pending");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -477,6 +568,189 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_oauth_device_poll_for_kiro_social_callback() {
|
||||
let token_requests = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
|
||||
let token_requests_clone = Arc::clone(&token_requests);
|
||||
let access_token = sample_kiro_device_access_token("social@example.com");
|
||||
let expected_access_token = access_token.clone();
|
||||
let token_server = Router::new().route(
|
||||
"/oauth/token",
|
||||
post(move |request: Request| {
|
||||
let token_requests_inner = Arc::clone(&token_requests_clone);
|
||||
let access_token_inner = access_token.clone();
|
||||
async move {
|
||||
let user_agent = request
|
||||
.headers()
|
||||
.get("user-agent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
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((user_agent, raw_body));
|
||||
Json(json!({
|
||||
"accessToken": access_token_inner,
|
||||
"refreshToken": "kiro-social-refresh-token",
|
||||
"profileArn": "arn:aws:kiro:profile/social",
|
||||
"idToken": "id-token-123",
|
||||
"tokenType": "Bearer",
|
||||
"expiresIn": 1800,
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
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 (token_url, token_handle) = start_server(token_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-social",
|
||||
json!({
|
||||
"provider_id": "provider-kiro",
|
||||
"region": "us-east-1",
|
||||
"client_id": "",
|
||||
"client_secret": "",
|
||||
"device_code": "",
|
||||
"auth_type": "social",
|
||||
"social_provider": "Github",
|
||||
"code_verifier": "verifier-123",
|
||||
"redirect_uri": "http://localhost:49153",
|
||||
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
|
||||
"interval": 5,
|
||||
"expires_at_unix_secs": 4_102_444_800u64,
|
||||
"status": "pending",
|
||||
"proxy_node_id": null,
|
||||
"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_social_token",
|
||||
format!("{token_url}/oauth/token"),
|
||||
);
|
||||
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-social",
|
||||
"callback_url": "http://localhost:49153/signin/callback?login_option=github&code=social-code-123&state=session-social"
|
||||
}))
|
||||
.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"], "social@example.com");
|
||||
assert_eq!(payload["replaced"], false);
|
||||
|
||||
{
|
||||
let requests = token_requests.lock().expect("mutex should lock");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].0,
|
||||
"KiroIDE-0.6.18-123e4567-e89b-12d3-a456-426614174000"
|
||||
);
|
||||
assert!(requests[0].1.contains("\"code\":\"social-code-123\""));
|
||||
assert!(requests[0].1.contains("\"code_verifier\":\"verifier-123\""));
|
||||
assert!(requests[0].1.contains(
|
||||
"\"redirect_uri\":\"http://localhost:49153/signin/callback?login_option=github\""
|
||||
));
|
||||
}
|
||||
|
||||
let stored = state
|
||||
.load_provider_oauth_device_session_for_tests("device_auth_session:session-social")
|
||||
.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"], "social@example.com");
|
||||
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.name, "social@example.com (Github)");
|
||||
assert_eq!(persisted.auth_type, "oauth");
|
||||
let decrypted_api_key = decrypt_python_fernet_ciphertext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
persisted
|
||||
.encrypted_api_key
|
||||
.as_deref()
|
||||
.expect("api key should be present"),
|
||||
)
|
||||
.expect("api key should decrypt");
|
||||
assert_eq!(decrypted_api_key, expected_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["provider"], "Github");
|
||||
assert_eq!(auth_config["auth_method"], "social");
|
||||
assert_eq!(auth_config["refresh_token"], "kiro-social-refresh-token");
|
||||
assert_eq!(auth_config["profile_arn"], "arn:aws:kiro:profile/social");
|
||||
assert_eq!(auth_config["email"], "social@example.com");
|
||||
assert_eq!(
|
||||
auth_config["machine_id"],
|
||||
"123e4567-e89b-12d3-a456-426614174000"
|
||||
);
|
||||
assert_eq!(auth_config["kiro_version"], "0.6.18");
|
||||
|
||||
gateway_handle.abort();
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_keeps_admin_provider_oauth_device_poll_pending_for_authorization_pending_error() {
|
||||
let token_server = Router::new().route(
|
||||
|
||||
Reference in New Issue
Block a user