olai-http 0.0.5

use std::sync::Arc;

use futures::future::BoxFuture;

use crate::service::make_service;
use crate::token::TemporaryToken;
use crate::{ClientOptions, RequestSigner, Result, RetryConfig, TokenProvider};

use self::credential::{AssumeRoleProvider, AwsAuthorizer, CredentialExt};
use crate::CredentialProvider;

mod builder;
pub(crate) mod credential;

pub use builder::*;
pub use credential::AwsCredential;

pub type AwsCredentialProvider = Arc<dyn CredentialProvider<Credential = AwsCredential>>;

#[derive(Debug, Clone)]
pub struct AmazonConfig {
    pub region: String,
    pub credentials: AwsCredentialProvider,
    pub retry_config: RetryConfig,
    pub client_options: ClientOptions,
    pub skip_signature: bool,
}

impl AmazonConfig {
    pub(crate) async fn get_credential(&self) -> Result<Option<Arc<AwsCredential>>> {
        Ok(match self.skip_signature {
            false => Some(self.credentials.get_credential().await?),
            true => None,
        })
    }
}

impl RequestSigner for AmazonConfig {
    fn sign<'a>(
        &'a self,
        req: reqwest::RequestBuilder,
    ) -> BoxFuture<'a, Result<reqwest::RequestBuilder>> {
        Box::pin(async move {
            if let Some(cred) = self.get_credential().await? {
                let authorizer = AwsAuthorizer::new(&cred, "execute-api", &self.region);
                Ok(req.with_aws_sigv4(Some(authorizer), None))
            } else {
                Ok(req)
            }
        })
    }
}

/// Assume an AWS IAM role using explicit base credentials.
///
/// Like [`assume_role`] but accepts an explicit `base_credentials` provider
/// instead of falling back to environment-based credential discovery.
/// Use this when the base identity (the caller's access key + secret) is
/// stored in a credential registry rather than in the server environment.
pub async fn assume_role_with_base(
    role_arn: &str,
    region: &str,
    sts_endpoint: Option<&str>,
    policy: Option<String>,
    base_credentials: AwsCredentialProvider,
) -> Result<TemporaryToken<Arc<AwsCredential>>> {
    let endpoint = sts_endpoint
        .map(|s| s.to_owned())
        .unwrap_or_else(|| format!("https://sts.{region}.amazonaws.com"));

    let provider = AssumeRoleProvider {
        role_arn: role_arn.to_owned(),
        session_name: "TrestleVending".to_owned(),
        endpoint,
        base_credentials,
        region: region.to_owned(),
        policy,
    };

    let client = ClientOptions::default().client()?;
    let service = make_service(client.clone(), None);
    provider
        .fetch_token(&client, &service, &RetryConfig::default())
        .await
}

/// Assume an AWS IAM role using ambient server credentials.
///
/// Uses the environment credential chain (`AWS_*` env vars, instance profile,
/// ECS task role, WebIdentity, etc.) as the base identity for the STS
/// `AssumeRole` call. Prefer [`assume_role_with_base`] when the base
/// credentials are stored in a credential registry.
pub async fn assume_role(
    role_arn: &str,
    region: &str,
    sts_endpoint: Option<&str>,
    policy: Option<String>,
) -> Result<TemporaryToken<Arc<AwsCredential>>> {
    let base_credentials = AmazonBuilder::from_env()
        .with_region(region)
        .build(None)?
        .credentials;
    assume_role_with_base(role_arn, region, sts_endpoint, policy, base_credentials).await
}