Improve tests for remote session credentials

trinodb · Aug 6, 2024 · 4c83263 · 4c83263
1 parent d956120
commit 4c83263
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/...roxy/src/test/java/io/trino/aws/proxy/server/testing/TestingCredentialsRolesProvider.java b/...roxy/src/test/java/io/trino/aws/proxy/server/testing/TestingCredentialsRolesProvider.java
@@ -69,7 +69,9 @@ public Optional<Credentials> credentials(String emulatedAccessKey, Optional<Stri
                 checkState(emulatedAccessKey.equals(session.sessionCredential.accessKey()), "emulatedAccessKey and session accessKey mismatch");
 
                 Credentials originalCredentials = requireNonNull(credentials.get(session.originalEmulatedAccessKey), "original credentials missing for: " + session.originalEmulatedAccessKey);
-                return Optional.of(Credentials.build(session.sessionCredential, originalCredentials.remote()));
+                return Optional.of(originalCredentials.remoteSessionRole()
+                        .map(remoteSessionRole -> Credentials.build(session.sessionCredential, originalCredentials.requiredRemoteCredential(), remoteSessionRole))
+                        .orElseGet(() -> Credentials.build(session.sessionCredential, originalCredentials.requiredRemoteCredential())));
             });
         }
 

diff --git a/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/testing/containers/S3Container.java b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/testing/containers/S3Container.java
@@ -75,12 +75,13 @@ public class S3Container
                "Statement": [
                   {
                      "Effect": "Allow",
-                     "Action": [
-                        "s3:*"
-                     ],
-                     "Resource": [
-                        "arn:aws:s3:::*"
-                     ]
+                     "Action": "s3:*",
+                     "Resource": "arn:aws:s3:::*",
+                     "Condition": {
+                       "StringEquals": {
+                         "aws:principaltype": "AssumedRole"
+                       }
+                     }
                   }
                ]
             }