1 use crate::Buf;
2 
3 use core::cmp;
4 
5 /// A `Buf` adapter which limits the bytes read from an underlying buffer.
6 ///
7 /// This struct is generally created by calling `take()` on `Buf`. See
8 /// documentation of [`take()`](trait.Buf.html#method.take) for more details.
9 #[derive(Debug)]
10 pub struct Take<T> {
11     inner: T,
12     limit: usize,
13 }
14 
new<T>(inner: T, limit: usize) -> Take<T>15 pub fn new<T>(inner: T, limit: usize) -> Take<T> {
16     Take {
17         inner,
18         limit,
19     }
20 }
21 
22 impl<T> Take<T> {
23     /// Consumes this `Take`, returning the underlying value.
24     ///
25     /// # Examples
26     ///
27     /// ```rust
28     /// use bytes::buf::{Buf, BufMut, BufExt};
29     ///
30     /// let mut buf = b"hello world".take(2);
31     /// let mut dst = vec![];
32     ///
33     /// dst.put(&mut buf);
34     /// assert_eq!(*dst, b"he"[..]);
35     ///
36     /// let mut buf = buf.into_inner();
37     ///
38     /// dst.clear();
39     /// dst.put(&mut buf);
40     /// assert_eq!(*dst, b"llo world"[..]);
41     /// ```
into_inner(self) -> T42     pub fn into_inner(self) -> T {
43         self.inner
44     }
45 
46     /// Gets a reference to the underlying `Buf`.
47     ///
48     /// It is inadvisable to directly read from the underlying `Buf`.
49     ///
50     /// # Examples
51     ///
52     /// ```rust
53     /// use bytes::{Buf, buf::BufExt};
54     ///
55     /// let mut buf = b"hello world".take(2);
56     ///
57     /// assert_eq!(11, buf.get_ref().remaining());
58     /// ```
get_ref(&self) -> &T59     pub fn get_ref(&self) -> &T {
60         &self.inner
61     }
62 
63     /// Gets a mutable reference to the underlying `Buf`.
64     ///
65     /// It is inadvisable to directly read from the underlying `Buf`.
66     ///
67     /// # Examples
68     ///
69     /// ```rust
70     /// use bytes::{Buf, BufMut, buf::BufExt};
71     ///
72     /// let mut buf = b"hello world".take(2);
73     /// let mut dst = vec![];
74     ///
75     /// buf.get_mut().advance(2);
76     ///
77     /// dst.put(&mut buf);
78     /// assert_eq!(*dst, b"ll"[..]);
79     /// ```
get_mut(&mut self) -> &mut T80     pub fn get_mut(&mut self) -> &mut T {
81         &mut self.inner
82     }
83 
84     /// Returns the maximum number of bytes that can be read.
85     ///
86     /// # Note
87     ///
88     /// If the inner `Buf` has fewer bytes than indicated by this method then
89     /// that is the actual number of available bytes.
90     ///
91     /// # Examples
92     ///
93     /// ```rust
94     /// use bytes::{Buf, buf::BufExt};
95     ///
96     /// let mut buf = b"hello world".take(2);
97     ///
98     /// assert_eq!(2, buf.limit());
99     /// assert_eq!(b'h', buf.get_u8());
100     /// assert_eq!(1, buf.limit());
101     /// ```
limit(&self) -> usize102     pub fn limit(&self) -> usize {
103         self.limit
104     }
105 
106     /// Sets the maximum number of bytes that can be read.
107     ///
108     /// # Note
109     ///
110     /// If the inner `Buf` has fewer bytes than `lim` then that is the actual
111     /// number of available bytes.
112     ///
113     /// # Examples
114     ///
115     /// ```rust
116     /// use bytes::{Buf, BufMut, buf::BufExt};
117     ///
118     /// let mut buf = b"hello world".take(2);
119     /// let mut dst = vec![];
120     ///
121     /// dst.put(&mut buf);
122     /// assert_eq!(*dst, b"he"[..]);
123     ///
124     /// dst.clear();
125     ///
126     /// buf.set_limit(3);
127     /// dst.put(&mut buf);
128     /// assert_eq!(*dst, b"llo"[..]);
129     /// ```
set_limit(&mut self, lim: usize)130     pub fn set_limit(&mut self, lim: usize) {
131         self.limit = lim
132     }
133 }
134 
135 impl<T: Buf> Buf for Take<T> {
remaining(&self) -> usize136     fn remaining(&self) -> usize {
137         cmp::min(self.inner.remaining(), self.limit)
138     }
139 
bytes(&self) -> &[u8]140     fn bytes(&self) -> &[u8] {
141         let bytes = self.inner.bytes();
142         &bytes[..cmp::min(bytes.len(), self.limit)]
143     }
144 
advance(&mut self, cnt: usize)145     fn advance(&mut self, cnt: usize) {
146         assert!(cnt <= self.limit);
147         self.inner.advance(cnt);
148         self.limit -= cnt;
149     }
150 }
151