|
| 1 | +//! HTTP body utilities. |
| 2 | +
|
| 3 | +use crate::{BoxError, Error}; |
| 4 | +use bytes::Bytes; |
| 5 | +use bytes::{Buf, BufMut}; |
| 6 | +use http_body::Body; |
| 7 | + |
| 8 | +/// A boxed [`Body`] trait object. |
| 9 | +/// |
| 10 | +/// This is used in axum as the response body type for applications. It's |
| 11 | +/// necessary to unify multiple response bodies types into one. |
| 12 | +pub type BoxBody = http_body::combinators::UnsyncBoxBody<Bytes, Error>; |
| 13 | + |
| 14 | +/// Convert a [`http_body::Body`] into a [`BoxBody`]. |
| 15 | +pub fn boxed<B>(body: B) -> BoxBody |
| 16 | +where |
| 17 | + B: http_body::Body<Data = Bytes> + Send + 'static, |
| 18 | + B::Error: Into<BoxError>, |
| 19 | +{ |
| 20 | + try_downcast(body).unwrap_or_else(|body| body.map_err(Error::new).boxed_unsync()) |
| 21 | +} |
| 22 | + |
| 23 | +pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K> |
| 24 | +where |
| 25 | + T: 'static, |
| 26 | + K: Send + 'static, |
| 27 | +{ |
| 28 | + let mut k = Some(k); |
| 29 | + if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) { |
| 30 | + Ok(k.take().unwrap()) |
| 31 | + } else { |
| 32 | + Err(k.unwrap()) |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +// copied from hyper under the following license: |
| 37 | +// Copyright (c) 2014-2021 Sean McArthur |
| 38 | + |
| 39 | +// Permission is hereby granted, free of charge, to any person obtaining a copy |
| 40 | +// of this software and associated documentation files (the "Software"), to deal |
| 41 | +// in the Software without restriction, including without limitation the rights |
| 42 | +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 43 | +// copies of the Software, and to permit persons to whom the Software is |
| 44 | +// furnished to do so, subject to the following conditions: |
| 45 | + |
| 46 | +// The above copyright notice and this permission notice shall be included in |
| 47 | +// all copies or substantial portions of the Software. |
| 48 | + |
| 49 | +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 50 | +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 51 | +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 52 | +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 53 | +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 54 | +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN |
| 55 | +// THE SOFTWARE. |
| 56 | +pub(crate) async fn to_bytes<T>(body: T) -> Result<Bytes, T::Error> |
| 57 | +where |
| 58 | + T: Body, |
| 59 | +{ |
| 60 | + futures_util::pin_mut!(body); |
| 61 | + |
| 62 | + // If there's only 1 chunk, we can just return Buf::to_bytes() |
| 63 | + let mut first = if let Some(buf) = body.data().await { |
| 64 | + buf? |
| 65 | + } else { |
| 66 | + return Ok(Bytes::new()); |
| 67 | + }; |
| 68 | + |
| 69 | + let second = if let Some(buf) = body.data().await { |
| 70 | + buf? |
| 71 | + } else { |
| 72 | + return Ok(first.copy_to_bytes(first.remaining())); |
| 73 | + }; |
| 74 | + |
| 75 | + // With more than 1 buf, we gotta flatten into a Vec first. |
| 76 | + let cap = first.remaining() + second.remaining() + body.size_hint().lower() as usize; |
| 77 | + let mut vec = Vec::with_capacity(cap); |
| 78 | + vec.put(first); |
| 79 | + vec.put(second); |
| 80 | + |
| 81 | + while let Some(buf) = body.data().await { |
| 82 | + vec.put(buf?); |
| 83 | + } |
| 84 | + |
| 85 | + Ok(vec.into()) |
| 86 | +} |
| 87 | + |
| 88 | +#[test] |
| 89 | +fn test_try_downcast() { |
| 90 | + assert_eq!(try_downcast::<i32, _>(5_u32), Err(5_u32)); |
| 91 | + assert_eq!(try_downcast::<i32, _>(5_i32), Ok(5_i32)); |
| 92 | +} |
0 commit comments