rustfs/crates/e2e_test/src/checksum_upload_test.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! E2E tests for PutObject and MultipartUpload with checksums (Content-MD5, x-amz-checksum-*).
//! Verifies that uploads with Content-MD5 and x-amz-checksum-sha256 succeed and content is correct.

#[cfg(test)]
mod tests {
    use crate::common::{RustFSTestEnvironment, init_logging};
    use aws_sdk_s3::Client;
    use aws_sdk_s3::primitives::{ByteStream, SdkBody};
    use aws_sdk_s3::types::{ChecksumAlgorithm, ChecksumMode, CompletedMultipartUpload, CompletedPart};
    use base64::Engine;
    use bytes::Bytes;
    use futures::StreamExt;
    use http_body::Frame;
    use http_body_util::StreamBody;
    use rustfs_rio::{Checksum, ChecksumType as RioChecksumType};
    use serial_test::serial;
    use sha2::{Digest, Sha256};
    use tracing::info;

    fn create_s3_client(env: &RustFSTestEnvironment) -> Client {
        env.create_s3_client()
    }

    async fn create_bucket(client: &Client, bucket: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        match client.create_bucket().bucket(bucket).send().await {
            Ok(_) => {
                info!("Bucket {} created successfully", bucket);
                Ok(())
            }
            Err(e) => {
                if e.to_string().contains("BucketAlreadyOwnedByYou") || e.to_string().contains("BucketAlreadyExists") {
                    info!("Bucket {} already exists", bucket);
                    Ok(())
                } else {
                    Err(Box::new(e))
                }
            }
        }
    }

    fn content_md5_base64(body: &[u8]) -> String {
        let digest = md5::compute(body);
        base64::engine::general_purpose::STANDARD.encode(digest.as_slice())
    }

    fn checksum_sha256_base64(body: &[u8]) -> String {
        let digest = Sha256::digest(body);
        base64::engine::general_purpose::STANDARD.encode(digest.as_slice())
    }

    fn checksum_crc64nvme_base64(body: &[u8]) -> String {
        Checksum::new_from_data(RioChecksumType::CRC64_NVME, body)
            .expect("crc64nvme checksum")
            .encoded
    }

    fn streamed_body_70kib_of_a() -> ByteStream {
        let bytes = Bytes::from_static(&[b'a'; 1024]);
        let stream = futures::stream::repeat_with(move || {
            let frame = Frame::data(bytes.clone());
            Ok::<_, std::io::Error>(frame)
        });
        let body = WithSizeHint::new(StreamBody::new(stream.take(70)), 70 * 1024);
        ByteStream::new(SdkBody::from_body_1_x(body))
    }

    struct WithSizeHint<T> {
        inner: T,
        size_hint: usize,
    }

    impl<T> WithSizeHint<T> {
        fn new(inner: T, size_hint: usize) -> Self {
            Self { inner, size_hint }
        }
    }

    impl<T> http_body::Body for WithSizeHint<T>
    where
        T: http_body::Body + Unpin,
    {
        type Data = T::Data;
        type Error = T::Error;

        fn poll_frame(
            self: std::pin::Pin<&mut Self>,
            cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
            let this = self.get_mut();
            std::pin::Pin::new(&mut this.inner).poll_frame(cx)
        }

        fn is_end_stream(&self) -> bool {
            self.inner.is_end_stream()
        }

        fn size_hint(&self) -> http_body::SizeHint {
            let mut hint = self.inner.size_hint();
            hint.set_exact(self.size_hint as u64);
            hint
        }
    }

    /// PutObject with Content-MD5: upload succeeds and GetObject returns same content.
    #[tokio::test]
    #[serial]
    async fn test_put_object_with_content_md5() {
        init_logging();
        info!("TEST: PutObject with Content-MD5");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-checksum-md5";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        let key = "obj-with-md5.txt";
        let content = b"Hello world with Content-MD5 checksum";
        let content_md5 = content_md5_base64(content);

        let result = client
            .put_object()
            .bucket(bucket)
            .key(key)
            .body(ByteStream::from_static(content))
            .content_md5(&content_md5)
            .send()
            .await;

        assert!(result.is_ok(), "PutObject with Content-MD5 failed: {:?}", result.err());

        let get_result = client.get_object().bucket(bucket).key(key).send().await;
        assert!(get_result.is_ok(), "GetObject failed: {:?}", get_result.err());
        let body_bytes = get_result.unwrap().body.collect().await.expect("collect body").into_bytes();
        assert_eq!(body_bytes.as_ref(), content, "GetObject body must match uploaded content");
        info!("PASSED: PutObject with Content-MD5 and GetObject content match");
    }

    /// PutObject with x-amz-checksum-sha256: upload succeeds and GetObject returns same content.
    #[tokio::test]
    #[serial]
    async fn test_put_object_with_checksum_sha256() {
        init_logging();
        info!("TEST: PutObject with x-amz-checksum-sha256");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-checksum-sha256";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        let key = "obj-with-sha256.txt";
        let content = b"Hello world with x-amz-checksum-sha256";
        let checksum = checksum_sha256_base64(content);

        let result = client
            .put_object()
            .bucket(bucket)
            .key(key)
            .body(ByteStream::from_static(content))
            .checksum_sha256(&checksum)
            .send()
            .await;

        assert!(result.is_ok(), "PutObject with checksum_sha256 failed: {:?}", result.err());

        let get_result = client.get_object().bucket(bucket).key(key).send().await;
        assert!(get_result.is_ok(), "GetObject failed: {:?}", get_result.err());
        let body_bytes = get_result.unwrap().body.collect().await.expect("collect body").into_bytes();
        assert_eq!(body_bytes.as_ref(), content, "GetObject body must match uploaded content");
        info!("PASSED: PutObject with checksum_sha256 and GetObject content match");
    }

    /// Mirrors `s3s-e2e` behavior: only request `checksum_algorithm`, then expect
    /// both PutObject and GetObject(checksum_mode=enabled) to expose the same checksum.
    #[tokio::test]
    #[serial]
    async fn test_put_object_with_checksum_algorithm_only() {
        init_logging();
        info!("TEST: PutObject with checksum_algorithm only");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-checksum-algorithm-only";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        let key = "obj-with-checksum-algorithm-only.txt";
        let content = vec![b'a'; 70 * 1024];

        let put_resp = client
            .put_object()
            .bucket(bucket)
            .key(key)
            .checksum_algorithm(ChecksumAlgorithm::Crc32)
            .body(ByteStream::from(content.clone()))
            .send()
            .await
            .expect("PutObject with checksum_algorithm should succeed");

        let put_checksum = put_resp
            .checksum_crc32()
            .expect("PutObject should return checksum_crc32 when checksum_algorithm is used")
            .to_string();

        let mut get_resp = client
            .get_object()
            .bucket(bucket)
            .key(key)
            .checksum_mode(ChecksumMode::Enabled)
            .send()
            .await
            .expect("GetObject should succeed");

        let body_bytes = std::mem::replace(&mut get_resp.body, ByteStream::new(aws_sdk_s3::primitives::SdkBody::empty()))
            .collect()
            .await
            .expect("collect body")
            .into_bytes();

        assert_eq!(body_bytes.as_ref(), content.as_slice(), "GetObject body must match uploaded content");
        assert_eq!(
            get_resp.checksum_crc32().map(str::to_string),
            Some(put_checksum),
            "GetObject(checksum_mode=enabled) should expose the stored CRC32 checksum"
        );
    }

    /// Matches the `s3s-e2e` streaming upload shape more closely than `ByteStream::from(Vec<u8>)`.
    #[tokio::test]
    #[serial]
    async fn test_put_object_with_checksum_algorithm_only_streaming_body() {
        init_logging();
        info!("TEST: PutObject with checksum_algorithm only using streaming body");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-checksum-algorithm-streaming";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        let key = "obj-with-checksum-algorithm-streaming.txt";
        let expected_content = vec![b'a'; 70 * 1024];

        let put_resp = client
            .put_object()
            .bucket(bucket)
            .key(key)
            .checksum_algorithm(ChecksumAlgorithm::Crc32)
            .body(streamed_body_70kib_of_a())
            .send()
            .await
            .expect("PutObject with streaming checksum_algorithm should succeed");

        let put_checksum = put_resp
            .checksum_crc32()
            .expect("PutObject should return checksum_crc32 for streaming checksum_algorithm uploads")
            .to_string();

        let mut get_resp = client
            .get_object()
            .bucket(bucket)
            .key(key)
            .checksum_mode(ChecksumMode::Enabled)
            .send()
            .await
            .expect("GetObject should succeed");

        let body_bytes = std::mem::replace(&mut get_resp.body, ByteStream::new(SdkBody::empty()))
            .collect()
            .await
            .expect("collect body")
            .into_bytes();

        assert_eq!(
            body_bytes.as_ref(),
            expected_content.as_slice(),
            "GetObject body must match uploaded content"
        );
        assert_eq!(
            get_resp.checksum_crc32().map(str::to_string),
            Some(put_checksum),
            "GetObject(checksum_mode=enabled) should expose the stored CRC32 checksum for streaming uploads"
        );
    }

    /// Multipart upload with checksum: CreateMultipartUpload, UploadPart(s) with checksum_sha256, CompleteMultipartUpload; then GetObject verifies content.
    /// Uses part size >= 5MB (server minimum) for two parts.
    #[tokio::test]
    #[serial]
    async fn test_multipart_upload_with_checksum() {
        init_logging();
        info!("TEST: MultipartUpload with checksum (checksum_sha256 on parts)");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-multipart-checksum";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        let key = "multipart-with-checksum.bin";
        const PART_SIZE: usize = 6 * 1024 * 1024; // 6 MB per part (>= 5MB minimum)
        let part1: Vec<u8> = (0..PART_SIZE).map(|i| (i % 256) as u8).collect();
        let part2: Vec<u8> = (0..PART_SIZE).map(|i| ((i + 1) % 256) as u8).collect();
        let full_content: Vec<u8> = part1.iter().chain(part2.iter()).copied().collect();

        let create_result = client
            .create_multipart_upload()
            .bucket(bucket)
            .key(key)
            .checksum_algorithm(aws_sdk_s3::types::ChecksumAlgorithm::Sha256)
            .send()
            .await
            .expect("Failed to create multipart upload");

        let upload_id = create_result.upload_id().expect("No upload_id").to_string();

        let checksum1 = checksum_sha256_base64(&part1);
        let upload1 = client
            .upload_part()
            .bucket(bucket)
            .key(key)
            .upload_id(&upload_id)
            .part_number(1)
            .body(ByteStream::from(part1.clone()))
            .checksum_sha256(&checksum1)
            .send()
            .await
            .expect("Failed to upload part 1");

        let etag1 = upload1.e_tag().expect("No etag part 1").to_string();
        let checksum_sha256_1 = upload1.checksum_sha256().map(|s| s.to_string());

        let checksum2 = checksum_sha256_base64(&part2);
        let upload2 = client
            .upload_part()
            .bucket(bucket)
            .key(key)
            .upload_id(&upload_id)
            .part_number(2)
            .body(ByteStream::from(part2.clone()))
            .checksum_sha256(&checksum2)
            .send()
            .await
            .expect("Failed to upload part 2");

        let etag2 = upload2.e_tag().expect("No etag part 2").to_string();
        let checksum_sha256_2 = upload2.checksum_sha256().map(|s| s.to_string());

        let mut part1_builder = CompletedPart::builder().part_number(1).e_tag(etag1);
        if let Some(ref cs) = checksum_sha256_1 {
            part1_builder = part1_builder.checksum_sha256(cs);
        }
        let mut part2_builder = CompletedPart::builder().part_number(2).e_tag(etag2);
        if let Some(ref cs) = checksum_sha256_2 {
            part2_builder = part2_builder.checksum_sha256(cs);
        }

        let completed_parts = vec![part1_builder.build(), part2_builder.build()];
        let completed_upload = CompletedMultipartUpload::builder().set_parts(Some(completed_parts)).build();

        let complete_result = client
            .complete_multipart_upload()
            .bucket(bucket)
            .key(key)
            .upload_id(&upload_id)
            .multipart_upload(completed_upload)
            .send()
            .await;

        assert!(complete_result.is_ok(), "CompleteMultipartUpload failed: {:?}", complete_result.err());

        let get_result = client.get_object().bucket(bucket).key(key).send().await;
        assert!(get_result.is_ok(), "GetObject failed: {:?}", get_result.err());
        let body_bytes = get_result.unwrap().body.collect().await.expect("collect body").into_bytes();
        assert_eq!(
            body_bytes.as_ref(),
            full_content.as_slice(),
            "GetObject body must match concatenated parts"
        );
        info!("PASSED: MultipartUpload with checksum and GetObject content match");
    }

    /// Mirrors `s3s-e2e` multipart behavior: request checksum algorithm at MPU creation,
    /// rely on auto checksum handling during UploadPart, and expect CompleteMultipartUpload to succeed.
    #[tokio::test]
    #[serial]
    async fn test_multipart_upload_with_crc32_algorithm_only() {
        init_logging();
        info!("TEST: MultipartUpload with checksum_algorithm only (CRC32)");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-multipart-checksum-crc32-auto";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        let key = "multipart-with-crc32-auto.bin";
        let part1_content = "a".repeat(5 * 1024 * 1024 + 1);
        let part2_content = "b".repeat(1024);

        let create_resp = client
            .create_multipart_upload()
            .bucket(bucket)
            .key(key)
            .checksum_algorithm(ChecksumAlgorithm::Crc32)
            .send()
            .await
            .expect("CreateMultipartUpload should succeed");

        let upload_id = create_resp.upload_id().expect("upload_id should be present");

        let part1_resp = client
            .upload_part()
            .bucket(bucket)
            .key(key)
            .upload_id(upload_id)
            .part_number(1)
            .body(ByteStream::from(part1_content.clone().into_bytes()))
            .send()
            .await
            .expect("UploadPart 1 should succeed");
        let part1_checksum = part1_resp
            .checksum_crc32()
            .expect("UploadPart 1 should return checksum_crc32")
            .to_string();

        let part2_resp = client
            .upload_part()
            .bucket(bucket)
            .key(key)
            .upload_id(upload_id)
            .part_number(2)
            .body(ByteStream::from(part2_content.clone().into_bytes()))
            .send()
            .await
            .expect("UploadPart 2 should succeed");
        let part2_checksum = part2_resp
            .checksum_crc32()
            .expect("UploadPart 2 should return checksum_crc32")
            .to_string();

        let completed_upload = CompletedMultipartUpload::builder()
            .parts(
                CompletedPart::builder()
                    .part_number(1)
                    .e_tag(part1_resp.e_tag().expect("etag part 1"))
                    .checksum_crc32(part1_checksum)
                    .build(),
            )
            .parts(
                CompletedPart::builder()
                    .part_number(2)
                    .e_tag(part2_resp.e_tag().expect("etag part 2"))
                    .checksum_crc32(part2_checksum)
                    .build(),
            )
            .build();

        client
            .complete_multipart_upload()
            .bucket(bucket)
            .key(key)
            .upload_id(upload_id)
            .multipart_upload(completed_upload)
            .send()
            .await
            .expect("CompleteMultipartUpload should succeed");

        let body_bytes = client
            .get_object()
            .bucket(bucket)
            .key(key)
            .send()
            .await
            .expect("GetObject should succeed")
            .body
            .collect()
            .await
            .expect("collect body")
            .into_bytes();

        let expected_content = format!("{part1_content}{part2_content}");
        assert_eq!(
            body_bytes.as_ref(),
            expected_content.as_bytes(),
            "completed multipart object must match concatenated parts"
        );
    }

    /// Regression test for issue #2282:
    /// CRC64NVME full-object checksum should match between direct PutObject and multipart upload.
    #[tokio::test]
    #[serial]
    async fn test_crc64nvme_matches_between_put_object_and_multipart_upload() {
        init_logging();
        info!("TEST: CRC64NVME matches between direct PutObject and multipart upload");

        let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
        env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");

        let client = create_s3_client(&env);
        let bucket = "test-crc64nvme-multipart-match";
        create_bucket(&client, bucket).await.expect("Failed to create bucket");

        const PART_SIZE: usize = 6 * 1024 * 1024;
        let part1: Vec<u8> = (0..PART_SIZE).map(|i| (i % 251) as u8).collect();
        let part2: Vec<u8> = (0..PART_SIZE).map(|i| ((i + 17) % 251) as u8).collect();
        let content: Vec<u8> = part1.iter().chain(part2.iter()).copied().collect();

        let direct_key = "crc64nvme-direct.bin";
        let multipart_key = "crc64nvme-multipart.bin";
        let full_checksum = checksum_crc64nvme_base64(&content);
        let part1_checksum = checksum_crc64nvme_base64(&part1);
        let part2_checksum = checksum_crc64nvme_base64(&part2);

        client
            .put_object()
            .bucket(bucket)
            .key(direct_key)
            .body(ByteStream::from(content.clone()))
            .checksum_algorithm(ChecksumAlgorithm::Crc64Nvme)
            .checksum_crc64_nvme(full_checksum.clone())
            .send()
            .await
            .expect("Failed to put direct object with CRC64NVME");

        let create_result = client
            .create_multipart_upload()
            .bucket(bucket)
            .key(multipart_key)
            .checksum_algorithm(ChecksumAlgorithm::Crc64Nvme)
            .send()
            .await
            .expect("Failed to create multipart upload");

        let upload_id = create_result.upload_id().expect("No upload_id").to_string();

        let upload1 = client
            .upload_part()
            .bucket(bucket)
            .key(multipart_key)
            .upload_id(&upload_id)
            .part_number(1)
            .body(ByteStream::from(part1.clone()))
            .checksum_algorithm(ChecksumAlgorithm::Crc64Nvme)
            .checksum_crc64_nvme(part1_checksum)
            .send()
            .await
            .expect("Failed to upload multipart part 1");

        let upload2 = client
            .upload_part()
            .bucket(bucket)
            .key(multipart_key)
            .upload_id(&upload_id)
            .part_number(2)
            .body(ByteStream::from(part2.clone()))
            .checksum_algorithm(ChecksumAlgorithm::Crc64Nvme)
            .checksum_crc64_nvme(part2_checksum)
            .send()
            .await
            .expect("Failed to upload multipart part 2");

        let completed_upload = CompletedMultipartUpload::builder()
            .parts(
                CompletedPart::builder()
                    .part_number(1)
                    .e_tag(upload1.e_tag().expect("No etag for part 1"))
                    .checksum_crc64_nvme(upload1.checksum_crc64_nvme().expect("No CRC64NVME for part 1"))
                    .build(),
            )
            .parts(
                CompletedPart::builder()
                    .part_number(2)
                    .e_tag(upload2.e_tag().expect("No etag for part 2"))
                    .checksum_crc64_nvme(upload2.checksum_crc64_nvme().expect("No CRC64NVME for part 2"))
                    .build(),
            )
            .build();

        client
            .complete_multipart_upload()
            .bucket(bucket)
            .key(multipart_key)
            .upload_id(&upload_id)
            .multipart_upload(completed_upload)
            .send()
            .await
            .expect("Failed to complete multipart upload");

        let direct_head = client
            .head_object()
            .bucket(bucket)
            .key(direct_key)
            .checksum_mode(ChecksumMode::Enabled)
            .send()
            .await
            .expect("Failed to head direct object");

        let multipart_head = client
            .head_object()
            .bucket(bucket)
            .key(multipart_key)
            .checksum_mode(ChecksumMode::Enabled)
            .send()
            .await
            .expect("Failed to head multipart object");

        assert_eq!(
            direct_head.checksum_crc64_nvme(),
            Some(full_checksum.as_str()),
            "Direct object should report the uploaded full-object CRC64NVME"
        );
        assert_eq!(
            multipart_head.checksum_crc64_nvme(),
            Some(full_checksum.as_str()),
            "Multipart object should report the same full-object CRC64NVME as direct upload"
        );
    }
}