rocket/data/
peekable.rs

Help
1use tokio::io::{AsyncRead, AsyncReadExt};
2
3pub struct Peekable<const N: usize, R> {
4    pub(crate) buffer: Vec<u8>,
5    pub(crate) complete: bool,
6    pub(crate) reader: R,
7}
8
9impl<const N: usize, R: AsyncRead + Unpin> Peekable<N, R> {
10    pub fn new(reader: R) -> Self {
11        Self { buffer: Vec::new(), complete: false, reader }
12    }
13
14    pub fn with_buffer(buffer: Vec<u8>, complete: bool, reader: R) -> Self {
15        Self { buffer, complete, reader }
16    }
17
18    pub async fn peek(&mut self, num: usize) -> &[u8] {
19        if self.complete {
20            return self.buffer.as_slice();
21        }
22
23        let to_read = std::cmp::min(N, num);
24        if self.buffer.len() >= to_read {
25            return self.buffer.as_slice();
26        }
27
28        if self.buffer.capacity() == 0 {
29            self.buffer.reserve(N);
30        }
31
32        while self.buffer.len() < to_read {
33            match self.reader.read_buf::<Vec<u8>>(&mut self.buffer).await {
34                Ok(0) => {
35                    self.complete = self.buffer.capacity() > self.buffer.len();
36                    break;
37                },
38                Ok(_) => { /* continue */ },
39                Err(e) => {
40                    error!("failed to read into peek buffer: {:?}.", e);
41                    break;
42                }
43            }
44        }
45
46        self.buffer.as_slice()
47    }
48}
rocket/data/peekable.rs

rocket/data/
peekable.rs
