use ::prelude::*; use super::*; use std::path::{Path, PathBuf}; use std::fs::File; use std::io::{self, Write, BufWriter}; use std::sync::{Arc, Mutex}; use chrono::prelude::*; quick_error!{ #[derive(Debug)] pub enum BundleWriterError { CompressionSetup(err: CompressionError) { cause(err) description("Failed to setup compression") display("Bundle writer error: failed to setup compression\n\tcaused by: {}", err) } Compression(err: CompressionError) { cause(err) description("Failed to compress data") display("Bundle writer error: failed to compress data\n\tcaused by: {}", err) } Encryption(err: EncryptionError) { from() cause(err) description("Encryption failed") display("Bundle writer error: failed to encrypt data\n\tcaused by: {}", err) } Encode(err: msgpack::EncodeError, path: PathBuf) { cause(err) context(path: &'a Path, err: msgpack::EncodeError) -> (err, path.to_path_buf()) description("Failed to encode bundle header to file") display("Bundle writer error: failed to encode bundle header to file {:?}\n\tcaused by: {}", path, err) } Write(err: io::Error, path: PathBuf) { cause(err) context(path: &'a Path, err: io::Error) -> (err, path.to_path_buf()) description("Failed to write data to file") display("Bundle writer error: failed to write data to file {:?}\n\tcaused by: {}", path, err) } } } pub struct BundleWriter { mode: BundleMode, hash_method: HashMethod, data: Vec, compression: Option, compression_stream: Option, encryption: Option, crypto: Arc>, raw_size: usize, chunk_count: usize, chunks: ChunkList, } impl BundleWriter { pub fn new(mode: BundleMode, hash_method: HashMethod, compression: Option, encryption: Option, crypto: Arc>) -> Result { let compression_stream = match compression { Some(ref compression) => Some(try!(compression.compress_stream().map_err(BundleWriterError::CompressionSetup))), None => None }; Ok(BundleWriter { mode: mode, hash_method: hash_method, data: vec![], compression: compression, compression_stream: compression_stream, encryption: encryption, crypto: crypto, raw_size: 0, chunk_count: 0, chunks: ChunkList::new() }) } pub fn add(&mut self, chunk: &[u8], hash: Hash) -> Result { if let Some(ref mut stream) = self.compression_stream { try!(stream.process(chunk, &mut self.data).map_err(BundleWriterError::Compression)) } else { self.data.extend_from_slice(chunk) } self.raw_size += chunk.len(); self.chunk_count += 1; self.chunks.push((hash, chunk.len() as u32)); Ok(self.chunk_count-1) } pub fn finish(mut self, db: &BundleDb) -> Result { if let Some(stream) = self.compression_stream { try!(stream.finish(&mut self.data).map_err(BundleWriterError::Compression)) } if let Some(ref encryption) = self.encryption { self.data = try!(self.crypto.lock().unwrap().encrypt(&encryption, &self.data)); } let encoded_size = self.data.len(); let mut chunk_data = Vec::with_capacity(self.chunks.encoded_size()); self.chunks.write_to(&mut chunk_data).unwrap(); let id = BundleId(self.hash_method.hash(&chunk_data)); if let Some(ref encryption) = self.encryption { chunk_data = try!(self.crypto.lock().unwrap().encrypt(&encryption, &chunk_data)); } let mut path = db.layout.temp_bundle_path(); let mut file = BufWriter::new(try!(File::create(&path).context(&path as &Path))); try!(file.write_all(&HEADER_STRING).context(&path as &Path)); try!(file.write_all(&[HEADER_VERSION]).context(&path as &Path)); let info = BundleInfo { mode: self.mode, hash_method: self.hash_method, compression: self.compression, encryption: self.encryption.clone(), chunk_count: self.chunk_count, id: id.clone(), raw_size: self.raw_size, encoded_size: encoded_size, chunk_list_size: chunk_data.len(), timestamp: Local::now().timestamp() }; let mut info_data = try!(msgpack::encode(&info).context(&path as &Path)); if let Some(ref encryption) = self.encryption { info_data = try!(self.crypto.lock().unwrap().encrypt(&encryption, &info_data)); } let header = BundleHeader { encryption: self.encryption, info_size: info_data.len() }; try!(msgpack::encode_to_stream(&header, &mut file).context(&path as &Path)); try!(file.write_all(&info_data).context(&path as &Path)); try!(file.write_all(&chunk_data).context(&path as &Path)); try!(file.write_all(&self.data).context(&path as &Path)); path = path.strip_prefix(db.layout.base_path()).unwrap().to_path_buf(); Ok(StoredBundle { path: path, info: info }) } #[inline] pub fn raw_size(&self) -> usize { self.raw_size } #[inline] pub fn estimate_final_size(&self) -> usize { self.data.len() + self.chunk_count * 20 + 500 } }