use super::error::{Error, Result};
use super::utils::collect_file_paths;
#[cfg(feature = "progress_bar")]
use super::utils::progress_bar;
use chrono::{offset::Utc, DateTime};
#[cfg(feature = "progress_bar")]
use indicatif::ProgressIterator;
use lazy_static::lazy_static;
use regex::Regex;
use sha2::{Digest, Sha256};
use std::sync::mpsc;
use std::thread;
use std::{
fs::{self, File},
io,
path::Path,
path::PathBuf,
};
use tree_magic_mini as magic;
#[derive(Debug, Clone)]
pub struct BlobRef {
value: String,
}
#[derive(Debug)]
pub struct BlobMetadata {
pub filename: String,
pub mime_type: String,
pub size: u64,
pub created: DateTime<Utc>,
}
impl From<Sha256> for BlobRef {
fn from(hasher: Sha256) -> Self {
BlobRef::new(&format!("{:x}", hasher.finalize())[..]).unwrap()
}
}
impl BlobRef {
pub fn new(value: &str) -> Result<BlobRef> {
lazy_static! {
static ref VALID_HASH_REGEX: Regex = Regex::new(r"^[a-z0-9]{64}$").unwrap();
}
if VALID_HASH_REGEX.is_match(value) {
Ok(BlobRef {
value: String::from(value),
})
} else {
Err(Error::InvalidRef)
}
}
pub fn to_path(&self) -> PathBuf {
PathBuf::from(&self.value[0..2])
.join(&self.value[2..4])
.join(&self.value[4..6])
.join(&self.value[6..])
}
pub fn reference(&self) -> &str {
&self.value
}
}
impl std::fmt::Display for BlobRef {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "BlobRef({})", &self.value[..10])
}
}
#[derive(Clone, Debug)]
pub struct BlobStore {
root: PathBuf,
}
type BlobRefAndPath = (PathBuf, BlobRef);
impl BlobStore {
pub fn new<P: AsRef<Path>>(path: P) -> Result<BlobStore> {
let path = path.as_ref();
if !path.exists() {
fs::create_dir_all(path)?
} else if !path.is_dir() {
return Err(io::Error::from(io::ErrorKind::Other).into());
}
Ok(BlobStore { root: path.into() })
}
pub fn hasher() -> Sha256 {
Sha256::new()
}
fn get_blob_path(&self, blob_ref: &BlobRef) -> PathBuf {
self.root.join(blob_ref.to_path())
}
fn get_blob_file_path(&self, blob_ref: &BlobRef) -> Result<PathBuf> {
let mut entries = self.get_blob_path(blob_ref).read_dir()?;
if let Some(Ok(entry)) = entries.next() {
return Ok(entry.path());
};
Err(Error::BlobNotFound)
}
pub fn add<P: AsRef<Path>>(&self, path: P) -> Result<BlobRef> {
let mut file = File::open(&path)?;
let mut hasher = BlobStore::hasher();
io::copy(&mut file, &mut hasher)?;
let blob_ref = BlobRef::from(hasher);
if !self.exists(&blob_ref) {
let save_path = self.get_blob_path(&blob_ref);
fs::create_dir_all(&save_path)?;
let filename = path.as_ref().file_name().unwrap();
let save_path = save_path.join(&filename);
fs::copy(path, save_path)?;
};
Ok(blob_ref)
}
pub fn add_files<P: AsRef<Path>>(
&self,
paths: &[P],
threads: u8,
) -> (Vec<BlobRefAndPath>, Vec<(PathBuf, Error)>) {
let paths: Vec<PathBuf> = paths
.iter()
.flat_map(|p| collect_file_paths(p.as_ref()))
.collect();
let (tx, rx) = mpsc::channel();
let chunk_size = std::cmp::max(paths.len() / threads as usize, 1_usize);
let chunks = paths.chunks(chunk_size);
for chunk in chunks {
let tx = tx.clone();
let chunk = chunk.to_owned();
let blob_store = self.clone();
thread::spawn(move || {
for path in chunk {
let blob_ref = blob_store.add(&path);
tx.send((path, blob_ref)).expect("err")
}
});
}
drop(tx);
let rx_iter = rx.iter();
#[cfg(feature = "progress_bar")]
let rx_iter = rx_iter.progress_with(progress_bar(paths.len() as u64));
let (success, errors): (Vec<_>, Vec<_>) = rx_iter.partition(|(_, b)| b.is_ok());
let success = success.into_iter().map(|(p, b)| (p, b.unwrap())).collect();
let errors = errors
.into_iter()
.map(|(p, b)| (p, b.unwrap_err()))
.collect();
(success, errors)
}
pub fn get(&self, blob_ref: &BlobRef) -> Result<Vec<u8>> {
Ok(fs::read(&self.get_blob_file_path(blob_ref)?)?)
}
pub fn exists(&self, blob_ref: &BlobRef) -> bool {
let dir = self.get_blob_path(blob_ref);
dir.exists() && dir.read_dir().unwrap().next().is_some()
}
pub fn delete(&self, blob_ref: &BlobRef) -> Result<()> {
Ok(fs::remove_dir_all(self.get_blob_path(blob_ref))?)
}
pub fn metadata(&self, blob_ref: &BlobRef) -> Result<BlobMetadata> {
let file_path = self.get_blob_file_path(blob_ref)?;
let mime = magic::from_filepath(&file_path).unwrap_or("application/octet-stream");
let filename = file_path.file_name().unwrap().to_str().unwrap().to_string();
let metadata = fs::metadata(file_path)?;
Ok(BlobMetadata {
mime_type: String::from(mime),
filename,
size: metadata.len(),
created: metadata.created()?.into(),
})
}
}
impl BlobMetadata {
pub fn created_str(&self) -> String {
self.created
.to_rfc3339_opts(chrono::SecondsFormat::Secs, false)
}
}