1use crate::Buf;
2
3use core::cmp;
4
5#[cfg(feature = "std")]
6use std::io::IoSlice;
7
8/// A `Buf` adapter which limits the bytes read from an underlying buffer.
9///
10/// This struct is generally created by calling `take()` on `Buf`. See
11/// documentation of [`take()`](Buf::take) for more details.
12#[derive(Debug)]
13pub struct Take<T> {
14 inner: T,
15 limit: usize,
16}
17
18pub fn new<T>(inner: T, limit: usize) -> Take<T> {
19 Take { inner, limit }
20}
21
22impl<T> Take<T> {
23 /// Consumes this `Take`, returning the underlying value.
24 ///
25 /// # Examples
26 ///
27 /// ```rust
28 /// use bytes::{Buf, BufMut};
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 /// ```
42 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;
54 ///
55 /// let buf = b"hello world".take(2);
56 ///
57 /// assert_eq!(11, buf.get_ref().remaining());
58 /// ```
59 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};
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 /// ```
80 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;
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 /// ```
102 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};
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 /// ```
130 pub fn set_limit(&mut self, lim: usize) {
131 self.limit = lim
132 }
133}
134
135impl<T: Buf> Buf for Take<T> {
136 fn remaining(&self) -> usize {
137 cmp::min(self.inner.remaining(), self.limit)
138 }
139
140 fn chunk(&self) -> &[u8] {
141 let bytes = self.inner.chunk();
142 &bytes[..cmp::min(bytes.len(), self.limit)]
143 }
144
145 fn advance(&mut self, cnt: usize) {
146 assert!(cnt <= self.limit);
147 self.inner.advance(cnt);
148 self.limit -= cnt;
149 }
150
151 fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes {
152 assert!(len <= self.remaining(), "`len` greater than remaining");
153
154 let r = self.inner.copy_to_bytes(len);
155 self.limit -= len;
156 r
157 }
158
159 #[cfg(feature = "std")]
160 fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
161 if self.limit == 0 {
162 return 0;
163 }
164
165 const LEN: usize = 16;
166 let mut slices: [IoSlice<'a>; LEN] = [
167 IoSlice::new(&[]),
168 IoSlice::new(&[]),
169 IoSlice::new(&[]),
170 IoSlice::new(&[]),
171 IoSlice::new(&[]),
172 IoSlice::new(&[]),
173 IoSlice::new(&[]),
174 IoSlice::new(&[]),
175 IoSlice::new(&[]),
176 IoSlice::new(&[]),
177 IoSlice::new(&[]),
178 IoSlice::new(&[]),
179 IoSlice::new(&[]),
180 IoSlice::new(&[]),
181 IoSlice::new(&[]),
182 IoSlice::new(&[]),
183 ];
184
185 let cnt = self
186 .inner
187 .chunks_vectored(&mut slices[..dst.len().min(LEN)]);
188 let mut limit = self.limit;
189 for (i, (dst, slice)) in dst[..cnt].iter_mut().zip(slices.iter()).enumerate() {
190 if let Some(buf) = slice.get(..limit) {
191 // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV.
192 let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(buf) };
193 *dst = IoSlice::new(buf);
194 return i + 1;
195 } else {
196 // SAFETY: We could do this safely with `IoSlice::advance` if we had a larger MSRV.
197 let buf = unsafe { std::mem::transmute::<&[u8], &'a [u8]>(slice) };
198 *dst = IoSlice::new(buf);
199 limit -= slice.len();
200 }
201 }
202 cnt
203 }
204}
205