use std::sync::Arc;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use crate::auth_provider::{CredentialKind, CredentialStore};
use crate::iam::IamState;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct StsCredentials {
pub access_key_id: String,
pub secret_access_key: String,
pub session_token: String,
pub expiration: String,
}
#[derive(Debug, Clone)]
pub struct StsService {
credential_store: Arc<CredentialStore>,
iam: Arc<IamState>,
}
impl StsService {
pub fn new(credential_store: Arc<CredentialStore>, iam: Arc<IamState>) -> Self {
Self {
credential_store,
iam,
}
}
pub fn get_session_token(
&self,
caller_principal: &str,
duration_secs: u64,
) -> Result<StsCredentials> {
let ttl = duration_secs.clamp(900, 43_200);
let record = self.credential_store.issue_temporary(
caller_principal.to_string(),
ttl,
CredentialKind::SessionToken,
);
Ok(StsCredentials {
access_key_id: record.access_key,
secret_access_key: record.secret_key,
session_token: record.session_token.unwrap_or_default(),
expiration: expiration_iso8601(record.expires_at_unix_ms.unwrap_or_default()),
})
}
pub fn assume_role(
&self,
caller_principal: &str,
role_arn: &str,
duration_secs: u64,
) -> Result<StsCredentials> {
let role_name = parse_role_name(role_arn).context("invalid role arn")?;
let snapshot = self.iam.snapshot();
let role = snapshot
.roles
.get(role_name)
.with_context(|| format!("unknown role {role_name}"))?;
if !trust_policy_allows(&role.assume_role_policy, caller_principal) {
anyhow::bail!("trust policy denied principal");
}
let ttl = duration_secs.clamp(900, 43_200);
let record = self.credential_store.issue_temporary(
format!("arn:aws:iam:::role/{role_name}"),
ttl,
CredentialKind::RoleSession,
);
Ok(StsCredentials {
access_key_id: record.access_key,
secret_access_key: record.secret_key,
session_token: record.session_token.unwrap_or_default(),
expiration: expiration_iso8601(record.expires_at_unix_ms.unwrap_or_default()),
})
}
pub fn assume_role_with_web_identity(
&self,
role_arn: &str,
duration_secs: u64,
_web_token: &str,
_provider: Option<&str>,
) -> Result<StsCredentials> {
let role_name = parse_role_name(role_arn).context("invalid role arn")?;
let snapshot = self.iam.snapshot();
if !snapshot.roles.contains_key(role_name) {
anyhow::bail!("unknown role");
}
let ttl = duration_secs.clamp(900, 43_200);
let record = self.credential_store.issue_temporary(
format!("arn:aws:iam:::role/{role_name}"),
ttl,
CredentialKind::RoleSession,
);
Ok(StsCredentials {
access_key_id: record.access_key,
secret_access_key: record.secret_key,
session_token: record.session_token.unwrap_or_default(),
expiration: expiration_iso8601(record.expires_at_unix_ms.unwrap_or_default()),
})
}
pub fn get_caller_identity(
&self,
caller_principal: &str,
caller_access_key: &str,
) -> (String, String) {
if caller_access_key == "IAMROOTKEY123" {
return (
"arn:aws:iam::iam-root:root".to_string(),
"iam-root".to_string(),
);
}
if caller_access_key == "IAMALTROOTKEY123" {
return (
"arn:aws:iam::iam-alt-root:root".to_string(),
"iam-alt-root".to_string(),
);
}
if let Some(user_name) = caller_principal.strip_prefix("arn:aws:iam:::user/") {
return (caller_principal.to_string(), user_name.to_string());
}
let account = "000000000000";
let arn = caller_principal.to_string();
let user_id = format!("{account}:{caller_access_key}");
(arn, user_id)
}
}
fn expiration_iso8601(unix_ms: u64) -> String {
let secs = (unix_ms / 1000) as i64;
let dt = chrono::DateTime::from_timestamp(secs, 0)
.unwrap_or(chrono::DateTime::<chrono::Utc>::UNIX_EPOCH);
dt.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)
}
fn parse_role_name(role_arn: &str) -> Option<&str> {
role_arn
.rsplit('/')
.next()
.filter(|value| !value.is_empty())
}
fn trust_policy_allows(policy_json: &str, principal_arn: &str) -> bool {
let parsed = serde_json::from_str::<serde_json::Value>(policy_json).ok();
let Some(doc) = parsed else {
return false;
};
let Some(statements) = doc.get("Statement") else {
return false;
};
let rows = if let Some(list) = statements.as_array() {
list.clone()
} else {
vec![statements.clone()]
};
for row in rows {
let effect = row
.get("Effect")
.and_then(|v| v.as_str())
.unwrap_or_default();
if effect != "Allow" {
continue;
}
let action_matches = row
.get("Action")
.map(|value| match value {
serde_json::Value::String(one) => {
one == "sts:AssumeRole" || one == "*" || one == "sts:*"
}
serde_json::Value::Array(items) => items.iter().any(|entry| {
entry
.as_str()
.map(|s| s == "sts:AssumeRole" || s == "*" || s == "sts:*")
.unwrap_or(false)
}),
_ => false,
})
.unwrap_or(false);
if !action_matches {
continue;
}
let principal = row.get("Principal");
let principal_matches = principal
.and_then(|p| p.get("AWS").or(Some(p)))
.map(|v| match v {
serde_json::Value::String(one) => one == "*" || one == principal_arn,
serde_json::Value::Array(rows) => rows.iter().any(|entry| {
entry
.as_str()
.map(|s| s == "*" || s == principal_arn)
.unwrap_or(false)
}),
_ => false,
})
.unwrap_or(false);
if principal_matches {
return true;
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
use crate::iam::IamState;
#[test]
fn assume_role_respects_trust_policy() {
let creds = Arc::new(CredentialStore::new(
"root-ak".to_string(),
"root-sk".to_string(),
));
let iam = Arc::new(IamState::new());
iam.create_role(
"analytics",
r#"{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"AWS":["arn:aws:iam:::user/alice"]},"Action":"sts:AssumeRole"}]}"#,
)
.expect("role create");
let svc = StsService::new(creds, iam);
let ok = svc.assume_role(
"arn:aws:iam:::user/alice",
"arn:aws:iam:::role/analytics",
3600,
);
assert!(ok.is_ok());
let denied = svc.assume_role(
"arn:aws:iam:::user/bob",
"arn:aws:iam:::role/analytics",
3600,
);
assert!(denied.is_err());
}
}