Expose configuration needed for testing profile authorization (#344)

IsaevIlya · web-flow · commit 56d632d4df51 · 2025-07-04T13:35:57.000+01:00
Expose configuration needed for testing profile authorization
diff --git a/.github/workflows/python-integration.yml b/.github/workflows/python-integration.yml
@@ -32,11 +32,15 @@ jobs:
             region: ${{ vars.S3_REGION }}
             storage-class: ""
             endpoint-url: ${{ vars.S3_CUSTOM_ENDPOINT_URL }}
+            profile-role: ${{ vars.PROFILE_IAM_ROLE }}
+            profile-bucket: ${{ vars.S3_PROFILE_BUCKET }}
           - name: "S3 Express"
             bucket: ${{ vars.S3_EXPRESS_BUCKET }}
             region: ${{ vars.S3_EXPRESS_REGION }}
             storage-class: "EXPRESS_ONEZONE"
             endpoint-url: ""
+            profile-role: ${{ vars.PROFILE_IAM_ROLE }}
+            profile-bucket: ${{ vars.S3_EXPRESS_PROFILE_BUCKET }}
     permissions:
       id-token: write
       contents: read
@@ -88,6 +92,8 @@ jobs:
           CI_REGION=${{ matrix.test-run.region }} \
           CI_BUCKET=${{ matrix.test-run.bucket }} \
           CI_STORAGE_CLASS=${{ matrix.test-run.storage-class }} \
+          CI_PROFILE_ROLE=${{ matrix.test-run.profile-role }} \
+          CI_PROFILE_BUCKET=${{ matrix.test-run.profile-bucket }} \
           pytest s3torchconnector/tst/e2e --ignore-glob '*/**/test_e2e_s3_lightning_checkpoint.py' --ignore-glob '*/**/dcp' --ignore-glob '*/**/test_distributed_training.py' -n auto
 
       - name: s3torchconnector ${{ matrix.test-run.name }} distributed training integration tests
@@ -125,6 +131,8 @@ jobs:
           CI_BUCKET=${{ matrix.test-run.bucket }} \
           CI_STORAGE_CLASS=${{ matrix.test-run.storage-class }} \
           CI_CUSTOM_ENDPOINT_URL=${{ matrix.test-run.endpoint-url }} \
+          CI_PROFILE_ROLE=${{ matrix.test-run.profile-role }} \
+          CI_PROFILE_BUCKET=${{ matrix.test-run.profile-bucket }} \
           pytest s3torchconnectorclient/python/tst/integration -n auto
 
       - name: Save Cargo cache
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
@@ -12,6 +12,9 @@ env:
   S3_PREFIX: ${{ vars.S3_PREFIX }}
   S3_EXPRESS_BUCKET: ${{ vars.S3_EXPRESS_BUCKET }}
   S3_EXPRESS_REGION: ${{ vars.S3_EXPRESS_REGION }}
+  PROFILE_IAM_ROLE: ${{vars.PROFILE_IAM_ROLE}}
+  S3_PROFILE_BUCKET: ${{vars.S3_PROFILE_BUCKET}}
+  S3_EXPRESS_PROFILE_BUCKET: ${{vars.S3_EXPRESS_PROFILE_BUCKET}}
 
 jobs:
   generate_third_party_licenses:
diff --git a/run_cibuildwheel_on_ec2.sh b/run_cibuildwheel_on_ec2.sh
@@ -1,6 +1,6 @@
-if [ $# -ne 7 ]; then
-    echo "Invalid number of parameters, you need to provide role name, region name, bucket name, prefix, express region name and express bucket name, custom endpoint for s3 standard"
-    echo "Usage: $0 S3RoleName us-west-2 s3torchconnector-test-bucket-name prefix-name/ us-east-1 s3torchconnectorclient-express-bucket-name https://s3.amazon.com"
+if [ $# -ne 10 ]; then
+    echo "Invalid number of parameters, you need to provide role name, region name, bucket name, prefix, express region name and express bucket name, custom endpoint for s3 standard, auth profile arn and buckets names for testing auth profile"
+    echo "Usage: $0 S3RoleName us-west-2 s3torchconnector-test-bucket-name prefix-name/ us-east-1 s3torchconnectorclient-express-bucket-name https://s3.amazon.com arn:aws:iam::XXXXXXXXXXX:role/RoleName profile-test-bucket-name profile-test-express-bucket-name "
     exit 1
 fi
 
@@ -11,6 +11,9 @@ PREFIX=$4
 EXPRESS_REGION_NAME=$5
 EXPRESS_BUCKET_NAME=$6
 S3_CUSTOM_ENDPOINT_URL=$7
+PROFILE_IAM_ROLE=$8
+S3_PROFILE_BUCKET=$9
+S3_EXPRESS_PROFILE_BUCKET=${10}
 
 FILE_NAME="tmp_cred.json"
 # Set metadata token TTL to 6 hours
@@ -30,5 +33,8 @@ export S3_PREFIX=${PREFIX}
 export S3_EXPRESS_REGION=${EXPRESS_REGION_NAME}
 export S3_EXPRESS_BUCKET=${EXPRESS_BUCKET_NAME}
 export S3_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL}
+export PROFILE_IAM_ROLE=${PROFILE_IAM_ROLE}
+export S3_PROFILE_BUCKET=${S3_PROFILE_BUCKET}
+export S3_EXPRESS_PROFILE_BUCKET=${S3_EXPRESS_PROFILE_BUCKET}
 
 CIBW_MANYLINUX_X86_64_IMAGE=manylinux2014 CIBW_MANYLINUX_AARCH64_IMAGE=manylinux2014 cibuildwheel --output-dir wheelhouse --platform linux s3torchconnectorclient
diff --git a/s3torchconnector/tst/e2e/conftest.py b/s3torchconnector/tst/e2e/conftest.py
@@ -27,6 +27,8 @@ class BucketPrefixData(object):
     prefix: str
     storage_class: str = None
     contents: dict
+    profile_arn: str = None
+    profile_bucket: str = None
 
     def __init__(
         self,
@@ -35,12 +37,16 @@ def __init__(
         prefix: str,
         storage_class: str = None,
         contents: dict = None,
+        profile_arn: str = None,
+        profile_bucket: str = None,
     ):
         self.bucket = bucket
         self.prefix = prefix
         self.region = region
         self.storage_class = storage_class
         self.contents = contents or {}
+        self.profile_arn = profile_arn
+        self.profile_bucket = profile_bucket
 
     @property
     def s3_uri(self):
@@ -61,9 +67,22 @@ class BucketPrefixFixture(BucketPrefixData):
     to this instance, so other concurrent tests won't affect its state."""
 
     def __init__(
-        self, region: str, bucket: str, prefix: str, storage_class: str = None
+        self,
+        region: str,
+        bucket: str,
+        prefix: str,
+        storage_class: str = None,
+        profile_arn: str = None,
+        profile_bucket: str = None,
     ):
-        super().__init__(region, bucket, prefix, storage_class)
+        super().__init__(
+            region,
+            bucket,
+            prefix,
+            storage_class,
+            profile_arn=profile_arn,
+            profile_bucket=profile_bucket,
+        )
         session = boto3.Session(region_name=region)
         self.s3 = session.client("s3")
 
@@ -80,7 +99,13 @@ def get_data_snapshot(self):
         Useful when passing data to another process without serializing s3 client
         """
         return BucketPrefixData(
-            self.region, self.bucket, self.prefix, self.storage_class, self.contents
+            self.region,
+            self.bucket,
+            self.prefix,
+            self.storage_class,
+            self.contents,
+            self.profile_arn,
+            self.profile_bucket,
         )
 
 
@@ -90,12 +115,16 @@ def get_test_bucket_prefix(name: str) -> BucketPrefixFixture:
     prefix = getenv("CI_PREFIX")
     region = getenv("CI_REGION")
     storage_class = getenv("CI_STORAGE_CLASS", optional=True)
+    profile_arn = getenv("CI_PROFILE_ROLE", optional=True)
+    profile_bucket = getenv("CI_PROFILE_BUCKET", optional=True)
     assert prefix == "" or prefix.endswith("/")
 
     nonce = random.randrange(2**64)
     prefix = f"{prefix}{name}/{nonce}/"
 
-    return BucketPrefixFixture(region, bucket, prefix, storage_class)
+    return BucketPrefixFixture(
+        region, bucket, prefix, storage_class, profile_arn, profile_bucket
+    )
 
 
 @pytest.fixture
@@ -139,3 +168,8 @@ def _create_image_directory_fixture(num_image: int, image_size: int, node_name:
 @pytest.fixture
 def checkpoint_directory(request) -> BucketPrefixFixture:
     return get_test_bucket_prefix(f"{request.node.name}/checkpoint_directory")
+
+
+@pytest.fixture
+def empty_directory(request) -> BucketPrefixFixture:
+    return get_test_bucket_prefix(f"{request.node.name}/empty_directory")
diff --git a/s3torchconnector/tst/e2e/test_s3_client.py b/s3torchconnector/tst/e2e/test_s3_client.py
@@ -0,0 +1,26 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+import pytest
+from s3torchconnectorclient import S3Exception
+
+from s3torchconnector._s3client import S3Client
+
+HELLO_WORLD_DATA = b"Hello, World!\n"
+
+
+def test_no_access_objects_without_profile(empty_directory):
+    if empty_directory.profile_bucket is None:
+        pytest.skip("No profile bucket configured")
+
+    client = S3Client(
+        empty_directory.region,
+    )
+    filename = f"{empty_directory.prefix}hello_world.txt"
+
+    with pytest.raises(S3Exception):
+        put_stream = client.put_object(
+            empty_directory.profile_bucket,
+            filename,
+        )
+        put_stream.write(HELLO_WORLD_DATA)
diff --git a/s3torchconnectorclient/pyproject.toml b/s3torchconnectorclient/pyproject.toml
@@ -32,7 +32,7 @@ test = [
     "flake8",
     "black",
     "mypy",
-    "Pillow"
+    "Pillow<=11.2.1" # installation of the newer versions fails in manylinux2014 images
 ]
 
 [tool.setuptools.packages]
@@ -51,10 +51,10 @@ test-extras = "test"
 test-command = [
     "pytest {package}/python/tst/unit",
     "pytest {package}/../s3torchconnector/tst/unit --ignore {package}/../s3torchconnector/tst/unit/lightning --ignore {package}/../s3torchconnector/tst/unit/dcp",
-    "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} pytest {package}/python/tst/integration",
-    "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} pytest {package}/../s3torchconnector/tst/e2e --ignore {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py --ignore {package}/../s3torchconnector/tst/e2e/dcp",
-    "CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' pytest {package}/python/tst/integration",
-    "CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' pytest {package}/../s3torchconnector/tst/e2e --ignore {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py --ignore {package}/../s3torchconnector/tst/e2e/dcp",
+    "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} CI_PROFILE_ROLE=${PROFILE_IAM_ROLE} CI_PROFILE_BUCKET=${S3_PROFILE_BUCKET} pytest {package}/python/tst/integration",
+    "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} CI_PROFILE_ROLE=${PROFILE_IAM_ROLE} CI_PROFILE_BUCKET=${S3_PROFILE_BUCKET} pytest {package}/../s3torchconnector/tst/e2e --ignore {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py --ignore {package}/../s3torchconnector/tst/e2e/dcp",
+    "CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' CI_PROFILE_ROLE=${PROFILE_IAM_ROLE} CI_PROFILE_BUCKET=${S3_EXPRESS_PROFILE_BUCKET} pytest {package}/python/tst/integration",
+    "CI_STORAGE_CLASS=EXPRESS_ONEZONE CI_REGION=${S3_EXPRESS_REGION} CI_BUCKET=${S3_EXPRESS_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL='' CI_PROFILE_ROLE=${PROFILE_IAM_ROLE} CI_PROFILE_BUCKET=${S3_EXPRESS_PROFILE_BUCKET} pytest {package}/../s3torchconnector/tst/e2e --ignore {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py --ignore {package}/../s3torchconnector/tst/e2e/dcp",
     "python -m pip install -e '{package}/../s3torchconnector[lightning-tests]'",
     "pytest {package}/../s3torchconnector/tst/unit/lightning",
     "CI_STORAGE_CLASS='' CI_REGION=${S3_REGION} CI_BUCKET=${S3_BUCKET} CI_PREFIX=${S3_PREFIX} CI_CUSTOM_ENDPOINT_URL=${S3_CUSTOM_ENDPOINT_URL} pytest {package}/../s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py",
@@ -69,7 +69,10 @@ environment-pass = [
     "S3_PREFIX",
     "S3_EXPRESS_BUCKET",
     "S3_EXPRESS_REGION",
-    "S3_CUSTOM_ENDPOINT_URL"
+    "S3_CUSTOM_ENDPOINT_URL",
+    "PROFILE_IAM_ROLE",
+    "S3_PROFILE_BUCKET",
+    "S3_EXPRESS_PROFILE_BUCKET"
 ]
 before-build = "cp README.md s3torchconnectorclient; cp LICENSE s3torchconnectorclient/; cp THIRD-PARTY-LICENSES s3torchconnectorclient/; cp NOTICE s3torchconnectorclient/"
 build = ["cp38*", "cp39*", "cp310*", "cp311*", "cp312*"]
diff --git a/s3torchconnectorclient/python/tst/integration/conftest.py b/s3torchconnectorclient/python/tst/integration/conftest.py
@@ -33,6 +33,8 @@ class BucketPrefixFixture:
     storage_class: Optional[str] = getenv("CI_STORAGE_CLASS", optional=True)
     endpoint_url: Optional[str] = getenv("CI_CUSTOM_ENDPOINT_URL", optional=True)
     contents: dict = field(default_factory=dict)
+    profile_arn: Optional[str] = getenv("CI_PROFILE_ROLE", optional=True)
+    profile_bucket: Optional[str] = getenv("CI_PROFILE_BUCKET", optional=True)
 
     def __post_init__(self):
         assert self.prefix == "" or self.prefix.endswith("/")
diff --git a/s3torchconnectorclient/python/tst/integration/test_mountpoint_s3_integration.py b/s3torchconnectorclient/python/tst/integration/test_mountpoint_s3_integration.py
@@ -6,6 +6,7 @@
 import os
 import pickle
 import sys
+import tempfile
 import uuid
 import random
 import pytest
@@ -31,6 +32,7 @@
 
 HELLO_WORLD_DATA = b"Hello, World!\n"
 TEST_USER_AGENT_PREFIX = "integration-tests"
+TEST_PROFILE_NAME = "test-profile"
 
 
 @pytest.mark.parametrize("force_path_style", [False, True])
@@ -489,6 +491,67 @@ def test_copy_object_raises_when_source_key_does_not_exist(
         )
 
 
+def test_no_access_objects_without_profile(sample_directory: BucketPrefixFixture):
+    if sample_directory.profile_bucket is None:
+        pytest.skip("No profile bucket configured")
+
+    client = MountpointS3Client(
+        sample_directory.region,
+        TEST_USER_AGENT_PREFIX,
+    )
+    filename = f"{sample_directory.prefix}hello_world.txt"
+
+    with pytest.raises(S3Exception):
+        put_stream = client.put_object(
+            sample_directory.profile_bucket,
+            filename,
+        )
+        put_stream.write(HELLO_WORLD_DATA)
+
+
+def test_access_objects_with_profile(sample_directory: BucketPrefixFixture):
+    if sample_directory.profile_bucket is None:
+        pytest.skip("No profile bucket configured")
+
+    try:
+        tmp_file = tempfile.NamedTemporaryFile()
+        tmp_file.write(
+            f"""[profile default]
+aws_access_key_id = {os.getenv("AWS_ACCESS_KEY_ID")}
+aws_secret_access_key = {os.getenv("AWS_SECRET_ACCESS_KEY")}
+aws_session_token = {os.getenv("AWS_SESSION_TOKEN")}
+ 
+[profile {TEST_PROFILE_NAME}]
+role_arn = {sample_directory.profile_arn}
+region = {sample_directory.region}
+source_profile = default""".encode()
+        )
+        tmp_file.flush()
+        os.environ["AWS_CONFIG_FILE"] = tmp_file.name
+
+        client = MountpointS3Client(
+            sample_directory.region,
+            TEST_USER_AGENT_PREFIX,
+            profile=TEST_PROFILE_NAME,
+        )
+        filename = f"{sample_directory.prefix}hello_world.txt"
+        put_stream = client.put_object(
+            sample_directory.profile_bucket,
+            filename,
+        )
+
+        put_stream.write(HELLO_WORLD_DATA)
+        put_stream.close()
+
+        get_stream = client.get_object(
+            sample_directory.profile_bucket,
+            filename,
+        )
+        assert b"".join(get_stream) == HELLO_WORLD_DATA
+    finally:
+        os.environ["AWS_CONFIG_FILE"] = ""
+
+
 def _parse_list_result(stream: ListObjectStream, max_keys: int):
     object_infos = []
     i = 0
