diff --git a/src/backend/decrypt.rs b/src/backend/decrypt.rs index ebb2a7e..0c01410 100644 --- a/src/backend/decrypt.rs +++ b/src/backend/decrypt.rs @@ -1,4 +1,5 @@ use std::fs::File; +use std::num::NonZeroU32; use anyhow::{bail, Result}; use async_trait::async_trait; @@ -23,6 +24,7 @@ pub trait DecryptReadBackend: ReadBackend { cacheable: bool, offset: u32, length: u32, + uncompressed_length: Option, ) -> Result>; async fn get_file(&self, id: &Id) -> Result { @@ -176,13 +178,21 @@ impl DecryptReadBackend for DecryptBackend { cacheable: bool, offset: u32, length: u32, + uncompressed_length: Option, ) -> Result> { - Ok(self.key.decrypt_data( + let mut data = self.key.decrypt_data( &self .backend .read_partial(tpe, id, cacheable, offset, length) .await?, - )?) + )?; + if let Some(length) = uncompressed_length { + data = decode_all(&*data).unwrap(); + if data.len() != length.get() as usize { + bail!("length of uncompressed data does not match!"); + } + } + Ok(data) } } diff --git a/src/backend/dry_run.rs b/src/backend/dry_run.rs index 0d0cc83..1371fe6 100644 --- a/src/backend/dry_run.rs +++ b/src/backend/dry_run.rs @@ -1,4 +1,5 @@ use std::fs::File; +use std::num::NonZeroU32; use anyhow::Result; use async_trait::async_trait; @@ -32,9 +33,10 @@ impl DecryptReadBackend for DryRunBackend { cacheable: bool, offset: u32, length: u32, + uncompressed_length: Option, ) -> Result> { self.be - .read_encrypted_partial(tpe, id, cacheable, offset, length) + .read_encrypted_partial(tpe, id, cacheable, offset, length, uncompressed_length) .await } } diff --git a/src/blob/packer.rs b/src/blob/packer.rs index 90b7b33..65eb2bb 100644 --- a/src/blob/packer.rs +++ b/src/blob/packer.rs @@ -359,6 +359,7 @@ impl Repacker { blob.tpe.is_cacheable(), blob.offset, blob.length, + blob.uncompressed_length, ) .await?; self.packer diff --git a/src/commands/restore.rs b/src/commands/restore.rs index 317884d..ca4ffd4 100644 --- a/src/commands/restore.rs +++ b/src/commands/restore.rs @@ -8,7 +8,6 @@ use derive_getters::Dissolve; use futures::{stream::FuturesUnordered, TryStreamExt}; use tokio::spawn; use vlog::*; -use zstd::decode_all; use super::{progress_bytes, progress_counter}; use crate::backend::{DecryptReadBackend, FileType, LocalBackend}; @@ -131,17 +130,17 @@ async fn restore_contents( stream.push(spawn(async move { // read pack at blob_offset with length blob_length let data = be - .read_encrypted_partial(FileType::Pack, &pack, false, bl.offset, bl.length) + .read_encrypted_partial( + FileType::Pack, + &pack, + false, + bl.offset, + bl.length, + bl.uncompressed_length, + ) .await .unwrap(); - let data = match bl.uncompressed_length.is_some() { - false => data, - true => decode_all(&*data).unwrap(), - }; - - // TODO: check length of uncomressed data - if !dry_run { // save into needed files in parallel for (name, start) in name_dests { diff --git a/src/index/mod.rs b/src/index/mod.rs index a01780a..ba822b7 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -9,7 +9,6 @@ use derive_more::Constructor; use futures::StreamExt; use indicatif::ProgressBar; use vlog::*; -use zstd::decode_all; use crate::backend::{DecryptReadBackend, FileType}; use crate::blob::BlobType; @@ -51,12 +50,10 @@ impl IndexEntry { self.blob_type.is_cacheable(), self.offset, self.length, + self.uncompressed_length, ) .await?; - Ok(match self.uncompressed_length { - None => data, - Some(_) => decode_all(&*data)?, - }) + Ok(data) } pub fn data_length(&self) -> u32 {