1use crate::{Buf, Bytes};
2
3use 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)]
10pub struct Take<T> {
11 inner: T,
12 limit: usize,
13}
14
15pub fn new<T>(inner: T, limit: usize) -> Take<T> {
16 Take { inner, limit }
17}
18
19impl<T> Take<T> {
20 /// Consumes this `Take`, returning the underlying value.
21 ///
22 /// # Examples
23 ///
24 /// ```rust
25 /// use bytes::{Buf, BufMut};
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 /// ```
39 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;
51 ///
52 /// let buf = b"hello world".take(2);
53 ///
54 /// assert_eq!(11, buf.get_ref().remaining());
55 /// ```
56 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};
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 /// ```
77 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;
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 /// ```
99 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::{Buf, BufMut};
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 /// ```
127 pub fn set_limit(&mut self, lim: usize) {
128 self.limit = lim
129 }
130}
131
132impl<T: Buf> Buf for Take<T> {
133 fn remaining(&self) -> usize {
134 cmp::min(self.inner.remaining(), self.limit)
135 }
136
137 fn chunk(&self) -> &[u8] {
138 let bytes = self.inner.chunk();
139 &bytes[..cmp::min(bytes.len(), self.limit)]
140 }
141
142 fn advance(&mut self, cnt: usize) {
143 assert!(cnt <= self.limit);
144 self.inner.advance(cnt);
145 self.limit -= cnt;
146 }
147
148 fn copy_to_bytes(&mut self, len: usize) -> Bytes {
149 assert!(len <= self.remaining(), "`len` greater than remaining");
150
151 let r = self.inner.copy_to_bytes(len);
152 self.limit -= len;
153 r
154 }
155}
156