Skip to content

Commit

Permalink
Fix IAM token expiration (#234)
Browse files Browse the repository at this point in the history
* Switch to custom session object
  • Loading branch information
yankovs authored Jan 15, 2024
1 parent bb8aba6 commit 88bfcbc
Showing 1 changed file with 23 additions and 18 deletions.
41 changes: 23 additions & 18 deletions karton/core/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
InstanceMetadataFetcher,
InstanceMetadataProvider,
)
from botocore.session import get_session
from redis import AuthenticationError, StrictRedis
from redis.client import Pipeline
from urllib3.response import HTTPResponse
Expand Down Expand Up @@ -120,7 +121,6 @@ def __init__(
config, identity=identity, service_info=service_info
)

session_token = None
endpoint = config.get("s3", "address")
access_key = config.get("s3", "access_key")
secret_key = config.get("s3", "secret_key")
Expand All @@ -136,22 +136,10 @@ def __init__(
)

if iam_auth:
iam_providers = [
ContainerProvider(),
InstanceMetadataProvider(
iam_role_fetcher=InstanceMetadataFetcher(
timeout=1000, num_attempts=2
)
),
]

for provider in iam_providers:
creds = provider.load()
if creds:
access_key = creds.access_key
secret_key = creds.secret_key
session_token = creds.token
break
s3_client = self.iam_auth_s3(endpoint)
if s3_client:
self.s3 = s3_client
return

if access_key is None or secret_key is None:
raise RuntimeError(
Expand All @@ -163,9 +151,26 @@ def __init__(
endpoint_url=endpoint,
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
aws_session_token=session_token,
)

def iam_auth_s3(self, endpoint: str):
boto_session = get_session()
iam_providers = [
ContainerProvider(),
InstanceMetadataProvider(
iam_role_fetcher=InstanceMetadataFetcher(timeout=1000, num_attempts=2)
),
]

for provider in iam_providers:
creds = provider.load()
if creds:
boto_session._credentials = creds # type: ignore
return boto3.Session(botocore_session=boto_session).client(
"s3",
endpoint_url=endpoint,
)

@staticmethod
def _validate_identity(identity: str):
disallowed_chars = [" ", "?"]
Expand Down

0 comments on commit 88bfcbc

Please sign in to comment.