Merge pull request #551 from voxel51/web-identity-refresh

Always reload web identity token file
voxel51 · Feb 9, 2022 · 83dbc79 · 83dbc79
2 parents e8cba32 + f4368c4
commit 83dbc79
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/eta/core/storage.py b/eta/core/storage.py
@@ -688,7 +688,7 @@ def __init__(
 
         self._prefixes = tuple(prefixes)
         self._role_arn = None
-        self._web_identity_token = None
+        self._web_identity_token_file = None
         self._duration_seconds = None
         self._sts_client = None
 
@@ -716,13 +716,11 @@ def _make_session(self, credentials):
             if "role_arn" not in credentials:
                 return boto3.Session(**credentials)
 
-            # Create session with autorefreshing temporary credentials
+            # Create session with auto-refreshing temporary credentials
             role_arn = credentials["role_arn"]
             web_identity_token_file = credentials["web_identity_token_file"]
             region_name = credentials.get("region_name", None)
 
-            web_identity_token = etau.read_file(web_identity_token_file)
-
             sts_client = boto3.client("sts", region_name=region_name)
 
             try:
@@ -732,7 +730,7 @@ def _make_session(self, credentials):
                 duration_seconds = 3600
 
             self._role_arn = role_arn
-            self._web_identity_token = web_identity_token
+            self._web_identity_token_file = web_identity_token_file
             self._duration_seconds = duration_seconds
             self._sts_client = sts_client
 
@@ -754,10 +752,14 @@ def _make_session(self, credentials):
                 os.environ["AWS_PROFILE"] = aws_profile
 
     def _refresh_temporary_credentials(self):
+        # This token is refreshed periodically, so we re-read it just in case
+        # https://kubernetes.io/docs/tasks/configure-pod-container/configure-service-account/#service-account-token-volume-projection
+        web_identity_token = etau.read_file(self._web_identity_token_file)
+
         response = self._sts_client.assume_role_with_web_identity(
             RoleArn=self._role_arn,
             RoleSessionName="voxel51",
-            WebIdentityToken=self._web_identity_token,
+            WebIdentityToken=web_identity_token,
             DurationSeconds=self._duration_seconds,
         )