retry InvalidIdentityToken errors (#1858)

aws · Aug 8, 2019 · 2f0d5af · 2f0d5af
1 parent f0c2725
commit 2f0d5af
Show file tree

Hide file tree

Showing 3 changed files with 214 additions and 26 deletions.
diff --git a/.changes/nextrelease/invalid_identity_token_retry.json b/.changes/nextrelease/invalid_identity_token_retry.json
@@ -0,0 +1,7 @@
+[
+    {
+      "type": "enhancement",
+      "category": "",
+      "description": "Retry InvalidIdentityToken errors for AssumeRoleWithWebIdentityCredentialProvider"
+    }
+]
diff --git a/src/Credentials/AssumeRoleWithWebIdentityCredentialProvider.php b/src/Credentials/AssumeRoleWithWebIdentityCredentialProvider.php
@@ -26,6 +26,11 @@ class AssumeRoleWithWebIdentityCredentialProvider
     /** @var StsClient */
     private $client;
 
+    /** @var integer */
+    private $retries;
+
+    /** @var integer */
+    private $attempts;
 
     /**
      * The constructor attempts to load config from environment variables.
@@ -53,6 +58,9 @@ public function __construct(array $config = [])
             throw new \InvalidArgumentException("'WebIdentityTokenFile' must be an absolute path.");
         }
 
+        $this->retries = isset($config['retries']) ? $config['retries'] : 3;
+        $this->attempts = 0;
+
         $this->session = isset($config['SessionName'])
             ? $config['SessionName']
             : 'aws-sdk-php-' . round(microtime(true) * 1000);
@@ -75,32 +83,49 @@ public function __construct(array $config = [])
      */
     public function __invoke()
     {
-        $client = $this->client;
-        try {
-            $token = file_get_contents($this->tokenFile);
-        } catch (\Exception $exception) {
-            throw new CredentialsException(
-                "Error reading WebIdentityTokenFile from " . $this->tokenFile,
-                0,
-                $exception
-            );
-        }
+        return Promise\coroutine(function () {
+            $client = $this->client;
+            $result = null;
+            while ($result == null) {
+                try {
+                    $token = file_get_contents($this->tokenFile);
+                } catch (\Exception $exception) {
+                    throw new CredentialsException(
+                        "Error reading WebIdentityTokenFile from " . $this->tokenFile,
+                        0,
+                        $exception
+                    );
+                }
+
+                $assumeParams = [
+                    'RoleArn' => $this->arn,
+                    'RoleSessionName' => $this->session,
+                    'WebIdentityToken' => $token
+                ];
+
+                try {
+                    $result = $client->assumeRoleWithWebIdentity($assumeParams);
+                } catch (\Exception $e) {
+                    if ($e->getAwsErrorCode() == 'InvalidIdentityToken') {
+                        if ($this->attempts < $this->retries) {
+                            sleep(pow(1.2, $this->attempts));
+                        } else {
+                            throw new CredentialsException(
+                                "InvalidIdentityToken, retries exhausted"
+                            );
+                        }
+                    } else {
+                        throw new CredentialsException(
+                            "Error assuming role from web identity credentials",
+                            0,
+                            $e
+                        );
+                    }
+                }
+                $this->attempts++;
+            }
 
-        $assumeParams = [
-            'RoleArn' => $this->arn,
-            'RoleSessionName' => $this->session,
-            'WebIdentityToken' => $token
-        ];
-
-        return $client->assumeRoleWithWebIdentityAsync($assumeParams)
-            ->then(function (Result $result) {
-                return $this->client->createCredentials($result);
-            })->otherwise(function (\Exception $exception) {
-                throw new CredentialsException(
-                    "Error assuming role from web identity credentials",
-                    0,
-                    $exception
-                );
-            });
+            yield $this->client->createCredentials($result);
+        });
     }
 }
diff --git a/tests/Credentials/AssumeRoleWithWebIdentityCredentialProviderTest.php b/tests/Credentials/AssumeRoleWithWebIdentityCredentialProviderTest.php
@@ -1,11 +1,13 @@
 <?php
 namespace Aws\Test\Credentials;
 
+use Aws\Command;
 use Aws\Credentials\AssumeRoleWithWebIdentityCredentialProvider;
 use Aws\Credentials\Credentials;
 use Aws\Exception\AwsException;
 use Aws\Result;
 use Aws\Sts\StsClient;
+use Aws\Sts\Exception\StsException;
 use Aws\Api\DateTimeResult;
 use GuzzleHttp\Promise;
 use GuzzleHttp\Promise\RejectedPromise;
@@ -193,4 +195,158 @@ public function testThrowsExceptionWhenRetrievingAssumeRoleCredentialFails()
             unlink($tokenPath);
         }
     }
+
+    public function testRetryInvalidIdentityToken()
+    {
+        $dir = $this->clearEnv();
+        $result = [
+            'Credentials' => [
+                'AccessKeyId'     => 'foo',
+                'SecretAccessKey' => 'bar',
+                'SessionToken'    => 'baz',
+                'Expiration'      => DateTimeResult::fromEpoch(time() + 10)
+            ],
+        ];
+        $retries = 1;
+
+        $sts = new StsClient([
+            'region' => 'us-west-2',
+            'version' => 'latest',
+            'credentials' => false,
+            'handler' => function () use (&$retries, $result) {
+                if (0 === $retries--) {
+                    return Promise\promise_for(new Result($result));
+                }
+
+                return new StsException(
+                    "foo",
+                    new Command("foo"),
+                    ['code' => 'InvalidIdentityToken']
+                );
+            }
+        ]);
+
+        $tokenPath = $dir . '/my-token.jwt';
+        file_put_contents($tokenPath, 'token');
+
+        $args['client'] = $sts;
+        $args['RoleArn'] = self::SAMPLE_ROLE_ARN;
+        $args['WebIdentityTokenFile'] = $tokenPath;
+        $provider = new AssumeRoleWithWebIdentityCredentialProvider($args);
+        $creds = $provider()->wait();
+        try {
+            $this->assertEquals('foo', $creds->getAccessKeyId());
+            $this->assertEquals('bar', $creds->getSecretKey());
+            $this->assertEquals('baz', $creds->getSecurityToken());
+            $this->assertInternalType('int', $creds->getExpiration());
+            $this->assertFalse($creds->isExpired());
+        } catch (\Exception $e) {
+            throw $e;
+        } finally {
+            unlink($tokenPath);
+        }
+    }
+
+    /**
+     * @expectedException \Aws\Exception\CredentialsException
+     * @expectedExceptionMessage InvalidIdentityToken, retries exhausted
+     */
+    public function testThrowsExceptionWhenInvalidIdentityTokenRetriesExhausted()
+    {
+        $dir = $this->clearEnv();
+        $result = [
+            'Credentials' => [
+                'AccessKeyId'     => 'foo',
+                'SecretAccessKey' => 'bar',
+                'SessionToken'    => 'baz',
+                'Expiration'      => DateTimeResult::fromEpoch(time() + 10)
+            ],
+        ];
+        $retries = 4;
+
+        $sts = new StsClient([
+            'region' => 'us-west-2',
+            'version' => 'latest',
+            'credentials' => false,
+            'handler' => function () use (&$retries, $result) {
+                if (0 === $retries--) {
+                    return Promise\promise_for(new Result($result));
+                }
+
+                return new StsException(
+                    "foo",
+                    new Command("foo"),
+                    ['code' => 'InvalidIdentityToken']
+                );
+            }
+        ]);
+
+        $tokenPath = $dir . '/my-token.jwt';
+        file_put_contents($tokenPath, 'token');
+
+        $args['client'] = $sts;
+        $args['RoleArn'] = self::SAMPLE_ROLE_ARN;
+        $args['WebIdentityTokenFile'] = $tokenPath;
+        $provider = new AssumeRoleWithWebIdentityCredentialProvider($args);
+        try {
+            $provider()->wait();
+        } catch (\Exception $e) {
+            throw $e;
+        } finally {
+            unlink($tokenPath);
+        }
+    }
+
+    /**
+     * @expectedException \Aws\Exception\CredentialsException
+     * @expectedExceptionMessage InvalidIdentityToken, retries exhausted
+     */
+    public function testCanDisableInvalidIdentityTokenRetries()
+    {
+        $dir = $this->clearEnv();
+        $result = [
+            'Credentials' => [
+                'AccessKeyId'     => 'foo',
+                'SecretAccessKey' => 'bar',
+                'SessionToken'    => 'baz',
+                'Expiration'      => DateTimeResult::fromEpoch(time() + 10)
+            ],
+        ];
+        $retries = 1;
+
+        $sts = new StsClient([
+            'region' => 'us-west-2',
+            'version' => 'latest',
+            'credentials' => false,
+            'handler' => function () use (&$retries, $result) {
+                if (0 === $retries--) {
+                    return Promise\promise_for(new Result($result));
+                }
+
+                return new StsException(
+                    "foo",
+                    new Command("foo"),
+                    ['code' => 'InvalidIdentityToken']
+                );
+            }
+        ]);
+
+        $tokenPath = $dir . '/my-token.jwt';
+        file_put_contents($tokenPath, 'token');
+
+        $args = [
+            'client' => $sts,
+            'RoleArn' => self::SAMPLE_ROLE_ARN,
+            'WebIdentityTokenFile' => $tokenPath,
+            'retries' => 0
+        ];
+        $provider = new AssumeRoleWithWebIdentityCredentialProvider($args);
+        try {
+            $provider()->wait();
+        } catch (\Exception $e) {
+            throw $e;
+        } finally {
+            unlink($tokenPath);
+        }
+    }
 }