1use aws_sdk_sts::error::{DisplayErrorContext, ProvideErrorMetadata, SdkError};
2use aws_sdk_sts::operation::assume_role::{AssumeRoleError, AssumeRoleOutput};
3use regex::Regex;
4use uuid::Uuid;
5
6use crate::errors::StsError;
7
8pub struct AwsTempCreds {
9 pub access_key: String,
10 pub secret_key: String,
11 pub session_token: String,
12 pub region: String,
13}
14
15struct Hop<'a> {
16 role_arn: &'a str,
17 external_id: Option<&'a str>,
18 policy: Option<&'static str>,
19}
20
21const ONLY_ASSUME_ROLE: &str = r#"{
25 "Version": "2012-10-17",
26 "Statement": [{
27 "Effect": "Allow",
28 "Action": "sts:AssumeRole",
29 "Resource": "*"
30 }]
31}"#;
32
33const fn chain<'a>(
35 arn: &'a str,
36 external_id: &'a str,
37 via: Option<&'a str>,
38) -> (Hop<'a>, Option<Hop<'a>>) {
39 let client = Hop {
40 role_arn: arn,
41 external_id: Some(external_id),
42 policy: None,
43 };
44
45 match via {
46 Some(intermediate) => (
47 Hop {
48 role_arn: intermediate,
49 external_id: None,
50 policy: Some(ONLY_ASSUME_ROLE),
51 },
52 Some(client),
53 ),
54 None => (client, None),
55 }
56}
57
58pub async fn assume_role(
59 arn: &str,
60 external_id: &str,
61 url: &str,
62 via: Option<&str>,
63) -> Result<AwsTempCreds, StsError> {
64 let region = extract_region(url);
65 let (first, then) = chain(arn, external_id, via);
66
67 let mut creds = assume(&first, ®ion, None).await?;
68 if let Some(hop) = then {
69 creds = assume(&hop, ®ion, Some(&creds)).await?;
70 }
71
72 Ok(creds)
73}
74
75async fn assume(
76 hop: &Hop<'_>,
77 region: &str,
78 from: Option<&AwsTempCreds>,
79) -> Result<AwsTempCreds, StsError> {
80 let mut config = aws_config::defaults(aws_config::BehaviorVersion::latest())
81 .region(aws_config::Region::new(region.to_owned()));
82
83 if let Some(session) = from {
84 config = config.credentials_provider(aws_sdk_sts::config::Credentials::new(
85 session.access_key.clone(),
86 session.secret_key.clone(),
87 Some(session.session_token.clone()),
88 None,
89 "fluidattacks-core-chain",
90 ));
91 }
92
93 let sts = aws_sdk_sts::Client::new(&config.load().await);
94
95 let mut request = sts
96 .assume_role()
97 .role_arn(hop.role_arn)
98 .role_session_name(format!("session-{}", Uuid::new_v4()));
99
100 if let Some(external_id) = hop.external_id {
101 request = request.external_id(external_id);
102 }
103
104 if let Some(policy) = hop.policy {
105 request = request.policy(policy);
106 }
107
108 let resp = request
109 .send()
110 .await
111 .map_err(|error| refusal(hop.role_arn, error))?;
112
113 session(hop.role_arn, &resp, region)
114}
115
116fn refusal<R: std::fmt::Debug>(arn: &str, error: SdkError<AssumeRoleError, R>) -> StsError {
118 match error {
119 SdkError::ServiceError(refused) => StsError::Refused {
120 arn: arn.to_owned(),
121 code: refused.err().code().unwrap_or("Unknown").to_owned(),
122 message: refused.err().message().unwrap_or("Unknown").to_owned(),
123 },
124 unanswered => StsError::Unreachable {
125 arn: arn.to_owned(),
126 reason: DisplayErrorContext(&unanswered).to_string(),
127 },
128 }
129}
130
131fn session(arn: &str, resp: &AssumeRoleOutput, region: &str) -> Result<AwsTempCreds, StsError> {
132 let no_credentials = || StsError::NoCredentials {
133 arn: arn.to_owned(),
134 };
135 let creds = resp.credentials().ok_or_else(no_credentials)?;
136
137 let access_key = creds.access_key_id().to_owned();
138 let secret_key = creds.secret_access_key().to_owned();
139 let session_token = creds.session_token().to_owned();
140
141 if access_key.is_empty() || secret_key.is_empty() {
142 return Err(no_credentials());
143 }
144
145 Ok(AwsTempCreds {
146 access_key,
147 secret_key,
148 session_token,
149 region: region.to_owned(),
150 })
151}
152
153fn extract_region(url: &str) -> String {
154 [
155 r"codecommit::([a-z0-9-]+)://",
156 r"git-codecommit\.([a-z0-9-]+)\.amazonaws\.com",
157 ]
158 .into_iter()
159 .find_map(|pattern| {
160 let region = Regex::new(pattern).ok()?.captures(url)?.get(1)?;
161 Some(region.as_str().to_owned())
162 })
163 .unwrap_or_else(|| "us-east-1".to_owned())
164}
165
166#[cfg(test)]
167mod tests {
168 use aws_sdk_sts::error::{ErrorMetadata, SdkError};
169 use aws_sdk_sts::operation::assume_role::{AssumeRoleError, AssumeRoleOutput};
170 use aws_sdk_sts::primitives::DateTime;
171 use aws_sdk_sts::types::Credentials;
172
173 use super::{chain, extract_region, refusal, session};
174 use crate::errors::StsError;
175
176 const ROLE: &str = "arn:aws:iam::1:role/client";
177
178 #[test]
179 fn a_refusal_carries_the_code_and_message_sts_answered() {
180 let refused = AssumeRoleError::generic(
181 ErrorMetadata::builder()
182 .code("AccessDenied")
183 .message("not authorized to perform sts:AssumeRole")
184 .build(),
185 );
186
187 let error = refusal(ROLE, SdkError::service_error(refused, ()));
188
189 assert_eq!(
190 error.to_string(),
191 "STS refused the role arn:aws:iam::1:role/client: AccessDenied: not authorized to perform sts:AssumeRole"
192 );
193 }
194
195 #[test]
196 fn an_sts_that_never_answered_is_not_a_refusal() {
197 let error = refusal(
198 ROLE,
199 SdkError::<AssumeRoleError, ()>::timeout_error("timed out"),
200 );
201
202 assert!(matches!(error, StsError::Unreachable { .. }), "{error:?}");
203 }
204
205 #[test]
206 fn an_answer_without_a_usable_key_is_no_session() {
207 let keys = |access: &str, secret: &str| {
208 Credentials::builder()
209 .access_key_id(access)
210 .secret_access_key(secret)
211 .session_token("a-token")
212 .expiration(DateTime::from_secs(0))
213 .build()
214 .unwrap()
215 };
216
217 for answer in [
218 AssumeRoleOutput::builder().build(),
219 AssumeRoleOutput::builder()
220 .credentials(keys("", "a-secret-key"))
221 .build(),
222 AssumeRoleOutput::builder()
223 .credentials(keys("an-access-key", ""))
224 .build(),
225 ] {
226 let outcome = session(ROLE, &answer, "us-east-1");
227
228 assert!(
229 matches!(outcome, Err(StsError::NoCredentials { .. })),
230 "{answer:?}"
231 );
232 }
233 }
234
235 #[test]
236 fn the_region_comes_from_the_url_or_falls_back_to_us_east_1() {
237 for (url, region) in [
238 ("codecommit::eu-west-1://a-repo", "eu-west-1"),
239 (
240 "https://git-codecommit.sa-east-1.amazonaws.com/v1/repos/a-repo",
241 "sa-east-1",
242 ),
243 ("codecommit://a-repo", "us-east-1"),
244 ] {
245 assert_eq!(extract_region(url), region);
246 }
247 }
248
249 #[test]
250 fn one_call_takes_the_client_role_with_its_external_id() {
251 let (first, then) = chain("arn:aws:iam::1:role/client", "an-external-id", None);
252
253 assert_eq!(first.role_arn, "arn:aws:iam::1:role/client");
254 assert_eq!(first.external_id, Some("an-external-id"));
255 assert_eq!(first.policy, None);
256 assert!(then.is_none());
257 }
258
259 #[test]
260 fn a_chain_takes_the_intermediate_role_first_and_without_the_external_id() {
261 let (first, then) = chain(
262 "arn:aws:iam::1:role/client",
263 "an-external-id",
264 Some("arn:aws:iam::2:role/intermediate"),
265 );
266
267 assert_eq!(first.role_arn, "arn:aws:iam::2:role/intermediate");
268 assert_eq!(first.external_id, None);
269
270 let client = then.expect("the client role is still taken");
271 assert_eq!(client.role_arn, "arn:aws:iam::1:role/client");
272 assert_eq!(client.external_id, Some("an-external-id"));
273 assert_eq!(client.policy, None);
274 }
275
276 #[test]
277 fn a_chain_lets_the_intermediate_session_only_assume_a_role() {
278 let (first, _) = chain(
279 "arn:aws:iam::1:role/client",
280 "an-external-id",
281 Some("arn:aws:iam::2:role/intermediate"),
282 );
283
284 let policy: serde_json::Value =
285 serde_json::from_str(first.policy.expect("the intermediate session is narrowed"))
286 .expect("the policy is json");
287 assert_eq!(
288 policy.get("Statement"),
289 Some(&serde_json::json!([{
290 "Effect": "Allow",
291 "Action": "sts:AssumeRole",
292 "Resource": "*",
293 }]))
294 );
295 }
296}