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
use std::io::Read;
use std::io;
use std::cmp;
use std::collections::VecDeque;
/// This trait specifies rustls's precise requirements doing writes with
/// vectored IO.
///
/// The purpose of vectored IO is to pass contigious output in many blocks
/// to the kernel without either coalescing it in user-mode (by allocating
/// and copying) or making many system calls.
///
/// We don't directly use types from the vecio crate because the traits
/// don't compose well: the most useful trait (`Rawv`) is hard to test
/// with (it can't be implemented without an FD) and implies a readable
/// source too. You will have to write a trivial adaptor struct which
/// glues either `vecio::Rawv` or `vecio::Writev` to this trait. See
/// the rustls examples.
pub trait WriteV {
/// Writes as much data from `vbytes` as possible, returning
/// the number of bytes written.
fn writev(&mut self, vbytes: &[&[u8]]) -> io::Result<usize>;
}
/// This is a byte buffer that is built from a vector
/// of byte vectors. This avoids extra copies when
/// appending a new byte vector, at the expense of
/// more complexity when reading out.
pub struct ChunkVecBuffer {
chunks: VecDeque<Vec<u8>>,
limit: usize,
}
impl ChunkVecBuffer {
pub fn new() -> ChunkVecBuffer {
ChunkVecBuffer { chunks: VecDeque::new(), limit: 0 }
}
/// Sets the upper limit on how many bytes this
/// object can store.
///
/// Setting a lower limit than the currently stored
/// data is not an error.
///
/// A zero limit is interpreted as no limit.
pub fn set_limit(&mut self, new_limit: usize) {
self.limit = new_limit;
}
/// If we're empty
pub fn is_empty(&self) -> bool {
self.chunks.is_empty()
}
/// How many bytes we're storing
pub fn len(&self) -> usize {
let mut len = 0;
for ch in &self.chunks {
len += ch.len();
}
len
}
/// For a proposed append of `len` bytes, how many
/// bytes should we actually append to adhere to the
/// currently set `limit`?
pub fn apply_limit(&self, len: usize) -> usize {
if self.limit == 0 {
len
} else {
let space =self.limit.saturating_sub(self.len());
cmp::min(len, space)
}
}
/// Append a copy of `bytes`, perhaps a prefix if
/// we're near the limit.
pub fn append_limited_copy(&mut self, bytes: &[u8]) -> usize {
let take = self.apply_limit(bytes.len());
self.append(bytes[..take].to_vec());
take
}
/// Take and append the given `bytes`.
pub fn append(&mut self, bytes: Vec<u8>) -> usize {
let len = bytes.len();
if !bytes.is_empty() {
self.chunks.push_back(bytes);
}
len
}
/// Take one of the chunks from this object. This
/// function panics if the object `is_empty`.
pub fn take_one(&mut self) -> Vec<u8> {
self.chunks.pop_front().unwrap()
}
/// Read data out of this object, writing it into `buf`
/// and returning how many bytes were written there.
pub fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let mut offs = 0;
while offs < buf.len() && !self.is_empty() {
let used = self.chunks[0].as_slice().read(&mut buf[offs..])?;
self.consume(used);
offs += used;
}
Ok(offs)
}
fn consume(&mut self, mut used: usize) {
while used > 0 && !self.is_empty() {
if used >= self.chunks[0].len() {
used -= self.chunks[0].len();
self.take_one();
} else {
self.chunks[0] = self.chunks[0].split_off(used);
used = 0;
}
}
}
/// Read data out of this object, passing it `wr`
pub fn write_to(&mut self, wr: &mut io::Write) -> io::Result<usize> {
if self.is_empty() {
return Ok(0);
}
let used = wr.write(&self.chunks[0])?;
self.consume(used);
Ok(used)
}
pub fn writev_to(&mut self, wr: &mut WriteV) -> io::Result<usize> {
if self.is_empty() {
return Ok(0);
}
let used = {
let chunks = self.chunks.iter()
.map(|ch| ch.as_ref())
.collect::<Vec<&[u8]>>();
wr.writev(&chunks)?
};
self.consume(used);
Ok(used)
}
}
#[cfg(test)]
mod test {
use super::ChunkVecBuffer;
#[test]
fn short_append_copy_with_limit()
{
let mut cvb = ChunkVecBuffer::new();
cvb.set_limit(12);
assert_eq!(cvb.append_limited_copy(b"hello"), 5);
assert_eq!(cvb.append_limited_copy(b"world"), 5);
assert_eq!(cvb.append_limited_copy(b"hello"), 2);
assert_eq!(cvb.append_limited_copy(b"world"), 0);
let mut buf = [0u8; 12];
assert_eq!(cvb.read(&mut buf).unwrap(), 12);
assert_eq!(buf.to_vec(),
b"helloworldhe".to_vec());
}
}