// Copyright 2026 PingCAP, Inc.

use std::{
    fs,
    sync::{Arc, Mutex},
    time::Duration,
};

use async_trait::async_trait;
use aws::ActiveRefreshingProvider;
use chrono::{DateTime, TimeZone, Utc};
use hmac::{Hmac, Mac};
use http::{HeaderValue, Uri};
use hyper::{Body, Client, Request, client::HttpConnector};
use hyper_tls::HttpsConnector;
use rusoto_credential::{AwsCredentials, CredentialsError, ProvideAwsCredentials};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use slog_global::{info, warn};
use thiserror::Error;

pub const DOMAIN_STRING: &str = "myqcloud.com";

// Ref https://cloud.tencent.com/document/product/457/81989#.E6.AD.A5.E9.AA.A410.EF.BC.9A.E9.83.A8.E7.BD.B2.E7.A4.BA.E4.BE.8B.E5.BA.94.E7.94.A8.E7.A8.8B.E5.BA.8F
// for the envs will be injected.
// Ref https://cloud.tencent.com/document/product/1312/73070 for how to get credentials by this
// envs.
pub const ENV_TKE_DEFAULT_REGION: &str = "TKE_DEFAULT_REGION";
pub const ENV_TKE_REGION: &str = "TKE_REGION";
pub const ENV_TKE_PROVIDER_ID: &str = "TKE_PROVIDER_ID";
pub const ENV_TKE_ROLE_ARN: &str = "TKE_ROLE_ARN";
pub const ENV_TKE_WEB_IDENTITY_TOKEN_FILE: &str = "TKE_WEB_IDENTITY_TOKEN_FILE";

const STS_VERSION: &str = "2018-08-13";
const STS_ACTION_ASSUME_ROLE_WITH_WEB_IDENTITY: &str = "AssumeRoleWithWebIdentity";

const CVM_ROLE_METADATA_URL: &str =
    "http://metadata.tencentyun.com/latest/meta-data/cam/security-credentials";
const CVM_METADATA_TIMEOUT: Duration = Duration::from_secs(3);

const KMS_HOST: &str = "kms.tencentcloudapi.com";
const KMS_ENDPOINT: &str = "https://kms.tencentcloudapi.com/";
const KMS_VERSION: &str = "2019-01-18";
const KMS_SERVICE: &str = "kms";
const KMS_ACTION_DECRYPT: &str = "Decrypt";

const DEFAULT_SESSION_NAME: &str = "TiKV";

const TKE_CREDENTIAL_PROVIDER: &str = "TKE workload identity";
const CVM_CREDENTIAL_PROVIDER: &str = "CVM instance role";

#[derive(Debug, Deserialize)]
struct TencentApiError {
    #[serde(rename = "Code")]
    code: String,
    #[serde(rename = "Message")]
    message: String,
}

#[derive(Debug, Deserialize)]
struct TencentResponse<T> {
    #[serde(rename = "Response")]
    response: TencentResponseInner<T>,
}

#[derive(Debug, Deserialize)]
struct TencentResponseInner<T> {
    #[serde(default, rename = "Error")]
    error: Option<TencentApiError>,
    #[serde(rename = "RequestId")]
    request_id: Option<String>,
    #[serde(flatten)]
    data: T,
}

#[derive(Debug, Serialize)]
struct AssumeRoleWithWebIdentityRequest<'a> {
    #[serde(rename = "ProviderId")]
    provider_id: &'a str,
    #[serde(rename = "WebIdentityToken")]
    web_identity_token: &'a str,
    #[serde(rename = "RoleArn")]
    role_arn: &'a str,
    #[serde(rename = "RoleSessionName")]
    role_session_name: &'a str,
}

#[derive(Debug, Deserialize)]
struct AssumeRoleWithWebIdentityResult {
    #[serde(rename = "Credentials")]
    credentials: Option<StsCredentials>,
    #[serde(rename = "ExpiredTime")]
    expired_time: Option<i64>,
    #[serde(rename = "Expiration")]
    expiration: Option<String>,
}

#[derive(Debug, Deserialize)]
struct StsCredentials {
    #[serde(rename = "TmpSecretId")]
    tmp_secret_id: String,
    #[serde(rename = "TmpSecretKey")]
    tmp_secret_key: String,
    #[serde(rename = "Token")]
    token: String,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
struct CvmRoleCredentials {
    tmp_secret_id: String,
    tmp_secret_key: String,
    token: String,
    code: String,
    #[serde(default)]
    expired_time: Option<i64>,
    #[serde(default)]
    expiration: Option<String>,
}

#[derive(Debug, Serialize)]
struct DecryptRequest<'a> {
    #[serde(rename = "CiphertextBlob")]
    ciphertext_blob: &'a str,
}

#[derive(Debug, Deserialize)]
struct DecryptResult {
    #[serde(rename = "Plaintext")]
    plaintext: Option<String>,
}

#[derive(Debug, Error)]
enum TencentError {
    #[error("http error: {0}")]
    Http(#[from] hyper::Error),
    #[error("invalid endpoint: {0}")]
    InvalidEndpoint(#[from] http::uri::InvalidUri),
    #[error("invalid response: {0}")]
    InvalidResponse(String),
    #[error("serde error: {0}")]
    Serde(#[from] serde_json::Error),
    #[error("base64 error: {0}")]
    Base64(#[from] base64::DecodeError),
}

pub struct AssumeRoleWithWebIdentityProvider {
    region: String,
    provider_id: String,
    role_arn: String,
    token_file: String,
    client: Client<HttpsConnector<HttpConnector>, Body>,
}

struct ChainProvider {
    providers: Vec<NamedProvider>,
    initialization_errors: Vec<String>,
    selected_provider: Mutex<Option<&'static str>>,
}

struct NamedProvider {
    name: &'static str,
    provider: Arc<dyn ProvideAwsCredentials + Send + Sync>,
}

impl Default for ChainProvider {
    fn default() -> Self {
        Self {
            providers: Vec::new(),
            initialization_errors: Vec::new(),
            selected_provider: Mutex::new(None),
        }
    }
}

impl ChainProvider {
    fn push<P>(&mut self, name: &'static str, provider: P)
    where
        P: ProvideAwsCredentials + Send + Sync + 'static,
    {
        self.providers.push(NamedProvider {
            name,
            provider: Arc::new(provider),
        });
    }

    fn push_initialization_error(&mut self, name: &'static str, error: CredentialsError) {
        self.initialization_errors.push(format!("{name}: {error}"));
    }

    fn record_selected_provider(&self, provider: &'static str) {
        let previous_provider = {
            let mut selected_provider = self
                .selected_provider
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            let previous_provider = *selected_provider;
            *selected_provider = Some(provider);
            previous_provider
        };

        match previous_provider {
            None => info!(
                "selected Tencent Cloud credential provider";
                "provider" => provider,
            ),
            Some(previous_provider) if previous_provider != provider => warn!(
                "switched Tencent Cloud credential provider";
                "previous_provider" => previous_provider,
                "provider" => provider,
            ),
            _ => {}
        }
    }
}

#[async_trait]
impl ProvideAwsCredentials for ChainProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let mut errors = self.initialization_errors.clone();
        for named_provider in &self.providers {
            match named_provider.provider.credentials().await {
                Ok(credentials) => {
                    self.record_selected_provider(named_provider.name);
                    return Ok(credentials);
                }
                Err(error) => {
                    errors.push(format!("{}: {}", named_provider.name, error));
                }
            }
        }
        if errors.is_empty() {
            return Err(CredentialsError::new("no credential provider found"));
        }
        Err(CredentialsError::new(format!(
            "all credential providers failed: {}",
            errors.join("; ")
        )))
    }
}

struct CvmRoleCredentialsProvider {
    metadata_url: String,
    timeout: Duration,
    client: Client<HttpConnector, Body>,
}

impl CvmRoleCredentialsProvider {
    fn new() -> Self {
        Self::with_metadata_url(CVM_ROLE_METADATA_URL.to_owned())
    }

    fn with_metadata_url(metadata_url: String) -> Self {
        Self::with_metadata_url_and_timeout(metadata_url, CVM_METADATA_TIMEOUT)
    }

    fn with_metadata_url_and_timeout(metadata_url: String, timeout: Duration) -> Self {
        Self {
            metadata_url,
            timeout,
            client: Client::new(),
        }
    }

    async fn get(&self, url: String) -> Result<Vec<u8>, CredentialsError> {
        let request = Request::get(url).body(Body::empty()).map_err(|e| {
            CredentialsError::new(format!("build CVM metadata request failed: {e}"))
        })?;
        let response = self
            .client
            .request(request)
            .await
            .map_err(|e| CredentialsError::new(format!("CVM metadata request failed: {e}")))?;
        let status = response.status();
        if !status.is_success() {
            return Err(CredentialsError::new(format!(
                "CVM metadata request returned HTTP {status}"
            )));
        }
        let body = hyper::body::to_bytes(response.into_body())
            .await
            .map_err(|e| {
                CredentialsError::new(format!("read CVM metadata response failed: {e}"))
            })?;
        Ok(body.to_vec())
    }

    async fn get_role_name(&self) -> Result<String, CredentialsError> {
        let body = self.get(format!("{}/", self.metadata_url)).await?;
        let role_name = String::from_utf8(body)
            .map_err(|e| CredentialsError::new(format!("CVM role name is not UTF-8: {e}")))?;
        let role_name = role_name.trim();
        if role_name.is_empty() {
            return Err(CredentialsError::new(
                "CVM metadata returned an empty role name",
            ));
        }
        Ok(role_name.to_owned())
    }

    async fn load_credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let role_name = self.get_role_name().await?;
        let body = self
            .get(format!("{}/{}", self.metadata_url, role_name))
            .await?;
        let response: CvmRoleCredentials = serde_json::from_slice(&body).map_err(|e| {
            CredentialsError::new(format!("parse CVM metadata response failed: {e}"))
        })?;
        if response.code != "Success" {
            return Err(CredentialsError::new(format!(
                "CVM metadata credential request failed with code {}",
                response.code
            )));
        }
        if response.tmp_secret_id.is_empty()
            || response.tmp_secret_key.is_empty()
            || response.token.is_empty()
        {
            return Err(CredentialsError::new(
                "CVM metadata response contains empty credentials",
            ));
        }
        let expiration = parse_expiration(
            response.expired_time,
            response.expiration,
            "CVM metadata response",
            "",
        )?;
        Ok(AwsCredentials::new(
            response.tmp_secret_id,
            response.tmp_secret_key,
            Some(response.token),
            Some(expiration),
        ))
    }
}

#[async_trait]
impl ProvideAwsCredentials for CvmRoleCredentialsProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        tokio::time::timeout(self.timeout, self.load_credentials())
            .await
            .map_err(|_| CredentialsError::new("CVM metadata credential lookup timed out"))?
    }
}

impl AssumeRoleWithWebIdentityProvider {
    pub fn new() -> Result<Self, CredentialsError> {
        let region = read_region()?;
        let provider_id = read_env_required(ENV_TKE_PROVIDER_ID)?;
        let role_arn = read_env_required(ENV_TKE_ROLE_ARN)?;
        let token_file = read_env_required(ENV_TKE_WEB_IDENTITY_TOKEN_FILE)?;
        let client = Client::builder().build(HttpsConnector::new());
        Ok(Self {
            region,
            provider_id,
            role_arn,
            token_file,
            client,
        })
    }
}

#[async_trait]
impl ProvideAwsCredentials for AssumeRoleWithWebIdentityProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let token = fs::read_to_string(&self.token_file)
            .map_err(|e| CredentialsError::new(format!("read oidc token file failed: {}", e)))?;
        let token = token.trim();
        if token.is_empty() {
            return Err(CredentialsError::new("oidc token is empty"));
        }

        let payload = AssumeRoleWithWebIdentityRequest {
            provider_id: &self.provider_id,
            web_identity_token: token,
            role_arn: &self.role_arn,
            role_session_name: DEFAULT_SESSION_NAME,
        };
        let payload = serde_json::to_string(&payload)
            .map_err(|e| CredentialsError::new(format!("serialize sts request failed: {}", e)))?;

        let timestamp = Utc::now().timestamp();
        let (sts_endpoint, sts_host) = build_sts_endpoint(&self.region);
        let req = Request::post(sts_endpoint.as_str())
            .header("Content-Type", "application/json")
            .header("Host", sts_host.as_str())
            .header("X-TC-Action", STS_ACTION_ASSUME_ROLE_WITH_WEB_IDENTITY)
            .header("X-TC-Version", STS_VERSION)
            .header("X-TC-Region", &self.region)
            .header("X-TC-Timestamp", timestamp.to_string())
            .header("Authorization", "SKIP")
            .body(Body::from(payload))
            .map_err(|e| CredentialsError::new(format!("build sts request failed: {}", e)))?;

        let resp = self
            .client
            .request(req)
            .await
            .map_err(|e| CredentialsError::new(format!("sts request failed: {}", e)))?;
        let status = resp.status();
        let body = hyper::body::to_bytes(resp.into_body())
            .await
            .map_err(|e| CredentialsError::new(format!("read sts response failed: {}", e)))?;

        if !status.is_success() {
            return Err(CredentialsError::new(format!(
                "sts http error {}: {}{}",
                status,
                String::from_utf8_lossy(&body),
                format_request_id(None)
            )));
        }

        let resp: TencentResponse<AssumeRoleWithWebIdentityResult> = serde_json::from_slice(&body)
            .map_err(|e| CredentialsError::new(format!("parse sts response failed: {}", e)))?;
        let response = resp.response;
        let request_id_suffix = format_request_id(response.request_id.as_deref());
        if let Some(err) = response.error {
            return Err(CredentialsError::new(format!(
                "sts error {}: {}{}",
                err.code, err.message, request_id_suffix
            )));
        }
        let creds = response.data.credentials.ok_or_else(|| {
            CredentialsError::new(format!(
                "sts response missing credentials field{}",
                request_id_suffix
            ))
        })?;
        let expire_at = parse_expiration(
            response.data.expired_time,
            response.data.expiration,
            "sts response",
            &request_id_suffix,
        )?;

        Ok(AwsCredentials::new(
            creds.tmp_secret_id,
            creds.tmp_secret_key,
            Some(creds.token),
            Some(expire_at),
        ))
    }
}

pub fn new_credential_provider() -> Result<ActiveRefreshingProvider, CredentialsError> {
    let mut provider = ChainProvider::default();
    match AssumeRoleWithWebIdentityProvider::new() {
        Ok(web_identity_provider) => provider.push(TKE_CREDENTIAL_PROVIDER, web_identity_provider),
        Err(error) => provider.push_initialization_error(TKE_CREDENTIAL_PROVIDER, error),
    }
    provider.push(CVM_CREDENTIAL_PROVIDER, CvmRoleCredentialsProvider::new());
    Ok(ActiveRefreshingProvider::new(Arc::new(provider)))
}

pub async fn decrypt_master_key(
    cipher_text: &str,
    region: &str,
    endpoint: &str,
) -> Result<Vec<u8>, CredentialsError> {
    let provider = new_credential_provider()?;
    let region = if region.is_empty() {
        read_region()?
    } else {
        region.to_owned()
    };
    let creds = provider.credentials().await?;
    let payload = DecryptRequest {
        ciphertext_blob: cipher_text,
    };
    let payload = serde_json::to_string(&payload)
        .map_err(|e| CredentialsError::new(format!("serialize kms request failed: {}", e)))?;
    let resp = call_kms(&region, endpoint, &creds, KMS_ACTION_DECRYPT, payload)
        .await
        .map_err(|e| CredentialsError::new(format!("kms decrypt failed: {}", e)))?;
    let request_id_suffix = format_request_id(resp.request_id.as_deref());
    let plaintext = resp.data.plaintext.ok_or_else(|| {
        CredentialsError::new(format!(
            "kms response missing plaintext field{}",
            request_id_suffix
        ))
    })?;
    let plaintext = base64::decode(plaintext).map_err(|e| {
        CredentialsError::new(format!(
            "decode plaintext failed: {}{}",
            e, request_id_suffix
        ))
    })?;
    Ok(plaintext)
}

async fn call_kms(
    region: &str,
    endpoint_override: &str,
    creds: &AwsCredentials,
    action: &str,
    payload: String,
) -> Result<TencentResponseInner<DecryptResult>, TencentError> {
    let endpoint = if endpoint_override.is_empty() {
        KMS_ENDPOINT.to_owned()
    } else {
        endpoint_override.to_owned()
    };
    let uri: Uri = endpoint.parse()?;
    let host = uri
        .host()
        .map(|s| s.to_owned())
        .unwrap_or_else(|| KMS_HOST.to_owned());
    let timestamp = Utc::now().timestamp();
    let authorization = build_authorization(
        creds.aws_access_key_id(),
        creds.aws_secret_access_key(),
        timestamp,
        &host,
        &payload,
    )?;

    // Ref https://cloud.tencent.com/document/api/573/34429 for Decrypt API.
    let client = Client::builder().build(HttpsConnector::new());
    let mut req = Request::post(endpoint)
        .header("Content-Type", "application/json")
        .header("Host", &host)
        .header("X-TC-Action", action)
        .header("X-TC-Version", KMS_VERSION)
        .header("X-TC-Region", region)
        .header("X-TC-Timestamp", timestamp.to_string())
        .header("Authorization", authorization)
        .body(Body::from(payload))
        .map_err(|e| TencentError::InvalidResponse(e.to_string()))?;

    if let Some(token) = creds.token() {
        let header = HeaderValue::from_str(token)
            .map_err(|e| TencentError::InvalidResponse(e.to_string()))?;
        req.headers_mut().insert("X-TC-Token", header);
    }

    let resp = client.request(req).await?;
    let status = resp.status();
    let body = hyper::body::to_bytes(resp.into_body()).await?;
    if !status.is_success() {
        return Err(TencentError::InvalidResponse(format!(
            "kms http error {}: {}{}",
            status,
            String::from_utf8_lossy(&body),
            format_request_id(None)
        )));
    }
    let resp: TencentResponse<DecryptResult> = serde_json::from_slice(&body)?;
    let response = resp.response;
    let request_id_suffix = format_request_id(response.request_id.as_deref());
    if let Some(err) = response.error {
        return Err(TencentError::InvalidResponse(format!(
            "kms error {}: {}{}",
            err.code, err.message, request_id_suffix
        )));
    }
    Ok(response)
}

fn parse_expiration(
    expired_time: Option<i64>,
    expiration: Option<String>,
    source: &str,
    request_id_suffix: &str,
) -> Result<DateTime<Utc>, CredentialsError> {
    let mut last_error = None;
    if let Some(ts) = expired_time {
        match Utc.timestamp_opt(ts, 0).single() {
            Some(parsed) => return Ok(parsed),
            None => {
                last_error = Some(format!(
                    "{source} has invalid ExpiredTime {ts}{request_id_suffix}"
                ));
            }
        }
    }
    if let Some(exp) = expiration {
        match DateTime::parse_from_rfc3339(&exp) {
            Ok(parsed) => return Ok(parsed.with_timezone(&Utc)),
            Err(err) => {
                last_error = Some(format!(
                    "{source} has invalid Expiration {exp}: {err}{request_id_suffix}"
                ));
            }
        }
    }
    Err(CredentialsError::new(last_error.unwrap_or_else(|| {
        format!("{source} is missing ExpiredTime/Expiration fields{request_id_suffix}")
    })))
}

fn read_env_required(name: &str) -> Result<String, CredentialsError> {
    match std::env::var(name) {
        Ok(v) if !v.is_empty() => Ok(v),
        Ok(_) => Err(CredentialsError::new(format!("{} is empty", name))),
        Err(_) => Err(CredentialsError::new(format!("{} is missing", name))),
    }
}

fn read_region() -> Result<String, CredentialsError> {
    if let Ok(region) = std::env::var(ENV_TKE_REGION) {
        if !region.is_empty() {
            return Ok(region);
        }
    }
    if let Ok(region) = std::env::var(ENV_TKE_DEFAULT_REGION) {
        if !region.is_empty() {
            return Ok(region);
        }
    }
    Err(CredentialsError::new("region is missing"))
}

fn build_sts_endpoint(region: &str) -> (String, String) {
    let host = format!("sts.{}.tencentcloudapi.com", region);
    let endpoint = format!("https://{}/", host);
    (endpoint, host)
}

fn build_authorization(
    secret_id: &str,
    secret_key: &str,
    timestamp: i64,
    host: &str,
    payload: &str,
) -> Result<String, TencentError> {
    let date = Utc
        .timestamp_opt(timestamp, 0)
        .single()
        .ok_or_else(|| TencentError::InvalidResponse("invalid timestamp".to_owned()))?
        .format("%Y-%m-%d")
        .to_string();
    let hashed_payload = sha256_hex(payload);
    let canonical_headers = format!("content-type:application/json\nhost:{}\n", host);
    let signed_headers = "content-type;host";
    let canonical_request = format!(
        "POST\n/\n\n{}\n{}\n{}",
        canonical_headers, signed_headers, hashed_payload
    );
    let hashed_canonical_request = sha256_hex(&canonical_request);
    let credential_scope = format!("{}/{}/tc3_request", date, KMS_SERVICE);
    let string_to_sign = format!(
        "TC3-HMAC-SHA256\n{}\n{}\n{}",
        timestamp, credential_scope, hashed_canonical_request
    );

    let secret_date = hmac_sha256(format!("TC3{}", secret_key).as_bytes(), &date);
    let secret_service = hmac_sha256(&secret_date, KMS_SERVICE);
    let secret_signing = hmac_sha256(&secret_service, "tc3_request");
    let signature = hex::encode(hmac_sha256(&secret_signing, &string_to_sign));

    Ok(format!(
        "TC3-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}",
        secret_id, credential_scope, signed_headers, signature
    ))
}

fn format_request_id(request_id: Option<&str>) -> String {
    match request_id {
        Some(id) if !id.is_empty() => format!(", RequestId: {}", id),
        _ => ", RequestId: unknown".to_owned(),
    }
}

fn sha256_hex(input: &str) -> String {
    hex::encode(Sha256::digest(input.as_bytes()))
}

fn hmac_sha256(key: &[u8], input: &str) -> Vec<u8> {
    let mut mac = Hmac::<Sha256>::new_from_slice(key).expect("hmac can take key of any size");
    mac.update(input.as_bytes());
    mac.finalize().into_bytes().to_vec()
}

#[cfg(test)]
mod tests {
    use std::{
        collections::VecDeque, convert::Infallible, net::TcpListener, sync::Mutex, time::Duration,
    };

    use async_trait::async_trait;
    use hyper::{
        Body, Response, Server,
        service::{make_service_fn, service_fn},
    };
    use rusoto_credential::{AwsCredentials, CredentialsError, ProvideAwsCredentials};

    use super::{
        CVM_CREDENTIAL_PROVIDER, ChainProvider, CvmRoleCredentialsProvider,
        TKE_CREDENTIAL_PROVIDER, build_authorization,
    };

    enum MockCredentialResult {
        Credentials(&'static str),
        Error(&'static str),
    }

    struct MockCredentialProvider {
        results: Mutex<VecDeque<MockCredentialResult>>,
    }

    impl MockCredentialProvider {
        fn new(results: Vec<MockCredentialResult>) -> Self {
            Self {
                results: Mutex::new(results.into()),
            }
        }
    }

    #[async_trait]
    impl ProvideAwsCredentials for MockCredentialProvider {
        async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
            let result = self
                .results
                .lock()
                .unwrap()
                .pop_front()
                .expect("mock credential provider called too many times");
            match result {
                MockCredentialResult::Credentials(access_key) => {
                    Ok(AwsCredentials::new(access_key, "secret", None, None))
                }
                MockCredentialResult::Error(error) => Err(CredentialsError::new(error)),
            }
        }
    }

    #[test]
    fn test_build_authorization() {
        let auth = build_authorization(
            "TENCENTCLOUD_SECRET_ID",
            "TENCENTCLOUD_SECRET_KEY",
            1551113065,
            "kms.ap-beijing.tencentcloudapi.com",
            r#"{"CiphertextBlob":"TEST"}"#,
        )
        .expect("build authorization");

        assert_eq!(
            auth,
            "TC3-HMAC-SHA256 Credential=TENCENTCLOUD_SECRET_ID/2019-02-25/kms/tc3_request, SignedHeaders=content-type;host, Signature=1ae5c203f5993836b3566b9f034554c71444fc248f5b21eea280a0490538398b"
        );
    }

    #[tokio::test]
    async fn test_chain_provider_tracks_provider_switches() {
        let mut provider = ChainProvider::default();
        provider.push(
            TKE_CREDENTIAL_PROVIDER,
            MockCredentialProvider::new(vec![
                MockCredentialResult::Error("TKE unavailable"),
                MockCredentialResult::Credentials("tke-key"),
                MockCredentialResult::Error("TKE unavailable again"),
            ]),
        );
        provider.push(
            CVM_CREDENTIAL_PROVIDER,
            MockCredentialProvider::new(vec![
                MockCredentialResult::Credentials("cvm-key-1"),
                MockCredentialResult::Credentials("cvm-key-2"),
            ]),
        );

        let credentials = provider.credentials().await.unwrap();
        assert_eq!(credentials.aws_access_key_id(), "cvm-key-1");
        assert_eq!(
            *provider.selected_provider.lock().unwrap(),
            Some(CVM_CREDENTIAL_PROVIDER)
        );

        let credentials = provider.credentials().await.unwrap();
        assert_eq!(credentials.aws_access_key_id(), "tke-key");
        assert_eq!(
            *provider.selected_provider.lock().unwrap(),
            Some(TKE_CREDENTIAL_PROVIDER)
        );

        let credentials = provider.credentials().await.unwrap();
        assert_eq!(credentials.aws_access_key_id(), "cvm-key-2");
        assert_eq!(
            *provider.selected_provider.lock().unwrap(),
            Some(CVM_CREDENTIAL_PROVIDER)
        );
    }

    #[tokio::test]
    async fn test_chain_provider_aggregates_runtime_errors() {
        let mut provider = ChainProvider::default();
        provider.push(
            TKE_CREDENTIAL_PROVIDER,
            MockCredentialProvider::new(vec![MockCredentialResult::Error("STS failed")]),
        );
        provider.push(
            CVM_CREDENTIAL_PROVIDER,
            MockCredentialProvider::new(vec![MockCredentialResult::Error("metadata failed")]),
        );

        let error = provider.credentials().await.unwrap_err();
        assert_eq!(
            error.to_string(),
            "all credential providers failed: TKE workload identity: STS failed; CVM instance role: metadata failed"
        );
    }

    #[tokio::test]
    async fn test_chain_provider_includes_initialization_errors() {
        let mut provider = ChainProvider::default();
        provider.push_initialization_error(
            TKE_CREDENTIAL_PROVIDER,
            CredentialsError::new("TKE_ROLE_ARN is missing"),
        );
        provider.push(
            CVM_CREDENTIAL_PROVIDER,
            MockCredentialProvider::new(vec![MockCredentialResult::Error("metadata failed")]),
        );

        let error = provider.credentials().await.unwrap_err();
        assert_eq!(
            error.to_string(),
            "all credential providers failed: TKE workload identity: TKE_ROLE_ARN is missing; CVM instance role: metadata failed"
        );
    }

    #[tokio::test]
    async fn test_cvm_role_provider_gets_role_and_credentials() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        listener.set_nonblocking(true).unwrap();
        let service = make_service_fn(|_| async {
            Ok::<_, Infallible>(service_fn(|request| async move {
                let body = match request.uri().path() {
                    "/latest/meta-data/cam/security-credentials/" => "test-role\n".to_owned(),
                    "/latest/meta-data/cam/security-credentials/test-role" => serde_json::json!({
                        "TmpSecretId": "temporary-secret-id",
                        "TmpSecretKey": "temporary-secret-key",
                        "ExpiredTime": 4_102_444_800_i64,
                        "Expiration": "2100-01-01T00:00:00Z",
                        "Token": "temporary-token",
                        "Code": "Success"
                    })
                    .to_string(),
                    path => panic!("unexpected metadata path: {path}"),
                };
                Ok::<_, Infallible>(Response::new(Body::from(body)))
            }))
        });
        let server = Server::from_tcp(listener).unwrap().serve(service);
        let server_task = tokio::spawn(server);

        let provider = CvmRoleCredentialsProvider::with_metadata_url(format!(
            "http://{address}/latest/meta-data/cam/security-credentials"
        ));
        let credentials = provider.credentials().await.unwrap();

        assert_eq!(credentials.aws_access_key_id(), "temporary-secret-id");
        assert_eq!(credentials.aws_secret_access_key(), "temporary-secret-key");
        assert_eq!(credentials.token().as_deref(), Some("temporary-token"));
        assert_eq!(credentials.expires_at().unwrap().timestamp(), 4_102_444_800);

        server_task.abort();
    }

    #[tokio::test]
    async fn test_cvm_role_provider_applies_total_timeout() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        listener.set_nonblocking(true).unwrap();
        let service = make_service_fn(|_| async {
            Ok::<_, Infallible>(service_fn(|request| async move {
                let body = match request.uri().path() {
                    "/latest/meta-data/cam/security-credentials/" => {
                        tokio::time::sleep(Duration::from_millis(75)).await;
                        "test-role".to_owned()
                    }
                    "/latest/meta-data/cam/security-credentials/test-role" => {
                        tokio::time::sleep(Duration::from_secs(1)).await;
                        unreachable!("the total metadata timeout should expire first")
                    }
                    path => panic!("unexpected metadata path: {path}"),
                };
                Ok::<_, Infallible>(Response::new(Body::from(body)))
            }))
        });
        let server = Server::from_tcp(listener).unwrap().serve(service);
        let server_task = tokio::spawn(server);

        let provider = CvmRoleCredentialsProvider::with_metadata_url_and_timeout(
            format!("http://{address}/latest/meta-data/cam/security-credentials"),
            Duration::from_millis(100),
        );
        let error = provider.credentials().await.unwrap_err();

        assert_eq!(
            error.to_string(),
            "CVM metadata credential lookup timed out"
        );

        server_task.abort();
    }
}
