google-cloud-storage 1.17.0

// Copyright 2025 Google LLC
//
// 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
//
//     https://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.

use super::active_read::ActiveRead;
use crate::Error;
use crate::model::Object;
use crate::read_object::dynamic::ReadObjectResponse;
use crate::{error::ReadError, model_ext::ObjectHighlights};
use std::sync::Arc;
use tokio::sync::mpsc::{Receiver, Sender};

/// Read the data from a [ObjectDescriptor][super::ObjectDescriptor] range.
///
/// This type is used to stream an open object descriptor range.
#[derive(Debug)]
pub struct RangeReader {
    inner: Receiver<Result<bytes::Bytes, ReadError>>,
    object: Arc<Object>,
    // Unused, holding to a copy prevents the worker task from terminating
    // early.
    _tx: Sender<ActiveRead>,
    checksum: crate::storage::checksum::details::Checksum,
    offset: u64,
    exhausted: bool,
}

impl RangeReader {
    /// Create a new instance.
    ///
    /// This constructor is useful when mocking `ObjectDescriptor`.
    pub fn new(
        inner: Receiver<Result<bytes::Bytes, ReadError>>,
        object: Arc<Object>,
        tx: Sender<ActiveRead>,
        checksum: crate::storage::checksum::details::Checksum,
    ) -> Self {
        // If the server didn't return any checksums in the metadata (e.g. for composite objects),
        // we disable local checksum computation to save CPU cycles.
        let checksum = if object.checksums.is_none() {
            crate::storage::checksum::details::Checksum::default()
        } else {
            checksum
        };

        Self {
            inner,
            object,
            _tx: tx,
            checksum,
            offset: 0,
            exhausted: false,
        }
    }
}

#[async_trait::async_trait]
impl ReadObjectResponse for RangeReader {
    fn object(&self) -> ObjectHighlights {
        ObjectHighlights {
            generation: self.object.generation,
            metageneration: self.object.metageneration,
            size: self.object.size,
            content_encoding: self.object.content_encoding.clone(),
            checksums: self.object.checksums.clone(),
            storage_class: self.object.storage_class.clone(),
            content_language: self.object.content_language.clone(),
            content_type: self.object.content_type.clone(),
            content_disposition: self.object.content_disposition.clone(),
            etag: self.object.etag.clone(),
            metadata: self.object.metadata.clone(),
        }
    }

    /// Fetches the next chunk of data from the underlying stream.
    ///
    /// This function reads from the gRPC stream channel, updating the running
    /// checksum as chunks are successfully received. If the stream reaches its
    /// end (`None`), this function will finalize and validate the running checksum
    /// against the object's expected checksum.
    async fn next(&mut self) -> Option<crate::Result<bytes::Bytes>> {
        if self.exhausted {
            return None;
        }
        let msg = self.inner.recv().await;
        match msg {
            // A valid chunk was received from the stream.
            // Update the running checksum and advance our offset tracker.
            Some(Ok(b)) => {
                self.checksum.update(self.offset, &b);
                self.offset += b.len() as u64;
                Some(Ok(b))
            }
            // A stream error occurred during the RPC.
            // We mark the stream as exhausted and propagate the error.
            Some(Err(e)) => {
                self.exhausted = true;
                Some(Err(Error::io(e)))
            }
            // The stream has successfully completed.
            // Before returning `None` to indicate EOF, we validate the running
            // checksum against the expected checksum returned in the object metadata.
            None => {
                self.exhausted = true;
                if let Some(expected) = &self.object.checksums {
                    let computed = self.checksum.finalize();
                    let res =
                        crate::storage::checksum::details::validate(expected, &Some(computed));
                    if let Err(e) = res {
                        return Some(Err(Error::deser(ReadError::ChecksumMismatch(e))));
                    }
                }
                None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::tests::permanent_error;
    use super::*;
    use crate::model::{Object, ObjectChecksums};
    use std::error::Error as _;

    #[tokio::test]
    async fn object() -> anyhow::Result<()> {
        let object = Object::new()
            .set_generation(123456)
            .set_metageneration(234567)
            .set_size(1024)
            .set_checksums(ObjectChecksums::new().set_crc32c(456789_u32))
            .set_etag("test-etag")
            .set_storage_class("STANDARD")
            .set_content_encoding("content-encoding")
            .set_content_language("content-language")
            .set_content_type("content-type")
            .set_content_disposition("content-disposition");
        let object = Arc::new(object);

        let (tx, inner) = tokio::sync::mpsc::channel(1);
        let (pending_tx, mut pending_rx) = tokio::sync::mpsc::channel(100);

        let mut reader = RangeReader::new(
            inner,
            object.clone(),
            pending_tx,
            crate::storage::checksum::details::Checksum::default(),
        );
        let got = reader.object();
        let want = ObjectHighlights {
            generation: 123456,
            metageneration: 234567,
            size: 1024,
            checksums: Some(ObjectChecksums::new().set_crc32c(456789_u32)),
            etag: "test-etag".into(),
            storage_class: "STANDARD".into(),
            content_encoding: "content-encoding".into(),
            content_language: "content-language".into(),
            content_type: "content-type".into(),
            content_disposition: "content-disposition".into(),
            metadata: std::collections::HashMap::new(),
        };
        assert_eq!(got, want);

        let data = bytes::Bytes::from_static(b"the quick brown fox jumps over the lazy dog");
        tx.send(Ok(data.clone())).await?;
        let got = reader.next().await;
        assert!(matches!(got, Some(Ok(ref d)) if *d == data), "{got:?}");

        let error = ReadError::UnrecoverableBidiReadInterrupt(Arc::new(permanent_error()));
        tx.send(Err(error)).await?;
        let got = reader.next().await;
        assert!(matches!(got, Some(Err(_))), "{got:?}");
        let got = got.unwrap().unwrap_err();
        assert!(got.is_io(), "{got:?}");
        let source = got.source().and_then(|e| e.downcast_ref::<ReadError>());
        assert!(
            matches!(source, Some(ReadError::UnrecoverableBidiReadInterrupt(_))),
            "{source:?}"
        );

        drop(tx);

        let got = reader.next().await;
        assert!(got.is_none(), "{got:?}");

        drop(reader);
        let done = pending_rx.recv().await;
        assert!(done.is_none(), "{done:?}");

        Ok(())
    }

    #[tokio::test]
    async fn object_with_checksum_success() -> anyhow::Result<()> {
        use crate::storage::checksum::details::{Checksum, Crc32c, Md5};
        let data = bytes::Bytes::from_static(b"the quick brown fox jumps over the lazy dog");
        let crc32c = crc32c::crc32c(&data);
        let md5 = bytes::Bytes::from_owner(Vec::from_iter(md5::compute(&data).0));

        let object = Object::new()
            .set_checksums(ObjectChecksums::new().set_crc32c(crc32c).set_md5_hash(md5));
        let object = Arc::new(object);

        let (tx, inner) = tokio::sync::mpsc::channel(1);
        let (pending_tx, _pending_rx) = tokio::sync::mpsc::channel(100);

        let mut reader = RangeReader::new(
            inner,
            object.clone(),
            pending_tx,
            Checksum {
                crc32c: Some(Crc32c::default()),
                md5_hash: Some(Md5::default()),
            },
        );

        tx.send(Ok(data.clone())).await?;
        let got = reader.next().await;
        assert!(matches!(got, Some(Ok(ref d)) if *d == data), "{got:?}");

        drop(tx);

        let got = reader.next().await;
        assert!(got.is_none(), "{got:?}"); // Successfully matched checksums

        Ok(())
    }

    #[tokio::test]
    async fn object_with_checksum_mismatch() -> anyhow::Result<()> {
        use crate::storage::checksum::details::{Checksum, Crc32c};
        let data = bytes::Bytes::from_static(b"the quick brown fox jumps over the lazy dog");
        let wrong_crc32c = crc32c::crc32c(&data) + 1; // Incorrect checksum

        let object = Object::new().set_checksums(ObjectChecksums::new().set_crc32c(wrong_crc32c));
        let object = Arc::new(object);

        let (tx, inner) = tokio::sync::mpsc::channel(1);
        let (pending_tx, _pending_rx) = tokio::sync::mpsc::channel(100);

        let mut reader = RangeReader::new(
            inner,
            object.clone(),
            pending_tx,
            Checksum {
                crc32c: Some(Crc32c::default()),
                md5_hash: None,
            },
        );

        tx.send(Ok(data.clone())).await?;
        let got = reader.next().await;
        assert!(matches!(got, Some(Ok(ref d)) if *d == data), "{got:?}");

        drop(tx);

        let got = reader.next().await;
        // Should yield a ChecksumMismatch error
        assert!(matches!(got, Some(Err(_))), "{got:?}");
        let err = got.unwrap().unwrap_err();

        let source = err.source().and_then(|e| e.downcast_ref::<ReadError>());
        assert!(
            matches!(source, Some(ReadError::ChecksumMismatch(_))),
            "{source:?}"
        );

        // Subsequent calls should yield None
        let got_none = reader.next().await;
        assert!(got_none.is_none(), "{got_none:?}");

        Ok(())
    }

    #[tokio::test]
    async fn object_with_default_checksum_skips_all_validation() -> anyhow::Result<()> {
        let data = bytes::Bytes::from_static(b"the quick brown fox jumps over the lazy dog");
        let wrong_crc32c = crc32c::crc32c(&data) + 1; // Incorrect checksum for the whole object

        let object = Object::new().set_checksums(ObjectChecksums::new().set_crc32c(wrong_crc32c));
        let object = Arc::new(object);

        let (tx, inner) = tokio::sync::mpsc::channel(1);
        let (pending_tx, _pending_rx) = tokio::sync::mpsc::channel(100);

        // Using Checksum::default() disables local validation
        let mut reader = RangeReader::new(
            inner,
            object.clone(),
            pending_tx,
            crate::storage::checksum::details::Checksum::default(),
        );

        tx.send(Ok(data.clone())).await?;
        let got = reader.next().await;
        assert!(matches!(got, Some(Ok(ref d)) if *d == data), "{got:?}");

        drop(tx);

        let got = reader.next().await;
        assert!(got.is_none(), "{got:?}"); // Successfully bypassed checksums, no error returned

        Ok(())
    }

    #[tokio::test]
    async fn object_with_md5_disabled_skips_md5_validation() -> anyhow::Result<()> {
        use crate::storage::checksum::details::{Checksum, Crc32c};
        let data = bytes::Bytes::from_static(b"the quick brown fox jumps over the lazy dog");
        let wrong_md5 = bytes::Bytes::from_static(b"wrongmd5hashhere");

        let object = Object::new().set_checksums(ObjectChecksums::new().set_md5_hash(wrong_md5));
        let object = Arc::new(object);

        let (tx, inner) = tokio::sync::mpsc::channel(1);
        let (pending_tx, _pending_rx) = tokio::sync::mpsc::channel(100);

        // Using Checksum with md5_hash=None simulates compute_md5 not called
        let mut reader = RangeReader::new(
            inner,
            object.clone(),
            pending_tx,
            Checksum {
                crc32c: Some(Crc32c::default()),
                md5_hash: None,
            },
        );

        tx.send(Ok(data.clone())).await?;
        let got = reader.next().await;
        assert!(matches!(got, Some(Ok(ref d)) if *d == data), "{got:?}");

        drop(tx);

        let got = reader.next().await;
        assert!(got.is_none(), "{got:?}"); // Successfully bypassed md5 checksums, no error returned

        Ok(())
    }
}