krabka_object_store/
read.rs

1//! Capped buffered reads.
2//!
3//! The module does a `head()` on an object, rejects the object if it exceeds a
4//! byte cap, then does a conditional, streaming GET. It centralises the OOM guard that runs before
5//! the code buffers a whole object, such as an index snapshot, into memory.
6
7use {
8    super::{error::ObjectStoreError, object_store_api as object_store},
9    bytes::{Bytes, BytesMut},
10    futures_util::TryStreamExt as _,
11    object_store::{GetOptions, ObjectStoreExt as _, path::Path},
12    std::sync::Arc,
13};
14
15/// Read a whole object, and reject it if its size exceeds `max_bytes`.
16///
17/// The function does a `head()` first, so it refuses an oversized object
18/// *before* it buffers any bytes. This is the guard against OOM on a corrupt or
19/// malicious object.
20/// The GET uses the HEAD version and etag when available. Its response size
21/// and streamed chunks are also checked, so an overwrite cannot bypass the cap.
22///
23/// # Errors
24///
25/// - [`ObjectStoreError::TooLarge`] if the object is larger than `max_bytes`.
26/// - [`ObjectStoreError::NotFound`] if the object does not exist.
27/// - [`ObjectStoreError::Precondition`] if its etag changed after the HEAD.
28/// - [`ObjectStoreError::Backend`] for any other backend failure.
29pub async fn read_capped(
30    store: &Arc<dyn object_store::ObjectStore>,
31    key: &Path,
32    max_bytes: u64,
33) -> Result<Bytes, ObjectStoreError> {
34    let meta = store.head(key).await?;
35    if meta.size > max_bytes {
36        return Err(ObjectStoreError::TooLarge {
37            key: key.clone(),
38            size: meta.size,
39            max_bytes,
40        });
41    }
42    let result = store
43        .get_opts(key, GetOptions {
44            version: meta.version,
45            if_match: meta.e_tag,
46            ..GetOptions::default()
47        })
48        .await?;
49    if result.meta.size > max_bytes {
50        return Err(ObjectStoreError::TooLarge {
51            key: key.clone(),
52            size: result.meta.size,
53            max_bytes,
54        });
55    }
56    let mut stream = result.into_stream();
57    let mut bytes = BytesMut::new();
58    while let Some(chunk) = stream.try_next().await? {
59        let size = u64::try_from(bytes.len())
60            .unwrap_or(u64::MAX)
61            .saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX));
62        if size > max_bytes {
63            return Err(ObjectStoreError::TooLarge {
64                key: key.clone(),
65                size,
66                max_bytes,
67            });
68        }
69        bytes.extend_from_slice(&chunk);
70    }
71    Ok(bytes.freeze())
72}
73
74#[cfg(test)]
75mod tests {
76    use {
77        super::{object_store, *},
78        assert2::assert,
79        futures_util::{StreamExt as _, stream},
80        object_store::{PutPayload, memory::InMemory, path::Path},
81        std::{
82            fmt, iter,
83            sync::{
84                Arc,
85                atomic::{AtomicUsize, Ordering::SeqCst},
86            },
87        },
88    };
89
90    async fn store_with(key: &str, bytes: &'static [u8]) -> Arc<dyn object_store::ObjectStore> {
91        let store: Arc<dyn object_store::ObjectStore> = Arc::new(InMemory::new());
92        store
93            .put(&Path::from(key), PutPayload::from(bytes))
94            .await
95            .unwrap();
96        store
97    }
98
99    #[tokio::test]
100    async fn under_cap_returns_bytes() {
101        let store = store_with("k", b"hello").await;
102        let got = read_capped(&store, &Path::from("k"), 1024).await.unwrap();
103        assert!(&got[..] == b"hello");
104    }
105
106    #[tokio::test]
107    async fn over_cap_returns_too_large() {
108        let store = store_with("k", b"hello world").await;
109        let err = read_capped(&store, &Path::from("k"), 4).await.unwrap_err();
110        assert!(matches!(err, ObjectStoreError::TooLarge {
111            size: 11,
112            max_bytes: 4,
113            ..
114        }));
115    }
116
117    #[tokio::test]
118    async fn missing_object_maps_to_not_found() {
119        let store: Arc<dyn object_store::ObjectStore> = Arc::new(InMemory::new());
120        let err = read_capped(&store, &Path::from("absent"), 1024)
121            .await
122            .unwrap_err();
123        assert!(matches!(err, ObjectStoreError::NotFound(_)));
124    }
125
126    #[derive(Debug, Clone, Copy)]
127    enum Race {
128        Etag,
129        Version,
130        Metadata,
131        Body,
132    }
133
134    #[derive(Debug)]
135    struct ChangingStore {
136        inner: InMemory,
137        race: Race,
138        chunks: Arc<AtomicUsize>,
139    }
140
141    impl fmt::Display for ChangingStore {
142        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("ChangingStore") }
143    }
144
145    #[krabka_macros::object_store_delegate(
146        put_multipart_opts,
147        delete_stream,
148        list,
149        list_with_delimiter,
150        copy_opts
151    )]
152    #[async_trait::async_trait]
153    impl object_store::ObjectStore for ChangingStore {
154        async fn put_opts(
155            &self,
156            location: &Path,
157            payload: PutPayload,
158            opts: object_store::PutOptions,
159        ) -> object_store::Result<object_store::PutResult> {
160            self.inner.put_opts(location, payload, opts).await
161        }
162
163        async fn get_opts(
164            &self,
165            location: &Path,
166            options: GetOptions,
167        ) -> object_store::Result<object_store::GetResult> {
168            if options.head {
169                let mut result = self.inner.get_opts(location, options).await?;
170                if matches!(self.race, Race::Version) {
171                    result.meta.version = Some("original".into());
172                } else if !matches!(self.race, Race::Etag) {
173                    result.meta.e_tag = None;
174                }
175                return Ok(result);
176            }
177            if matches!(self.race, Race::Version) {
178                assert!(options.version.as_deref() == Some("original"));
179                let original = self
180                    .inner
181                    .get_opts(location, GetOptions {
182                        version: None,
183                        ..options
184                    })
185                    .await?;
186                self.inner
187                    .put(location, PutPayload::from_static(b"123456789"))
188                    .await?;
189                return Ok(original);
190            }
191            self.inner
192                .put(location, PutPayload::from_static(b"123456789"))
193                .await?;
194            let mut result = self.inner.get_opts(location, options).await?;
195            if matches!(self.race, Race::Body) {
196                result.meta.size = 3;
197                let chunks = Arc::clone(&self.chunks);
198                result.payload = object_store::GetResultPayload::Stream(
199                    stream::iter(iter::repeat_with(|| Ok(Bytes::from_static(b"123"))).take(3))
200                        .inspect(move |_| {
201                            chunks.fetch_add(1, SeqCst);
202                        })
203                        .boxed(),
204                );
205            }
206            Ok(result)
207        }
208    }
209
210    #[yare::parameterized(
211        etag = { Race::Etag },
212        metadata = { Race::Metadata },
213        body = { Race::Body },
214    )]
215    #[test_macro(tokio::test)]
216    async fn overwrite_between_head_and_get_cannot_bypass_the_cap(race: Race) {
217        let changing = Arc::new(ChangingStore {
218            inner: InMemory::new(),
219            race,
220            chunks: Arc::default(),
221        });
222        let key = Path::from("changing");
223        changing
224            .put(&key, PutPayload::from_static(b"123"))
225            .await
226            .unwrap();
227        let store: Arc<dyn object_store::ObjectStore> = Arc::<ChangingStore>::clone(&changing);
228        let error = read_capped(&store, &key, 5).await.unwrap_err();
229        match race {
230            Race::Etag => assert2::assert!(matches!(error, ObjectStoreError::Precondition { .. })),
231            Race::Version => unreachable!("the versioned case returns the original bytes"),
232            Race::Metadata => {
233                assert2::assert!(matches!(error, ObjectStoreError::TooLarge { size: 9, .. }));
234            }
235            Race::Body => {
236                assert2::assert!(matches!(error, ObjectStoreError::TooLarge { size: 6, .. }));
237                assert2::assert!(changing.chunks.load(SeqCst) == 2);
238            }
239        }
240    }
241
242    #[tokio::test]
243    async fn versioned_get_reads_the_head_version_after_an_overwrite() {
244        let changing = Arc::new(ChangingStore {
245            inner: InMemory::new(),
246            race: Race::Version,
247            chunks: Arc::default(),
248        });
249        let key = Path::from("versioned");
250        changing
251            .put(&key, PutPayload::from_static(b"123"))
252            .await
253            .unwrap();
254        let store: Arc<dyn object_store::ObjectStore> = Arc::<ChangingStore>::clone(&changing);
255        assert!(read_capped(&store, &key, 5).await.unwrap().as_ref() == b"123");
256        assert!(changing.inner.head(&key).await.unwrap().size == 9);
257    }
258}
krabka_object_store/read.rs

krabka_object_store/
read.rs
