mkit_server/http_objects/
body.rs

1//! Response bodies: the length-enforcing wrapper and the end-of-body hook.
2
3use core::pin::Pin;
4use core::task::{Context, Poll};
5
6use bytes::Bytes;
7use futures_core::Stream;
8
9use crate::store::StoreError;
10use crate::{BoxStream, Code, ServerError};
11
12/// A response body. `Stream` carries its exact length: the stream errors
13/// instead of yielding more or fewer bytes.
14pub enum HttpBody {
15    /// No body (HEAD, 204, 304 and every error but the canned 404).
16    Empty,
17    /// The whole body in memory.
18    Bytes(Bytes),
19    /// A body of exactly `len` bytes, streamed in pieces.
20    Stream {
21        /// The exact number of body bytes.
22        len: u64,
23        /// The pieces, in order.
24        stream: BoxStream<'static, Result<Bytes, ServerError>>,
25    },
26}
27
28impl core::fmt::Debug for HttpBody {
29    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
30        match self {
31            Self::Empty => f.write_str("Empty"),
32            Self::Bytes(b) => f.debug_tuple("Bytes").field(&b.len()).finish(),
33            Self::Stream { len, .. } => f.debug_struct("Stream").field("len", len).finish(),
34        }
35    }
36}
37
38/// Called once when a body ends, with the bytes sent and how it ended
39/// (WP-4.13's `ReadServed{bytes}` and `Aborted(INTERNAL)` hook). A body
40/// dropped before its end reports `canceled`.
41#[cfg(not(target_arch = "wasm32"))]
42pub type EndHook = Box<dyn FnOnce(u64, Result<(), &ServerError>) + Send>;
43/// Called once when a body ends; see the native definition.
44#[cfg(target_arch = "wasm32")]
45pub type EndHook = Box<dyn FnOnce(u64, Result<(), &ServerError>)>;
46
47fn overrun() -> ServerError {
48    ServerError::unavailable("object body length changed")
49}
50
51/// Enforces `remaining` bytes and fires the hook exactly once.
52struct Exact {
53    inner: BoxStream<'static, Result<Bytes, ServerError>>,
54    remaining: u64,
55    sent: u64,
56    hook: Option<EndHook>,
57    done: bool,
58}
59
60impl Exact {
61    fn end(&mut self, result: Result<(), &ServerError>) {
62        self.done = true;
63        if let Some(hook) = self.hook.take() {
64            hook(self.sent, result);
65        }
66    }
67
68    fn fail(&mut self, error: ServerError) -> Poll<Option<Result<Bytes, ServerError>>> {
69        self.end(Err(&error));
70        Poll::Ready(Some(Err(error)))
71    }
72}
73
74impl Stream for Exact {
75    type Item = Result<Bytes, ServerError>;
76
77    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
78        if self.done {
79            return Poll::Ready(None);
80        }
81        match self.inner.as_mut().poll_next(cx) {
82            Poll::Pending => Poll::Pending,
83            Poll::Ready(Some(Ok(piece))) => {
84                let n = piece.len() as u64;
85                if n > self.remaining {
86                    return self.fail(overrun());
87                }
88                self.remaining -= n;
89                self.sent += n;
90                Poll::Ready(Some(Ok(piece)))
91            }
92            Poll::Ready(Some(Err(error))) => self.fail(error),
93            Poll::Ready(None) if self.remaining != 0 => self.fail(overrun()),
94            Poll::Ready(None) => {
95                self.end(Ok(()));
96                Poll::Ready(None)
97            }
98        }
99    }
100}
101
102impl Drop for Exact {
103    fn drop(&mut self) {
104        if !self.done {
105            self.end(Err(&ServerError::new(Code::Canceled, "request canceled")));
106        }
107    }
108}
109
110fn wrap(
111    stream: BoxStream<'static, Result<Bytes, ServerError>>,
112    len: u64,
113    hook: Option<EndHook>,
114) -> HttpBody {
115    HttpBody::Stream {
116        len,
117        stream: Box::pin(Exact {
118            inner: stream,
119            remaining: len,
120            sent: 0,
121            hook,
122            done: false,
123        }),
124    }
125}
126
127/// Wrap a backend `stream` so it yields exactly `len` bytes or fails,
128/// reporting its end to `hook`. Backend errors are logged and redacted to a
129/// fixed message.
130#[must_use]
131pub fn exact(
132    stream: BoxStream<'static, Result<Bytes, StoreError>>,
133    len: u64,
134    hook: Option<EndHook>,
135) -> HttpBody {
136    let stream = futures::StreamExt::map(stream, |piece| {
137        piece.map_err(|error| {
138            tracing::warn!(detail = %error, "object body read failed");
139            ServerError::unavailable("object storage request failed")
140        })
141    });
142    wrap(Box::pin(stream), len, hook)
143}
144
145/// `body` with `hook` attached: a stream is wrapped, an in-memory body
146/// becomes a one-piece stream so the hook still fires when it is consumed.
147#[must_use]
148pub fn with_hook(body: HttpBody, hook: Option<EndHook>) -> HttpBody {
149    let Some(hook) = hook else {
150        return body;
151    };
152    match body {
153        HttpBody::Empty => {
154            hook(0, Ok(()));
155            HttpBody::Empty
156        }
157        HttpBody::Bytes(bytes) => {
158            let len = bytes.len() as u64;
159            wrap(
160                Box::pin(futures::stream::once(core::future::ready(Ok(bytes)))),
161                len,
162                Some(hook),
163            )
164        }
165        HttpBody::Stream { len, stream } => wrap(stream, len, Some(hook)),
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    #![allow(clippy::unwrap_used)]
172
173    use std::sync::{Arc, Mutex};
174
175    use futures::StreamExt as _;
176    use futures_executor::block_on;
177
178    use super::*;
179
180    type Ended = Arc<Mutex<Vec<(u64, Option<Code>)>>>;
181
182    fn hook(ended: &Ended) -> EndHook {
183        let ended = ended.clone();
184        Box::new(move |sent, result| {
185            ended
186                .lock()
187                .unwrap()
188                .push((sent, result.err().map(ServerError::code)));
189        })
190    }
191
192    fn source(
193        pieces: Vec<Result<&'static [u8], StoreError>>,
194    ) -> BoxStream<'static, Result<Bytes, StoreError>> {
195        Box::pin(futures::stream::iter(
196            pieces
197                .into_iter()
198                .map(|piece| piece.map(Bytes::from_static))
199                .collect::<Vec<_>>(),
200        ))
201    }
202
203    fn collect(body: HttpBody) -> Vec<Result<usize, Code>> {
204        let HttpBody::Stream { stream, .. } = body else {
205            panic!("expected a stream");
206        };
207        block_on(
208            stream
209                .map(|p| p.map(|b| b.len()).map_err(|e| e.code()))
210                .collect(),
211        )
212    }
213
214    #[test]
215    fn an_exact_body_reports_its_end_once() {
216        let ended = Ended::default();
217        let body = exact(source(vec![Ok(b"abc"), Ok(b"de")]), 5, Some(hook(&ended)));
218        assert_eq!(collect(body), vec![Ok(3), Ok(2)]);
219        assert_eq!(*ended.lock().unwrap(), [(5, None)]);
220    }
221
222    #[test]
223    fn a_wrong_length_or_a_backend_failure_fails_the_body_and_the_hook() {
224        for (pieces, len, sent, want) in [
225            (vec![Ok(&b"abcd"[..])], 3, 0, vec![Err(Code::Unavailable)]),
226            (
227                vec![Ok(&b"ab"[..])],
228                3,
229                2,
230                vec![Ok(2), Err(Code::Unavailable)],
231            ),
232            (
233                vec![
234                    Ok(&b"ab"[..]),
235                    Err(StoreError::unavailable("SECRET backend detail")),
236                ],
237                4,
238                2,
239                vec![Ok(2), Err(Code::Unavailable)],
240            ),
241        ] {
242            let ended = Ended::default();
243            let got = collect(exact(source(pieces), len, Some(hook(&ended))));
244            assert_eq!(got, want);
245            assert_eq!(*ended.lock().unwrap(), [(sent, Some(Code::Unavailable))]);
246        }
247    }
248
249    #[test]
250    fn a_dropped_body_reports_canceled_with_what_was_sent() {
251        let ended = Ended::default();
252        let HttpBody::Stream { mut stream, .. } =
253            exact(source(vec![Ok(b"abc"), Ok(b"def")]), 6, Some(hook(&ended)))
254        else {
255            panic!();
256        };
257        assert_eq!(block_on(stream.next()).unwrap().unwrap().len(), 3);
258        drop(stream);
259        assert_eq!(*ended.lock().unwrap(), [(3, Some(Code::Canceled))]);
260    }
261
262    #[test]
263    fn a_backend_error_is_redacted_to_a_fixed_message() {
264        let body = exact(
265            source(vec![Err(StoreError::unavailable("SECRET backend detail"))]),
266            1,
267            None,
268        );
269        let HttpBody::Stream { mut stream, .. } = body else {
270            panic!();
271        };
272        let error = block_on(stream.next()).unwrap().unwrap_err();
273        assert!(!format!("{error:?} {error}").contains("SECRET"));
274        assert_eq!(error.public_message(), "object storage request failed");
275    }
276
277    #[test]
278    fn a_hook_also_covers_in_memory_and_empty_bodies() {
279        let ended = Ended::default();
280        let body = with_hook(
281            HttpBody::Bytes(Bytes::from_static(b"xyz")),
282            Some(hook(&ended)),
283        );
284        assert!(ended.lock().unwrap().is_empty(), "fires when consumed");
285        assert_eq!(collect(body), vec![Ok(3)]);
286        assert_eq!(*ended.lock().unwrap(), [(3, None)]);
287        assert!(matches!(
288            with_hook(HttpBody::Empty, Some(hook(&ended))),
289            HttpBody::Empty
290        ));
291        assert_eq!(ended.lock().unwrap()[1], (0, None));
292        // Without a hook the body is untouched.
293        assert!(matches!(
294            with_hook(HttpBody::Bytes(Bytes::new()), None),
295            HttpBody::Bytes(_)
296        ));
297    }
298}
mkit_server/http_objects/body.rs

mkit_server/http_objects/
body.rs
