feat(aws): add a fallback for expiration key (#4455)

* feat(aws): add a fallback for `expiration` * fix(aws): intermittent test failures - extend the time range from `-2s,0s` to `-5s,+2s` * fix: `docs/config/README.md` readability Co-authored-by: David Knaack <davidkna@users.noreply.github.com> Co-authored-by: David Knaack <davidkna@users.noreply.github.com>
starship · Oct 15, 2022 · 5a2c85d · 5a2c85d
1 parent 865e68d
commit 5a2c85d
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 44 deletions.
diff --git a/docs/config/README.md b/docs/config/README.md
@@ -328,6 +328,9 @@ When using [AWSume](https://awsu.me) the profile
 is read from the `AWSUME_PROFILE` env var and the credentials expiration
 date is read from the `AWSUME_EXPIRATION` env var.
 
+When using [saml2aws](https://github.com/Versent/saml2aws) the expiration information obtained from `~/.aws/credentials`
+falls back to the `x_security_token_expires` key.
+
 ### Options
 
 | Option              | Default                                                           | Description                                                                                                 |

diff --git a/src/modules/aws.rs b/src/modules/aws.rs
@@ -131,8 +131,10 @@ fn get_credentials_duration(
         let creds = get_creds(context, aws_creds)?;
         let section = get_profile_creds(creds, aws_profile)?;
 
-        section
-            .get("expiration")
+        let expiration_keys = ["expiration", "x_security_token_expires"];
+        expiration_keys
+            .iter()
+            .find_map(|expiration_key| section.get(expiration_key))
             .and_then(|expiration| DateTime::parse_from_rfc3339(expiration).ok())
     }?;
 
@@ -655,54 +657,63 @@ credential_process = /opt/bin/awscreds-retriever
 
         let expiration_date = now_plus_half_hour.to_rfc3339_opts(chrono::SecondsFormat::Secs, true);
 
-        file.write_all(
-            format!(
-                "[astronauts]
+        let expiration_keys = ["expiration", "x_security_token_expires"];
+        expiration_keys.iter().for_each(|key| {
+            file.write_all(
+                format!(
+                    "[astronauts]
 aws_access_key_id=dummy
 aws_secret_access_key=dummy
-expiration={}
+{}={}
 ",
-                expiration_date
-            )
-            .as_bytes(),
-        )?;
-
-        let actual = ModuleRenderer::new("aws")
-            .env("AWS_PROFILE", "astronauts")
-            .env("AWS_REGION", "ap-northeast-2")
-            .env(
-                "AWS_SHARED_CREDENTIALS_FILE",
-                credentials_path.to_string_lossy().as_ref(),
+                    key, expiration_date
+                )
+                .as_bytes(),
             )
-            .collect();
-
-        let actual_variant = ModuleRenderer::new("aws")
-            .env("AWS_PROFILE", "astronauts")
-            .env("AWS_REGION", "ap-northeast-2")
-            .env(
-                "AWS_CREDENTIALS_FILE",
-                credentials_path.to_string_lossy().as_ref(),
-            )
-            .collect();
-
-        assert_eq!(
-            actual, actual_variant,
-            "both AWS_SHARED_CREDENTIALS_FILE and AWS_CREDENTIALS_FILE should work"
-        );
-
-        // In principle, "30m" should be correct. However, bad luck in scheduling
-        // on shared runners may delay it. Allow for up to 2 seconds of delay.
-        let possible_values = ["30m", "29m59s", "29m58s"];
-        let possible_values = possible_values.map(|duration| {
-            let segment_colored = format!("☁️  astronauts (ap-northeast-2) [{}] ", duration);
-            Some(format!(
-                "on {}",
-                Color::Yellow.bold().paint(segment_colored)
-            ))
+            .unwrap();
+
+            let actual = ModuleRenderer::new("aws")
+                .env("AWS_PROFILE", "astronauts")
+                .env("AWS_REGION", "ap-northeast-2")
+                .env(
+                    "AWS_SHARED_CREDENTIALS_FILE",
+                    credentials_path.to_string_lossy().as_ref(),
+                )
+                .collect();
+
+            let actual_variant = ModuleRenderer::new("aws")
+                .env("AWS_PROFILE", "astronauts")
+                .env("AWS_REGION", "ap-northeast-2")
+                .env(
+                    "AWS_CREDENTIALS_FILE",
+                    credentials_path.to_string_lossy().as_ref(),
+                )
+                .collect();
+
+            assert_eq!(
+                actual, actual_variant,
+                "both AWS_SHARED_CREDENTIALS_FILE and AWS_CREDENTIALS_FILE should work"
+            );
+
+            // In principle, "30m" should be correct. However, bad luck in scheduling
+            // on shared runners may delay it.
+            let possible_values = [
+                "30m2s", "30m1s", "30m", "29m59s", "29m58s", "29m57s", "29m56s", "29m55s",
+            ];
+            let possible_values = possible_values.map(|duration| {
+                let segment_colored = format!("☁️  astronauts (ap-northeast-2) [{}] ", duration);
+                Some(format!(
+                    "on {}",
+                    Color::Yellow.bold().paint(segment_colored)
+                ))
+            });
+
+            assert!(
+                possible_values.contains(&actual),
+                "time is not in range: {actual:?}"
+            );
         });
 
-        assert!(possible_values.contains(&actual));
-
         dir.close()
     }