Check content before zip/unzip the txhashset (#1174)

* Check txhashset content before zip/unzip
* Add header in txhashset verification
* Add copy function and test
* Add file util
* Now check and remove unexpected files instead of just crashing

chain/src/chain.rs

@@ -529,7 +529,7 @@ impl Chain {
 		}
 
 		// prepares the zip and return the corresponding Read
-		let txhashset_reader = txhashset::zip_read(self.db_root.clone())?;
+		let txhashset_reader = txhashset::zip_read(self.db_root.clone(), &header)?;
 		Ok((
 			header.output_mmr_size,
 			header.kernel_mmr_size,
@@ -558,7 +558,7 @@ impl Chain {
 		}
 
 		let header = self.store.get_block_header(&h)?;
-		txhashset::zip_write(self.db_root.clone(), txhashset_data)?;
+		txhashset::zip_write(self.db_root.clone(), txhashset_data, &header)?;
 
 		let mut txhashset =
 			txhashset::TxHashSet::open(self.db_root.clone(), self.store.clone(), Some(&header))?;

chain/src/txhashset.rs

@@ -15,9 +15,8 @@
 //! Utility structs to handle the 3 hashtrees (output, range proof,
 //! kernel) more conveniently and transactionally.
 
-use std::collections::HashMap;
-use std::fs;
-use std::fs::File;
+use std::collections::{HashMap, HashSet};
+use std::fs::{self, File};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::Instant;
@@ -38,11 +37,11 @@ use core::ser::{PMMRIndexHashable, PMMRable};
 
 use error::{Error, ErrorKind};
 use grin_store;
-use grin_store::pmmr::PMMRBackend;
+use grin_store::pmmr::{PMMRBackend, PMMR_FILES};
 use grin_store::types::prune_noop;
 use store::{Batch, ChainStore};
 use types::{TxHashSetRoots, TxHashsetWriteStatus};
-use util::{secp_static, zip, LOGGER};
+use util::{file, secp_static, zip, LOGGER};
 
 const TXHASHSET_SUBDIR: &'static str = "txhashset";
 const OUTPUT_SUBDIR: &'static str = "output";
@@ -1055,13 +1054,23 @@ impl<'a> Extension<'a> {
 
 /// Packages the txhashset data files into a zip and returns a Read to the
 /// resulting file
-pub fn zip_read(root_dir: String) -> Result<File, Error> {
+pub fn zip_read(root_dir: String, header: &BlockHeader) -> Result<File, Error> {
 	let txhashset_path = Path::new(&root_dir).join(TXHASHSET_SUBDIR);
 	let zip_path = Path::new(&root_dir).join(TXHASHSET_ZIP);
-
 	// create the zip archive
 	{
-		zip::compress(&txhashset_path, &File::create(zip_path.clone())?)
+		// Temp txhashset directory
+		let temp_txhashset_path = Path::new(&root_dir).join(TXHASHSET_SUBDIR.to_string() + "_zip");
+		// Remove temp dir if it exist
+		if temp_txhashset_path.exists() {
+			fs::remove_dir_all(&temp_txhashset_path)?;
+		}
+		// Copy file to another dir
+		file::copy_dir_to(&txhashset_path,&temp_txhashset_path)?;
+		// Check and remove file that are not supposed to be there
+		check_and_remove_files(&temp_txhashset_path, header)?;
+		// Compress zip
+		zip::compress(&temp_txhashset_path, &File::create(zip_path.clone())?)
 			.map_err(|ze| ErrorKind::Other(ze.to_string()))?;
 	}
 
@@ -1072,12 +1081,82 @@ pub fn zip_read(root_dir: String) -> Result<File, Error> {
 
 /// Extract the txhashset data from a zip file and writes the content into the
 /// txhashset storage dir
-pub fn zip_write(root_dir: String, txhashset_data: File) -> Result<(), Error> {
+pub fn zip_write(root_dir: String, txhashset_data: File, header: &BlockHeader) -> Result<(), Error> {
 	let txhashset_path = Path::new(&root_dir).join(TXHASHSET_SUBDIR);
-
 	fs::create_dir_all(txhashset_path.clone())?;
 	zip::decompress(txhashset_data, &txhashset_path)
-		.map_err(|ze| ErrorKind::Other(ze.to_string()).into())
+		.map_err(|ze| ErrorKind::Other(ze.to_string()))?;
+	check_and_remove_files(&txhashset_path, header)
+}
+
+/// Check a txhashset directory and remove any unexpected
+fn check_and_remove_files(txhashset_path: &PathBuf, header: &BlockHeader) -> Result<(), Error> {
+	// First compare the subdirectories
+	let subdirectories_expected: HashSet<_> = [OUTPUT_SUBDIR, KERNEL_SUBDIR, RANGE_PROOF_SUBDIR]
+		.iter()
+		.cloned()
+		.map(|s| String::from(s))
+		.collect();
+
+	let subdirectories_found: HashSet<_> = fs::read_dir(txhashset_path)?
+		.filter_map(|entry| {
+			entry.ok().and_then(|e| {
+				e.path()
+					.file_name()
+					.and_then(|n| n.to_str().map(|s| String::from(s)))
+			})
+		})
+		.collect();
+
+	let dir_difference: Vec<String> = subdirectories_found
+		.difference(&subdirectories_expected)
+		.cloned()
+		.collect();
+
+	// Removing unexpected directories if needed
+	if !dir_difference.is_empty() {
+		debug!(LOGGER, "Unexpected folder(s) found in txhashset folder, removing.");
+		for diff in dir_difference {
+			let diff_path = txhashset_path.join(diff);
+			file::delete(diff_path)?;
+		}
+	}
+
+	// Then compare the files found in the subdirectories
+	let pmmr_files_expected: HashSet<_> = PMMR_FILES
+		.iter()
+		.cloned()
+		.map(|s| if s.contains("pmmr_leaf.bin") {
+		format!("{}.{}", s, header.hash())}
+		else {String::from(s) })
+		.collect();
+
+	let subdirectories = fs::read_dir(txhashset_path)?;
+	for subdirectory in subdirectories {
+		let subdirectory_path = subdirectory?.path();
+		let pmmr_files = fs::read_dir(&subdirectory_path)?;
+		let pmmr_files_found: HashSet<_> = pmmr_files
+			.filter_map(|entry| {
+				entry.ok().and_then(|e| {
+					e.path()
+						.file_name()
+						.and_then(|n| n.to_str().map(|s| String::from(s)))
+				})
+			})
+			.collect();
+		let difference: Vec<String> = pmmr_files_found
+			.difference(&pmmr_files_expected)
+			.cloned()
+			.collect();
+		if !difference.is_empty() {
+			debug!(LOGGER, "Unexpected file(s) found in txhashset subfolder {:?}, removing.", &subdirectory_path);
+			for diff in difference {
+				let diff_path = subdirectory_path.join(diff);
+				file::delete(diff_path)?;
+			}
+		}
+	}
+	Ok(())
 }
 
 /// Given a block header to rewind to and the block header at the