1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

use bytes::{Buf, BufMut, Bytes};

use super::HttpBody;

/// Concatenate the buffers from a body into a single `Bytes` asynchronously.
///
/// This may require copying the data into a single buffer. If you don't need
/// a contiguous buffer, prefer the [`aggregate`](crate::body::aggregate())
/// function.
///
/// # Note
///
/// Care needs to be taken if the remote is untrusted. The function doesn't implement any length
/// checks and an malicious peer might make it consume arbitrary amounts of memory. Checking the
/// `Content-Length` is a possibility, but it is not strictly mandated to be present.
///
/// # Example
///
/// ```
/// # #[cfg(all(feature = "client", any(feature = "http1", feature = "http2")))]
/// # async fn doc() -> hyper::Result<()> {
/// use hyper::{body::HttpBody};
///
/// # let request = hyper::Request::builder()
/// #        .method(hyper::Method::POST)
/// #        .uri("http://httpbin.org/post")
/// #        .header("content-type", "application/json")
/// #        .body(hyper::Body::from(r#"{"library":"hyper"}"#)).unwrap();
/// # let client = hyper::Client::new();
/// let response = client.request(request).await?;
///
/// const MAX_ALLOWED_RESPONSE_SIZE: u64 = 1024;
///
/// let response_content_length = match response.body().size_hint().upper() {
///     Some(v) => v,
///     None => MAX_ALLOWED_RESPONSE_SIZE + 1 // Just to protect ourselves from a malicious response
/// };
///
/// if response_content_length < MAX_ALLOWED_RESPONSE_SIZE {
///     let body_bytes = hyper::body::to_bytes(response.into_body()).await?;
///     println!("body: {:?}", body_bytes);
/// }
///
/// # Ok(())
/// # }
/// ```
pub async fn to_bytes<T>(body: T) -> Result<Bytes, T::Error>
where
    T: HttpBody,
{
    futures_util::pin_mut!(body);

    // If there's only 1 chunk, we can just return Buf::to_bytes()
    let mut first = if let Some(buf) = body.data().await {
        buf?
    } else {
        return Ok(Bytes::new());
    };

    let second = if let Some(buf) = body.data().await {
        buf?
    } else {
        return Ok(first.copy_to_bytes(first.remaining()));
    };

    // With more than 1 buf, we gotta flatten into a Vec first.
    let cap = first.remaining() + second.remaining() + body.size_hint().lower() as usize;
    let mut vec = Vec::with_capacity(cap);
    vec.put(first);
    vec.put(second);

    while let Some(buf) = body.data().await {
        vec.put(buf?);
    }

    Ok(vec.into())
}