From d7308203a92bc067a3cb5177a5c6c32981c40959 Mon Sep 17 00:00:00 2001
From: David Knaack <davidkna@users.noreply.github.com>
Date: Wed, 16 Mar 2022 23:14:52 +0100
Subject: [PATCH] fix(aws): accept sso credentials (#3718)

---
 docs/config/README.md |  2 +-
 src/modules/aws.rs    | 51 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index 3499b261..727a3198 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -272,7 +272,7 @@ format = "$all$directory$character"
 ## AWS
 
 The `aws` module shows the current AWS region and profile when
-credentials or a `credential_process` have been setup. This is based on
+credentials, a `credential_process` or a `sso_start_url` have been setup. This is based on
 `AWS_REGION`, `AWS_DEFAULT_REGION`, and `AWS_PROFILE` env var with
 `~/.aws/config` file. This module also shows an expiration timer when using temporary
 credentials.
diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index e63a41c8..c8fc9d38 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -119,8 +119,15 @@ fn alias_name(name: Option<String>, aliases: &HashMap<String, &str>) -> Option<S
         .or(name)
 }
 
-fn get_credential_process(context: &Context, aws_profile: Option<&Profile>) -> Option<String> {
-    let contents = read_file(get_config_file_path(context)?).ok()?;
+fn has_credential_process_or_sso(context: &Context, aws_profile: Option<&Profile>) -> bool {
+    let fp = match get_config_file_path(context) {
+        Some(fp) => fp,
+        None => return false,
+    };
+    let contents = match read_file(fp) {
+        Ok(contents) => contents,
+        Err(_) => return false,
+    };
 
     let profile_line = if let Some(aws_profile) = aws_profile {
         format!("[profile {}]", aws_profile)
@@ -128,15 +135,12 @@ fn get_credential_process(context: &Context, aws_profile: Option<&Profile>) -> O
         "[default]".to_string()
     };
 
-    let cred_proc_line = contents
+    contents
         .lines()
         .skip_while(|line| line != &profile_line)
         .skip(1)
         .take_while(|line| !line.starts_with('['))
-        .find(|line| line.starts_with("credential_process"))?;
-
-    let cred_proc = cred_proc_line.split('=').nth(1)?.trim();
-    Some(cred_proc.to_string())
+        .any(|line| line.starts_with("credential_process") || line.starts_with("sso_start_url"))
 }
 
 fn get_defined_credentials(context: &Context, aws_profile: Option<&Profile>) -> Option<String> {
@@ -182,7 +186,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     }
 
     // only display if credential_process is defined or has valid credentials
-    if get_credential_process(context, aws_profile.as_ref()).is_none()
+    if !has_credential_process_or_sso(context, aws_profile.as_ref())
         && get_defined_credentials(context, aws_profile.as_ref()).is_none()
     {
         return None;
@@ -844,6 +848,37 @@ credential_process = /opt/bin/awscreds-retriever
         dir.close()
     }
 
+    #[test]
+    fn sso_set() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let config_path = dir.path().join("config");
+        let mut file = File::create(&config_path)?;
+
+        file.write_all(
+            "[default]
+region = ap-northeast-2
+sso_start_url = https://starship.rs/sso
+sso_region = <SSO-Default-Region>
+sso_account_id = <AWS ACCOUNT ID>
+sso_role_name = <AWS-ROLE-NAME>
+"
+            .as_bytes(),
+        )?;
+
+        file.sync_all()?;
+
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_CONFIG_FILE", config_path.to_string_lossy().as_ref())
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  (ap-northeast-2) ")
+        ));
+
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
     #[test]
     fn access_key_env_var_set() {
         let actual = ModuleRenderer::new("aws")