jiachengdb commented on code in PR #6230:
URL: https://github.com/apache/arrow-rs/pull/6230#discussion_r1717370262
##########
object_store/src/aws/builder.rs:
##########
@@ -1067,35 +1100,89 @@ pub struct S3EncryptionHeaders(HeaderMap);
impl S3EncryptionHeaders {
fn try_new(
encryption_type: &S3EncryptionType,
- key_id: Option<String>,
+ encryption_kms_key_id: Option<String>,
bucket_key_enabled: Option<bool>,
+ encryption_customer_key_base64: Option<String>,
) -> Result<Self> {
let mut headers = HeaderMap::new();
- // Note: if we later add support for SSE-C, we should be sure to use
- // HeaderValue::set_sensitive to prevent the key from being logged.
- 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(
- "x-amz-server-side-encryption-bucket-key-enabled",
- HeaderValue::from_static(if bucket_key_enabled { "true" } else
{ "false" }),
- );
+ match encryption_type {
+ S3EncryptionType::S3 | S3EncryptionType::SseKms |
S3EncryptionType::DsseKms => {
+ headers.insert(
+ "x-amz-server-side-encryption",
+ HeaderValue::from_static(encryption_type.into()),
+ );
+ if let Some(key_id) = encryption_kms_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(
+ "x-amz-server-side-encryption-bucket-key-enabled",
+ HeaderValue::from_static(if bucket_key_enabled {
"true" } else { "false" }),
+ );
+ }
+ }
+ S3EncryptionType::SseC => {
+ headers.insert(
+ "x-amz-server-side-encryption-customer-algorithm",
+ HeaderValue::from_static("AES256"),
+ );
+ if let Some(key) = encryption_customer_key_base64 {
+ let mut header_value: HeaderValue =
+ key.clone()
+ .try_into()
+ .map_err(|err| Error::InvalidEncryptionHeader {
+ header:
"x-amz-server-side-encryption-customer-key",
+ source: Box::new(err),
+ })?;
+ header_value.set_sensitive(true);
+
headers.insert("x-amz-server-side-encryption-customer-key", header_value);
+
+ let decoded_key =
BASE64_STANDARD.decode(key.as_bytes()).map_err(|err| {
+ Error::InvalidEncryptionHeader {
+ header:
"x-amz-server-side-encryption-customer-key",
+ source: Box::new(err),
+ }
+ })?;
+ let mut hasher = Md5::new();
+ hasher.update(decoded_key);
+ let md5 = BASE64_STANDARD.encode(hasher.finalize());
+ let mut md5_header_value: HeaderValue =
+ md5.try_into()
+ .map_err(|err| Error::InvalidEncryptionHeader {
+ header:
"x-amz-server-side-encryption-customer-key-MD5",
+ source: Box::new(err),
+ })?;
+ md5_header_value.set_sensitive(true);
+ headers.insert(
+ "x-amz-server-side-encryption-customer-key-MD5",
+ md5_header_value,
+ );
+ } else {
+ return Err(Error::InvalidEncryptionHeader {
+ header: "x-amz-server-side-encryption-customer-key",
+ source: Box::new(std::io::Error::new(
+ std::io::ErrorKind::InvalidInput,
+ "Missing customer key",
+ )),
+ }
+ .into());
+ }
+ }
}
Ok(Self(headers))
}
+
+ pub fn header_map(&self) -> &HeaderMap {
Review Comment:
Thanks for the suggestions. Made the changes.
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
