pbs-datastore/src/crypt_reader.rs

   1 use std::sync::Arc;
   2 use std::io::{Read, BufRead};
   3
   4 use anyhow::{bail, Error};
   5
   6 use super::CryptConfig;
   7
   8 pub struct CryptReader<R> {
   9     reader: R,
  10     small_read_buf: Vec<u8>,
  11     block_size: usize,
  12     crypter: openssl::symm::Crypter,
  13     finalized: bool,
  14 }
  15
  16 impl <R: BufRead> CryptReader<R> {
  17
  18     pub fn new(reader: R, iv: [u8; 16], tag: [u8; 16], config: Arc<CryptConfig>) -> Result<Self, Error> {
  19         let block_size = config.cipher().block_size(); // Note: block size is normally 1 byte for stream ciphers
  20         if block_size.count_ones() != 1 || block_size > 512 {
  21             bail!("unexpected Cipher block size {}", block_size);
  22         }
  23         let mut crypter = config.data_crypter(&iv, openssl::symm::Mode::Decrypt)?;
  24         crypter.set_tag(&tag)?;
  25
  26         Ok(Self { reader, crypter, block_size, finalized: false, small_read_buf: Vec::new() })
  27     }
  28
  29     pub fn finish(self) -> Result<R, Error> {
  30         if !self.finalized {
  31             bail!("CryptReader not successfully finalized.");
  32         }
  33         Ok(self.reader)
  34     }
  35 }
  36
  37 impl <R: BufRead> Read for CryptReader<R> {
  38
  39     fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
  40         if !self.small_read_buf.is_empty() {
  41             let max = if self.small_read_buf.len() > buf.len() {  buf.len() } else { self.small_read_buf.len() };
  42             let rest = self.small_read_buf.split_off(max);
  43             buf[..max].copy_from_slice(&self.small_read_buf);
  44             self.small_read_buf = rest;
  45             return Ok(max);
  46         }
  47
  48         let data = self.reader.fill_buf()?;
  49
  50         // handle small read buffers
  51         if buf.len() <= 2*self.block_size {
  52             let mut outbuf = [0u8; 1024];
  53
  54             let count = if data.is_empty() { // EOF
  55                 let written = self.crypter.finalize(&mut outbuf)?;
  56                 self.finalized = true;
  57                 written
  58             } else {
  59                 let mut read_size = outbuf.len() - self.block_size;
  60                 if read_size > data.len() {
  61                     read_size = data.len();
  62                 }
  63                 let written = self.crypter.update(&data[..read_size], &mut outbuf)?;
  64                 self.reader.consume(read_size);
  65                 written
  66             };
  67
  68             if count > buf.len() {
  69                 buf.copy_from_slice(&outbuf[..buf.len()]);
  70                 self.small_read_buf = outbuf[buf.len()..count].to_vec();
  71                 Ok(buf.len())
  72             } else {
  73                 buf[..count].copy_from_slice(&outbuf[..count]);
  74                 Ok(count)
  75             }
  76         } else if data.is_empty() { // EOF
  77             let rest = self.crypter.finalize(buf)?;
  78             self.finalized = true;
  79             Ok(rest)
  80         } else {
  81             let mut read_size = buf.len() - self.block_size;
  82             if read_size > data.len() {
  83                 read_size = data.len();
  84             }
  85             let count = self.crypter.update(&data[..read_size], buf)?;
  86             self.reader.consume(read_size);
  87             Ok(count)
  88         }
  89     }
  90 }