Xuanwo commented on code in PR #5402:
URL: https://github.com/apache/arrow-rs/pull/5402#discussion_r1495228710


##########
object_store/src/aws/builder.rs:
##########
@@ -912,6 +1002,120 @@ fn parse_bucket_az(bucket: &str) -> Option<&str> {
     Some(bucket.strip_suffix("--x-s3")?.rsplit_once("--")?.1)
 }
 
+/// Encryption configuration options for S3.
+///
+/// These options are used to configure server-side encryption for S3 objects.
+/// To configure them, pass them to [`AmazonS3Builder::with_config`].
+///
+/// Both [SSE-KMS] and [DSSE-KMS] are supported. [SSE-C] is not yet supported.
+///
+/// [SSE-KMS]: 
https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingKMSEncryption.html
+/// [DSSE-KMS]: 
https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingDSSEncryption.html
+/// [SSE-C]: 
https://docs.aws.amazon.com/AmazonS3/latest/userguide/ServerSideEncryptionCustomerKeys.html
+#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy, Serialize, Deserialize)]
+#[non_exhaustive]
+pub enum S3EncryptionConfigKey {
+    /// Type of encryption to use. If set, must be one of "AES256", "aws:kms", 
or "aws:kms:dsse".
+    ServerSideEncryption,
+    /// The KMS key ID to use for server-side encryption. If set, 
ServerSideEncryption
+    /// must be "aws:kms" or "aws:kms:dsse".
+    KmsKeyId,
+    /// If set to true, will use the bucket's default KMS key for server-side 
encryption.
+    /// If set to false, will disable the use of the bucket's default KMS key 
for server-side encryption.
+    BucketKeyEnabled,
+}
+
+impl AsRef<str> for S3EncryptionConfigKey {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::ServerSideEncryption => "aws_server_side_encryption",
+            Self::KmsKeyId => "aws_sse_kms_key_id",
+            Self::BucketKeyEnabled => "aws_sse_bucket_key_enabled",
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+enum S3EncryptionType {
+    S3,
+    SseKms,
+    DsseKms,
+}
+
+impl crate::config::Parse for S3EncryptionType {
+    fn parse(s: &str) -> Result<Self> {
+        match s {
+            "AES256" => Ok(Self::S3),
+            "aws:kms" => Ok(Self::SseKms),
+            "aws:kms:dsse" => Ok(Self::DsseKms),
+            _ => Err(Error::InvalidEncryptionType { passed: s.into() }.into()),
+        }
+    }
+}
+
+impl From<&S3EncryptionType> for &'static str {
+    fn from(value: &S3EncryptionType) -> Self {
+        match value {
+            S3EncryptionType::S3 => "AES256",
+            S3EncryptionType::SseKms => "aws:kms",
+            S3EncryptionType::DsseKms => "aws:kms:dsse",
+        }
+    }
+}
+
+impl std::fmt::Display for S3EncryptionType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.into())
+    }
+}
+
+/// A sequence of headers to be sent for write requests that specify 
server-side
+/// encryption.
+///
+/// Whether these headers are sent depends on both the kind of encryption set
+/// and the kind of request being made.
+#[derive(Default, Clone)]
+pub struct S3EncryptionHeaders(pub HeaderMap);
+
+impl std::fmt::Debug for S3EncryptionHeaders {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // TODO: if we take a user-provided key, hide the key from debug 
output.
+        f.debug_map().entries(self.0.iter()).finish()
+    }
+}
+
+impl S3EncryptionHeaders {
+    fn try_new(
+        encryption_type: &S3EncryptionType,
+        key_id: Option<String>,
+        bucket_key_enabled: Option<bool>,
+    ) -> Result<Self> {
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            "x-amz-server-side-encryption",
+            HeaderValue::from_static(encryption_type.into()),

Review Comment:
   HeaderValue has an API called 
[set_sensitive](https://docs.rs/http/latest/http/header/struct.HeaderValue.html#method.set_sensitive).
 We can adopt this to prevent accidental leakage of our header value.



##########
object_store/src/aws/builder.rs:
##########
@@ -912,6 +1002,120 @@ fn parse_bucket_az(bucket: &str) -> Option<&str> {
     Some(bucket.strip_suffix("--x-s3")?.rsplit_once("--")?.1)
 }
 
+/// Encryption configuration options for S3.
+///
+/// These options are used to configure server-side encryption for S3 objects.
+/// To configure them, pass them to [`AmazonS3Builder::with_config`].
+///
+/// Both [SSE-KMS] and [DSSE-KMS] are supported. [SSE-C] is not yet supported.
+///
+/// [SSE-KMS]: 
https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingKMSEncryption.html
+/// [DSSE-KMS]: 
https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingDSSEncryption.html
+/// [SSE-C]: 
https://docs.aws.amazon.com/AmazonS3/latest/userguide/ServerSideEncryptionCustomerKeys.html
+#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy, Serialize, Deserialize)]
+#[non_exhaustive]
+pub enum S3EncryptionConfigKey {
+    /// Type of encryption to use. If set, must be one of "AES256", "aws:kms", 
or "aws:kms:dsse".
+    ServerSideEncryption,
+    /// The KMS key ID to use for server-side encryption. If set, 
ServerSideEncryption
+    /// must be "aws:kms" or "aws:kms:dsse".
+    KmsKeyId,
+    /// If set to true, will use the bucket's default KMS key for server-side 
encryption.
+    /// If set to false, will disable the use of the bucket's default KMS key 
for server-side encryption.
+    BucketKeyEnabled,
+}
+
+impl AsRef<str> for S3EncryptionConfigKey {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::ServerSideEncryption => "aws_server_side_encryption",
+            Self::KmsKeyId => "aws_sse_kms_key_id",
+            Self::BucketKeyEnabled => "aws_sse_bucket_key_enabled",
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+enum S3EncryptionType {
+    S3,
+    SseKms,
+    DsseKms,
+}
+
+impl crate::config::Parse for S3EncryptionType {
+    fn parse(s: &str) -> Result<Self> {
+        match s {
+            "AES256" => Ok(Self::S3),
+            "aws:kms" => Ok(Self::SseKms),
+            "aws:kms:dsse" => Ok(Self::DsseKms),
+            _ => Err(Error::InvalidEncryptionType { passed: s.into() }.into()),
+        }
+    }
+}
+
+impl From<&S3EncryptionType> for &'static str {
+    fn from(value: &S3EncryptionType) -> Self {
+        match value {
+            S3EncryptionType::S3 => "AES256",
+            S3EncryptionType::SseKms => "aws:kms",
+            S3EncryptionType::DsseKms => "aws:kms:dsse",
+        }
+    }
+}
+
+impl std::fmt::Display for S3EncryptionType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.into())
+    }
+}
+
+/// A sequence of headers to be sent for write requests that specify 
server-side
+/// encryption.
+///
+/// Whether these headers are sent depends on both the kind of encryption set
+/// and the kind of request being made.
+#[derive(Default, Clone)]
+pub struct S3EncryptionHeaders(pub HeaderMap);
+
+impl std::fmt::Debug for S3EncryptionHeaders {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // TODO: if we take a user-provided key, hide the key from debug 
output.
+        f.debug_map().entries(self.0.iter()).finish()
+    }
+}
+
+impl S3EncryptionHeaders {
+    fn try_new(
+        encryption_type: &S3EncryptionType,
+        key_id: Option<String>,
+        bucket_key_enabled: Option<bool>,
+    ) -> Result<Self> {
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            "x-amz-server-side-encryption",
+            HeaderValue::from_static(encryption_type.into()),
+        );
+        if let Some(key_id) = key_id {
+            headers.insert(
+                "x-amz-server-side-encryption-aws-kms-key-id",
+                key_id
+                    .try_into()
+                    .map_err(|err| Error::InvalidEncryptionHeader {
+                        header: "kms-key-id",
+                        source: Box::new(err),
+                    })?,
+            );
+        }
+        if let Some(bucket_key_enabled) = bucket_key_enabled {
+            headers.insert(

Review Comment:
   Maybe we can only insert this key while `bucket_key_enabled` is `true`? This 
can save one extra long header.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to