mountpoint_s3_client/s3_crt_client/
get_object.rs

1use std::ops::Deref;
2use std::os::unix::prelude::OsStrExt;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::task::{Context, Poll};
7
8use async_trait::async_trait;
9use bytes::Bytes;
10use futures::channel::mpsc::UnboundedReceiver;
11use futures::stream::FusedStream;
12use futures::{Stream, StreamExt};
13use mountpoint_s3_crt::http::request_response::{Header, Headers};
14use mountpoint_s3_crt::s3::client::{MetaRequest, MetaRequestResult};
15use pin_project::pin_project;
16use tracing::trace;
17
18use crate::error_metadata::ClientErrorMetadata;
19use crate::object_client::{
20    Checksum, ChecksumMode, ClientBackpressureHandle, GetBodyPart, GetObjectError, GetObjectParams, GetObjectResponse,
21    ObjectChecksumError, ObjectClientError, ObjectClientResult, ObjectMetadata,
22};
23
24use super::{CancellingMetaRequest, ResponseHeadersError, S3CrtClient, S3Operation, S3RequestError, parse_checksum};
25
26impl S3CrtClient {
27    /// Create and begin a new GetObject request. The returned [S3GetObjectResponse] is a [Stream] of
28    /// body parts of the object, which will be delivered in order.
29    pub(super) async fn get_object(
30        &self,
31        bucket: &str,
32        key: &str,
33        params: &GetObjectParams,
34    ) -> Result<S3GetObjectResponse, ObjectClientError<GetObjectError, S3RequestError>> {
35        let requested_checksums = params.checksum_mode.as_ref() == Some(&ChecksumMode::Enabled);
36        let next_offset = params.range.as_ref().map(|r| r.start).unwrap_or(0);
37        let (event_sender, mut event_receiver) = futures::channel::mpsc::unbounded();
38        let meta_request = {
39            let span =
40                request_span!(self.inner, "get_object", bucket, key, range=?params.range, if_match=?params.if_match);
41
42            let mut message = self
43                .inner
44                .new_request_template("GET", bucket)
45                .map_err(S3RequestError::construction_failure)?;
46
47            // Overwrite "accept" header since this returns raw object data.
48            message
49                .set_header(&Header::new("accept", "*/*"))
50                .map_err(S3RequestError::construction_failure)?;
51
52            if requested_checksums {
53                // Add checksum header to receive object checksums.
54                message
55                    .set_header(&Header::new("x-amz-checksum-mode", "enabled"))
56                    .map_err(S3RequestError::construction_failure)?;
57            }
58
59            if let Some(etag) = params.if_match.as_ref() {
60                // Return the object only if its entity tag (ETag) is matched
61                message
62                    .set_header(&Header::new("If-Match", etag.as_str()))
63                    .map_err(S3RequestError::construction_failure)?;
64            }
65
66            if let Some(range) = params.range.as_ref() {
67                // Range HTTP header is bounded below *inclusive*
68                let range_value = format!("bytes={}-{}", range.start, range.end.saturating_sub(1));
69                message
70                    .set_header(&Header::new("Range", range_value))
71                    .map_err(S3RequestError::construction_failure)?;
72            }
73
74            let key = format!("/{key}");
75            message
76                .set_request_path(key)
77                .map_err(S3RequestError::construction_failure)?;
78
79            let mut options = message.into_options(S3Operation::GetObject);
80            options.part_size(self.inner.read_part_size as u64);
81
82            let mut headers_sender = Some(event_sender.clone());
83            let part_sender = event_sender.clone();
84
85            self.inner.meta_request_with_callbacks(
86                options,
87                span,
88                |_| (),
89                move |headers, status| {
90                    // Only send headers if we have a 2xx status code. If we only get other status codes,
91                    // then on_meta_request_result will send an error.
92                    if (200..300).contains(&status) {
93                        // Headers can be returned multiple times, but the metadata/checksums don't change.
94                        // We only send the first occurence to the channel.
95                        if let Some(headers_sender) = headers_sender.take() {
96                            _ = headers_sender.unbounded_send(S3GetObjectEvent::Headers(headers.clone()));
97                        }
98                    }
99                },
100                move |offset, data| {
101                    let owned_buffer = data
102                        .to_owned_buffer()
103                        .expect("buffers returned from GetObject can always be acquired");
104                    let bytes = Bytes::from_owner(owned_buffer);
105                    let body_part = GetBodyPart { offset, data: bytes };
106                    _ = part_sender.unbounded_send(S3GetObjectEvent::BodyPart(body_part));
107                },
108                parse_get_object_error,
109                move |result| {
110                    if let Err(e) = result {
111                        _ = event_sender.unbounded_send(S3GetObjectEvent::Error(e));
112                    }
113                    event_sender.close_channel();
114                },
115            )?
116        };
117
118        let headers = match event_receiver.next().await {
119            Some(S3GetObjectEvent::Headers(headers)) => headers,
120            Some(S3GetObjectEvent::Error(e)) => {
121                return Err(e);
122            }
123            event => {
124                // If we did not received the headers first, the request must have failed.
125                trace!(?event, "unexpected GetObject event while waiting for headers");
126                return Err(S3RequestError::internal_failure(ResponseHeadersError::MissingHeaders).into());
127            }
128        };
129
130        let backpressure_handle = if self.inner.enable_backpressure {
131            let read_window_end_offset =
132                Arc::new(AtomicU64::new(next_offset + self.inner.initial_read_window_size as u64));
133            Some(S3BackpressureHandle {
134                read_window_end_offset,
135                meta_request: meta_request.clone(),
136            })
137        } else {
138            None
139        };
140        Ok(S3GetObjectResponse {
141            meta_request,
142            event_receiver,
143            requested_checksums,
144            backpressure_handle,
145            headers,
146            next_offset,
147        })
148    }
149}
150
151#[derive(Debug)]
152enum S3GetObjectEvent {
153    Headers(Headers),
154    BodyPart(GetBodyPart),
155    Error(ObjectClientError<GetObjectError, S3RequestError>),
156}
157
158#[derive(Clone, Debug)]
159pub struct S3BackpressureHandle {
160    /// Upper bound of the current read window. When backpressure is enabled, [S3GetObjectRequest]
161    /// can return data up to this offset *exclusively*.
162    read_window_end_offset: Arc<AtomicU64>,
163    meta_request: MetaRequest,
164}
165
166impl ClientBackpressureHandle for S3BackpressureHandle {
167    fn increment_read_window(&mut self, len: usize) {
168        self.read_window_end_offset.fetch_add(len as u64, Ordering::SeqCst);
169        self.meta_request.increment_read_window(len as u64);
170    }
171
172    fn ensure_read_window(&mut self, desired_end_offset: u64) {
173        let diff = desired_end_offset.saturating_sub(self.read_window_end_offset()) as usize;
174        self.increment_read_window(diff);
175    }
176
177    fn read_window_end_offset(&self) -> u64 {
178        self.read_window_end_offset.load(Ordering::SeqCst)
179    }
180}
181
182/// A streaming response to a GetObject request.
183///
184/// This struct implements [`futures::Stream`], which you can use to read the body of the object.
185/// Each item of the stream is a part of the object body together with the part's offset within the
186/// object.
187#[derive(Debug)]
188#[pin_project]
189pub struct S3GetObjectResponse {
190    meta_request: CancellingMetaRequest,
191    #[pin]
192    event_receiver: UnboundedReceiver<S3GetObjectEvent>,
193    requested_checksums: bool,
194    backpressure_handle: Option<S3BackpressureHandle>,
195    headers: Headers,
196    /// Next offset of the data to be polled from [poll_next]
197    next_offset: u64,
198}
199
200#[cfg_attr(not(docsrs), async_trait)]
201impl GetObjectResponse for S3GetObjectResponse {
202    type BackpressureHandle = S3BackpressureHandle;
203    type ClientError = S3RequestError;
204
205    fn backpressure_handle(&mut self) -> Option<&mut Self::BackpressureHandle> {
206        self.backpressure_handle.as_mut()
207    }
208
209    fn get_object_metadata(&self) -> ObjectMetadata {
210        self.headers
211            .iter()
212            .filter_map(|(key, value)| {
213                let metadata_header = key.to_str()?.strip_prefix("x-amz-meta-")?;
214                let value = value.to_str()?;
215                Some((metadata_header.to_string(), value.to_string()))
216            })
217            .collect()
218    }
219
220    fn get_object_checksum(&self) -> Result<Checksum, ObjectChecksumError> {
221        if !self.requested_checksums {
222            return Err(ObjectChecksumError::DidNotRequestChecksums);
223        }
224
225        parse_checksum(&self.headers).map_err(|e| ObjectChecksumError::HeadersError(Box::new(e)))
226    }
227}
228
229impl Stream for S3GetObjectResponse {
230    type Item = ObjectClientResult<GetBodyPart, GetObjectError, S3RequestError>;
231
232    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
233        if self.event_receiver.is_terminated() {
234            return Poll::Ready(None);
235        }
236
237        let this = self.project();
238        match this.event_receiver.poll_next(cx) {
239            Poll::Ready(None) => Poll::Ready(None),
240            Poll::Ready(Some(S3GetObjectEvent::BodyPart(part))) => {
241                *this.next_offset = part.offset + part.data.len() as u64;
242                Poll::Ready(Some(Ok(part)))
243            }
244            Poll::Ready(Some(S3GetObjectEvent::Headers(_))) => {
245                unreachable!("headers are only sent once and received before returning the stream")
246            }
247            Poll::Ready(Some(S3GetObjectEvent::Error(e))) => Poll::Ready(Some(Err(e))),
248            Poll::Pending => {
249                // If the request is still not finished but the read window is not enough to poll
250                // the next chunk we want to return error instead of keeping the request blocked.
251                // This prevents a risk of deadlock from using the [S3CrtClient], users must implement
252                // their own logic to block the request if they really want to block a [S3GetObjectResponse].
253                if let Some(handle) = &this.backpressure_handle
254                    && *this.next_offset >= handle.read_window_end_offset()
255                {
256                    let err = ObjectClientError::from(S3RequestError::EmptyReadWindow);
257                    return Poll::Ready(Some(Err(err)));
258                }
259                Poll::Pending
260            }
261        }
262    }
263}
264
265fn parse_get_object_error(result: &MetaRequestResult) -> Option<GetObjectError> {
266    let client_error_metadata = ClientErrorMetadata::from_meta_request_result(result);
267    match result.response_status {
268        404 => {
269            let body = result.error_response_body.as_ref()?;
270            let root = xmltree::Element::parse(body.as_bytes()).ok()?;
271            let error_code = root.get_child("Code")?;
272            let error_str = error_code.get_text()?;
273            match error_str.deref() {
274                "NoSuchBucket" => Some(GetObjectError::NoSuchBucket(client_error_metadata)),
275                "NoSuchKey" => Some(GetObjectError::NoSuchKey(client_error_metadata)),
276                _ => None,
277            }
278        }
279        412 => Some(GetObjectError::PreconditionFailed(client_error_metadata)),
280        _ => None,
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use std::ffi::{OsStr, OsString};
287
288    use super::*;
289
290    fn make_result(response_status: i32, body: impl Into<OsString>) -> MetaRequestResult {
291        MetaRequestResult {
292            response_status,
293            crt_error: 1i32.into(),
294            error_response_headers: None,
295            error_response_body: Some(body.into()),
296        }
297    }
298
299    #[test]
300    fn parse_404_no_such_key() {
301        let body = br#"<?xml version="1.0" encoding="UTF-8"?><Error><Code>NoSuchKey</Code><Message>The specified key does not exist.</Message><Key>not-a-real-key</Key><RequestId>NTKJWKHQBYNS73A9</RequestId><HostId>Nc9kWNrf4kGoq5NIUnQ4t7u04ZZXGm/i463v+jwCI8sIrZBqeYI8uffLHQ+/qusdMWNuUwqeXHU=</HostId></Error>"#;
302        let result = make_result(404, OsStr::from_bytes(&body[..]));
303        let result = parse_get_object_error(&result);
304        assert!(matches!(result, Some(GetObjectError::NoSuchKey(_))));
305    }
306
307    #[test]
308    fn parse_404_no_such_bucket() {
309        let body = br#"<?xml version="1.0" encoding="UTF-8"?><Error><Code>NoSuchBucket</Code><Message>The specified bucket does not exist</Message><BucketName>amzn-s3-demo-bucket</BucketName><RequestId>4VAGDP5HMYTDNB3Y</RequestId><HostId>JMgGqpVKIaaTieG68IODiV2piWw/q9VCTowGvWP36BEz6oIVEXiesn8cDE5ph7if0gpY5WU1Wc8=</HostId></Error>"#;
310        let result = make_result(404, OsStr::from_bytes(&body[..]));
311        let result = parse_get_object_error(&result);
312        assert!(matches!(result, Some(GetObjectError::NoSuchBucket(_))));
313    }
314
315    #[test]
316    fn parse_403_glacier_storage_class() {
317        let body = br#"<?xml version="1.0" encoding="UTF-8"?><Error><Code>InvalidObjectState</Code><Message>The action is not valid for the object's storage class</Message><RequestId>9FEFFF118E15B86F</RequestId><HostId>WVQ5kzhiT+oiUfDCOiOYv8W4Tk9eNcxWi/MK+hTS/av34Xy4rBU3zsavf0aaaaa</HostId></Error>"#;
318        let result = make_result(403, OsStr::from_bytes(&body[..]));
319        let result = parse_get_object_error(&result);
320        assert_eq!(result, None);
321    }
322}