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
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use SizeHint;

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

use std::error::Error;
use std::fmt;
use std::usize;

/// Conversion from a `BufStream`.
///
/// By implementing `FromBufStream` for a type, you define how it will be
/// created from a buf stream. This is common for types which describe byte
/// storage of some kind.
///
/// `FromBufStream` is rarely called explicitly, and it is instead used through
/// `BufStream`'s `collect` method.
pub trait FromBufStream<T: Buf>: Sized {
    /// Type that is used to build `Self` while the `BufStream` is being
    /// consumed.
    type Builder;

    /// Error that might happen on conversion.
    type Error;

    /// Create a new, empty, builder. The provided `hint` can be used to inform
    /// reserving capacity.
    fn builder(hint: &SizeHint) -> Self::Builder;

    /// Extend the builder with the `Buf`.
    ///
    /// This method is called whenever a new `Buf` value is obtained from the
    /// buf stream.
    ///
    /// The provided size hint represents the state of the stream **after**
    /// `buf` has been yielded. The lower bound represents the minimum amount of
    /// data that will be provided after this call to `extend` returns.
    fn extend(builder: &mut Self::Builder, buf: &mut T, hint: &SizeHint)
        -> Result<(), Self::Error>;

    /// Finalize the building of `Self`.
    ///
    /// Called once the buf stream is fully consumed.
    fn build(builder: Self::Builder) -> Result<Self, Self::Error>;
}

/// Error returned from collecting into a `Vec<u8>`
#[derive(Debug)]
pub struct CollectVecError {
    _p: (),
}

/// Error returned from collecting into a `Bytes`
#[derive(Debug)]
pub struct CollectBytesError {
    _p: (),
}

impl<T: Buf> FromBufStream<T> for Vec<u8> {
    type Builder = Vec<u8>;
    type Error = CollectVecError;

    fn builder(hint: &SizeHint) -> Vec<u8> {
        Vec::with_capacity(hint.lower() as usize)
    }

    fn extend(builder: &mut Self, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> {
        let lower = hint.lower();

        // If the lower bound is greater than `usize::MAX` then we have a
        // problem
        if lower > usize::MAX as u64 {
            return Err(CollectVecError { _p: () });
        }

        let mut reserve = lower as usize;

        // If `upper` is set, use this value if it is less than or equal to 64.
        // This only really impacts the first iteration.
        match hint.upper() {
            Some(upper) if upper <= 64 => {
                reserve = upper as usize;
            }
            _ => {}
        }

        // hint.lower() represents the minimum amount of data that will be
        // received *after* this function call. We reserve this amount on top of
        // the amount of data in `buf`.
        reserve = match reserve.checked_add(buf.remaining()) {
            Some(n) => n,
            None => return Err(CollectVecError { _p: () }),
        };

        // Always reserve 64 bytes the first time, unless `upper` is set and is
        // less than 64.
        if builder.is_empty() {
            reserve = reserve.max(match hint.upper() {
                Some(upper) if upper < 64 => upper as usize,
                _ => 64,
            });
        }

        // Make sure overflow won't happen when reserving
        if reserve.checked_add(builder.len()).is_none() {
            return Err(CollectVecError { _p: () });
        }

        // Reserve space
        builder.reserve(reserve);

        // Copy the data
        builder.put(buf);

        Ok(())
    }

    fn build(builder: Self) -> Result<Self, Self::Error> {
        Ok(builder)
    }
}

impl<T: Buf> FromBufStream<T> for Bytes {
    type Builder = Vec<u8>;
    type Error = CollectBytesError;

    fn builder(hint: &SizeHint) -> Vec<u8> {
        <Vec<u8> as FromBufStream<T>>::builder(hint)
    }

    fn extend(builder: &mut Vec<u8>, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> {
        <Vec<u8> as FromBufStream<T>>::extend(builder, buf, hint)
            .map_err(|_| CollectBytesError { _p: () })
    }

    fn build(builder: Vec<u8>) -> Result<Self, Self::Error> {
        Ok(builder.into())
    }
}

impl fmt::Display for CollectVecError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "BufStream is too big")
    }
}

impl Error for CollectVecError {
    fn description(&self) -> &str {
        "BufStream too big"
    }
}

impl fmt::Display for CollectBytesError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "BufStream too big")
    }
}

impl Error for CollectBytesError {
    fn description(&self) -> &str {
        "BufStream too big"
    }
}