Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 37 additions & 10 deletions crates/aws_s3/src/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ use aws_sdk_s3::{
Client,
};
use aws_utils::{
are_checksums_disabled,
is_sse_disabled,
must_s3_config_from_env,
s3::S3Client,
};
Expand Down Expand Up @@ -216,15 +218,25 @@ impl<RT: Runtime> Storage for S3Storage<RT> {
async fn start_upload(&self) -> anyhow::Result<Box<BufferedUpload>> {
let key: ObjectKey = self.runtime.new_uuid_v4().to_string().try_into()?;
let s3_key = S3Key(self.key_prefix.clone() + &key);
let output = self
let mut upload_builder = self
.client
.create_multipart_upload()
.bucket(self.bucket.clone())
.key(&s3_key.0)
.server_side_encryption(ServerSideEncryption::Aes256)
.key(&s3_key.0);

// Add server-side encryption if not disabled for S3 compatibility
if !is_sse_disabled() {
upload_builder = upload_builder.server_side_encryption(ServerSideEncryption::Aes256);
}

// Add checksum algorithm if not disabled for S3 compatibility
if !are_checksums_disabled() {
// Because we're using multipart uploads, we're really specifying the part checksum
// algorithm here, so it needs to match what we use for each part.
.checksum_algorithm(ChecksumAlgorithm::Crc32)
upload_builder = upload_builder.checksum_algorithm(ChecksumAlgorithm::Crc32);
}

let output = upload_builder
.send()
.await
.context("Failed to create multipart upload")?;
Expand Down Expand Up @@ -254,15 +266,25 @@ impl<RT: Runtime> Storage for S3Storage<RT> {
async fn start_client_driven_upload(&self) -> anyhow::Result<ClientDrivenUploadToken> {
let key: ObjectKey = self.runtime.new_uuid_v4().to_string().try_into()?;
let s3_key = S3Key(self.key_prefix.clone() + &key);
let output = self
let mut upload_builder = self
.client
.create_multipart_upload()
.bucket(self.bucket.clone())
.key(&s3_key.0)
.server_side_encryption(ServerSideEncryption::Aes256)
.key(&s3_key.0);

// Add server-side encryption if not disabled for S3 compatibility
if !is_sse_disabled() {
upload_builder = upload_builder.server_side_encryption(ServerSideEncryption::Aes256);
}

// Add checksum algorithm if not disabled for S3 compatibility
if !are_checksums_disabled() {
// Because we're using multipart uploads, we're really specifying the part checksum
// algorithm here, so it needs to match what we use for each part.
.checksum_algorithm(ChecksumAlgorithm::Crc32)
upload_builder = upload_builder.checksum_algorithm(ChecksumAlgorithm::Crc32);
}

let output = upload_builder
.send()
.await
.context("Failed to create multipart upload")?;
Expand Down Expand Up @@ -565,15 +587,20 @@ impl<RT: Runtime> S3Upload<RT> {
let part_number = self.next_part_number()?;
crate::metrics::log_aws_s3_part_upload_size_bytes(data.len());

let builder = self
let mut builder = self
.client
.upload_part()
.checksum_algorithm(ChecksumAlgorithm::Crc32)
.body(ByteStream::from(data))
.bucket(self.bucket.clone())
.key(&self.s3_key.0)
.part_number(Into::<u16>::into(part_number) as i32)
.upload_id(self.upload_id.to_string());

// Add checksum algorithm if not disabled for S3 compatibility
if !are_checksums_disabled() {
builder = builder.checksum_algorithm(ChecksumAlgorithm::Crc32);
}

Ok(UploadPart {
part_number,
builder,
Expand Down
24 changes: 24 additions & 0 deletions crates/aws_utils/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,20 @@ static AWS_S3_FORCE_PATH_STYLE: LazyLock<bool> = LazyLock::new(|| {
.unwrap_or_default()
});

static AWS_S3_DISABLE_SSE: LazyLock<bool> = LazyLock::new(|| {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you'll also want to add them to self-hosted/docker/docker-compose.yml

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll take care of it

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alright. Let me know if there is anything else you want me to fix or change.

env::var("AWS_S3_DISABLE_SSE")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or_default()
});

static AWS_S3_DISABLE_CHECKSUMS: LazyLock<bool> = LazyLock::new(|| {
env::var("AWS_S3_DISABLE_CHECKSUMS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or_default()
});

/// Similar aws_config::from_env but returns an error if credentials or
/// region is are not. It also doesn't spew out log lines every time
/// credentials are accessed.
Expand Down Expand Up @@ -62,3 +76,13 @@ pub async fn must_s3_config_from_env() -> anyhow::Result<S3ConfigBuilder> {
s3_config_builder = s3_config_builder.force_path_style(*AWS_S3_FORCE_PATH_STYLE);
Ok(s3_config_builder)
}

/// Returns true if server-side encryption headers should be disabled
pub fn is_sse_disabled() -> bool {
*AWS_S3_DISABLE_SSE
}

/// Returns true if checksum headers should be disabled
pub fn are_checksums_disabled() -> bool {
*AWS_S3_DISABLE_CHECKSUMS
}