1use std::sync::Arc;
2use std::time::Duration;
3
4use chrono::{DateTime, Utc};
5
6use futures::future::BoxFuture;
7
8use crate::service::make_service;
9use crate::token::TemporaryToken;
10use crate::{ClientOptions, CloudClient, RequestSigner, Result, RetryConfig, TokenProvider};
11
12use self::credential::{AssumeRoleParams, AssumeRoleProvider, AwsAuthorizer, CredentialExt};
13use crate::CredentialProvider;
14
15mod builder;
16pub(crate) mod credential;
17
18pub use builder::*;
19pub use credential::AwsCredential;
20
21pub type AwsCredentialProvider = Arc<dyn CredentialProvider<Credential = AwsCredential>>;
22
23#[derive(Debug, Clone)]
24pub struct AmazonConfig {
25 pub region: String,
26 pub credentials: AwsCredentialProvider,
27 pub retry_config: RetryConfig,
28 pub client_options: ClientOptions,
29 pub skip_signature: bool,
30}
31
32impl AmazonConfig {
33 pub(crate) async fn get_credential(&self) -> Result<Option<Arc<AwsCredential>>> {
34 Ok(match self.skip_signature {
35 false => Some(self.credentials.get_credential().await?),
36 true => None,
37 })
38 }
39}
40
41impl RequestSigner for AmazonConfig {
42 fn sign<'a>(
43 &'a self,
44 req: reqwest::RequestBuilder,
45 ) -> BoxFuture<'a, Result<reqwest::RequestBuilder>> {
46 Box::pin(async move {
47 if let Some(cred) = self.get_credential().await? {
48 let authorizer = AwsAuthorizer::new(&cred, "execute-api", &self.region);
49 Ok(req.with_aws_sigv4(Some(authorizer), None))
50 } else {
51 Ok(req)
52 }
53 })
54 }
55}
56
57pub const S3_GLOBAL_ENDPOINT: &str = "https://s3.amazonaws.com";
59
60pub async fn bucket_region(
74 client: &CloudClient,
75 endpoint: &str,
76 bucket: &str,
77) -> Result<Option<String>> {
78 let mut url = reqwest::Url::parse(endpoint).map_err(|source| crate::Error::Generic {
79 source: Box::new(source),
80 })?;
81 url.path_segments_mut()
82 .map_err(|()| crate::Error::Generic {
83 source: format!("endpoint `{endpoint}` cannot carry a path").into(),
84 })?
85 .pop_if_empty()
86 .push(bucket);
87 let request = client.reqwest_client.head(url).build()?;
90 let response = client.service.call(request).await?;
91 Ok(response
92 .headers()
93 .get("x-amz-bucket-region")
94 .and_then(|region| region.to_str().ok())
95 .filter(|region| !region.is_empty())
96 .map(str::to_owned))
97}
98
99pub async fn assume_role_with_base(
106 role_arn: &str,
107 region: &str,
108 sts_endpoint: Option<&str>,
109 policy: Option<String>,
110 base_credentials: AwsCredentialProvider,
111) -> Result<TemporaryToken<Arc<AwsCredential>>> {
112 assume_role_with_base_and_options(
113 role_arn,
114 region,
115 sts_endpoint,
116 policy,
117 base_credentials,
118 &ClientOptions::default(),
119 )
120 .await
121}
122
123pub async fn assume_role_with_base_and_options(
129 role_arn: &str,
130 region: &str,
131 sts_endpoint: Option<&str>,
132 policy: Option<String>,
133 base_credentials: AwsCredentialProvider,
134 client_options: &ClientOptions,
135) -> Result<TemporaryToken<Arc<AwsCredential>>> {
136 let endpoint = sts_endpoint
137 .map(|s| s.to_owned())
138 .unwrap_or_else(|| format!("https://sts.{region}.amazonaws.com"));
139
140 let provider = AssumeRoleProvider {
141 role_arn: role_arn.to_owned(),
142 session_name: DEFAULT_SESSION_NAME.to_owned(),
143 endpoint,
144 base_credentials,
145 region: region.to_owned(),
146 policy,
147 };
148
149 let client = client_options.client()?;
150 let service = make_service(client.clone(), None);
151 provider
152 .fetch_token(&client, &service, &RetryConfig::default())
153 .await
154}
155
156pub const DEFAULT_SESSION_NAME: &str = "TrestleVending";
158
159#[derive(Debug, Clone, Default)]
163pub struct AssumeRoleOptions {
164 pub duration: Option<Duration>,
167 pub session_name: Option<String>,
169 pub policy: Option<String>,
171 pub external_id: Option<String>,
173 pub source_identity: Option<String>,
175 pub tags: Vec<(String, String)>,
177 pub sts_endpoint: Option<String>,
179 pub client_options: ClientOptions,
181 pub retry: RetryConfig,
183}
184
185#[derive(Debug)]
187pub struct AssumedRole {
188 pub credential: AwsCredential,
190 pub expiration: DateTime<Utc>,
192}
193
194pub async fn assume_role_with_options(
200 role_arn: &str,
201 region: &str,
202 base_credentials: AwsCredentialProvider,
203 options: &AssumeRoleOptions,
204) -> Result<AssumedRole> {
205 let endpoint = options
206 .sts_endpoint
207 .clone()
208 .unwrap_or_else(|| format!("https://sts.{region}.amazonaws.com"));
209 let base = base_credentials
210 .get_credential()
211 .await
212 .map_err(|source| crate::Error::Generic {
213 source: Box::new(source),
214 })?;
215 let params = AssumeRoleParams {
216 role_arn,
217 session_name: options
218 .session_name
219 .as_deref()
220 .unwrap_or(DEFAULT_SESSION_NAME),
221 duration_secs: options.duration.map(|d| d.as_secs()),
222 policy: options.policy.as_deref(),
223 external_id: options.external_id.as_deref(),
224 source_identity: options.source_identity.as_deref(),
225 tags: &options.tags,
226 };
227 let client = options.client_options.client()?;
228 let service = make_service(client.clone(), None);
229 let (credential, expiration) = credential::assume_role_request(
230 &client,
231 &service,
232 &options.retry,
233 &base,
234 region,
235 &endpoint,
236 ¶ms,
237 )
238 .await
239 .map_err(|source| crate::Error::Generic { source })?;
240 Ok(AssumedRole {
241 credential,
242 expiration,
243 })
244}
245
246pub async fn assume_role(
253 role_arn: &str,
254 region: &str,
255 sts_endpoint: Option<&str>,
256 policy: Option<String>,
257) -> Result<TemporaryToken<Arc<AwsCredential>>> {
258 let base_credentials = AmazonBuilder::from_env()
259 .with_region(region)
260 .build(None)?
261 .credentials;
262 assume_role_with_base(role_arn, region, sts_endpoint, policy, base_credentials).await
263}
264
265#[cfg(test)]
266mod tests {
267 use mockito::Matcher;
268
269 use super::*;
270
271 fn client() -> CloudClient {
272 CloudClient::new_unauthenticated_with_options(
273 &ClientOptions::new().with_allow_http(true),
274 RetryConfig::default(),
275 None,
276 )
277 .unwrap()
278 }
279
280 #[tokio::test]
281 async fn bucket_region_is_read_from_any_status() {
282 let mut server = mockito::Server::new_async().await;
283 for (bucket, status, region) in [
284 ("open", 200, "us-east-1"),
285 ("elsewhere", 301, "eu-west-1"),
286 ("private.with.dots", 403, "ap-south-1"),
287 ] {
288 let mock = server
289 .mock("HEAD", format!("/{bucket}").as_str())
290 .match_header("authorization", Matcher::Missing)
291 .with_status(status)
292 .with_header("x-amz-bucket-region", region)
293 .create_async()
294 .await;
295 let found = bucket_region(&client(), &server.url(), bucket)
296 .await
297 .unwrap();
298 assert_eq!(found.as_deref(), Some(region), "{bucket}");
299 mock.assert_async().await;
300 }
301 }
302
303 #[tokio::test]
304 async fn bucket_region_is_none_without_the_header() {
305 let mut server = mockito::Server::new_async().await;
306 let _mock = server
307 .mock("HEAD", "/missing")
308 .with_status(404)
309 .create_async()
310 .await;
311 let found = bucket_region(&client(), &server.url(), "missing")
312 .await
313 .unwrap();
314 assert_eq!(found, None);
315 }
316
317 const ASSUME_ROLE_RESPONSE: &str = r#"<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
318 <AssumeRoleResult>
319 <Credentials>
320 <AccessKeyId>ASIASESSION</AccessKeyId>
321 <SecretAccessKey>session-secret</SecretAccessKey>
322 <SessionToken>session-token</SessionToken>
323 <Expiration>2030-01-01T00:15:00Z</Expiration>
324 </Credentials>
325 </AssumeRoleResult>
326</AssumeRoleResponse>"#;
327
328 fn base() -> AwsCredentialProvider {
329 Arc::new(crate::StaticCredentialProvider::new(AwsCredential {
330 key_id: "AKIDBASE".into(),
331 secret_key: "base-secret".into(),
332 token: None,
333 }))
334 }
335
336 fn options(server: &mockito::Server) -> AssumeRoleOptions {
337 AssumeRoleOptions {
338 sts_endpoint: Some(server.url()),
339 client_options: ClientOptions::new().with_allow_http(true),
340 ..AssumeRoleOptions::default()
341 }
342 }
343
344 #[tokio::test]
345 async fn assume_role_sends_every_option() {
346 let mut server = mockito::Server::new_async().await;
347 let sts = server
348 .mock("POST", "/")
349 .match_header(
350 "authorization",
351 Matcher::Regex("^AWS4-HMAC-SHA256 Credential=AKIDBASE/".into()),
352 )
353 .match_body(Matcher::AllOf(vec![
354 Matcher::UrlEncoded("Action".into(), "AssumeRole".into()),
355 Matcher::UrlEncoded("RoleArn".into(), "arn:aws:iam::1:role/vend".into()),
356 Matcher::UrlEncoded("DurationSeconds".into(), "900".into()),
357 Matcher::UrlEncoded("RoleSessionName".into(), "alice@table".into()),
358 Matcher::UrlEncoded("ExternalId".into(), "ext-42".into()),
359 Matcher::UrlEncoded("SourceIdentity".into(), "alice".into()),
360 Matcher::UrlEncoded("Policy".into(), r#"{"Version":"2012-10-17"}"#.into()),
361 Matcher::UrlEncoded("Tags.member.1.Key".into(), "table".into()),
362 Matcher::UrlEncoded("Tags.member.1.Value".into(), "main.default.t".into()),
363 Matcher::UrlEncoded("Tags.member.2.Key".into(), "env".into()),
364 ]))
365 .with_body(ASSUME_ROLE_RESPONSE)
366 .create_async()
367 .await;
368 let options = AssumeRoleOptions {
369 duration: Some(Duration::from_secs(900)),
370 session_name: Some("alice@table".into()),
371 policy: Some(r#"{"Version":"2012-10-17"}"#.into()),
372 external_id: Some("ext-42".into()),
373 source_identity: Some("alice".into()),
374 tags: vec![
375 ("table".into(), "main.default.t".into()),
376 ("env".into(), "dev".into()),
377 ],
378 ..options(&server)
379 };
380 let role =
381 assume_role_with_options("arn:aws:iam::1:role/vend", "us-east-1", base(), &options)
382 .await
383 .unwrap();
384 assert_eq!(role.credential.key_id, "ASIASESSION");
385 assert_eq!(role.credential.token.as_deref(), Some("session-token"));
386 assert_eq!(role.expiration.to_rfc3339(), "2030-01-01T00:15:00+00:00");
387 sts.assert_async().await;
388 }
389
390 #[tokio::test]
391 async fn assume_role_defaults() {
392 let mut server = mockito::Server::new_async().await;
393 let body = Arc::new(std::sync::Mutex::new(String::new()));
394 let seen = Arc::clone(&body);
395 let sts = server
396 .mock("POST", "/")
397 .with_body_from_request(move |request| {
398 *seen.lock().unwrap() =
399 String::from_utf8_lossy(request.body().unwrap()).into_owned();
400 ASSUME_ROLE_RESPONSE.into()
401 })
402 .create_async()
403 .await;
404 assume_role_with_options(
405 "arn:aws:iam::1:role/vend",
406 "us-east-1",
407 base(),
408 &options(&server),
409 )
410 .await
411 .unwrap();
412 sts.assert_async().await;
413 let body = body.lock().unwrap();
414 assert!(body.contains("DurationSeconds=3600"), "{body}");
415 assert!(
416 body.contains(&format!("RoleSessionName={DEFAULT_SESSION_NAME}")),
417 "{body}"
418 );
419 for absent in ["ExternalId", "SourceIdentity", "Tags", "Policy"] {
421 assert!(!body.contains(absent), "{absent} in {body}");
422 }
423 }
424}