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
10pub trait BodyExtractExt: private::Sealed {
12 fn try_into_json<T: serde::de::DeserializeOwned + Send + 'static>(
17 self,
18 ) -> impl Future<Output = Result<T, BoxError>> + Send;
19
20 fn try_into_json_streaming<T: serde::de::DeserializeOwned + Send + 'static>(
40 self,
41 ) -> impl Future<Output = Result<T, BoxError>> + Send;
42
43 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 fn try_into_string(self) -> impl Future<Output = Result<String, BoxError>> + Send;
56
57 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 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
268async 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 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
293async 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
306async 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 #[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}