Skip to content

feat(cosh-ng): [core,shell] automate auth #3298

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kongche-jbw merged 3 commits into from
Sep 17, 2026
Merged

feat(cosh-ng): [core,shell] automate auth #3298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d83c1cd
chore(cosh-ng): [shell] align validation
SunnyQjm Sep 16, 2026
1ad355d
feat(cosh-ng): [core,shell] automate auth
SunnyQjm Sep 16, 2026
85ecfbf
docs(cosh-ng): [core,shell] document ECS auth protocol
SunnyQjm Sep 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 48 additions & 8 deletions src/cosh-ng/crates/cosh-core/src/auth/preflight.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,16 @@ use sha2::{Digest, Sha256};
use uuid::Uuid;

use crate::config::ResolvedProvider;
use crate::provider::sysom::{CredentialStatus, ProbeError};

fn ecs_preflight_result(
result: Result<CredentialStatus, ProbeError>,
) -> Result<(), AuthPreflightError> {
match result.map_err(AuthPreflightError::MetadataProbe)? {
CredentialStatus::Ready => Ok(()),
CredentialStatus::NotReady(_) => Err(AuthPreflightError::CredentialSourceUnavailable),
}
}

const CONNECT_TIMEOUT: Duration = Duration::from_secs(3);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(8);
Expand Down Expand Up @@ -46,6 +56,7 @@ pub(crate) enum AuthPreflightError {
ProviderUnavailable,
ServiceNotReady,
CredentialSourceUnavailable,
MetadataProbe(ProbeError),
UnsupportedResponse,
}

Expand All @@ -61,6 +72,7 @@ impl AuthPreflightError {
Self::ProviderUnavailable => "provider_unavailable",
Self::ServiceNotReady => "service_not_ready",
Self::CredentialSourceUnavailable => "credential_source_unavailable",
Self::MetadataProbe(error) => error.code(),
Self::UnsupportedResponse => "unsupported_response",
}
}
Expand Down Expand Up @@ -95,8 +107,9 @@ impl fmt::Display for AuthPreflightError {
"Aliyun SysOM is not authorized for this account. Complete service authorization and try again.",
),
Self::CredentialSourceUnavailable => formatter.write_str(
"ECS RAM Role credentials are not available yet. Authorize the instance role and try again.",
"ECS RAM Role credentials are not available yet. Check the instance role or wait for credential refresh.",
),
Self::MetadataProbe(error) => error.fmt(formatter),
Self::UnsupportedResponse => formatter.write_str(
"The endpoint returned an unsupported validation response. Check the endpoint configuration and provider compatibility.",
),
Expand Down Expand Up @@ -372,13 +385,7 @@ async fn response_explicitly_reports_missing_model(response: Response) -> bool {

async fn preflight_aliyun(provider: &ResolvedProvider) -> Result<(), AuthPreflightError> {
if provider.auth_source.as_deref() == Some("ecs_ram_role") {
return tokio::task::spawn_blocking(
crate::provider::sysom::ecs_ram_role_credentials_available,
)
.await
.map_err(|_| AuthPreflightError::EndpointUnreachable)?
.then_some(())
.ok_or(AuthPreflightError::CredentialSourceUnavailable);
return ecs_preflight_result(crate::provider::sysom::probe_ecs_ram_role().await);
}

let resolved = crate::provider::sysom::endpoint::resolve(&provider.sysom_endpoint).await;
Expand Down Expand Up @@ -740,6 +747,39 @@ mod tests {

use super::*;

#[test]
fn ecs_metadata_preflight_distinguishes_ready_from_not_ready() {
use crate::provider::sysom::NotReadyReason;
assert_eq!(ecs_preflight_result(Ok(CredentialStatus::Ready)), Ok(()));
for reason in [
NotReadyReason::RoleMissing,
NotReadyReason::CredentialsExpired,
] {
assert_eq!(
ecs_preflight_result(Ok(CredentialStatus::NotReady(reason))),
Err(AuthPreflightError::CredentialSourceUnavailable)
);
}
}

#[test]
fn ecs_metadata_preflight_preserves_probe_errors_through_configure() {
for probe_error in [
ProbeError::AccessDenied,
ProbeError::InvalidResponse,
ProbeError::Unreachable,
ProbeError::Timeout,
ProbeError::Http,
] {
let error = ecs_preflight_result(Err(probe_error)).expect_err("probe failure");
assert_eq!(error.code(), probe_error.code());
assert_eq!(error.to_string(), probe_error.to_string());
let configure_error = crate::auth::AuthConfigureError::from(error);
assert_eq!(configure_error.code(), probe_error.code());
assert_eq!(configure_error.to_string(), probe_error.to_string());
}
}

#[derive(Clone)]
struct Reply {
status: u16,
Expand Down
152 changes: 49 additions & 103 deletions src/cosh-ng/crates/cosh-core/src/provider/sysom.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@
//! Uses ACS3-HMAC-SHA256 signing and parses the cumulative SSE stream format
//! into incremental `GenerateEvent`s compatible with `ContentGenerator` trait.

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use std::time::Duration;
Expand All @@ -15,12 +13,16 @@ use futures::StreamExt;
use hmac::{Hmac, Mac};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::sync::OnceCell;
use uuid::Uuid;

use super::{ContentGenerator, GenerateConfig, GenerateStream, Message, ToolDeclaration};

use self::stream::sysom_event_stream;

pub use self::ecs_metadata::{probe_ecs_ram_role, CredentialStatus, NotReadyReason, ProbeError};

mod ecs_metadata;
pub mod endpoint;
mod stream;

Expand All @@ -38,8 +40,6 @@ const CONSOLE_URL_TEMPLATE: &str =
const INSTANCE_ID_CACHE_TTL_SECS: u64 = 3 * 3600;
/// Connect timeout for ECS metadata service.
const METADATA_CONNECT_TIMEOUT: Duration = Duration::from_secs(1);
/// Read timeout for ECS metadata service.
const METADATA_READ_TIMEOUT: Duration = Duration::from_secs(2);
/// Connect timeout for the SysOM API.
///
/// Bounds connection setup so an unreachable endpoint surfaces as a prompt
Expand Down Expand Up @@ -104,7 +104,9 @@ pub struct SysomProvider {
credentials: RwLock<SysomCredentials>,
is_sts: bool,
cancelled: Arc<AtomicBool>,
instance_id: Option<String>,
/// Resolved once per provider instance, including a negative result; only the
/// on-disk cache honours [`INSTANCE_ID_CACHE_TTL_SECS`].
instance_id: OnceCell<Option<String>>,
}

impl SysomProvider {
Expand All @@ -115,7 +117,6 @@ impl SysomProvider {
configured_endpoint: &str,
) -> Self {
let is_sts = security_token.is_some();
let instance_id = resolve_instance_id();
Self {
configured_endpoint: configured_endpoint.to_string(),
credentials: RwLock::new(SysomCredentials {
Expand All @@ -125,12 +126,11 @@ impl SysomProvider {
}),
is_sts,
cancelled: Arc::new(AtomicBool::new(false)),
instance_id,
instance_id: OnceCell::new(),
}
}

pub fn from_ecs_ram_role(configured_endpoint: &str) -> Self {
let instance_id = resolve_instance_id();
Self {
configured_endpoint: configured_endpoint.to_string(),
credentials: RwLock::new(SysomCredentials {
Expand All @@ -140,7 +140,7 @@ impl SysomProvider {
}),
is_sts: true,
cancelled: Arc::new(AtomicBool::new(false)),
instance_id,
instance_id: OnceCell::new(),
}
}

Expand Down Expand Up @@ -190,7 +190,7 @@ impl SysomProvider {
}
}

if let Some(ref id) = self.instance_id {
if let Some(Some(id)) = self.instance_id.get() {
inner["instance_id"] = serde_json::json!(id);
}

Expand Down Expand Up @@ -269,6 +269,7 @@ impl ContentGenerator for SysomProvider {
config: &GenerateConfig,
) -> Result<GenerateStream, String> {
self.cancelled.store(false, Ordering::SeqCst);
self.instance_id.get_or_init(resolve_instance_id).await;

let body = self.build_request_body(messages, tools, config);
let body_bytes = serde_json::to_vec(&body).map_err(|e| format!("JSON serialize: {e}"))?;
Expand Down Expand Up @@ -409,47 +410,21 @@ impl SysomProvider {
Ok(sysom_event_stream(Box::pin(byte_stream), cancelled))
}

/// Check if an error indicates STS credential expiration.
/// Refresh STS credentials from ECS metadata service.
/// Refresh only with credentials accepted by the same validator as auth preflight.
async fn refresh_sts_credentials(&self) -> bool {
let url = format!(
"{}/latest/meta-data/ram/security-credentials/{}",
ECS_METADATA_ENDPOINT, ECS_RAM_ROLE_NAME
);
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(2))
.timeout(std::time::Duration::from_secs(5))
.build()
.unwrap_or_default();

let resp = match client.get(&url).send().await {
Ok(r) => r,
Err(e) => {
tracing::warn!("STS refresh failed: {e}");
return false;
match ecs_metadata::load_role_credentials().await {
Ok(ecs_metadata::RoleCredentials::Ready(credentials)) => {
*self.credentials.write().unwrap() = credentials;
true
}
Ok(ecs_metadata::RoleCredentials::NotReady(reason)) => {
tracing::debug!(reason = reason.as_str(), "STS credentials not ready");
false
}
};
let body: Value = match resp.json().await {
Ok(v) => v,
Err(e) => {
tracing::warn!("STS refresh parse failed: {e}");
return false;
Err(error) => {
tracing::warn!(error_code = error.code(), "STS metadata refresh failed");
false
}
};

let ak = body.get("AccessKeyId").and_then(|v| v.as_str());
let sk = body.get("AccessKeySecret").and_then(|v| v.as_str());
let token = body.get("SecurityToken").and_then(|v| v.as_str());

if let (Some(ak), Some(sk), Some(token)) = (ak, sk, token) {
let mut creds = self.credentials.write().unwrap();
creds.access_key_id = ak.to_string();
creds.access_key_secret = sk.to_string();
creds.security_token = Some(token.to_string());
true
} else {
tracing::warn!("STS refresh: missing fields in response");
false
}
}
}
Expand All @@ -467,19 +442,26 @@ fn is_sts_error(error: &str) -> bool {

/// Resolve instance_id: read from local cache if valid, otherwise fetch from
/// ECS metadata service and update the cache.
fn resolve_instance_id() -> Option<String> {
async fn resolve_instance_id() -> Option<String> {
let config_dir = dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".copilot-shell");
let cache_path = config_dir.join("instance_id");
#[cfg(test)]
let cache_path = metadata_tests::TEST_CACHE_PATH
.try_with(Clone::clone)
.unwrap_or(cache_path);
resolve_instance_id_cached(&cache_path).await
}

async fn resolve_instance_id_cached(cache_path: &std::path::Path) -> Option<String> {
// Try reading from cache
if let Ok(metadata) = std::fs::metadata(&cache_path) {
if let Ok(metadata) = std::fs::metadata(cache_path) {
if let Ok(modified) = metadata.modified() {
let age = modified.elapsed().unwrap_or(Duration::from_secs(u64::MAX));
if age < Duration::from_secs(INSTANCE_ID_CACHE_TTL_SECS) {
// Cache is still valid
let content = std::fs::read_to_string(&cache_path).unwrap_or_default();
let content = std::fs::read_to_string(cache_path).unwrap_or_default();
let trimmed = content.trim();
if trimmed.is_empty() {
// Empty file = previously failed to fetch
Expand All @@ -491,55 +473,32 @@ fn resolve_instance_id() -> Option<String> {
}

// Cache miss or expired — fetch from metadata service
let instance_id = fetch_instance_id_from_metadata();
let instance_id = ecs_metadata::fetch_instance_id().await.ok();

// Write cache (create parent dir if needed)
if let Some(parent) = cache_path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let content = instance_id.as_deref().unwrap_or("");
if let Err(e) = std::fs::write(&cache_path, content) {
if let Err(e) = std::fs::write(cache_path, content) {
tracing::debug!("failed to write instance_id cache: {e}");
}

instance_id
}

/// Fetch instance-id from ECS metadata service via raw TCP.
/// Returns None if not running on ECS or if the request fails.
fn fetch_instance_id_from_metadata() -> Option<String> {
fetch_ecs_metadata_text("/latest/meta-data/instance-id").and_then(|body| {
let instance_id = body.trim();
if instance_id.starts_with("i-") {
Some(instance_id.to_string())
} else {
None
}
})
}

pub fn detect_ecs_auth_challenge() -> Option<EcsAuthChallenge> {
let instance_id = fetch_instance_id_from_metadata()?;
let region_id = fetch_ecs_region_id().unwrap_or_else(|| "cn-hangzhou".to_string());
Some(EcsAuthChallenge {
/// Detect ECS identity without treating denied or malformed metadata as non-ECS.
///
/// # Errors
/// Returns a safe error for HTTP rejection, invalid metadata, or any identity fetch
/// failure after acquiring a valid token. Only an unreachable or timed-out session
/// start is treated as non-ECS so manual authentication remains available.
pub async fn detect_ecs_auth_challenge() -> Result<Option<EcsAuthChallenge>, ProbeError> {
let identity = ecs_metadata::fetch_identity().await?;
Ok(identity.map(|(instance_id, region_id)| EcsAuthChallenge {
console_url: generate_console_url(&instance_id, &region_id),
instance_id,
})
}

pub fn ecs_ram_role_credentials_available() -> bool {
let path = format!(
"/latest/meta-data/ram/security-credentials/{}",
ECS_RAM_ROLE_NAME
);
fetch_ecs_metadata_text(&path)
.map(|body| body.contains("AccessKeyId") && body.contains("SecurityToken"))
.unwrap_or(false)
}

fn fetch_ecs_region_id() -> Option<String> {
let zone_id = fetch_ecs_metadata_text("/latest/meta-data/zone-id")?;
region_id_from_zone_id(zone_id.trim())
}))
}

fn region_id_from_zone_id(zone_id: &str) -> Option<String> {
Expand All @@ -561,23 +520,6 @@ fn generate_console_url(instance_id: &str, region_id: &str) -> String {
.replace("{instanceId}", instance_id)
}

fn fetch_ecs_metadata_text(path: &str) -> Option<String> {
let addr: SocketAddr = "100.100.100.200:80".parse().ok()?;
let mut stream = TcpStream::connect_timeout(&addr, METADATA_CONNECT_TIMEOUT).ok()?;
stream.set_read_timeout(Some(METADATA_READ_TIMEOUT)).ok()?;
stream
.set_write_timeout(Some(METADATA_CONNECT_TIMEOUT))
.ok()?;

let request = format!("GET {path} HTTP/1.0\r\nHost: 100.100.100.200\r\n\r\n");
stream.write_all(request.as_bytes()).ok()?;

let mut response = String::new();
stream.read_to_string(&mut response).ok()?;

response.split("\r\n\r\n").nth(1).map(str::to_string)
}

// ---------------------------------------------------------------------------
// Crypto helpers
// ---------------------------------------------------------------------------
Expand All @@ -597,3 +539,7 @@ fn hex_hmac_sha256(key: &[u8], data: &[u8]) -> String {
#[cfg(test)]
#[path = "sysom/tests.rs"]
mod tests;

#[cfg(test)]
#[path = "sysom/metadata_tests.rs"]
mod metadata_tests;
Loading
Loading