mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor: extract provider pool abstractions
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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> {
|
||||
|
||||
Reference in New Issue
Block a user