tlspuffin/tls/rustls/
vecbuf.rs

1use std::collections::VecDeque;
2use std::io::Read;
3use std::{cmp, io};
4
5/// This is a byte buffer that is built from a vector
6/// of byte vectors.  This avoids extra copies when
7/// appending a new byte vector, at the expense of
8/// more complexity when reading out.
9pub struct ChunkVecBuffer {
10    chunks: VecDeque<Vec<u8>>,
11    limit: Option<usize>,
12}
13
14impl ChunkVecBuffer {
15    pub fn new(limit: Option<usize>) -> Self {
16        Self {
17            chunks: VecDeque::new(),
18            limit,
19        }
20    }
21
22    /// Sets the upper limit on how many bytes this
23    /// object can store.
24    ///
25    /// Setting a lower limit than the currently stored
26    /// data is not an error.
27    ///
28    /// A [`None`] limit is interpreted as no limit.
29    pub fn set_limit(&mut self, new_limit: Option<usize>) {
30        self.limit = new_limit;
31    }
32
33    /// If we're empty
34    pub fn is_empty(&self) -> bool {
35        self.chunks.is_empty()
36    }
37
38    /// How many bytes we're storing
39    pub fn len(&self) -> usize {
40        let mut len = 0;
41        for ch in &self.chunks {
42            len += ch.len();
43        }
44        len
45    }
46
47    /// For a proposed append of `len` bytes, how many
48    /// bytes should we actually append to adhere to the
49    /// currently set `limit`?
50    pub fn apply_limit(&self, len: usize) -> usize {
51        if let Some(limit) = self.limit {
52            let space = limit.saturating_sub(self.len());
53            cmp::min(len, space)
54        } else {
55            len
56        }
57    }
58
59    /// Append a copy of `bytes`, perhaps a prefix if
60    /// we're near the limit.
61    pub fn append_limited_copy(&mut self, bytes: &[u8]) -> usize {
62        let take = self.apply_limit(bytes.len());
63        self.append(bytes[..take].to_vec());
64        take
65    }
66
67    /// Take and append the given `bytes`.
68    pub fn append(&mut self, bytes: Vec<u8>) -> usize {
69        let len = bytes.len();
70
71        if !bytes.is_empty() {
72            self.chunks.push_back(bytes);
73        }
74
75        len
76    }
77
78    /// Take one of the chunks from this object.  This
79    /// function panics if the object `is_empty`.
80    pub fn pop(&mut self) -> Option<Vec<u8>> {
81        self.chunks.pop_front()
82    }
83
84    /// Read data out of this object, writing it into `buf`
85    /// and returning how many bytes were written there.
86    pub fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
87        let mut offs = 0;
88
89        while offs < buf.len() && !self.is_empty() {
90            let used = self.chunks[0].as_slice().read(&mut buf[offs..])?;
91
92            self.consume(used);
93            offs += used;
94        }
95
96        Ok(offs)
97    }
98
99    fn consume(&mut self, mut used: usize) {
100        while let Some(mut buf) = self.chunks.pop_front() {
101            if used < buf.len() {
102                self.chunks.push_front(buf.split_off(used));
103                break;
104            } else {
105                used -= buf.len();
106            }
107        }
108    }
109
110    /// Read data out of this object, passing it `wr`
111    pub fn write_to(&mut self, wr: &mut dyn io::Write) -> io::Result<usize> {
112        if self.is_empty() {
113            return Ok(0);
114        }
115
116        let mut bufs = [io::IoSlice::new(&[]); 64];
117        for (iov, chunk) in bufs.iter_mut().zip(self.chunks.iter()) {
118            *iov = io::IoSlice::new(chunk);
119        }
120        let len = cmp::min(bufs.len(), self.chunks.len());
121        let used = wr.write_vectored(&bufs[..len])?;
122        self.consume(used);
123        Ok(used)
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::ChunkVecBuffer;
130
131    #[test_log::test]
132    fn short_append_copy_with_limit() {
133        let mut cvb = ChunkVecBuffer::new(Some(12));
134        assert_eq!(cvb.append_limited_copy(b"hello"), 5);
135        assert_eq!(cvb.append_limited_copy(b"world"), 5);
136        assert_eq!(cvb.append_limited_copy(b"hello"), 2);
137        assert_eq!(cvb.append_limited_copy(b"world"), 0);
138
139        let mut buf = [0u8; 12];
140        assert_eq!(cvb.read(&mut buf).unwrap(), 12);
141        assert_eq!(buf.to_vec(), b"helloworldhe".to_vec());
142    }
143}
tlspuffin/tls/rustls/vecbuf.rs

tlspuffin/tls/rustls/
vecbuf.rs
