Refactor/hardened unpack more generic unpack archive (#26677)

* Add entry_processor fn to unpack_archive - not used yet

* Add generic unpack_snapshot_with_processors to separate account paths handling from entry checking
This commit is contained in:
apfitzge 2022-07-19 14:48:57 -05:00 committed by GitHub
parent 7efe72a74f
commit 84c8cfe901
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 48 additions and 11 deletions

View File

@ -90,15 +90,18 @@ pub enum UnpackPath<'a> {
Invalid,
}
fn unpack_archive<'a, A: Read, C>(
fn unpack_archive<'a, A, C, D>(
archive: &mut Archive<A>,
apparent_limit_size: u64,
actual_limit_size: u64,
limit_count: u64,
mut entry_checker: C,
mut entry_checker: C, // checks if entry is valid
entry_processor: D, // processes entry after setting permissions
) -> Result<()>
where
A: Read,
C: FnMut(&[&str], tar::EntryType) -> UnpackPath<'a>,
D: Fn(PathBuf),
{
let mut apparent_total_size: u64 = 0;
let mut actual_total_size: u64 = 0;
@ -175,7 +178,11 @@ where
GNUSparse | Regular => 0o644,
_ => 0o755,
};
set_perms(&unpack_dir.join(entry.path()?), mode)?;
let entry_path_buf = unpack_dir.join(entry.path()?);
set_perms(&entry_path_buf, mode)?;
// Process entry after setting permissions
entry_processor(entry_path_buf);
total_entries += 1;
let now = Instant::now();
@ -299,8 +306,35 @@ pub fn unpack_snapshot<A: Read>(
account_paths: &[PathBuf],
parallel_selector: Option<ParallelSelector>,
) -> Result<UnpackedAppendVecMap> {
assert!(!account_paths.is_empty());
let mut unpacked_append_vec_map = UnpackedAppendVecMap::new();
unpack_snapshot_with_processors(
archive,
ledger_dir,
account_paths,
parallel_selector,
|file, path| {
unpacked_append_vec_map.insert(file.to_string(), path.join("accounts").join(file));
},
|_| {},
)
.map(|_| unpacked_append_vec_map)
}
fn unpack_snapshot_with_processors<A, F, G>(
archive: &mut Archive<A>,
ledger_dir: &Path,
account_paths: &[PathBuf],
parallel_selector: Option<ParallelSelector>,
mut accounts_path_processor: F,
entry_processor: G,
) -> Result<()>
where
A: Read,
F: FnMut(&str, &Path),
G: Fn(PathBuf),
{
assert!(!account_paths.is_empty());
let mut i = 0;
unpack_archive(
@ -322,12 +356,14 @@ pub fn unpack_snapshot<A: Read>(
if let ["accounts", file] = parts {
// Randomly distribute the accounts files about the available `account_paths`,
let path_index = thread_rng().gen_range(0, account_paths.len());
match account_paths.get(path_index).map(|path_buf| {
unpacked_append_vec_map
.insert(file.to_string(), path_buf.join("accounts").join(file));
path_buf.as_path()
}) {
Some(path) => UnpackPath::Valid(path),
match account_paths
.get(path_index)
.map(|path_buf| path_buf.as_path())
{
Some(path) => {
accounts_path_processor(*file, path);
UnpackPath::Valid(path)
}
None => UnpackPath::Invalid,
}
} else {
@ -337,8 +373,8 @@ pub fn unpack_snapshot<A: Read>(
UnpackPath::Invalid
}
},
entry_processor,
)
.map(|_| unpacked_append_vec_map)
}
fn all_digits(v: &str) -> bool {
@ -450,6 +486,7 @@ fn unpack_genesis<A: Read>(
max_genesis_archive_unpacked_size,
MAX_GENESIS_ARCHIVE_UNPACKED_COUNT,
|p, k| is_valid_genesis_archive_entry(unpack_dir, p, k),
|_| {},
)
}