pbs-datastore/src/data_blob_reader.rs

   1 use std::io::{BufReader, Read};
   2 use std::sync::Arc;
   3
   4 use anyhow::{bail, format_err, Error};
   5
   6 use proxmox_io::ReadExt;
   7
   8 use pbs_tools::crypt_config::CryptConfig;
   9
  10 use crate::checksum_reader::ChecksumReader;
  11 use crate::crypt_reader::CryptReader;
  12 use crate::file_formats::{self, DataBlobHeader};
  13
  14 enum BlobReaderState<'reader, R: Read> {
  15     Uncompressed {
  16         expected_crc: u32,
  17         csum_reader: ChecksumReader<R>,
  18     },
  19     Compressed {
  20         expected_crc: u32,
  21         decompr: zstd::stream::read::Decoder<'reader, BufReader<ChecksumReader<R>>>,
  22     },
  23     Encrypted {
  24         expected_crc: u32,
  25         decrypt_reader: CryptReader<BufReader<ChecksumReader<R>>>,
  26     },
  27     EncryptedCompressed {
  28         expected_crc: u32,
  29         decompr: zstd::stream::read::Decoder<
  30             'reader,
  31             BufReader<CryptReader<BufReader<ChecksumReader<R>>>>,
  32         >,
  33     },
  34 }
  35
  36 /// Read data blobs
  37 pub struct DataBlobReader<'reader, R: Read> {
  38     state: BlobReaderState<'reader, R>,
  39 }
  40
  41 // zstd_safe::DCtx is not sync but we are, since
  42 // the only public interface is on mutable reference
  43 unsafe impl<R: Read> Sync for DataBlobReader<'_, R> {}
  44
  45 impl<R: Read> DataBlobReader<'_, R> {
  46     pub fn new(mut reader: R, config: Option<Arc<CryptConfig>>) -> Result<Self, Error> {
  47         let head: DataBlobHeader = unsafe { reader.read_le_value()? };
  48         match head.magic {
  49             file_formats::UNCOMPRESSED_BLOB_MAGIC_1_0 => {
  50                 let expected_crc = u32::from_le_bytes(head.crc);
  51                 let csum_reader = ChecksumReader::new(reader, None);
  52                 Ok(Self {
  53                     state: BlobReaderState::Uncompressed {
  54                         expected_crc,
  55                         csum_reader,
  56                     },
  57                 })
  58             }
  59             file_formats::COMPRESSED_BLOB_MAGIC_1_0 => {
  60                 let expected_crc = u32::from_le_bytes(head.crc);
  61                 let csum_reader = ChecksumReader::new(reader, None);
  62
  63                 let decompr = zstd::stream::read::Decoder::new(csum_reader)?;
  64                 Ok(Self {
  65                     state: BlobReaderState::Compressed {
  66                         expected_crc,
  67                         decompr,
  68                     },
  69                 })
  70             }
  71             file_formats::ENCRYPTED_BLOB_MAGIC_1_0 => {
  72                 let config = config
  73                     .ok_or_else(|| format_err!("unable to read encrypted blob without key"))?;
  74                 let expected_crc = u32::from_le_bytes(head.crc);
  75                 let mut iv = [0u8; 16];
  76                 let mut expected_tag = [0u8; 16];
  77                 reader.read_exact(&mut iv)?;
  78                 reader.read_exact(&mut expected_tag)?;
  79                 let csum_reader = ChecksumReader::new(reader, None);
  80                 let decrypt_reader = CryptReader::new(
  81                     BufReader::with_capacity(64 * 1024, csum_reader),
  82                     iv,
  83                     expected_tag,
  84                     config,
  85                 )?;
  86                 Ok(Self {
  87                     state: BlobReaderState::Encrypted {
  88                         expected_crc,
  89                         decrypt_reader,
  90                     },
  91                 })
  92             }
  93             file_formats::ENCR_COMPR_BLOB_MAGIC_1_0 => {
  94                 let config = config
  95                     .ok_or_else(|| format_err!("unable to read encrypted blob without key"))?;
  96                 let expected_crc = u32::from_le_bytes(head.crc);
  97                 let mut iv = [0u8; 16];
  98                 let mut expected_tag = [0u8; 16];
  99                 reader.read_exact(&mut iv)?;
 100                 reader.read_exact(&mut expected_tag)?;
 101                 let csum_reader = ChecksumReader::new(reader, None);
 102                 let decrypt_reader = CryptReader::new(
 103                     BufReader::with_capacity(64 * 1024, csum_reader),
 104                     iv,
 105                     expected_tag,
 106                     config,
 107                 )?;
 108                 let decompr = zstd::stream::read::Decoder::new(decrypt_reader)?;
 109                 Ok(Self {
 110                     state: BlobReaderState::EncryptedCompressed {
 111                         expected_crc,
 112                         decompr,
 113                     },
 114                 })
 115             }
 116             _ => bail!("got wrong magic number {:?}", head.magic),
 117         }
 118     }
 119
 120     pub fn finish(self) -> Result<R, Error> {
 121         match self.state {
 122             BlobReaderState::Uncompressed {
 123                 csum_reader,
 124                 expected_crc,
 125             } => {
 126                 let (reader, crc, _) = csum_reader.finish()?;
 127                 if crc != expected_crc {
 128                     bail!("blob crc check failed");
 129                 }
 130                 Ok(reader)
 131             }
 132             BlobReaderState::Compressed {
 133                 expected_crc,
 134                 decompr,
 135             } => {
 136                 let csum_reader = decompr.finish().into_inner();
 137                 let (reader, crc, _) = csum_reader.finish()?;
 138                 if crc != expected_crc {
 139                     bail!("blob crc check failed");
 140                 }
 141                 Ok(reader)
 142             }
 143             BlobReaderState::Encrypted {
 144                 expected_crc,
 145                 decrypt_reader,
 146             } => {
 147                 let csum_reader = decrypt_reader.finish()?.into_inner();
 148                 let (reader, crc, _) = csum_reader.finish()?;
 149                 if crc != expected_crc {
 150                     bail!("blob crc check failed");
 151                 }
 152                 Ok(reader)
 153             }
 154             BlobReaderState::EncryptedCompressed {
 155                 expected_crc,
 156                 decompr,
 157             } => {
 158                 let decrypt_reader = decompr.finish().into_inner();
 159                 let csum_reader = decrypt_reader.finish()?.into_inner();
 160                 let (reader, crc, _) = csum_reader.finish()?;
 161                 if crc != expected_crc {
 162                     bail!("blob crc check failed");
 163                 }
 164                 Ok(reader)
 165             }
 166         }
 167     }
 168 }
 169
 170 impl<R: Read> Read for DataBlobReader<'_, R> {
 171     fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
 172         match &mut self.state {
 173             BlobReaderState::Uncompressed { csum_reader, .. } => csum_reader.read(buf),
 174             BlobReaderState::Compressed { decompr, .. } => decompr.read(buf),
 175             BlobReaderState::Encrypted { decrypt_reader, .. } => decrypt_reader.read(buf),
 176             BlobReaderState::EncryptedCompressed { decompr, .. } => decompr.read(buf),
 177         }
 178     }
 179 }