Merge pull request #1707 from HubSpot/encrypted_uploads

Enable server side encryption params for uploads
HubSpot · Feb 13, 2018 · 49968a5 · 49968a5
2 parents 57ce9ba + 4f06570
commit 49968a5
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 3 deletions.
diff --git a/...rc/main/java/com/hubspot/singularity/executor/task/SingularityExecutorTaskLogManager.java b/...rc/main/java/com/hubspot/singularity/executor/task/SingularityExecutorTaskLogManager.java
@@ -255,7 +255,7 @@ private boolean writeS3MetadataFile(String filenameHint, Path pathToS3Directory,
     }
 
     S3UploadMetadata s3UploadMetadata = new S3UploadMetadata(pathToS3Directory.toString(), globForS3Files, s3UploaderBucket, getS3KeyPattern(s3KeyPattern.or(taskDefinition.getExecutorData().getS3UploaderKeyPattern())), finished, Optional.<String> absent(),
-        Optional. absent(), Optional. absent(), Optional. absent(), Optional. absent(), s3StorageClass, applyS3StorageClassAfterBytes, Optional.of(finished), Optional.of(checkSubdirectories), Optional.absent(), Collections.emptyMap(), Optional.absent());
+        Optional. absent(), Optional. absent(), Optional. absent(), Optional. absent(), s3StorageClass, applyS3StorageClassAfterBytes, Optional.of(finished), Optional.of(checkSubdirectories), Optional.absent(), Collections.emptyMap(), Optional.absent(), Optional.absent());
 
     String s3UploadMetadataFileName = String.format("%s-%s%s", taskDefinition.getTaskId(), filenameHint, baseConfiguration.getS3UploaderMetadataSuffix());
 

diff --git a/...RunnerBase/src/main/java/com/hubspot/singularity/runner/base/shared/S3UploadMetadata.java b/...RunnerBase/src/main/java/com/hubspot/singularity/runner/base/shared/S3UploadMetadata.java
@@ -58,6 +58,7 @@ public class S3UploadMetadata {
   private final SingularityUploaderType uploaderType;
   private final Map<String, Object> gcsCredentials;
   private final Optional<String> gcsStorageClass;
+  private final Optional<String> encryptionKey;
 
   @JsonCreator
   public S3UploadMetadata(@JsonProperty("directory") String directory,
@@ -76,7 +77,8 @@ public S3UploadMetadata(@JsonProperty("directory") String directory,
                           @JsonProperty("checkSubdirectories") Optional<Boolean> checkSubdirectories,
                           @JsonProperty("uploaderType") Optional<SingularityUploaderType> uploaderType,
                           @JsonProperty("gcsCredentials") Map<String, Object> gcsCredentials,
-                          @JsonProperty("gcsStorageClass") Optional<String> gcsStorageClass) {
+                          @JsonProperty("gcsStorageClass") Optional<String> gcsStorageClass,
+                          @JsonProperty("encryptionKey") Optional<String> encryptionKey) {
     Preconditions.checkNotNull(directory);
     Preconditions.checkNotNull(fileGlob);
     Preconditions.checkNotNull(s3Bucket);
@@ -99,6 +101,7 @@ public S3UploadMetadata(@JsonProperty("directory") String directory,
     this.uploaderType = uploaderType.or(SingularityUploaderType.S3);
     this.gcsCredentials = gcsCredentials != null ? gcsCredentials : Collections.emptyMap();
     this.gcsStorageClass = gcsStorageClass;
+    this.encryptionKey = encryptionKey;
   }
 
 
@@ -209,6 +212,10 @@ public Map<String, Object> getGcsCredentials() {
     return gcsCredentials;
   }
 
+  public Optional<String> getEncryptionKey() {
+    return encryptionKey;
+  }
+
   @Override
   public String toString() {
     return "S3UploadMetadata{" +
@@ -228,6 +235,7 @@ public String toString() {
         ", checkSubdirectories=" + checkSubdirectories +
         ", uploaderType=" + uploaderType +
         ", gcsStorageClass=" + gcsStorageClass +
+        ", encryptionKey=" + encryptionKey +
         '}';
   }
 }
diff --git a/...tyS3Uploader/src/main/java/com/hubspot/singularity/s3uploader/SingularityGCSUploader.java b/...tyS3Uploader/src/main/java/com/hubspot/singularity/s3uploader/SingularityGCSUploader.java
@@ -17,6 +17,7 @@
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.storage.BlobInfo;
 import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.Storage.BlobWriteOption;
 import com.google.cloud.storage.StorageClass;
 import com.google.cloud.storage.StorageException;
 import com.google.cloud.storage.StorageOptions;
@@ -107,7 +108,11 @@ protected  void uploadSingle(int sequence, Path file) throws Exception {
       }
 
       try (FileInputStream fileInputStream = new FileInputStream(file.toFile())){
-        storage.create(blobInfoBuilder.build(), fileInputStream);
+        if (uploadMetadata.getEncryptionKey().isPresent()) {
+          storage.create(blobInfoBuilder.build(), fileInputStream, BlobWriteOption.encryptionKey(uploadMetadata.getEncryptionKey().get()));
+        } else {
+          storage.create(blobInfoBuilder.build(), fileInputStream);
+        }
         LOG.info("{} Uploaded {} in {}", logIdentifier, key, JavaUtils.duration(start));
         return true;
       } catch (StorageException se) {

diff --git a/...ityS3Uploader/src/main/java/com/hubspot/singularity/s3uploader/SingularityS3Uploader.java b/...ityS3Uploader/src/main/java/com/hubspot/singularity/s3uploader/SingularityS3Uploader.java
@@ -20,6 +20,7 @@
 import com.amazonaws.services.s3.model.ObjectMetadata;
 import com.amazonaws.services.s3.model.PartETag;
 import com.amazonaws.services.s3.model.PutObjectRequest;
+import com.amazonaws.services.s3.model.SSEAwsKeyManagementParams;
 import com.amazonaws.services.s3.model.StorageClass;
 import com.amazonaws.services.s3.model.UploadPartRequest;
 import com.github.rholder.retry.Retryer;
@@ -115,6 +116,9 @@ protected void uploadSingle(int sequence, Path file) throws Exception {
           if (maybeStorageClass.isPresent()) {
             putObjectRequest.setStorageClass(maybeStorageClass.get());
           }
+          if (uploadMetadata.getEncryptionKey().isPresent()) {
+            putObjectRequest.withSSEAwsKeyManagementParams(new SSEAwsKeyManagementParams(uploadMetadata.getEncryptionKey().get()));
+          }
           s3Client.putObject(putObjectRequest);
         }
       } catch (AmazonS3Exception se) {