use crate::{Error, Result};
pub struct FileCheck {
path: std::path::PathBuf,
_file: Option<std::fs::File>,
}
impl FileCheck {
pub fn path(&self) -> &std::path::Path {
&self.path
}
}
pub fn file_check(
file_data: &[u8],
file_hash: &str,
file_name_prefix: &str,
file_name_ext: &str,
) -> Result<FileCheck> {
let file_name = format!("{file_name_prefix}-{file_hash}{file_name_ext}");
let mut pref_path =
dirs::data_local_dir().expect("failed to get data_local_dir");
pref_path.push(&file_name);
if let Ok(file) = validate(&pref_path, file_hash) {
return Ok(FileCheck {
path: pref_path,
_file: Some(file),
});
}
let tmp = write(file_data)?;
match tmp.persist_noclobber(&pref_path) {
Ok(mut file) => {
set_perms(&mut file)?;
drop(file);
let file = validate(&pref_path, file_hash)?;
Ok(FileCheck {
path: pref_path,
_file: Some(file),
})
}
Err(err) => {
let tempfile::PersistError { file: tmp, .. } = err;
if let Ok(file) = validate(&pref_path, file_hash) {
let _ = tmp.close();
return Ok(FileCheck {
path: pref_path,
_file: Some(file),
});
}
let path = tmp.path().to_owned();
let tmp = tmp.into_temp_path();
std::mem::forget(tmp);
let file = validate(&path, file_hash)?;
Ok(FileCheck {
path,
_file: Some(file),
})
}
}
}
fn validate(path: &std::path::Path, hash: &str) -> Result<std::fs::File> {
use std::io::Read;
let mut file = std::fs::OpenOptions::new().read(true).open(path)?;
let mut data = Vec::new();
file.read_to_end(&mut data).expect("failed to read lib");
use sha2::Digest;
let mut hasher = sha2::Sha256::new();
hasher.update(data);
let on_disk_hash =
base64::encode_config(hasher.finalize(), base64::URL_SAFE_NO_PAD);
if on_disk_hash != hash {
return Err(Error::err(format!("FileCheckHashMiss({path:?})")));
}
let perms = file
.metadata()
.expect("failed to get lib metadata")
.permissions();
if !perms.readonly() {
return Err(Error::err(format!("FileCheckNotReadonly({path:?})")));
}
tracing::trace!("success correct file_check: {path:?}");
Ok(file)
}
fn write(file_data: &[u8]) -> Result<tempfile::NamedTempFile> {
use std::io::Write;
let mut tmp = tempfile::NamedTempFile::new()?;
tmp.as_file_mut().write_all(file_data)?;
tmp.as_file_mut().flush()?;
set_perms(tmp.as_file_mut())?;
Ok(tmp)
}
fn set_perms(file: &mut std::fs::File) -> Result<()> {
let mut perms = file.metadata()?.permissions();
perms.set_readonly(true);
#[cfg(unix)]
std::os::unix::fs::PermissionsExt::set_mode(&mut perms, 0o500);
file.set_permissions(perms)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
#[tokio::test(flavor = "multi_thread")]
async fn file_check_stress() {
use rand::Rng;
let mut data = vec![0; 1024 * 1024 * 10]; rand::thread_rng().fill(&mut data[..]);
let data = Arc::new(data);
use sha2::Digest;
let mut hasher = sha2::Sha256::new();
hasher.update(&data[..]);
let hash =
base64::encode_config(hasher.finalize(), base64::URL_SAFE_NO_PAD);
let mut task_list = Vec::new();
const COUNT: usize = 3;
let barrier = Arc::new(std::sync::Barrier::new(COUNT));
for _ in 0..3 {
let data = data.clone();
let hash = hash.clone();
let barrier = barrier.clone();
task_list.push(tokio::task::spawn_blocking(move || {
barrier.wait();
file_check(
data.as_slice(),
&hash,
"tx5-core-file-check-test",
".data",
)
}));
}
let mut tmp = Vec::new();
for task in task_list {
tmp.push(task.await.unwrap().unwrap());
}
for tmp in tmp {
let path = tmp.path().to_owned();
drop(tmp);
let _ = std::fs::remove_file(&path);
}
}
}