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.BufExt.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 { inner, limit }
17 }
18 
19 impl<T> Take<T> {
20     /// Consumes this `Take`, returning the underlying value.
21     ///
22     /// # Examples
23     ///
24     /// ```rust
25     /// use bytes::buf::{BufMut, BufExt};
26     ///
27     /// let mut buf = b"hello world".take(2);
28     /// let mut dst = vec![];
29     ///
30     /// dst.put(&mut buf);
31     /// assert_eq!(*dst, b"he"[..]);
32     ///
33     /// let mut buf = buf.into_inner();
34     ///
35     /// dst.clear();
36     /// dst.put(&mut buf);
37     /// assert_eq!(*dst, b"llo world"[..]);
38     /// ```
into_inner(self) -> T39     pub fn into_inner(self) -> T {
40         self.inner
41     }
42 
43     /// Gets a reference to the underlying `Buf`.
44     ///
45     /// It is inadvisable to directly read from the underlying `Buf`.
46     ///
47     /// # Examples
48     ///
49     /// ```rust
50     /// use bytes::{Buf, buf::BufExt};
51     ///
52     /// let buf = b"hello world".take(2);
53     ///
54     /// assert_eq!(11, buf.get_ref().remaining());
55     /// ```
get_ref(&self) -> &T56     pub fn get_ref(&self) -> &T {
57         &self.inner
58     }
59 
60     /// Gets a mutable reference to the underlying `Buf`.
61     ///
62     /// It is inadvisable to directly read from the underlying `Buf`.
63     ///
64     /// # Examples
65     ///
66     /// ```rust
67     /// use bytes::{Buf, BufMut, buf::BufExt};
68     ///
69     /// let mut buf = b"hello world".take(2);
70     /// let mut dst = vec![];
71     ///
72     /// buf.get_mut().advance(2);
73     ///
74     /// dst.put(&mut buf);
75     /// assert_eq!(*dst, b"ll"[..]);
76     /// ```
get_mut(&mut self) -> &mut T77     pub fn get_mut(&mut self) -> &mut T {
78         &mut self.inner
79     }
80 
81     /// Returns the maximum number of bytes that can be read.
82     ///
83     /// # Note
84     ///
85     /// If the inner `Buf` has fewer bytes than indicated by this method then
86     /// that is the actual number of available bytes.
87     ///
88     /// # Examples
89     ///
90     /// ```rust
91     /// use bytes::{Buf, buf::BufExt};
92     ///
93     /// let mut buf = b"hello world".take(2);
94     ///
95     /// assert_eq!(2, buf.limit());
96     /// assert_eq!(b'h', buf.get_u8());
97     /// assert_eq!(1, buf.limit());
98     /// ```
limit(&self) -> usize99     pub fn limit(&self) -> usize {
100         self.limit
101     }
102 
103     /// Sets the maximum number of bytes that can be read.
104     ///
105     /// # Note
106     ///
107     /// If the inner `Buf` has fewer bytes than `lim` then that is the actual
108     /// number of available bytes.
109     ///
110     /// # Examples
111     ///
112     /// ```rust
113     /// use bytes::{BufMut, buf::BufExt};
114     ///
115     /// let mut buf = b"hello world".take(2);
116     /// let mut dst = vec![];
117     ///
118     /// dst.put(&mut buf);
119     /// assert_eq!(*dst, b"he"[..]);
120     ///
121     /// dst.clear();
122     ///
123     /// buf.set_limit(3);
124     /// dst.put(&mut buf);
125     /// assert_eq!(*dst, b"llo"[..]);
126     /// ```
set_limit(&mut self, lim: usize)127     pub fn set_limit(&mut self, lim: usize) {
128         self.limit = lim
129     }
130 }
131 
132 impl<T: Buf> Buf for Take<T> {
remaining(&self) -> usize133     fn remaining(&self) -> usize {
134         cmp::min(self.inner.remaining(), self.limit)
135     }
136 
bytes(&self) -> &[u8]137     fn bytes(&self) -> &[u8] {
138         let bytes = self.inner.bytes();
139         &bytes[..cmp::min(bytes.len(), self.limit)]
140     }
141 
advance(&mut self, cnt: usize)142     fn advance(&mut self, cnt: usize) {
143         assert!(cnt <= self.limit);
144         self.inner.advance(cnt);
145         self.limit -= cnt;
146     }
147 }
148