proxmox-backup/pbs-datastore/src/crypt_reader.rs

use std::sync::Arc;
use std::io::{Read, BufRead};

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)
        }
    }
}