From c1f0691fc72f7dbd727eb3fce7d105046303cc42 Mon Sep 17 00:00:00 2001
From: Daniel Salazar <dsalazar@fluidattacks.com>
Date: Tue, 6 Sep 2022 13:15:31 -0500
Subject: [PATCH] feat(back): #889 support retries

- Support retries in secrets for aws from gitlab
to deal with authentication errors
- Add documentation
---
 README.md                                     |  7 ++++
 .../default.nix                               |  2 ++
 .../template.sh                               | 32 ++++++++++++++++---
 .../secrets-for-aws-from-gitlab/default.nix   | 14 +++++---
 4 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 578a43a8..2f5786d8 100644
--- a/README.md
+++ b/README.md
@@ -2285,6 +2285,11 @@ Types:
       Duration in seconds of the session.
       Defaults to `3600`.
 
+    - retries (`ints.positive`): Optional.
+      Number of login retries before failing.
+      One retry per second.
+      Defaults to `15`.
+
 Example `makes.nix`:
 
 ```nix
@@ -2298,10 +2303,12 @@ Example `makes.nix`:
     makesDev = {
       roleArn = "arn:aws:iam::123456789012:role/dev";
       duration = 3600;
+      retries = 30;
     };
     makesProd = {
       roleArn = "arn:aws:iam::123456789012:role/prod";
       duration = 7200;
+      retries = 30;
     };
   };
   lintTerraform = {
diff --git a/src/args/make-secret-for-aws-from-gitlab/default.nix b/src/args/make-secret-for-aws-from-gitlab/default.nix
index 5d06b7f6..0108b076 100644
--- a/src/args/make-secret-for-aws-from-gitlab/default.nix
+++ b/src/args/make-secret-for-aws-from-gitlab/default.nix
@@ -6,12 +6,14 @@
 }: {
   duration,
   name,
+  retries,
   roleArn,
 }:
 makeTemplate {
   replace = {
     __argDuration__ = duration;
     __argName__ = toDerivationName name;
+    __argRetries__ = retries;
     __argRoleArn__ = roleArn;
   };
   name = "make-secret-for-aws-from-gitlab-for-${name}";
diff --git a/src/args/make-secret-for-aws-from-gitlab/template.sh b/src/args/make-secret-for-aws-from-gitlab/template.sh
index e24c6686..157a2c2e 100644
--- a/src/args/make-secret-for-aws-from-gitlab/template.sh
+++ b/src/args/make-secret-for-aws-from-gitlab/template.sh
@@ -8,22 +8,44 @@ function _get_credential {
 }
 
 function login {
+  # AWS STS args
   local args=(
     --role-arn "${1}"
     --role-session-name "gitlab-${CI_PROJECT_ID}-${CI_PIPELINE_ID}-${CI_JOB_ID}"
     --web-identity-token "${CI_JOB_JWT_V2}"
     --duration-seconds "${2}"
   )
+
+  # Retry logic
+  local retries="__argRetries__"
+  local wait="1"
+  local try="1"
+  local success="1"
+
+  # Session variables
   local session
   export AWS_ACCESS_KEY_ID
   export AWS_SECRET_ACCESS_KEY
   export AWS_SESSION_TOKEN
 
   : \
-    && session="$(aws sts assume-role-with-web-identity "${args[@]}")" \
-    && AWS_ACCESS_KEY_ID="$(_get_credential "AccessKeyId" "${session}")" \
-    && AWS_SECRET_ACCESS_KEY="$(_get_credential "SecretAccessKey" "${session}")" \
-    && AWS_SESSION_TOKEN="$(_get_credential "SessionToken" "${session}")"
+    && while [ "${try}" -le "${retries}" ]; do
+      if session="$(aws sts assume-role-with-web-identity "${args[@]}" 2> /dev/null)"; then
+        success="0" \
+          && break
+      else
+        info "Login failed. Attempt ${try} of ${retries}." \
+          && sleep "${wait}" \
+          && try=$((try + 1))
+      fi
+    done \
+    && if [ "${success}" == "0" ]; then
+      AWS_ACCESS_KEY_ID="$(_get_credential "AccessKeyId" "${session}")" \
+        && AWS_SECRET_ACCESS_KEY="$(_get_credential "SecretAccessKey" "${session}")" \
+        && AWS_SESSION_TOKEN="$(_get_credential "SessionToken" "${session}")"
+    else
+      error "Could not login to AWS."
+    fi
 }
 
 function main {
@@ -33,7 +55,7 @@ function main {
       info "Logging in as '__argName__' using GitLab OIDC." \
         && login "__argRoleArn__" "__argDuration__"
     else
-      warn "Looks like this job is not running on GitLab CI. Skipping."
+      warn "It looks like this job is not running on GitLab CI. Skipping."
     fi
 }
 
diff --git a/src/evaluator/modules/secrets-for-aws-from-gitlab/default.nix b/src/evaluator/modules/secrets-for-aws-from-gitlab/default.nix
index 76932615..47517d41 100644
--- a/src/evaluator/modules/secrets-for-aws-from-gitlab/default.nix
+++ b/src/evaluator/modules/secrets-for-aws-from-gitlab/default.nix
@@ -9,23 +9,29 @@
 }: let
   type = lib.types.submodule (_: {
     options = {
-      roleArn = lib.mkOption {
-        type = lib.types.str;
-      };
       duration = lib.mkOption {
         default = 3600;
         type = lib.types.ints.positive;
       };
+      retries = lib.mkOption {
+        default = 15;
+        type = lib.types.ints.positive;
+      };
+      roleArn = lib.mkOption {
+        type = lib.types.str;
+      };
     };
   });
   output = name: {
-    roleArn,
     duration,
+    retries,
+    roleArn,
   }: {
     name = "/secretsForAwsFromGitlab/${name}";
     value = makeSecretForAwsFromGitlab {
       inherit duration;
       inherit name;
+      inherit retries;
       inherit roleArn;
     };
   };