tokio_buf/util/
from.rs

1use SizeHint;
2
3use bytes::{Buf, BufMut, Bytes};
4
5use std::error::Error;
6use std::fmt;
7use std::usize;
8
9/// Conversion from a `BufStream`.
10///
11/// By implementing `FromBufStream` for a type, you define how it will be
12/// created from a buf stream. This is common for types which describe byte
13/// storage of some kind.
14///
15/// `FromBufStream` is rarely called explicitly, and it is instead used through
16/// `BufStream`'s `collect` method.
17pub trait FromBufStream<T: Buf>: Sized {
18    /// Type that is used to build `Self` while the `BufStream` is being
19    /// consumed.
20    type Builder;
21
22    /// Error that might happen on conversion.
23    type Error;
24
25    /// Create a new, empty, builder. The provided `hint` can be used to inform
26    /// reserving capacity.
27    fn builder(hint: &SizeHint) -> Self::Builder;
28
29    /// Extend the builder with the `Buf`.
30    ///
31    /// This method is called whenever a new `Buf` value is obtained from the
32    /// buf stream.
33    ///
34    /// The provided size hint represents the state of the stream **after**
35    /// `buf` has been yielded. The lower bound represents the minimum amount of
36    /// data that will be provided after this call to `extend` returns.
37    fn extend(builder: &mut Self::Builder, buf: &mut T, hint: &SizeHint)
38        -> Result<(), Self::Error>;
39
40    /// Finalize the building of `Self`.
41    ///
42    /// Called once the buf stream is fully consumed.
43    fn build(builder: Self::Builder) -> Result<Self, Self::Error>;
44}
45
46/// Error returned from collecting into a `Vec<u8>`
47#[derive(Debug)]
48pub struct CollectVecError {
49    _p: (),
50}
51
52/// Error returned from collecting into a `Bytes`
53#[derive(Debug)]
54pub struct CollectBytesError {
55    _p: (),
56}
57
58impl<T: Buf> FromBufStream<T> for Vec<u8> {
59    type Builder = Vec<u8>;
60    type Error = CollectVecError;
61
62    fn builder(hint: &SizeHint) -> Vec<u8> {
63        Vec::with_capacity(hint.lower() as usize)
64    }
65
66    fn extend(builder: &mut Self, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> {
67        let lower = hint.lower();
68
69        // If the lower bound is greater than `usize::MAX` then we have a
70        // problem
71        if lower > usize::MAX as u64 {
72            return Err(CollectVecError { _p: () });
73        }
74
75        let mut reserve = lower as usize;
76
77        // If `upper` is set, use this value if it is less than or equal to 64.
78        // This only really impacts the first iteration.
79        match hint.upper() {
80            Some(upper) if upper <= 64 => {
81                reserve = upper as usize;
82            }
83            _ => {}
84        }
85
86        // hint.lower() represents the minimum amount of data that will be
87        // received *after* this function call. We reserve this amount on top of
88        // the amount of data in `buf`.
89        reserve = match reserve.checked_add(buf.remaining()) {
90            Some(n) => n,
91            None => return Err(CollectVecError { _p: () }),
92        };
93
94        // Always reserve 64 bytes the first time, unless `upper` is set and is
95        // less than 64.
96        if builder.is_empty() {
97            reserve = reserve.max(match hint.upper() {
98                Some(upper) if upper < 64 => upper as usize,
99                _ => 64,
100            });
101        }
102
103        // Make sure overflow won't happen when reserving
104        if reserve.checked_add(builder.len()).is_none() {
105            return Err(CollectVecError { _p: () });
106        }
107
108        // Reserve space
109        builder.reserve(reserve);
110
111        // Copy the data
112        builder.put(buf);
113
114        Ok(())
115    }
116
117    fn build(builder: Self) -> Result<Self, Self::Error> {
118        Ok(builder)
119    }
120}
121
122impl<T: Buf> FromBufStream<T> for Bytes {
123    type Builder = Vec<u8>;
124    type Error = CollectBytesError;
125
126    fn builder(hint: &SizeHint) -> Vec<u8> {
127        <Vec<u8> as FromBufStream<T>>::builder(hint)
128    }
129
130    fn extend(builder: &mut Vec<u8>, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> {
131        <Vec<u8> as FromBufStream<T>>::extend(builder, buf, hint)
132            .map_err(|_| CollectBytesError { _p: () })
133    }
134
135    fn build(builder: Vec<u8>) -> Result<Self, Self::Error> {
136        Ok(builder.into())
137    }
138}
139
140impl fmt::Display for CollectVecError {
141    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
142        write!(fmt, "BufStream is too big")
143    }
144}
145
146impl Error for CollectVecError {
147    fn description(&self) -> &str {
148        "BufStream too big"
149    }
150}
151
152impl fmt::Display for CollectBytesError {
153    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
154        write!(fmt, "BufStream too big")
155    }
156}
157
158impl Error for CollectBytesError {
159    fn description(&self) -> &str {
160        "BufStream too big"
161    }
162}
tokio_buf/util/from.rs

tokio_buf/util/
from.rs
