From 93ce48c374ec9fc9f529ced26875967c1f908ef2 Mon Sep 17 00:00:00 2001
From: tpetr <tpetr@hubspot.com>
Date: Wed, 25 Nov 2015 09:41:56 -0500
Subject: [PATCH] support per-bucket creds for downloading artifacts

---
 .../singularity/s3/base/S3ArtifactDownloader.java | 10 +++++++++-
 .../base/config/SingularityS3Configuration.java   | 15 +++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/S3ArtifactDownloader.java b/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/S3ArtifactDownloader.java
index 45b881e593..09b4394860 100644
--- a/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/S3ArtifactDownloader.java
+++ b/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/S3ArtifactDownloader.java
@@ -57,13 +57,21 @@ public void download(S3Artifact s3Artifact, Path downloadTo) {
     }
   }
 
+  private AWSCredentials getCredentialsForBucket(String bucketName) {
+    if (configuration.getS3BucketCredentials().containsKey(bucketName)) {
+      return configuration.getS3BucketCredentials().get(bucketName).toAWSCredentials();
+    }
+
+    return new AWSCredentials(configuration.getS3AccessKey().get(), configuration.getS3SecretKey().get());
+  }
+
   private void downloadThrows(final S3Artifact s3Artifact, final Path downloadTo) throws Exception {
     log.info("Downloading {}", s3Artifact);
 
     Jets3tProperties jets3tProperties = Jets3tProperties.getInstance(Constants.JETS3T_PROPERTIES_FILENAME);
     jets3tProperties.setProperty("httpclient.socket-timeout-ms", Long.toString(configuration.getS3ChunkDownloadTimeoutMillis()));
 
-    final S3Service s3 = new RestS3Service(new AWSCredentials(configuration.getS3AccessKey().get(), configuration.getS3SecretKey().get()), null, null, jets3tProperties);
+    final S3Service s3 = new RestS3Service(getCredentialsForBucket(s3Artifact.getS3Bucket()), null, null, jets3tProperties);
 
     long length = 0;
 
diff --git a/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/config/SingularityS3Configuration.java b/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/config/SingularityS3Configuration.java
index c378dd1ebd..1d5e38edc4 100644
--- a/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/config/SingularityS3Configuration.java
+++ b/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/config/SingularityS3Configuration.java
@@ -2,6 +2,8 @@
 
 import static com.hubspot.mesos.JavaUtils.obfuscateValue;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import javax.validation.constraints.Min;
@@ -57,6 +59,10 @@ public class SingularityS3Configuration extends BaseRunnerConfiguration {
   @JsonProperty
   private String localDownloadPath = "/download";
 
+  @NotNull
+  @JsonProperty
+  private Map<String, SingularityS3Credentials> s3BucketCredentials = new HashMap<>();
+
   public SingularityS3Configuration() {
     super(Optional.<String>absent());
   }
@@ -133,6 +139,14 @@ public void setS3ChunkDownloadTimeoutMillis(long s3ChunkDownloadTimeoutMillis) {
     this.s3ChunkDownloadTimeoutMillis = s3ChunkDownloadTimeoutMillis;
   }
 
+  public Map<String, SingularityS3Credentials> getS3BucketCredentials() {
+    return s3BucketCredentials;
+  }
+
+  public void setS3BucketCredentials(Map<String, SingularityS3Credentials> s3BucketCredentials) {
+    this.s3BucketCredentials = s3BucketCredentials;
+  }
+
   @Override
   public String toString() {
     return "SingularityS3Configuration[" +
@@ -145,6 +159,7 @@ public String toString() {
             ", s3ChunkRetries=" + s3ChunkRetries +
             ", localDownloadHttpPort=" + localDownloadHttpPort +
             ", localDownloadPath='" + localDownloadPath + '\'' +
+            ", s3BucketCredentials=" + s3BucketCredentials +
             ']';
   }
 }