refactor: extract provider pool abstractions

This commit is contained in:
fawney19
2026-05-13 18:19:15 +08:00
parent 3c2497f019
commit 5d1460e051
55 changed files with 3469 additions and 2184 deletions

View File

@@ -1,333 +1,7 @@
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::time::{SystemTime, UNIX_EPOCH};
pub const DEFAULT_REGION: &str = "us-east-1";
pub const DEFAULT_KIRO_VERSION: &str = "0.3.210";
pub const DEFAULT_NODE_VERSION: &str = "22.21.1";
pub const DEFAULT_SYSTEM_VERSION: &str = "other#unknown";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KiroAuthConfig {
pub auth_method: Option<String>,
pub refresh_token: Option<String>,
pub expires_at: Option<u64>,
pub profile_arn: Option<String>,
pub region: Option<String>,
pub auth_region: Option<String>,
pub api_region: Option<String>,
pub client_id: Option<String>,
pub client_secret: Option<String>,
pub machine_id: Option<String>,
pub kiro_version: Option<String>,
pub system_version: Option<String>,
pub node_version: Option<String>,
pub access_token: Option<String>,
}
impl KiroAuthConfig {
pub fn from_raw_json(raw: Option<&str>) -> Option<Self> {
let raw = raw?.trim();
if raw.is_empty() {
return None;
}
let parsed: Value = serde_json::from_str(raw).ok()?;
Self::from_json_value(&parsed)
}
pub fn from_json_value(raw: &Value) -> Option<Self> {
let object = raw.as_object()?;
Some(Self {
auth_method: get_nonempty_string(
object,
&["auth_method", "authMethod", "auth_type", "authType"],
)
.map(|value| normalize_auth_method(&value)),
refresh_token: get_nonempty_string(object, &["refresh_token", "refreshToken"]),
expires_at: get_epoch_seconds(object.get("expires_at"))
.or_else(|| get_epoch_seconds(object.get("expiresAt"))),
profile_arn: get_nonempty_string(object, &["profile_arn", "profileArn"]),
region: get_nonempty_string(object, &["region"]),
auth_region: get_nonempty_string(object, &["auth_region", "authRegion"]),
api_region: get_nonempty_string(object, &["api_region", "apiRegion"]),
client_id: get_nonempty_string(object, &["client_id", "clientId"]),
client_secret: get_nonempty_string(object, &["client_secret", "clientSecret"]),
machine_id: get_nonempty_string(object, &["machine_id", "machineId"]),
kiro_version: get_nonempty_string(object, &["kiro_version", "kiroVersion"]),
system_version: get_nonempty_string(object, &["system_version", "systemVersion"]),
node_version: get_nonempty_string(object, &["node_version", "nodeVersion"]),
access_token: get_nonempty_string(object, &["access_token", "accessToken"]),
})
}
pub fn to_json_value(&self) -> Value {
let mut object = serde_json::Map::new();
insert_optional_string(&mut object, "auth_method", self.auth_method.as_deref());
insert_optional_string(&mut object, "refresh_token", self.refresh_token.as_deref());
if let Some(expires_at) = self.expires_at {
object.insert("expires_at".to_string(), Value::from(expires_at));
}
insert_optional_string(&mut object, "profile_arn", self.profile_arn.as_deref());
insert_optional_string(&mut object, "region", self.region.as_deref());
insert_optional_string(&mut object, "auth_region", self.auth_region.as_deref());
insert_optional_string(&mut object, "api_region", self.api_region.as_deref());
insert_optional_string(&mut object, "client_id", self.client_id.as_deref());
insert_optional_string(&mut object, "client_secret", self.client_secret.as_deref());
insert_optional_string(&mut object, "machine_id", self.machine_id.as_deref());
insert_optional_string(&mut object, "kiro_version", self.kiro_version.as_deref());
insert_optional_string(
&mut object,
"system_version",
self.system_version.as_deref(),
);
insert_optional_string(&mut object, "node_version", self.node_version.as_deref());
insert_optional_string(&mut object, "access_token", self.access_token.as_deref());
Value::Object(object)
}
pub fn effective_api_region(&self) -> &str {
self.api_region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_auth_region(&self) -> &str {
self.auth_region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
self.region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
})
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_kiro_version(&self) -> &str {
self.kiro_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_KIRO_VERSION)
}
pub fn effective_system_version(&self) -> &str {
self.system_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_SYSTEM_VERSION)
}
pub fn effective_node_version(&self) -> &str {
self.node_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_NODE_VERSION)
}
pub fn cached_access_token(&self) -> Option<&str> {
self.access_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn cached_access_token_requires_refresh(&self, skew_seconds: u64) -> bool {
let Some(expires_at) = self.expires_at else {
return self.can_refresh_access_token();
};
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default();
now >= expires_at.saturating_sub(skew_seconds)
}
pub fn is_idc_auth(&self) -> bool {
let explicit_method = self
.auth_method
.as_deref()
.map(normalize_auth_method)
.unwrap_or_else(|| "social".to_string());
if explicit_method != "social" {
return matches!(explicit_method.as_str(), "idc" | "external_idp");
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
pub fn uses_external_idp_token_type(&self) -> bool {
self.auth_method
.as_deref()
.map(normalize_auth_method)
.as_deref()
== Some("external_idp")
}
pub fn profile_arn_for_payload(&self) -> Option<&str> {
if self.is_idc_auth() {
return None;
}
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn profile_arn_for_mcp(&self) -> Option<&str> {
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn can_refresh_access_token(&self) -> bool {
let refresh_token = self
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| value.len() >= 100 && !value.contains("..."));
if refresh_token.is_none() {
return false;
}
if !self.is_idc_auth() {
return true;
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
}
pub fn normalize_machine_id(raw: &str) -> Option<String> {
let raw = raw.trim();
if raw.is_empty() {
return None;
}
if raw.len() == 64 && raw.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Some(raw.to_ascii_lowercase());
}
if raw.len() == 36
&& raw.chars().enumerate().all(|(idx, ch)| match idx {
8 | 13 | 18 | 23 => ch == '-',
_ => ch.is_ascii_hexdigit(),
})
{
let normalized = raw.replace('-', "").to_ascii_lowercase();
return Some(format!("{normalized}{normalized}"));
}
None
}
pub fn generate_machine_id(
auth_config: &KiroAuthConfig,
fallback_secret: Option<&str>,
) -> Option<String> {
if let Some(machine_id) = auth_config
.machine_id
.as_deref()
.and_then(normalize_machine_id)
{
return Some(machine_id);
}
let seed = auth_config
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
fallback_secret
.map(str::trim)
.filter(|value| !value.is_empty())
})?;
let mut hasher = Sha256::new();
hasher.update(b"KotlinNativeAPI/");
hasher.update(seed.as_bytes());
Some(format!("{:x}", hasher.finalize()))
}
fn get_nonempty_string(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<String> {
keys.iter()
.find_map(|key| object.get(*key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn insert_optional_string(
object: &mut serde_json::Map<String, Value>,
key: &str,
value: Option<&str>,
) {
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
object.insert(key.to_string(), Value::String(value.to_string()));
}
fn get_epoch_seconds(value: Option<&Value>) -> Option<u64> {
match value? {
Value::Number(number) => number.as_u64().or_else(|| {
number
.as_i64()
.and_then(|value| (value >= 0).then_some(value as u64))
}),
Value::String(text) => text.trim().parse::<u64>().ok(),
_ => None,
}
}
fn normalize_auth_method(raw: &str) -> String {
let value = raw.trim().to_ascii_lowercase();
match value.as_str() {
"" => "social".to_string(),
"builder-id"
| "builder_id"
| "builderid"
| "device"
| "device-auth"
| "device_authorization"
| "iam"
| "identity-center"
| "identity_center"
| "identitycenter"
| "idc" => "idc".to_string(),
"external-idp" | "external_idp" | "externalidp" => "external_idp".to_string(),
_ => value,
}
}
pub use aether_oauth::provider::providers::{
generate_kiro_machine_id as generate_machine_id,
normalize_kiro_machine_id as normalize_machine_id, KiroAuthConfig, DEFAULT_REGION,
};
#[cfg(test)]
mod tests {

View File

@@ -1,7 +1,4 @@
use std::collections::BTreeMap;
use aether_oauth::provider::providers::KiroProviderOAuthAdapter as CoreKiroProviderOAuthAdapter;
use aether_oauth::provider::{ProviderOAuthAccount, ProviderOAuthAdapter};
use async_trait::async_trait;
use super::super::oauth_refresh::{
@@ -48,26 +45,10 @@ impl KiroOAuthRefreshAdapter {
let oauth_executor =
ProviderOAuthLocalHttpExecutor::new(PROVIDER_TYPE, transport, executor);
let ctx = provider_oauth_transport_context_from_snapshot(transport);
let account = ProviderOAuthAccount {
provider_type: PROVIDER_TYPE.to_string(),
access_token: auth_config
.cached_access_token()
.map(ToOwned::to_owned)
.unwrap_or_default(),
auth_config: auth_config.to_json_value(),
expires_at_unix_secs: auth_config.expires_at,
identity: BTreeMap::new(),
};
let refreshed = adapter
.refresh(&oauth_executor, &ctx, &account)
adapter
.refresh_auth_config(&oauth_executor, &ctx, auth_config)
.await
.map_err(|error| oauth_error_to_local_refresh_error(PROVIDER_TYPE, error))?;
KiroAuthConfig::from_json_value(&refreshed.auth_config).ok_or_else(|| {
LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "kiro refresh returned invalid auth_config".to_string(),
}
})
.map_err(|error| oauth_error_to_local_refresh_error(PROVIDER_TYPE, error))
}
fn auth_config_from_entry(entry: &CachedOAuthEntry) -> Option<KiroAuthConfig> {