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