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
use std::io::{self, Read};
pub trait ThBufRead {
/// Returns the currently buffered data.
fn get_buf(&self) -> &[u8];
/// Request more data to be buffered.
/// Indicates that the buffer from `get_buf` is insufficient to read the next chunk.
/// Returns the number of extra bytes read, 0 indicates eof.
fn fill_buf(&mut self) -> std::io::Result<usize>;
/// Number of bytes to discard from the buffer.
/// Must be no more than the buffer len from `get_buf`.
fn consume(&mut self, amount: usize);
}
pub struct ThBufReader<R> {
inner: R,
buf: Vec<u8>,
pos: usize,
cap: usize,
}
impl<R: Read> ThBufReader<R> {
pub fn new(inner: R) -> Self {
Self::with_capacity(8192, inner)
}
pub fn with_capacity(capacity: usize, inner: R) -> Self {
// `next_power_of_to`, to follow the rules of the Vec capacity sizes.
// At least 8, as Vec allocates at minimum 8 elements.
let capacity = capacity.next_power_of_two().max(8);
Self {
inner,
buf: vec![0; capacity],
pos: 0,
cap: 0,
}
}
pub fn reuse(mut self, inner: R) -> Self {
self.inner = inner;
self.pos = 0;
self.cap = 0;
self
}
}
impl<R: Read> ThBufRead for ThBufReader<R> {
fn get_buf(&self) -> &[u8] {
&self.buf[self.pos..self.cap]
}
fn fill_buf(&mut self) -> std::io::Result<usize> {
// If we've reached the end of our internal buffer then we need to fetch
// some more data from the underlying reader.
// Branch using `>=` instead of the more correct `==`
// to tell the compiler that the pos..cap slice is always valid.
if self.pos >= self.cap {
debug_assert!(self.pos == self.cap);
self.pos = 0;
self.cap = 0;
}
let remaining_size = self.buf.capacity() - self.cap;
if self.pos == 0 && remaining_size == 0 {
// RESIZE
// The 'valid' data takes up the entire buffer
if self.buf.capacity() >= 4194304 {
// 2^22 = 4M buffer size not sufficient
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"Too large chunk, possible DOS",
));
}
self.buf.resize(self.buf.capacity() * 2, 0);
} else if self.pos != 0 && remaining_size < 1024 {
// MOVE TO FRONT
// The 'valid' data is almost the end of the buffer.
// We need to make room after it for the next `read` call.
self.buf.copy_within(self.pos..self.cap, 0);
self.cap -= self.pos;
self.pos = 0;
}
let num_read = self.inner.read(&mut self.buf[self.cap..])?;
self.cap += num_read;
Ok(num_read)
}
fn consume(&mut self, amount: usize) {
self.pos += amount;
assert!(self.pos <= self.cap);
}
}
impl ThBufRead for &[u8] {
fn get_buf(&self) -> &[u8] {
self
}
fn fill_buf(&mut self) -> std::io::Result<usize> {
Ok(0)
}
fn consume(&mut self, amount: usize) {
*self = &self[amount..];
}
}