106 lines
3.2 KiB
Rust
106 lines
3.2 KiB
Rust
use std::io::{BufRead, Read};
|
|
use std::sync::Arc;
|
|
|
|
use anyhow::{bail, Error};
|
|
|
|
use pbs_tools::crypt_config::CryptConfig;
|
|
|
|
pub struct CryptReader<R> {
|
|
reader: R,
|
|
small_read_buf: Vec<u8>,
|
|
block_size: usize,
|
|
crypter: openssl::symm::Crypter,
|
|
finalized: bool,
|
|
}
|
|
|
|
impl<R: BufRead> CryptReader<R> {
|
|
pub fn new(
|
|
reader: R,
|
|
iv: [u8; 16],
|
|
tag: [u8; 16],
|
|
config: Arc<CryptConfig>,
|
|
) -> Result<Self, Error> {
|
|
let block_size = config.cipher().block_size(); // Note: block size is normally 1 byte for stream ciphers
|
|
if block_size.count_ones() != 1 || block_size > 512 {
|
|
bail!("unexpected Cipher block size {}", block_size);
|
|
}
|
|
let mut crypter = config.data_crypter(&iv, openssl::symm::Mode::Decrypt)?;
|
|
crypter.set_tag(&tag)?;
|
|
|
|
Ok(Self {
|
|
reader,
|
|
crypter,
|
|
block_size,
|
|
finalized: false,
|
|
small_read_buf: Vec::new(),
|
|
})
|
|
}
|
|
|
|
pub fn finish(self) -> Result<R, Error> {
|
|
if !self.finalized {
|
|
bail!("CryptReader not successfully finalized.");
|
|
}
|
|
Ok(self.reader)
|
|
}
|
|
}
|
|
|
|
impl<R: BufRead> Read for CryptReader<R> {
|
|
fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
|
|
if !self.small_read_buf.is_empty() {
|
|
let max = if self.small_read_buf.len() > buf.len() {
|
|
buf.len()
|
|
} else {
|
|
self.small_read_buf.len()
|
|
};
|
|
let rest = self.small_read_buf.split_off(max);
|
|
buf[..max].copy_from_slice(&self.small_read_buf);
|
|
self.small_read_buf = rest;
|
|
return Ok(max);
|
|
}
|
|
|
|
let data = self.reader.fill_buf()?;
|
|
|
|
// handle small read buffers
|
|
if buf.len() <= 2 * self.block_size {
|
|
let mut outbuf = [0u8; 1024];
|
|
|
|
let count = if data.is_empty() {
|
|
// EOF
|
|
let written = self.crypter.finalize(&mut outbuf)?;
|
|
self.finalized = true;
|
|
written
|
|
} else {
|
|
let mut read_size = outbuf.len() - self.block_size;
|
|
if read_size > data.len() {
|
|
read_size = data.len();
|
|
}
|
|
let written = self.crypter.update(&data[..read_size], &mut outbuf)?;
|
|
self.reader.consume(read_size);
|
|
written
|
|
};
|
|
|
|
if count > buf.len() {
|
|
buf.copy_from_slice(&outbuf[..buf.len()]);
|
|
self.small_read_buf = outbuf[buf.len()..count].to_vec();
|
|
Ok(buf.len())
|
|
} else {
|
|
buf[..count].copy_from_slice(&outbuf[..count]);
|
|
Ok(count)
|
|
}
|
|
} else if data.is_empty() {
|
|
// EOF
|
|
let rest = self.crypter.finalize(buf)?;
|
|
self.finalized = true;
|
|
Ok(rest)
|
|
} else {
|
|
let mut read_size = buf.len() - self.block_size;
|
|
if read_size > data.len() {
|
|
read_size = data.len();
|
|
}
|
|
let count = self.crypter.update(&data[..read_size], buf)?;
|
|
self.reader.consume(read_size);
|
|
Ok(count)
|
|
}
|
|
}
|
|
}
|