spiral-auth 0.1.9

use std::fmt::{Debug, Formatter};
use std::time::SystemTime;

use aws_credential_types::Credentials;
use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider};
use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign};
use cloud_detect::ProviderId;
use reqwest::Client;
use serde::Deserialize;
use tracing::{debug, warn};
use url::Url;
use url::form_urlencoded::Serializer;

use super::{cloud_provider_id, jwt_bearer};
use crate::http::{Accepted as _, shared_client};
use crate::{CredentialResult, Token};

const HOST: &str = "host";
const CONTENT_TYPE: &str = "content-type";
const STS_FORM: &str = "application/x-www-form-urlencoded; charset=utf-8";
/// STS has no global endpoint for `GetWebIdentityToken`.
const DEFAULT_REGION: &str = "us-east-1";

/// Authenticates a workload with its AWS identity: STS `GetWebIdentityToken` signs a JWT for
/// the authorization server's token endpoint, presented there as a JWT-bearer grant. The AWS
/// account must enable outbound identity federation; its tokens' issuer is the account's, and
/// their `sub` is the caller's role ARN.
pub struct AwsAuthProvider {
    client: Client,
    token_endpoint: Url,
    workload_id: String,
    credentials_provider: Option<SharedCredentialsProvider>,
}

impl Debug for AwsAuthProvider {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AwsAuthProvider").finish()
    }
}

impl AwsAuthProvider {
    /// Obtain the workload's tokens at `token_endpoint`, the authorization server's, with the
    /// default AWS credential chain.
    pub fn new(token_endpoint: Url, workload_id: String) -> Self {
        AwsAuthProvider {
            client: shared_client(),
            token_endpoint,
            workload_id,
            credentials_provider: None,
        }
    }

    /// Create an AWS authentication provider backed by an explicit credentials provider.
    ///
    /// Unlike [`Self::new`], this bypasses cloud-platform detection and the default AWS SDK
    /// credential chain.
    pub fn new_with_credentials_provider(
        token_endpoint: Url,
        workload_id: String,
        credentials_provider: SharedCredentialsProvider,
    ) -> Self {
        AwsAuthProvider {
            client: shared_client(),
            token_endpoint,
            workload_id,
            credentials_provider: Some(credentials_provider),
        }
    }

    /// The caller's credentials and the region whose STS endpoint signs its token.
    async fn credentials(&self) -> CredentialResult<Option<(Credentials, String)>> {
        let (provider, region) = match &self.credentials_provider {
            Some(provider) => (provider.clone(), environment_region()),
            None => {
                let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
                let Some(provider) = config.credentials_provider() else {
                    return Ok(None);
                };
                let region = config
                    .region()
                    .map(ToString::to_string)
                    .unwrap_or_else(environment_region);
                (provider, region)
            }
        };
        match provider.provide_credentials().await {
            Ok(credentials) => Ok(Some((credentials, region))),
            Err(error) => {
                warn!("Failed to get AWS credentials: {error}");
                Ok(None)
            }
        }
    }

    /// Call STS `GetWebIdentityToken` for a JWT whose audience is the token endpoint.
    async fn web_identity_token(
        &self,
        credentials: &Credentials,
        region: &str,
    ) -> CredentialResult<String> {
        let host = format!("sts.{region}.amazonaws.com");
        let url = format!("https://{host}/");
        // RFC 7523 §3: the assertion's audience is the token endpoint it is presented to.
        let body = Serializer::new(String::new())
            .append_pair("Action", "GetWebIdentityToken")
            .append_pair("Version", "2011-06-15")
            .append_pair("Audience.member.1", self.token_endpoint.as_str())
            .append_pair("SigningAlgorithm", "RS256")
            .finish();

        let identity = credentials.clone().into();
        let signing_params = aws_sigv4::sign::v4::SigningParams::builder()
            .identity(&identity)
            .region(region)
            .name("sts")
            .time(SystemTime::now())
            .settings(SigningSettings::default())
            .build()
            .map_err(|error| invalid!("signing params were invalid: {error}"))?
            .into();
        let signable = SignableRequest::new(
            "POST",
            &url,
            [(HOST, host.as_str()), (CONTENT_TYPE, STS_FORM)].into_iter(),
            SignableBody::Bytes(body.as_bytes()),
        )
        .map_err(|error| invalid!("bad signable request: {error}"))?;
        let (instructions, _signature) = sign(signable, &signing_params)
            .map_err(|error| invalid!("sign failed: {error}"))?
            .into_parts();
        let (signed_headers, _params) = instructions.into_parts();

        let mut request = self
            .client
            .post(&url)
            .header(CONTENT_TYPE, STS_FORM)
            .body(body);
        for header in signed_headers {
            request = request.header(header.name(), header.value());
        }
        let response = request.send().await?.accepted().await?.text().await?;
        let GetWebIdentityTokenResponse {
            get_web_identity_token_result,
        } = quick_xml::de::from_str(&response)
            .map_err(|error| invalid!("STS returned an unexpected response: {error}"))?;
        Ok(get_web_identity_token_result.web_identity_token)
    }
}

impl AwsAuthProvider {
    /// A usable token, or `None` when this provider has no credentials to offer.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        if self.credentials_provider.is_none()
            && !matches!(cloud_provider_id().await, Some(ProviderId::AWS))
        {
            return Ok(None);
        }
        let Some((credentials, region)) = self.credentials().await? else {
            debug!("No AWS credentials found.");
            return Ok(None);
        };
        let assertion = self.web_identity_token(&credentials, &region).await?;
        let token = jwt_bearer(
            &self.client,
            &self.token_endpoint,
            self.workload_id.clone(),
            assertion,
            None,
        )
        .await?;
        debug!("AwsAuthProvider obtained a workload token");
        Ok(Some(token))
    }
}

impl_token_source!(AwsAuthProvider);

fn environment_region() -> String {
    std::env::var("AWS_REGION")
        .or_else(|_| std::env::var("AWS_DEFAULT_REGION"))
        .unwrap_or_else(|_| DEFAULT_REGION.to_owned())
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
struct GetWebIdentityTokenResponse {
    get_web_identity_token_result: GetWebIdentityTokenResult,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
struct GetWebIdentityTokenResult {
    web_identity_token: String,
}

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

    #[test]
    fn reads_the_token_from_an_sts_response() {
        let response = r#"<GetWebIdentityTokenResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
  <GetWebIdentityTokenResult>
    <WebIdentityToken>eyJ.payload.signature</WebIdentityToken>
    <Expiration>2025-11-11T22:44:21.703Z</Expiration>
  </GetWebIdentityTokenResult>
  <ResponseMetadata>
    <RequestId>c6104cbe-af31-11e0-8154-cbc7ccf896c7</RequestId>
  </ResponseMetadata>
</GetWebIdentityTokenResponse>"#;
        let parsed: GetWebIdentityTokenResponse = quick_xml::de::from_str(response).unwrap();
        assert_eq!(
            parsed.get_web_identity_token_result.web_identity_token,
            "eyJ.payload.signature"
        );
    }
}