use crate::custom::credential::NewAwsCredsForStsCreds;
use crate::{AssumeRoleWithWebIdentityRequest, Sts, StsClient, PolicyDescriptorType};
use rusoto_core::credential::{
AwsCredentials, CredentialsError, ProvideAwsCredentials, Secret, Variable,
};
use rusoto_core::request::HttpClient;
use rusoto_core::{Client, Region};
use async_trait::async_trait;
const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
const AWS_ROLE_SESSION_NAME: &str = "AWS_ROLE_SESSION_NAME";
#[derive(Debug, Clone)]
pub struct WebIdentityProvider {
pub web_identity_token: Variable<Secret, CredentialsError>,
pub role_arn: Variable<String, CredentialsError>,
pub role_session_name: Option<Variable<Option<String>, CredentialsError>>,
pub duration_seconds: Option<i64>,
pub policy: Option<String>,
pub policy_arns: Option<Vec<PolicyDescriptorType>>,
}
impl WebIdentityProvider {
pub fn new<A, B, C>(web_identity_token: A, role_arn: B, role_session_name: Option<C>) -> Self
where
A: Into<Variable<Secret, CredentialsError>>,
B: Into<Variable<String, CredentialsError>>,
C: Into<Variable<Option<String>, CredentialsError>>,
{
Self {
web_identity_token: web_identity_token.into(),
role_arn: role_arn.into(),
role_session_name: role_session_name.map(|v| v.into()),
duration_seconds: None,
policy: None,
policy_arns: None
}
}
pub fn from_k8s_env() -> Self {
Self::_from_k8s_env(
Variable::from_env_var(AWS_WEB_IDENTITY_TOKEN_FILE),
Variable::from_env_var(AWS_ROLE_ARN),
Variable::from_env_var_optional(AWS_ROLE_SESSION_NAME),
)
}
pub(crate) fn _from_k8s_env(
token_file: Variable<String, CredentialsError>,
role: Variable<String, CredentialsError>,
session_name: Variable<Option<String>, CredentialsError>,
) -> Self {
Self::new(
Variable::dynamic(move || Variable::from_text_file(token_file.resolve()?).resolve()),
role,
Some(session_name),
)
}
#[cfg(test)]
pub(crate) fn load_token(&self) -> Result<Secret, CredentialsError> {
self.web_identity_token.resolve()
}
fn create_session_name() -> String {
"WebIdentitySession".to_string()
}
}
#[async_trait]
impl ProvideAwsCredentials for WebIdentityProvider {
async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
let http_client = match HttpClient::new() {
Ok(c) => c,
Err(e) => return Err(CredentialsError::new(e)),
};
let client = Client::new_not_signing(http_client);
let sts = StsClient::new_with_client(client, Region::default());
let mut req = AssumeRoleWithWebIdentityRequest::default();
req.role_arn = self.role_arn.resolve()?;
req.web_identity_token = self.web_identity_token.resolve()?.as_ref().to_string();
req.policy = self.policy.to_owned();
req.duration_seconds = self.duration_seconds.to_owned();
req.policy_arns = self.policy_arns.to_owned();
req.role_session_name = match self.role_session_name {
Some(ref role_session_name) => match role_session_name.resolve()? {
Some(session_name) => session_name,
None => Self::create_session_name(),
},
None => Self::create_session_name(),
};
let assume_role = sts.assume_role_with_web_identity(req).await;
match assume_role {
Err(e) => Err(CredentialsError::new(e)),
Ok(role) => match role.credentials {
None => Err(CredentialsError::new(format!(
"No credentials found in AssumeRoleWithWebIdentityResponse: {:?}",
role
))),
Some(c) => AwsCredentials::new_for_credentials(c),
},
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[test]
fn api_ergonomy() {
WebIdentityProvider::new(Secret::from("".to_string()), "", Some(Some("".to_string())));
}
#[test]
fn from_k8s_env() -> Result<(), CredentialsError> {
const TOKEN_VALUE: &str = "secret";
const ROLE_ARN: &str = "role";
const SESSION_NAME: &str = "session";
let mut file = NamedTempFile::new()?;
writeln!(file, "{}", TOKEN_VALUE)?;
let p = WebIdentityProvider::_from_k8s_env(
Variable::with_value(file.path().to_string_lossy().to_string()),
Variable::with_value(ROLE_ARN.to_string()),
Variable::with_value(SESSION_NAME.to_string()),
);
let token = p.load_token()?;
assert_eq!(token.as_ref(), TOKEN_VALUE);
Ok(())
}
}