Skip to main content

rama_http_types/body/
ext.rs

1use super::StreamingBody;
2use super::util::{BodyExt, CollectError, CollectOptions};
3use rama_core::bytes::{Buf, Bytes};
4use rama_core::error::{BoxError, ErrorContext};
5use rama_json::capture::{
6    CaptureHandler, CaptureResult, CapturedValue, JsonCapturer, OwnedCapturedValue,
7};
8use rama_json::path::JsonPath;
9
10/// An extension trait for [`StreamingBody`] that provides methods to extract data from it.
11pub trait BodyExtractExt: private::Sealed {
12    /// Try to deserialize the (contained) body as a JSON object.
13    ///
14    /// Buffers the entire body in memory before deserializing. For large bodies prefer
15    /// [`BodyExtractExt::try_into_json_streaming`].
16    fn try_into_json<T: serde::de::DeserializeOwned + Send + 'static>(
17        self,
18    ) -> impl Future<Output = Result<T, BoxError>> + Send;
19
20    /// Try to deserialize the (contained) body as a JSON object, streaming bytes
21    /// from the body directly into the JSON parser instead of buffering the whole
22    /// body first.
23    ///
24    /// Preferable to [`BodyExtractExt::try_into_json`] for large bodies where peak
25    /// memory matters.
26    ///
27    /// # Note
28    ///
29    /// Internally this runs `serde_json::from_reader` inside a
30    /// [`tokio::task::spawn_blocking`] task using [`SyncIoBridge`] to bridge the
31    /// async body to serde's synchronous `io::Read` interface. The thread hop is
32    /// unavoidable today because [`serde::Deserialize`] is a pull-based,
33    /// synchronous trait and no production-ready async-first JSON crate currently
34    /// ships a drop-in `serde::Deserialize` integration. Contributions that
35    /// remove this hop — e.g. by building a `serde::Deserializer` on top of an
36    /// async event-based JSON parser — are welcome.
37    ///
38    /// [`SyncIoBridge`]: rama_core::stream::io::SyncIoBridge
39    fn try_into_json_streaming<T: serde::de::DeserializeOwned + Send + 'static>(
40        self,
41    ) -> impl Future<Output = Result<T, BoxError>> + Send;
42
43    /// Capture JSON values matching `selectors` while streaming over the body.
44    ///
45    /// Unlike [`try_into_json`](Self::try_into_json), this does not buffer the
46    /// whole body. Only selected values are copied, and each selected scalar or
47    /// object/array subtree is bounded by `max_capture_bytes`.
48    fn try_capture_json(
49        self,
50        selectors: impl IntoIterator<Item = JsonPath> + Send,
51        max_capture_bytes: usize,
52    ) -> impl Future<Output = Result<Vec<OwnedCapturedValue>, BoxError>> + Send;
53
54    /// Try to turn the (contained) body in an utf-8 string.
55    fn try_into_string(self) -> impl Future<Output = Result<String, BoxError>> + Send;
56
57    /// Like [`try_into_json`](Self::try_into_json), but bounded by `opts` (a size
58    /// cap and/or timeout).
59    ///
60    /// On success returns the decoded value. Otherwise a [`CollectError`] tells
61    /// you why — cap reached, timed out, stream failure, or decode failure — and
62    /// for everything but a stream failure [`CollectError::into_full_body`] hands
63    /// the body back so it can be forwarded on untouched (handy for proxies).
64    ///
65    /// Unlike wrapping the body in [`Limited`], hitting the cap here does not
66    /// destroy the body.
67    ///
68    /// [`Limited`]: crate::body::util::Limited
69    /// [`CollectError::into_full_body`]: crate::body::util::CollectError::into_full_body
70    fn try_into_json_with<T: serde::de::DeserializeOwned + Send + 'static>(
71        self,
72        opts: CollectOptions,
73    ) -> impl Future<Output = Result<T, CollectError>> + Send;
74
75    /// Like [`try_into_string`](Self::try_into_string), but bounded by `opts`.
76    /// See [`try_into_json_with`](Self::try_into_json_with) for the error semantics.
77    fn try_into_string_with(
78        self,
79        opts: CollectOptions,
80    ) -> impl Future<Output = Result<String, CollectError>> + Send;
81}
82
83impl<Body> BodyExtractExt for crate::Response<Body>
84where
85    Body: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
86{
87    async fn try_into_json<T: serde::de::DeserializeOwned + Send + 'static>(
88        self,
89    ) -> Result<T, BoxError> {
90        let body = self.into_body().collect().await.into_box_error()?;
91        serde_json::from_slice(body.to_bytes().as_ref())
92            .context("deserialize response body as JSON")
93    }
94
95    async fn try_into_json_streaming<T: serde::de::DeserializeOwned + Send + 'static>(
96        self,
97    ) -> Result<T, BoxError> {
98        body_into_json_streaming(self.into_body())
99            .await
100            .context("streaming-deserialize response body as JSON")
101    }
102
103    async fn try_into_string(self) -> Result<String, BoxError> {
104        let body = self.into_body().collect().await.into_box_error()?;
105        let bytes = body.to_bytes();
106        String::from_utf8(bytes.to_vec()).context("parse body as utf-8 string")
107    }
108
109    async fn try_capture_json(
110        self,
111        selectors: impl IntoIterator<Item = JsonPath> + Send,
112        max_capture_bytes: usize,
113    ) -> Result<Vec<OwnedCapturedValue>, BoxError> {
114        let capturer = JsonCapturer::new(selectors, max_capture_bytes, CaptureCollector::default());
115        body_capture_json(self.into_body(), capturer)
116            .await
117            .context("capture selected JSON values from response body")
118    }
119
120    async fn try_into_json_with<T: serde::de::DeserializeOwned + Send + 'static>(
121        self,
122        opts: CollectOptions,
123    ) -> Result<T, CollectError> {
124        body_into_json_with(crate::Body::new(self.into_body()), opts).await
125    }
126
127    async fn try_into_string_with(self, opts: CollectOptions) -> Result<String, CollectError> {
128        body_into_string_with(crate::Body::new(self.into_body()), opts).await
129    }
130}
131
132impl<Body> BodyExtractExt for crate::Request<Body>
133where
134    Body: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
135{
136    async fn try_into_json<T: serde::de::DeserializeOwned + Send + 'static>(
137        self,
138    ) -> Result<T, BoxError> {
139        let body = self.into_body().collect().await.into_box_error()?;
140        serde_json::from_slice(body.to_bytes().as_ref()).context("deserialize request body as JSON")
141    }
142
143    async fn try_into_json_streaming<T: serde::de::DeserializeOwned + Send + 'static>(
144        self,
145    ) -> Result<T, BoxError> {
146        body_into_json_streaming(self.into_body())
147            .await
148            .context("streaming-deserialize request body as JSON")
149    }
150
151    async fn try_into_string(self) -> Result<String, BoxError> {
152        let body = self.into_body().collect().await.into_box_error()?;
153        let bytes = body.to_bytes();
154        String::from_utf8(bytes.to_vec()).context("parse request body as utf-8 string")
155    }
156
157    async fn try_capture_json(
158        self,
159        selectors: impl IntoIterator<Item = JsonPath> + Send,
160        max_capture_bytes: usize,
161    ) -> Result<Vec<OwnedCapturedValue>, BoxError> {
162        let capturer = JsonCapturer::new(selectors, max_capture_bytes, CaptureCollector::default());
163        body_capture_json(self.into_body(), capturer)
164            .await
165            .context("capture selected JSON values from request body")
166    }
167
168    async fn try_into_json_with<T: serde::de::DeserializeOwned + Send + 'static>(
169        self,
170        opts: CollectOptions,
171    ) -> Result<T, CollectError> {
172        body_into_json_with(crate::Body::new(self.into_body()), opts).await
173    }
174
175    async fn try_into_string_with(self, opts: CollectOptions) -> Result<String, CollectError> {
176        body_into_string_with(crate::Body::new(self.into_body()), opts).await
177    }
178}
179
180impl<B: Into<crate::Body> + Send + 'static> BodyExtractExt for B {
181    async fn try_into_json<T: serde::de::DeserializeOwned + Send + 'static>(
182        self,
183    ) -> Result<T, BoxError> {
184        let body = self.into();
185        let collected_body = body.collect().await.context("collect body")?;
186        serde_json::from_slice(collected_body.to_bytes().as_ref())
187            .context("deserialize body as JSON")
188    }
189
190    async fn try_into_json_streaming<T: serde::de::DeserializeOwned + Send + 'static>(
191        self,
192    ) -> Result<T, BoxError> {
193        body_into_json_streaming(self.into())
194            .await
195            .context("streaming-deserialize body as JSON")
196    }
197
198    async fn try_into_string(self) -> Result<String, BoxError> {
199        let body = self.into();
200        let collected_body = body.collect().await.context("collect body")?;
201        let bytes = collected_body.to_bytes();
202        String::from_utf8(bytes.to_vec()).context("parse body as utf-8 string")
203    }
204
205    async fn try_capture_json(
206        self,
207        selectors: impl IntoIterator<Item = JsonPath> + Send,
208        max_capture_bytes: usize,
209    ) -> Result<Vec<OwnedCapturedValue>, BoxError> {
210        let capturer = JsonCapturer::new(selectors, max_capture_bytes, CaptureCollector::default());
211        body_capture_json(self.into(), capturer)
212            .await
213            .context("capture selected JSON values from body")
214    }
215
216    async fn try_into_json_with<T: serde::de::DeserializeOwned + Send + 'static>(
217        self,
218        opts: CollectOptions,
219    ) -> Result<T, CollectError> {
220        body_into_json_with(self.into(), opts).await
221    }
222
223    async fn try_into_string_with(self, opts: CollectOptions) -> Result<String, CollectError> {
224        body_into_string_with(self.into(), opts).await
225    }
226}
227
228#[derive(Debug, Default)]
229struct CaptureCollector {
230    values: Vec<OwnedCapturedValue>,
231}
232
233impl CaptureHandler for CaptureCollector {
234    fn handle_capture(&mut self, value: CapturedValue<'_>) -> CaptureResult {
235        self.values.push(value.into_owned());
236        Ok(())
237    }
238}
239
240async fn body_capture_json<B>(
241    body: B,
242    mut capturer: JsonCapturer<CaptureCollector>,
243) -> Result<Vec<OwnedCapturedValue>, BoxError>
244where
245    B: StreamingBody<Data: Send + 'static, Error: Into<BoxError>> + Send + 'static,
246{
247    use rama_core::futures::TryStreamExt;
248
249    let data_stream = crate::body::util::BodyDataStream::new(body);
250    let mut data_stream = std::pin::pin!(data_stream);
251    while let Some(mut data) = data_stream
252        .as_mut()
253        .try_next()
254        .await
255        .map_err(|err| err.into())?
256    {
257        while data.has_remaining() {
258            let chunk = data.chunk();
259            let len = chunk.len();
260            capturer.write(chunk)?;
261            data.advance(len);
262        }
263    }
264    capturer.end()?;
265    Ok(capturer.into_handler().values)
266}
267
268/// Drive a body's data frames through an `AsyncRead` and into
269/// `serde_json::from_reader`, running the sync deserializer on a blocking
270/// task. See [`BodyExtractExt::try_into_json_streaming`] for rationale.
271async fn body_into_json_streaming<B, T>(body: B) -> Result<T, BoxError>
272where
273    B: StreamingBody<Data: Send + 'static, Error: Into<BoxError>> + Send + 'static,
274    T: serde::de::DeserializeOwned + Send + 'static,
275{
276    use rama_core::futures::TryStreamExt;
277    use rama_core::stream::io::{StreamReader, SyncIoBridge};
278
279    // `http_body::Body::Data: Buf` is a supertrait bound, so the data-frame
280    // stream items implement `Buf` and can feed `StreamReader` directly.
281    let data_stream =
282        crate::body::util::BodyDataStream::new(body).map_err(|e| std::io::Error::other(e.into()));
283    let async_reader = StreamReader::new(Box::pin(data_stream));
284
285    tokio::task::spawn_blocking(move || {
286        let reader = SyncIoBridge::new(async_reader);
287        serde_json::from_reader::<_, T>(reader).map_err(BoxError::from)
288    })
289    .await
290    .map_err(BoxError::from)?
291}
292
293/// Collect `body` within `opts`, then JSON-decode the buffered bytes. A decode
294/// failure surfaces as a [`CollectError`] with the full body still recoverable.
295async fn body_into_json_with<T: serde::de::DeserializeOwned + Send + 'static>(
296    body: crate::Body,
297    opts: CollectOptions,
298) -> Result<T, CollectError> {
299    let bytes = body.collect_with(opts).await?.to_bytes();
300    match serde_json::from_slice::<T>(bytes.as_ref()) {
301        Ok(value) => Ok(value),
302        Err(err) => Err(CollectError::decode(bytes, err.into())),
303    }
304}
305
306/// Collect `body` within `opts`, then UTF-8 decode the buffered bytes. A decode
307/// failure surfaces as a [`CollectError`] with the full body still recoverable.
308async fn body_into_string_with(
309    body: crate::Body,
310    opts: CollectOptions,
311) -> Result<String, CollectError> {
312    let bytes = body.collect_with(opts).await?.to_bytes();
313    match std::str::from_utf8(bytes.as_ref()) {
314        Ok(s) => Ok(s.to_owned()),
315        Err(err) => Err(CollectError::decode(bytes, err.into())),
316    }
317}
318
319mod private {
320    pub trait Sealed {}
321
322    impl<Body> Sealed for crate::Response<Body> {}
323    impl<Body> Sealed for crate::Request<Body> {}
324    impl<B: Into<crate::Body> + Send + 'static> Sealed for B {}
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use crate::Body;
331    use rama_core::bytes::Bytes;
332    use rama_core::futures::stream;
333
334    #[derive(Debug, serde::Deserialize, PartialEq, Eq)]
335    struct Foo {
336        name: String,
337        age: u8,
338    }
339
340    /// Build a body from multiple `Bytes` chunks so we actually exercise the
341    /// streaming path (split-across-frames JSON).
342    ///
343    /// Uses the multi_thread flavor because `SyncIoBridge` calls
344    /// `Handle::block_on` on the blocking task — that needs runtime workers
345    /// to be available on another thread.
346    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
347    async fn streaming_json_across_frames() {
348        let chunks: Vec<Result<Bytes, std::io::Error>> = vec![
349            Ok(Bytes::from_static(b"{\"name\":")),
350            Ok(Bytes::from_static(b"\"alice\",\"age\"")),
351            Ok(Bytes::from_static(b":42}")),
352        ];
353        let body = Body::from_stream(stream::iter(chunks));
354
355        let foo: Foo = body.try_into_json_streaming().await.unwrap();
356        assert_eq!(
357            foo,
358            Foo {
359                name: "alice".to_owned(),
360                age: 42,
361            }
362        );
363    }
364
365    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
366    async fn streaming_json_invalid_payload() {
367        let body = Body::from("not actually json");
368        let result: Result<serde_json::Value, _> = body.try_into_json_streaming().await;
369        result.unwrap_err();
370    }
371
372    #[tokio::test]
373    async fn capture_json_selected_values_across_frames() {
374        let chunks: Vec<Result<Bytes, std::io::Error>> = vec![
375            Ok(Bytes::from_static(br#"{"items":[{"id":"#)),
376            Ok(Bytes::from_static(br#"1},{"id":2}],"ok":true}"#)),
377        ];
378        let body = Body::from_stream(stream::iter(chunks));
379        let selectors = [JsonPath::builder()
380            .member("items")
381            .wildcard()
382            .member("id")
383            .build()];
384
385        let values = body.try_capture_json(selectors, 32).await.unwrap();
386        assert_eq!(values.len(), 2);
387        assert_eq!(values[0].path().to_string(), "$.items[0].id");
388        assert_eq!(values[0].as_u8(), Some(1));
389        assert_eq!(values[1].path().to_string(), "$.items[1].id");
390        assert_eq!(values[1].deserialize::<u8>().unwrap(), 2);
391    }
392
393    #[tokio::test]
394    async fn request_capture_json_selected_values() {
395        let req = crate::Request::new(Body::from(br#"{"name":"Ada"}"#.as_slice()));
396        let selectors = [JsonPath::builder().member("name").build()];
397
398        let values = req.try_capture_json(selectors, 32).await.unwrap();
399        assert_eq!(values.len(), 1);
400        assert_eq!(values[0].path().to_string(), "$.name");
401        assert_eq!(values[0].as_str().as_deref(), Some("Ada"));
402    }
403
404    #[tokio::test]
405    async fn response_capture_json_selected_values() {
406        let res = crate::Response::new(Body::from(br#"{"name":"Ada"}"#.as_slice()));
407        let selectors = [JsonPath::builder().member("name").build()];
408
409        let values = res.try_capture_json(selectors, 32).await.unwrap();
410        assert_eq!(values.len(), 1);
411        assert_eq!(values[0].path().to_string(), "$.name");
412        assert_eq!(values[0].as_str().as_deref(), Some("Ada"));
413    }
414
415    #[tokio::test]
416    async fn capture_json_reports_capture_limit() {
417        let body = Body::from(br#"{"item":{"name":"Ada"}}"#.as_slice());
418        let selectors = [JsonPath::builder().member("item").build()];
419
420        let err = body.try_capture_json(selectors, 8).await.unwrap_err();
421        assert!(err.to_string().contains("capture limit"));
422    }
423
424    #[tokio::test]
425    async fn try_into_string_with_complete() {
426        let body = Body::from("hello");
427        let s = body
428            .try_into_string_with(CollectOptions::new().with_max_size(100))
429            .await
430            .unwrap();
431        assert_eq!(s, "hello");
432    }
433
434    #[tokio::test]
435    async fn try_into_string_with_cap_returns_passthrough_body() {
436        let body = Body::from("hello world");
437        let err = body
438            .try_into_string_with(CollectOptions::new().with_max_size(5))
439            .await
440            .unwrap_err();
441        assert!(err.is_cap_reached());
442        let full = err.into_full_body().unwrap();
443        assert_eq!(full.try_into_string().await.unwrap(), "hello world");
444    }
445
446    #[tokio::test]
447    async fn try_into_string_with_invalid_utf8_is_decode_error() {
448        let body = Body::from(vec![0xff_u8, 0xfe, 0xfd]);
449        let err = body
450            .try_into_string_with(CollectOptions::new().with_max_size(1024))
451            .await
452            .unwrap_err();
453        assert!(err.is_decode_error());
454        assert_eq!(err.bytes_read().to_vec(), vec![0xff, 0xfe, 0xfd]);
455    }
456
457    #[tokio::test]
458    async fn try_into_json_with_complete() {
459        let body = Body::from(r#"{"name":"alice","age":42}"#);
460        let foo: Foo = body
461            .try_into_json_with(CollectOptions::new().with_max_size(1024))
462            .await
463            .unwrap();
464        assert_eq!(
465            foo,
466            Foo {
467                name: "alice".to_owned(),
468                age: 42,
469            }
470        );
471    }
472
473    #[tokio::test]
474    async fn try_into_json_with_cap_returns_passthrough_body() {
475        let body = Body::from(r#"{"name":"alice","age":42}"#);
476        let err = body
477            .try_into_json_with::<Foo>(CollectOptions::new().with_max_size(5))
478            .await
479            .unwrap_err();
480        assert!(err.is_cap_reached());
481        let recovered = err
482            .into_full_body()
483            .unwrap()
484            .try_into_string()
485            .await
486            .unwrap();
487        assert_eq!(recovered, r#"{"name":"alice","age":42}"#);
488    }
489
490    #[tokio::test]
491    async fn try_into_json_with_decode_error_recovers_full_body() {
492        let body = Body::from("not json");
493        let err = body
494            .try_into_json_with::<Foo>(CollectOptions::new().with_max_size(1024))
495            .await
496            .unwrap_err();
497        assert!(err.is_decode_error());
498        let recovered = err
499            .into_full_body()
500            .unwrap()
501            .try_into_string()
502            .await
503            .unwrap();
504        assert_eq!(recovered, "not json");
505    }
506}