Add SSE-C option on native-filesystem security mapping

trinodb · Dec 27, 2024 · be26d6d · be26d6d
1 parent 975e062
commit be26d6d
Show file tree

Hide file tree

Showing 12 changed files with 274 additions and 10 deletions.
diff --git a/docs/src/main/sphinx/object-storage/file-system-s3.md b/docs/src/main/sphinx/object-storage/file-system-s3.md
@@ -262,6 +262,8 @@ Example JSON configuration:
   - The name of the *extra credential* used to provide the IAM role.
 * - `s3.security-mapping.kms-key-id-credential-name`
   - The name of the *extra credential* used to provide the KMS-managed key ID.
+* - `s3.security-mapping.customer-key-credential-name`
+  - The name of the *extra credential* used to provide the server-side encryption with customer-provided keys (SSE-C).
 * - `s3.security-mapping.refresh-period`
   - How often to refresh the security mapping configuration, specified as a
     {ref}`prop-type-duration`. By default, the configuration is not refreshed.

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Context.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Context.java
@@ -25,6 +25,7 @@
 import java.util.Optional;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.CUSTOMER;
 import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.KMS;
 import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY;
 import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY;
@@ -70,6 +71,11 @@ public S3Context withCredentials(ConnectorIdentity identity)
         return this;
     }
 
+    public S3Context withCustomerKey(String sseCustomerKey)
+    {
+        return new S3Context(partSize, requesterPays, S3SseContext.withCustomerKey(sseCustomerKey), credentialsProviderOverride, cannedAcl, exclusiveWriteSupported);
+    }
+
     public S3Context withCredentialsProviderOverride(AwsCredentialsProvider credentialsProviderOverride)
     {
         return new S3Context(
@@ -109,5 +115,10 @@ public static S3SseContext withKmsKeyId(String kmsKeyId)
         {
             return new S3SseContext(KMS, Optional.ofNullable(kmsKeyId), Optional.empty());
         }
+
+        public static S3SseContext withCustomerKey(String sseCustomerKey)
+        {
+            return new S3SseContext(CUSTOMER, Optional.empty(), Optional.ofNullable(sseCustomerKey).map(S3SseCustomerKey::onAes256));
+        }
     }
 }
diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemLoader.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemLoader.java
@@ -46,6 +46,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.function.Function;
 
+import static com.google.common.base.Preconditions.checkState;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static io.trino.filesystem.s3.S3FileSystemConfig.RetryMode.getRetryStrategy;
 import static java.lang.Math.toIntExact;
@@ -108,9 +109,14 @@ public TrinoFileSystemFactory apply(Location location)
             S3Context context = this.context.withCredentials(identity);
 
             if (mapping.isPresent() && mapping.get().kmsKeyId().isPresent()) {
+                checkState(mapping.get().customerKey().isEmpty(), "Customer key (SSE-C) must not be used when KMS-managed key is set");
                 context = context.withKmsKeyId(mapping.get().kmsKeyId().get());
             }
 
+            if (mapping.isPresent() && mapping.get().customerKey().isPresent()) {
+                context = context.withCustomerKey(mapping.get().customerKey().get());
+            }
+
             return new S3FileSystem(uploadExecutor, client, preSigner, context);
         };
     }

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMapping.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMapping.java
@@ -42,6 +42,8 @@ public final class S3SecurityMapping
     private final Set<String> allowedIamRoles;
     private final Optional<String> kmsKeyId;
     private final Set<String> allowedKmsKeyIds;
+    private final Optional<String> customerKey;
+    private final Set<String> allowedCustomerKeys;
     private final Optional<AwsCredentials> credentials;
     private final boolean useClusterDefault;
     private final Optional<String> endpoint;
@@ -57,6 +59,8 @@ public S3SecurityMapping(
             @JsonProperty("allowedIamRoles") Optional<List<String>> allowedIamRoles,
             @JsonProperty("kmsKeyId") Optional<String> kmsKeyId,
             @JsonProperty("allowedKmsKeyIds") Optional<List<String>> allowedKmsKeyIds,
+            @JsonProperty("customerKey") Optional<String> customerKey,
+            @JsonProperty("allowedCustomerKeys") Optional<List<String>> allowedCustomerKeys,
             @JsonProperty("accessKey") Optional<String> accessKey,
             @JsonProperty("secretKey") Optional<String> secretKey,
             @JsonProperty("useClusterDefault") Optional<Boolean> useClusterDefault,
@@ -86,6 +90,10 @@ public S3SecurityMapping(
 
         this.allowedKmsKeyIds = ImmutableSet.copyOf(allowedKmsKeyIds.orElse(ImmutableList.of()));
 
+        this.customerKey = requireNonNull(customerKey, "customerKey is null");
+
+        this.allowedCustomerKeys = ImmutableSet.copyOf(allowedCustomerKeys.orElse(ImmutableList.of()));
+
         requireNonNull(accessKey, "accessKey is null");
         requireNonNull(secretKey, "secretKey is null");
         checkArgument(accessKey.isPresent() == secretKey.isPresent(), "accessKey and secretKey must be provided together");
@@ -96,6 +104,8 @@ public S3SecurityMapping(
         checkArgument(this.useClusterDefault != roleOrCredentialsArePresent, "must either allow useClusterDefault role or provide role and/or credentials");
 
         checkArgument(!this.useClusterDefault || this.kmsKeyId.isEmpty(), "KMS key ID cannot be provided together with useClusterDefault");
+        checkArgument(!this.useClusterDefault || this.customerKey.isEmpty(), "SSE Customer key cannot be provided together with useClusterDefault");
+        checkArgument(this.kmsKeyId.isEmpty() || this.customerKey.isEmpty(), "SSE Customer key cannot be provided together with KMS key ID");
 
         this.endpoint = requireNonNull(endpoint, "endpoint is null");
         this.region = requireNonNull(region, "region is null");
@@ -133,6 +143,16 @@ public Set<String> allowedKmsKeyIds()
         return allowedKmsKeyIds;
     }
 
+    public Optional<String> customerKey()
+    {
+        return customerKey;
+    }
+
+    public Set<String> allowedCustomerKeys()
+    {
+        return allowedCustomerKeys;
+    }
+
     public Optional<AwsCredentials> credentials()
     {
         return credentials;

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingConfig.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingConfig.java
@@ -31,6 +31,7 @@ public class S3SecurityMappingConfig
     private String jsonPointer = "";
     private String roleCredentialName;
     private String kmsKeyIdCredentialName;
+    private String customerKeyCredentialName;
     private Duration refreshPeriod;
     private String colonReplacement;
 
@@ -100,6 +101,19 @@ public S3SecurityMappingConfig setKmsKeyIdCredentialName(String kmsKeyIdCredenti
         return this;
     }
 
+    public Optional<String> getCustomerKeyCredentialName()
+    {
+        return Optional.ofNullable(customerKeyCredentialName);
+    }
+
+    @Config("s3.security-mapping.customer-key-credential-name")
+    @ConfigDescription("Name of the extra credential used to provide SSE Customer key")
+    public S3SecurityMappingConfig setCustomerKeyCredentialName(String customerKeyCredentialName)
+    {
+        this.customerKeyCredentialName = customerKeyCredentialName;
+        return this;
+    }
+
     public Optional<Duration> getRefreshPeriod()
     {
         return Optional.ofNullable(refreshPeriod);

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingProvider.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingProvider.java
@@ -33,6 +33,7 @@ final class S3SecurityMappingProvider
     private final Supplier<S3SecurityMappings> mappingsProvider;
     private final Optional<String> roleCredentialName;
     private final Optional<String> kmsKeyIdCredentialName;
+    private final Optional<String> customerKeyCredentialName;
     private final Optional<String> colonReplacement;
 
     @Inject
@@ -41,18 +42,21 @@ public S3SecurityMappingProvider(S3SecurityMappingConfig config, Supplier<S3Secu
         this(mappingsProvider(mappingsProvider, config.getRefreshPeriod()),
                 config.getRoleCredentialName(),
                 config.getKmsKeyIdCredentialName(),
+                config.getCustomerKeyCredentialName(),
                 config.getColonReplacement());
     }
 
     public S3SecurityMappingProvider(
             Supplier<S3SecurityMappings> mappingsProvider,
             Optional<String> roleCredentialName,
             Optional<String> kmsKeyIdCredentialName,
+            Optional<String> customerKeyCredentialName,
             Optional<String> colonReplacement)
     {
         this.mappingsProvider = requireNonNull(mappingsProvider, "mappingsProvider is null");
         this.roleCredentialName = requireNonNull(roleCredentialName, "roleCredentialName is null");
         this.kmsKeyIdCredentialName = requireNonNull(kmsKeyIdCredentialName, "kmsKeyIdCredentialName is null");
+        this.customerKeyCredentialName = requireNonNull(customerKeyCredentialName, "customerKeyCredentialName is null");
         this.colonReplacement = requireNonNull(colonReplacement, "colonReplacement is null");
     }
 
@@ -70,6 +74,7 @@ public Optional<S3SecurityMappingResult> getMapping(ConnectorIdentity identity,
                 selectRole(mapping, identity),
                 mapping.roleSessionName().map(name -> name.replace("${USER}", identity.getUser())),
                 selectKmsKeyId(mapping, identity),
+                selectCustomerKey(mapping, identity),
                 mapping.endpoint(),
                 mapping.region()));
     }
@@ -131,6 +136,31 @@ private Optional<String> getKmsKeyIdFromExtraCredential(ConnectorIdentity identi
         return kmsKeyIdCredentialName.map(name -> identity.getExtraCredentials().get(name));
     }
 
+    private Optional<String> selectCustomerKey(S3SecurityMapping mapping, ConnectorIdentity identity)
+    {
+        Optional<String> userSelected = getCustomerKeyFromExtraCredential(identity);
+
+        if (userSelected.isEmpty()) {
+            return mapping.customerKey();
+        }
+
+        String selected = userSelected.get();
+
+        // selected Customer key must match default or be allowed
+        if (!selected.equals(mapping.customerKey().orElse(null)) &&
+                !mapping.allowedCustomerKeys().contains(selected) &&
+                !mapping.allowedCustomerKeys().contains("*")) {
+            throw new AccessDeniedException("Selected Customer key is not allowed");
+        }
+
+        return userSelected;
+    }
+
+    private Optional<String> getCustomerKeyFromExtraCredential(ConnectorIdentity identity)
+    {
+        return customerKeyCredentialName.map(name -> identity.getExtraCredentials().get(name));
+    }
+
     private static Supplier<S3SecurityMappings> mappingsProvider(Supplier<S3SecurityMappings> supplier, Optional<Duration> refreshPeriod)
     {
         return refreshPeriod

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingResult.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingResult.java
@@ -26,6 +26,7 @@ record S3SecurityMappingResult(
         Optional<String> iamRole,
         Optional<String> roleSessionName,
         Optional<String> kmsKeyId,
+        Optional<String> customerKey,
         Optional<String> endpoint,
         Optional<String> region)
 {
@@ -35,6 +36,7 @@ record S3SecurityMappingResult(
         requireNonNull(iamRole, "iamRole is null");
         requireNonNull(roleSessionName, "roleSessionName is null");
         requireNonNull(kmsKeyId, "kmsKeyId is null");
+        requireNonNull(customerKey, "customerKey is null");
         requireNonNull(endpoint, "endpoint is null");
         requireNonNull(region, "region is null");
     }