pbs-datastore/src/crypt_reader.rs

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