// Copyright 2025 PingCAP, Inc.

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

use async_trait::async_trait;
use aws::ActiveRefreshingProvider;
use chrono::{DateTime, Utc};
use cloud::kms::{DataKeyPair, EncryptedKey, PlainKey};
use hmac::{Hmac, Mac};
use hyper::{Body, Client, Method, Request, client::HttpConnector, header::HeaderValue};
use hyper_tls::HttpsConnector;
use rusoto_credential::{AwsCredentials, CredentialsError, ProvideAwsCredentials};
use serde_derive::{Deserialize, Serialize};
use sha1::Sha1;

pub const OIDC_PROVIDER_ARN: &str = "ALIBABA_CLOUD_OIDC_PROVIDER_ARN";
pub const ROLE_ARN: &str = "ALIBABA_CLOUD_ROLE_ARN";
pub const TOKEN_FILE: &str = "ALIBABA_CLOUD_OIDC_TOKEN_FILE";
pub const REGION_ID: &str = "ALIBABA_CLOUD_REGION_ID";
pub const DOMAIN_STRING: &str = "aliyuncs.com";
pub const ECS_METADATA_DISABLED: &str = "ALIBABA_CLOUD_ECS_METADATA_DISABLED";
pub const ECS_METADATA_DISABLED_V1: &str = "ALIBABA_CLOUD_IMDSV1_DISABLED";
pub const ECS_ROLE_NAME: &str = "ALIBABA_CLOUD_ECS_METADATA";

const IMDS_BASE_URL: &str = "http://100.100.100.200";
const DEFAULT_METADATA_TOKEN_TTL_SECS: u64 = 21600;

pub const TIMEOUT: Duration = Duration::from_secs(60);

#[derive(Default)]
pub struct ChainProvider {
    providers: Vec<Arc<dyn ProvideAwsCredentials + Send + Sync>>,
}

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

#[async_trait]
impl ProvideAwsCredentials for ChainProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        for p in self.providers.iter() {
            match p.credentials().await {
                Ok(c) => return Ok(c),
                Err(_) => {
                    // ignore this credential error, try next
                    continue;
                }
            }
        }
        Err(CredentialsError::new("all credential providers failed"))
    }
}

pub struct EcsRamRoleCredentialsProvider {
    role_name: Option<String>,
    disable_imdsv1: bool,
    client: hyper::Client<HttpsConnector<HttpConnector>>,
}
#[derive(Deserialize, Debug)]
struct EcsRamRoleCredentials {
    #[serde(rename = "Code")]
    code: Option<String>,
    #[serde(rename = "AccessKeyId")]
    access_key_id: Option<String>,
    #[serde(rename = "AccessKeySecret")]
    access_key_secret: Option<String>,
    #[serde(rename = "SecurityToken")]
    security_token: Option<String>,
    #[serde(rename = "Expiration")]
    expiration: Option<String>,
}

pub struct AssumeRoleWithOidcProvider {
    region: String,
    oidc_provider_arn: String,
    role_arn: String,
    oidc_token_file: String,
    client: hyper::Client<HttpsConnector<HttpConnector>>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
#[serde(rename_all = "PascalCase")]
struct AssumeRoleResponse {
    credentials: Credentials,
}

#[derive(Serialize, Deserialize, Default, Debug)]
#[serde(rename_all = "PascalCase")]
struct Credentials {
    access_key_id: String,
    access_key_secret: String,
    security_token: String,
    expiration: String,
}

impl AssumeRoleWithOidcProvider {
    pub fn new() -> Result<AssumeRoleWithOidcProvider, CredentialsError> {
        let region = std::env::var(REGION_ID)
            .map_err(|_| CredentialsError::new("failed to get region id from env"))?;
        let oidc_provider_arn = std::env::var(OIDC_PROVIDER_ARN)
            .map_err(|_| CredentialsError::new("failed to get oidc provider arn from env"))?;
        let role_arn = std::env::var(ROLE_ARN)
            .map_err(|_| CredentialsError::new("failed to get role arn from env"))?;
        let oidc_token_file = std::env::var(TOKEN_FILE)
            .map_err(|_| CredentialsError::new("failed to get oidc token file from env"))?;
        let client = hyper::Client::builder().build(hyper_tls::HttpsConnector::new());
        Ok(AssumeRoleWithOidcProvider {
            region,
            oidc_provider_arn,
            role_arn,
            oidc_token_file,
            client,
        })
    }
}

#[async_trait]
impl ProvideAwsCredentials for AssumeRoleWithOidcProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let token = String::from_utf8(fs::read(&self.oidc_token_file)?)?;
        let timestamp = Utc::now().to_rfc3339();
        let params = vec![
            ("Action", "AssumeRoleWithOIDC"),
            ("Format", "JSON"),
            ("Version", "2015-04-01"),
            ("RoleSessionName", "TiKV"),
            ("OIDCProviderArn", &self.oidc_provider_arn),
            ("RoleArn", &self.role_arn),
            ("Timestamp", &timestamp),
            ("OIDCToken", &token),
        ];
        let query_string = build_query(&params);
        let uri = format!(
            "https://sts.{}.{}/?{}",
            self.region, DOMAIN_STRING, query_string,
        );
        let request = hyper::Request::post(uri)
            .body(hyper::Body::empty())
            .unwrap();
        let resp = self.client.request(request).await?;
        if !resp.status().is_success() {
            return Err(CredentialsError::new(format!(
                "failed to AssumeRoleWithOIDC: {}",
                resp.status()
            )));
        }
        let body = hyper::body::to_bytes(resp.into_body()).await?;
        let resp: AssumeRoleResponse = serde_json::from_slice(&body)?;
        let expiration = chrono::DateTime::parse_from_rfc3339(&resp.credentials.expiration)?;
        Ok(AwsCredentials::new(
            resp.credentials.access_key_id,
            resp.credentials.access_key_secret,
            Some(resp.credentials.security_token),
            Some(expiration.with_timezone(&Utc)),
        ))
    }
}

fn is_truthy(v: &str) -> bool {
    matches!(v.to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on")
}

impl EcsRamRoleCredentialsProvider {
    pub fn new() -> Result<EcsRamRoleCredentialsProvider, CredentialsError> {
        if let Ok(v) = std::env::var(ECS_METADATA_DISABLED) {
            if is_truthy(&v) {
                return Err(CredentialsError::new("ECS metadata is disabled by env"));
            }
        }

        let disable_imdsv1: bool = std::env::var(ECS_METADATA_DISABLED_V1)
            .ok()
            .map(|v| is_truthy(&v))
            .unwrap_or(false);

        let role_name: Option<String> = std::env::var(ECS_ROLE_NAME).ok();

        let https = HttpsConnector::new();
        let client: Client<_, Body> = Client::builder().build(https);

        Ok(EcsRamRoleCredentialsProvider {
            role_name,
            disable_imdsv1,
            client,
        })
    }

    async fn get_metadata_token(&self) -> Result<Option<String>, CredentialsError> {
        let url = format!("{IMDS_BASE_URL}/latest/api/token");
        let mut req = Request::builder()
            .method(Method::PUT)
            .uri(url)
            .body(Body::empty())
            .map_err(|e| CredentialsError::new(format!("build IMDS token request: {e}")))?;
        // X-aliyun-ecs-metadata-token-ttl-seconds
        req.headers_mut().insert(
            "X-aliyun-ecs-metadata-token-ttl-seconds",
            HeaderValue::from_str(&DEFAULT_METADATA_TOKEN_TTL_SECS.to_string())
                .map_err(|e| CredentialsError::new(format!("invalid ttl header: {e}")))?,
        );

        // Try IMDSv2 first; if it fails, attempt IMDSv1 unless explicitly disabled
        let resp = match tokio::time::timeout(TIMEOUT, self.client.request(req)).await {
            Ok(Ok(r)) => r,
            Ok(Err(e)) => {
                if self.disable_imdsv1 {
                    return Err(CredentialsError::new(format!(
                        "get metadata token failed: {e}"
                    )));
                } else {
                    return Ok(None);
                }
            }
            Err(_) => {
                if self.disable_imdsv1 {
                    return Err(CredentialsError::new("get metadata token timeout"));
                } else {
                    return Ok(None);
                }
            }
        };

        let status = resp.status();
        let body =
            match tokio::time::timeout(TIMEOUT, hyper::body::to_bytes(resp.into_body())).await {
                Ok(Ok(b)) => b,
                Ok(Err(e)) => {
                    if self.disable_imdsv1 {
                        return Err(CredentialsError::new(format!(
                            "read token body failed: {e}"
                        )));
                    } else {
                        return Ok(None);
                    }
                }
                Err(_) => {
                    if self.disable_imdsv1 {
                        return Err(CredentialsError::new("read token body timeout"));
                    } else {
                        return Ok(None);
                    }
                }
            };

        if !status.is_success() {
            if self.disable_imdsv1 {
                return Err(CredentialsError::new(format!(
                    "refresh ECS sts token err (IMDSv2 token), httpStatus: {status}, body={}",
                    String::from_utf8_lossy(&body)
                )));
            } else {
                return Ok(None);
            }
        }

        let token = String::from_utf8(body.to_vec())
            .map_err(|e| CredentialsError::new(format!("token utf8 error: {e}")))?;
        Ok(Some(token.trim().to_string()))
    }

    async fn get_role_name(&self) -> Result<String, CredentialsError> {
        if let Some(name) = &self.role_name {
            return Ok(name.clone());
        }

        let url = format!("{IMDS_BASE_URL}/latest/meta-data/ram/security-credentials/");
        let mut req = Request::builder()
            .method(Method::GET)
            .uri(url)
            .body(Body::empty())
            .map_err(|e| CredentialsError::new(format!("build roleName request: {e}")))?;

        if let Some(token) = self.get_metadata_token().await? {
            req.headers_mut().insert(
                "X-aliyun-ecs-metadata-token",
                HeaderValue::from_str(&token)
                    .map_err(|e| CredentialsError::new(format!("set token header: {e}")))?,
            );
        }

        let resp = tokio::time::timeout(TIMEOUT, self.client.request(req))
            .await
            .map_err(|_| CredentialsError::new("get role name timeout"))?
            .map_err(|e| CredentialsError::new(format!("get role name failed: {e}")))?;

        if resp.status() != hyper::StatusCode::OK {
            return Err(CredentialsError::new(format!(
                "get role name failed: http {}",
                resp.status()
            )));
        }

        let body = tokio::time::timeout(TIMEOUT, hyper::body::to_bytes(resp.into_body()))
            .await
            .map_err(|_| CredentialsError::new("read role name body timeout"))?
            .map_err(|e| CredentialsError::new(format!("read role name body: {e}")))?;

        let name = String::from_utf8(body.to_vec())
            .map_err(|e| CredentialsError::new(format!("role name utf8 error: {e}")))?;
        Ok(name.trim().to_string())
    }

    async fn get_session_credentials(
        &self,
    ) -> Result<(String, String, String, DateTime<Utc>), CredentialsError> {
        let role_name = self.get_role_name().await?;
        let url = format!("{IMDS_BASE_URL}/latest/meta-data/ram/security-credentials/{role_name}");
        let mut req = Request::builder()
            .method(Method::GET)
            .uri(url)
            .body(Body::empty())
            .map_err(|e| CredentialsError::new(format!("build creds request: {e}")))?;

        if let Some(token) = self.get_metadata_token().await? {
            req.headers_mut().insert(
                "X-aliyun-ecs-metadata-token",
                HeaderValue::from_str(&token)
                    .map_err(|e| CredentialsError::new(format!("set token header: {e}")))?,
            );
        }

        let resp = tokio::time::timeout(TIMEOUT, self.client.request(req))
            .await
            .map_err(|_| CredentialsError::new("get creds timeout"))?
            .map_err(|e| CredentialsError::new(format!("refresh ECS sts token err: {e}")))?;

        let status = resp.status();
        let body = tokio::time::timeout(TIMEOUT, hyper::body::to_bytes(resp.into_body()))
            .await
            .map_err(|_| CredentialsError::new("read creds body timeout"))?
            .map_err(|e| CredentialsError::new(format!("read creds body: {e}")))?;

        if status != hyper::StatusCode::OK {
            return Err(CredentialsError::new(format!(
                "refresh ECS sts token err, httpStatus: {status}, message={}",
                String::from_utf8_lossy(&body)
            )));
        }

        let data: EcsRamRoleCredentials = serde_json::from_slice(&body)
            .map_err(|e| CredentialsError::new(format!("json unmarshal fail: {e}")))?;

        if data.code.as_deref() != Some("Success") {
            return Err(CredentialsError::new(format!(
                "refresh ECS sts token err, Code:{} is not Success",
                data.code.as_deref().unwrap_or("None")
            )));
        }

        let (ak, sk, token, exp) = match (
            data.access_key_id,
            data.access_key_secret,
            data.security_token,
            data.expiration,
        ) {
            (Some(ak), Some(sk), Some(tok), Some(exp)) => (ak, sk, tok, exp),
            _ => {
                return Err(CredentialsError::new(
                    "refresh ECS sts token err, fail to get credentials",
                ));
            }
        };

        let expiration = chrono::DateTime::parse_from_rfc3339(&exp)
            .map_err(|e| CredentialsError::new(format!("parse expiration: {e}")))?
            .with_timezone(&Utc);

        Ok((ak, sk, token, expiration))
    }
}

#[async_trait]
impl ProvideAwsCredentials for EcsRamRoleCredentialsProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let (ak, sk, token, expiration) = self.get_session_credentials().await?;
        Ok(AwsCredentials::new(ak, sk, Some(token), Some(expiration)))
    }
}

pub fn new_credential_provider() -> Result<ActiveRefreshingProvider, CredentialsError> {
    let mut chain_provider = ChainProvider::default();
    if let Ok(oidc_provider) = AssumeRoleWithOidcProvider::new() {
        chain_provider.push(oidc_provider);
    }
    if let Ok(ecs_ram_provider) = EcsRamRoleCredentialsProvider::new() {
        chain_provider.push(ecs_ram_provider);
    }
    let auto_refreshing_provider = ActiveRefreshingProvider::new(Arc::new(chain_provider));
    Ok(auto_refreshing_provider)
}

// Aliyun KMS: GenerateDataKey
const ALIYUN_KMS_VERSION: &str = "2016-01-20";
const ALIYUN_KMS_DATA_KEY_SPEC: &str = "AES_256";

fn build_aliyun_kms_uri(region: &str, signed_query: &str) -> Result<String, CredentialsError> {
    if !region
        .split('-')
        .all(|part| !part.is_empty() && part.bytes().all(|b| b.is_ascii_alphanumeric()))
    {
        return Err(CredentialsError::new(format!(
            "invalid aliyun kms region: {region}"
        )));
    }
    Ok(format!(
        "https://kms.{}.aliyuncs.com/?{}",
        region, signed_query
    ))
}

fn build_aliyun_kms_uri_with_endpoint(
    region: &str,
    endpoint_override: &str,
    signed_query: &str,
) -> Result<String, CredentialsError> {
    if endpoint_override.is_empty() {
        build_aliyun_kms_uri(region, signed_query)
    } else {
        Ok(format!("{}/?{}", endpoint_override, signed_query))
    }
}

pub async fn generate_data_key(
    key_id: &str,
    region: &str,
    endpoint_override: &str,
    encryption_context: &BTreeMap<String, String>,
) -> Result<DataKeyPair, CredentialsError> {
    let credential_provider = new_credential_provider()?;
    let cred = credential_provider.credentials().await?;
    let fallback_region;
    let region = if !region.is_empty() {
        region
    } else if endpoint_override.is_empty() {
        fallback_region = std::env::var(REGION_ID)
            .map_err(|_| CredentialsError::new("failed to get region id from env"))?;
        fallback_region.as_str()
    } else {
        ""
    };
    let resp = generate_data_key_with_credential(
        key_id,
        endpoint_override,
        &cred,
        region,
        encryption_context,
    )
    .await?;
    decode_generate_data_key_response(resp)
}

#[derive(Serialize, Deserialize, Default, Debug)]
#[serde(rename_all = "PascalCase")]
struct GenerateDataKeyResponse {
    ciphertext_blob: String,
    plaintext: String,
}

fn decode_generate_data_key_response(
    resp: GenerateDataKeyResponse,
) -> Result<DataKeyPair, CredentialsError> {
    // Aliyun KMS returns base64-encoded plaintext/ciphertext strings.
    let plaintext_key = base64::decode(resp.plaintext)
        .map_err(|e| CredentialsError::new(format!("failed to decode plaintext base64: {e}")))?;
    let ciphertext_key = base64::decode(resp.ciphertext_blob)
        .map_err(|e| CredentialsError::new(format!("failed to decode ciphertext base64: {e}")))?;

    Ok(DataKeyPair {
        encrypted: EncryptedKey::new(ciphertext_key)
            .map_err(|e| CredentialsError::new(format!("invalid ciphertext key: {e}")))?,
        plaintext: PlainKey::new(plaintext_key)
            .map_err(|e| CredentialsError::new(format!("invalid plaintext key: {e}")))?,
    })
}

async fn generate_data_key_with_credential(
    key_id: &str,
    endpoint_override: &str,
    cred: &AwsCredentials,
    region: &str,
    encryption_context: &BTreeMap<String, String>,
) -> Result<GenerateDataKeyResponse, CredentialsError> {
    let mut params = vec![
        ("Action".to_owned(), "GenerateDataKey".to_owned()),
        ("KeyId".to_owned(), key_id.to_owned()),
        ("KeySpec".to_owned(), ALIYUN_KMS_DATA_KEY_SPEC.to_owned()),
        ("Version".to_owned(), ALIYUN_KMS_VERSION.to_owned()),
    ];
    if !encryption_context.is_empty() {
        let encryption_context_json = serde_json::to_string(encryption_context)
            .map_err(|e| CredentialsError::new(format!("serialize EncryptionContext: {e}")))?;
        params.push(("EncryptionContext".to_owned(), encryption_context_json));
    }

    let signed_query = build_signed_query_owned("GET", params, cred);
    let uri = build_aliyun_kms_uri_with_endpoint(region, endpoint_override, &signed_query)?;

    let client = hyper::Client::builder().build(hyper_tls::HttpsConnector::new());
    let request = hyper::Request::get(uri)
        .body(hyper::Body::empty())
        .map_err(|e| CredentialsError::new(format!("build GenerateDataKey request: {e}")))?;

    let resp = client
        .request(request)
        .await
        .map_err(|e| CredentialsError::new(format!("send GenerateDataKey request: {e}")))?;

    let status = resp.status();
    let body = hyper::body::to_bytes(resp.into_body())
        .await
        .map_err(|e| CredentialsError::new(format!("read GenerateDataKey response body: {e}")))?;

    if !status.is_success() {
        return Err(CredentialsError::new(format!(
            "failed to GenerateDataKey: httpStatus={status}, body={}",
            String::from_utf8_lossy(&body)
        )));
    }

    serde_json::from_slice::<GenerateDataKeyResponse>(&body)
        .map_err(|e| CredentialsError::new(format!("parse GenerateDataKey response: {e}")))
}

pub async fn decrypt_master_key(
    cipher_text_blob: &str,
    region: &str,
    endpoint_override: &str,
    encryption_context: &BTreeMap<String, String>,
) -> Result<Vec<u8>, CredentialsError> {
    let credential_provider = new_credential_provider()?;
    let cred = credential_provider.credentials().await?;
    let fallback_region;
    let region = if !region.is_empty() {
        region
    } else {
        fallback_region = std::env::var(REGION_ID)
            .map_err(|_| CredentialsError::new("failed to get region id from env"))?;
        fallback_region.as_str()
    };
    let resp = decrypt_key_with_credential(
        cipher_text_blob,
        endpoint_override,
        &cred,
        region,
        encryption_context,
    )
    .await?;
    base64::decode(resp.plaintext).map_err(|_| CredentialsError::new("failed to decode base64"))
}

async fn decrypt_key_with_credential(
    cipher_text_blob: &str,
    endpoint_override: &str,
    cred: &AwsCredentials,
    region: &str,
    encryption_context: &BTreeMap<String, String>,
) -> Result<DecryptResponse, CredentialsError> {
    let mut params = vec![
        ("Action".to_owned(), "Decrypt".to_owned()),
        ("CiphertextBlob".to_owned(), cipher_text_blob.to_owned()),
        ("Version".to_owned(), ALIYUN_KMS_VERSION.to_owned()),
    ];
    if !encryption_context.is_empty() {
        let encryption_context_json = serde_json::to_string(encryption_context)
            .map_err(|e| CredentialsError::new(format!("serialize EncryptionContext: {e}")))?;
        params.push(("EncryptionContext".to_owned(), encryption_context_json));
    }
    let signed_query = build_signed_query_owned("GET", params, cred);
    let uri = build_aliyun_kms_uri_with_endpoint(region, endpoint_override, &signed_query)?;
    let client = hyper::Client::builder().build(hyper_tls::HttpsConnector::new());
    let request = hyper::Request::get(uri).body(hyper::Body::empty()).unwrap();
    let resp = client.request(request).await?;
    let status = resp.status();
    let body = match hyper::body::to_bytes(resp.into_body()).await {
        Ok(body) => body,
        Err(err) => {
            return Err(CredentialsError::new(format!(
                "failed to decrypt: httpStatus={status}, readBodyError={err}",
            )));
        }
    };
    if !status.is_success() {
        return Err(CredentialsError::new(format!(
            "failed to decrypt: httpStatus={status}, body={}",
            String::from_utf8_lossy(&body)
        )));
    }
    let resp: DecryptResponse = serde_json::from_slice(&body)?;
    Ok(resp)
}

#[derive(Serialize, Deserialize, Default, Debug)]
#[serde(rename_all = "PascalCase")]
struct DecryptResponse {
    plaintext: String,
}

fn build_signed_query_owned(
    method: &str,
    input_params: Vec<(String, String)>,
    cred: &AwsCredentials,
) -> String {
    let timestamp = Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
    build_signed_query_with_timestamp_owned(method, input_params, cred, &timestamp)
}

fn build_signed_query_with_timestamp_owned(
    method: &str,
    mut all_params: Vec<(String, String)>,
    cred: &AwsCredentials,
    timestamp: &str,
) -> String {
    all_params.push((
        "AccessKeyId".to_owned(),
        cred.aws_access_key_id().to_owned(),
    ));
    all_params.push(("Timestamp".to_owned(), timestamp.to_owned()));
    all_params.push(("Format".to_owned(), "JSON".to_owned()));
    all_params.push(("SignatureMethod".to_owned(), "HMAC-SHA1".to_owned()));
    all_params.push(("SignatureVersion".to_owned(), "1.0".to_owned()));
    if let Some(token) = cred.token().as_ref() {
        all_params.push(("SecurityToken".to_owned(), token.to_owned()));
    }
    all_params.sort_by(|a, b| a.0.cmp(&b.0));
    let canonical_query = build_query_owned(&all_params);
    let string_to_sign = format!("{}&%2F&{}", method, percent_encode(&canonical_query));
    let key_secret = format!("{}&", cred.aws_secret_access_key());
    let mut mac =
        Hmac::<Sha1>::new_from_slice(key_secret.as_bytes()).expect("HMAC accepts keys of any size");
    mac.update(string_to_sign.as_bytes());
    let signature = base64::encode(mac.finalize().into_bytes());
    all_params.push(("Signature".to_owned(), signature));
    build_query_owned(&all_params)
}

fn build_query_owned(params: &[(String, String)]) -> String {
    params
        .iter()
        .map(|(key, value)| format!("{}={}", percent_encode(key), percent_encode(value)))
        .collect::<Vec<_>>()
        .join("&")
}

fn build_query(params: &[(&str, &str)]) -> String {
    build_query_owned(
        &params
            .iter()
            .map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
            .collect::<Vec<_>>(),
    )
}

// Percent-encode for aliyun signature
fn percent_encode(input: &str) -> String {
    let mut encoded = String::new();
    for byte in input.as_bytes() {
        match byte {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                encoded.push(*byte as char);
            }
            _ => encoded.push_str(&format!("%{:02X}", byte)),
        }
    }
    encoded
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_percent_encode() {
        assert_eq!(percent_encode("abcABC123-_.~"), "abcABC123-_.~");
        assert_eq!(percent_encode("test/key"), "test%2Fkey");
        assert_eq!(
            percent_encode("2017-02-10T10:10:10Z"),
            "2017-02-10T10%3A10%3A10Z"
        );
        assert_eq!(percent_encode(" "), "%20");
        assert_eq!(percent_encode("*"), "%2A");
    }

    #[test]
    fn test_build_query() {
        let params = vec![
            ("KeyId".to_owned(), "test/key".to_owned()),
            ("Timestamp".to_owned(), "2017-02-10T10:10:10Z".to_owned()),
        ];
        assert_eq!(
            build_query_owned(&params),
            "KeyId=test%2Fkey&Timestamp=2017-02-10T10%3A10%3A10Z"
        );
    }

    #[test]
    fn test_build_signed_query_with_timestamp() {
        let cred = AwsCredentials::new(
            "AKIDEXAMPLE".to_owned(),
            "SECRETEXAMPLE".to_owned(),
            Some("TOKENEXAMPLE".to_owned()),
            None,
        );
        let params = vec![
            ("Action".to_owned(), "GenerateDataKey".to_owned()),
            ("KeyId".to_owned(), "test/key".to_owned()),
            ("KeySpec".to_owned(), ALIYUN_KMS_DATA_KEY_SPEC.to_owned()),
            ("Version".to_owned(), ALIYUN_KMS_VERSION.to_owned()),
        ];

        let query =
            build_signed_query_with_timestamp_owned("GET", params, &cred, "2017-02-10T10:10:10Z");
        assert_eq!(
            query,
            "AccessKeyId=AKIDEXAMPLE&Action=GenerateDataKey&Format=JSON&KeyId=test%2Fkey&KeySpec=AES_256&SecurityToken=TOKENEXAMPLE&SignatureMethod=HMAC-SHA1&SignatureVersion=1.0&Timestamp=2017-02-10T10%3A10%3A10Z&Version=2016-01-20&Signature=ROdlqmvSpSULNPsvOEBPy7EBhDQ%3D"
        );
    }

    #[test]
    fn test_build_signed_query_without_security_token() {
        let cred = AwsCredentials::new(
            "AKIDEXAMPLE".to_owned(),
            "SECRETEXAMPLE".to_owned(),
            None,
            None,
        );
        let params = vec![
            ("Action".to_owned(), "GenerateDataKey".to_owned()),
            ("KeyId".to_owned(), "test/key".to_owned()),
            ("KeySpec".to_owned(), ALIYUN_KMS_DATA_KEY_SPEC.to_owned()),
            ("Version".to_owned(), ALIYUN_KMS_VERSION.to_owned()),
        ];

        let query =
            build_signed_query_with_timestamp_owned("GET", params, &cred, "2017-02-10T10:10:10Z");
        assert_eq!(
            query,
            "AccessKeyId=AKIDEXAMPLE&Action=GenerateDataKey&Format=JSON&KeyId=test%2Fkey&KeySpec=AES_256&SignatureMethod=HMAC-SHA1&SignatureVersion=1.0&Timestamp=2017-02-10T10%3A10%3A10Z&Version=2016-01-20&Signature=0baQIMlZQK0aHAeO%2Be3UisRXQOc%3D"
        );
    }

    #[test]
    fn test_build_signed_query_with_encryption_context() {
        let cred = AwsCredentials::new(
            "AKIDEXAMPLE".to_owned(),
            "SECRETEXAMPLE".to_owned(),
            None,
            None,
        );
        let mut encryption_context = BTreeMap::new();
        encryption_context.insert("ExternalId".to_owned(), "customer-42".to_owned());
        encryption_context.insert("tenant".to_owned(), "acme".to_owned());
        let mut params = vec![
            ("Action".to_owned(), "GenerateDataKey".to_owned()),
            ("KeyId".to_owned(), "test/key".to_owned()),
            ("KeySpec".to_owned(), ALIYUN_KMS_DATA_KEY_SPEC.to_owned()),
            ("Version".to_owned(), ALIYUN_KMS_VERSION.to_owned()),
            (
                "EncryptionContext".to_owned(),
                serde_json::to_string(&encryption_context).unwrap(),
            ),
        ];

        let query = build_signed_query_with_timestamp_owned(
            "GET",
            std::mem::take(&mut params),
            &cred,
            "2017-02-10T10:10:10Z",
        );
        assert!(query.contains(
            "EncryptionContext=%7B%22ExternalId%22%3A%22customer-42%22%2C%22tenant%22%3A%22acme%22%7D"
        ));
    }

    #[test]
    fn test_build_aliyun_kms_uri_rejects_invalid_region() {
        for region in [
            "",
            "-ap-southeast-1",
            "ap-southeast-1-",
            "ap--southeast-1",
            "evil.com?",
            "cn/hangzhou",
        ] {
            build_aliyun_kms_uri(region, "Action=GenerateDataKey").unwrap_err();
        }
    }

    #[test]
    fn test_build_aliyun_kms_uri_accepts_valid_region() {
        let uri = build_aliyun_kms_uri("ap-southeast-5", "Action=GenerateDataKey").unwrap();
        assert_eq!(
            uri,
            "https://kms.ap-southeast-5.aliyuncs.com/?Action=GenerateDataKey"
        );
    }

    #[test]
    fn test_decode_generate_data_key_response() {
        let plaintext = vec![42u8; 32];
        let ciphertext = vec![1u8, 2, 3, 4];
        let resp = GenerateDataKeyResponse {
            ciphertext_blob: base64::encode(ciphertext.clone()),
            plaintext: base64::encode(plaintext.clone()),
        };

        let data_key = decode_generate_data_key_response(resp).unwrap();
        assert_eq!(data_key.encrypted.clone().into_inner(), ciphertext);
        assert_eq!(&*data_key.plaintext, &plaintext);

        let resp = GenerateDataKeyResponse {
            ciphertext_blob: "".to_owned(),
            plaintext: base64::encode(plaintext),
        };
        decode_generate_data_key_response(resp).unwrap_err();
    }
}