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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
use core::{
cmp,
io::{BorrowedBuf, BorrowedCursor},
};
use crate::{BufRead, Error, IoBuf, Read, Result, Seek, SeekFrom};
/// Reader adapter which limits the bytes read from an underlying reader.
///
/// This struct is generally created by calling [`take`] on a reader.
/// Please see the documentation of [`take`] for more details.
///
/// See [`std::io::Take`] for more details.
///
/// [`take`]: Read::take
#[derive(Debug)]
pub struct Take<T> {
inner: T,
len: u64,
limit: u64,
}
impl<T> Take<T> {
pub(crate) fn new(inner: T, limit: u64) -> Self {
Take {
inner,
len: limit,
limit,
}
}
/// Returns the number of bytes that can be read before this instance will
/// return EOF.
pub fn limit(&self) -> u64 {
self.limit
}
/// Returns the number of bytes read so far.
pub fn position(&self) -> u64 {
self.len - self.limit
}
/// Sets the number of bytes that can be read before this instance will
/// return EOF. This is the same as constructing a new `Take` instance, so
/// the amount of bytes read and the previous limit value don't matter when
/// calling this method.
pub fn set_limit(&mut self, limit: u64) {
self.len = limit;
self.limit = limit;
}
/// Consumes the `Take`, returning the wrapped reader.
pub fn into_inner(self) -> T {
self.inner
}
/// Gets a reference to the underlying reader.
///
/// Care should be taken to avoid modifying the internal I/O state of the
/// underlying reader as doing so may corrupt the internal limit of this
/// `Take`.
pub fn get_ref(&self) -> &T {
&self.inner
}
/// Gets a mutable reference to the underlying reader.
///
/// Care should be taken to avoid modifying the internal I/O state of the
/// underlying reader as doing so may corrupt the internal limit of this
/// `Take`.
pub fn get_mut(&mut self) -> &mut T {
&mut self.inner
}
}
impl<T: Read> Read for Take<T> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
// Don't call into inner reader at all at EOF because it may still block
if self.limit == 0 {
return Ok(0);
}
let max = cmp::min(buf.len() as u64, self.limit) as usize;
let n = self.inner.read(&mut buf[..max])?;
assert!(n as u64 <= self.limit, "number of read bytes exceeds limit");
self.limit -= n as u64;
Ok(n)
}
fn read_buf(&mut self, mut buf: BorrowedCursor<'_>) -> Result<()> {
// Don't call into inner reader at all at EOF because it may still block
if self.limit == 0 {
return Ok(());
}
if self.limit < buf.capacity() as u64 {
// The condition above guarantees that `self.limit` fits in `usize`.
let limit = self.limit as usize;
let is_init = buf.is_init();
// SAFETY: no uninit data is written to ibuf
let ibuf = unsafe { &mut buf.as_mut()[..limit] };
let mut sliced_buf: BorrowedBuf<'_> = ibuf.into();
// SAFETY: extra_init bytes of ibuf are known to be initialized
if is_init {
unsafe { sliced_buf.set_init() };
}
let mut cursor = sliced_buf.unfilled();
let result = self.inner.read_buf(cursor.reborrow());
let should_init = cursor.is_init();
let filled = sliced_buf.len();
// cursor / sliced_buf / ibuf must drop here
// Avoid accidentally quadratic behaviour by initializing the whole
// cursor if only part of it was initialized.
if should_init {
// SAFETY: no uninit data is written
let uninit = unsafe { &mut buf.as_mut()[limit..] };
uninit.write_filled(0);
// SAFETY: all bytes that were not initialized by `T::read_buf`
// have just been written to.
unsafe { buf.set_init() };
}
unsafe {
// SAFETY: filled bytes have been filled and therefore initialized
buf.advance(filled);
}
self.limit -= filled as u64;
result
} else {
let written = buf.written();
let result = self.inner.read_buf(buf.reborrow());
self.limit -= (buf.written() - written) as u64;
result
}
}
}
impl<T: BufRead> BufRead for Take<T> {
fn fill_buf(&mut self) -> Result<&[u8]> {
// Don't call into inner reader at all at EOF because it may still block
if self.limit == 0 {
return Ok(&[]);
}
let buf = self.inner.fill_buf()?;
let cap = cmp::min(buf.len() as u64, self.limit) as usize;
Ok(&buf[..cap])
}
fn consume(&mut self, amt: usize) {
// Don't let callers reset the limit by passing an overlarge value
let amt = cmp::min(amt as u64, self.limit) as usize;
self.limit -= amt as u64;
self.inner.consume(amt);
}
}
impl<T: Seek> Seek for Take<T> {
fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
let new_position = match pos {
SeekFrom::Start(v) => Some(v),
SeekFrom::Current(v) => self.position().checked_add_signed(v),
SeekFrom::End(v) => self.len.checked_add_signed(v),
};
let new_position = match new_position {
Some(v) if v <= self.len => v,
_ => return Err(Error::InvalidInput),
};
while new_position != self.position() {
if let Some(offset) = new_position.checked_signed_diff(self.position()) {
self.inner.seek_relative(offset)?;
self.limit = self.limit.wrapping_sub(offset as u64);
break;
}
let offset = if new_position > self.position() {
i64::MAX
} else {
i64::MIN
};
self.inner.seek_relative(offset)?;
self.limit = self.limit.wrapping_sub(offset as u64);
}
Ok(new_position)
}
fn stream_len(&mut self) -> Result<u64> {
Ok(self.len)
}
fn stream_position(&mut self) -> Result<u64> {
Ok(self.position())
}
fn seek_relative(&mut self, offset: i64) -> Result<()> {
if self
.position()
.checked_add_signed(offset)
.is_none_or(|p| p > self.len)
{
return Err(Error::InvalidInput);
}
self.inner.seek_relative(offset)?;
self.limit = self.limit.wrapping_sub(offset as u64);
Ok(())
}
}
impl<T: IoBuf> IoBuf for Take<T> {
fn remaining(&self) -> usize {
cmp::min(self.inner.remaining(), self.limit as usize)
}
}