93 lines
2.5 KiB
Rust
93 lines
2.5 KiB
Rust
// SPDX-License-Identifier: MIT-0
|
|
|
|
use std::io;
|
|
use std::io::{BufRead, Read};
|
|
use zeroize::Zeroizing;
|
|
|
|
pub enum ReadInfoChunk {
|
|
Normal(#[allow(dead_code)] usize),
|
|
Last(usize),
|
|
Empty,
|
|
}
|
|
|
|
pub struct AheadReader {
|
|
inner: Box<dyn BufRead + Send>,
|
|
buf: Zeroizing<Vec<u8>>,
|
|
bufsz: usize,
|
|
capacity: usize,
|
|
}
|
|
|
|
impl AheadReader {
|
|
pub fn from(reader: Box<dyn BufRead + Send>, capacity: usize) -> Self {
|
|
Self {
|
|
inner: reader,
|
|
buf: Zeroizing::new(vec![0; capacity]),
|
|
bufsz: 0,
|
|
capacity,
|
|
}
|
|
}
|
|
|
|
fn read_until_full(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
|
|
let mut total = 0;
|
|
loop {
|
|
match self.inner.read(buf) {
|
|
Ok(0) => break,
|
|
Ok(n) => {
|
|
total += n;
|
|
let tmp = buf;
|
|
buf = &mut tmp[n..];
|
|
}
|
|
Err(e) => match e.kind() {
|
|
io::ErrorKind::Interrupted => continue,
|
|
_ => return Err(e),
|
|
},
|
|
}
|
|
}
|
|
|
|
Ok(total)
|
|
}
|
|
|
|
pub fn read_ahead(&mut self, userbuf: &mut [u8]) -> io::Result<ReadInfoChunk> {
|
|
if self.bufsz == 0 {
|
|
return self.first_read(userbuf);
|
|
}
|
|
self.normal_read(userbuf)
|
|
}
|
|
|
|
fn first_read(&mut self, userbuf: &mut [u8]) -> io::Result<ReadInfoChunk> {
|
|
// 1st read directly to userbuf (we have no cached data yet)
|
|
let n = self.read_until_full(userbuf)?;
|
|
if n == 0 {
|
|
return Ok(ReadInfoChunk::Empty);
|
|
}
|
|
|
|
// 2nd read directly into our internal buf
|
|
let mut tmp = Zeroizing::new(vec![0u8; self.capacity]);
|
|
let n2 = self.read_until_full(&mut tmp)?;
|
|
self.buf = tmp;
|
|
self.bufsz = n2;
|
|
if n2 == 0 {
|
|
return Ok(ReadInfoChunk::Last(n));
|
|
}
|
|
|
|
Ok(ReadInfoChunk::Normal(n))
|
|
}
|
|
|
|
fn normal_read(&mut self, userbuf: &mut [u8]) -> io::Result<ReadInfoChunk> {
|
|
// copy internal buf to userbuf
|
|
userbuf.copy_from_slice(&self.buf);
|
|
let userbuf_sz = self.bufsz;
|
|
|
|
// 2nd read directly into our internal buf
|
|
let mut tmp = Zeroizing::new(vec![0u8; self.capacity]);
|
|
let n2 = self.read_until_full(&mut tmp)?;
|
|
self.buf = tmp;
|
|
self.bufsz = n2;
|
|
if n2 == 0 {
|
|
return Ok(ReadInfoChunk::Last(userbuf_sz));
|
|
}
|
|
|
|
Ok(ReadInfoChunk::Normal(userbuf_sz))
|
|
}
|
|
}
|