use base64::{
Engine,
engine::general_purpose::{URL_SAFE_NO_PAD as base64_engine},
};
use flate2::Compression;
use flate2::read::ZlibDecoder;
use flate2::write::ZlibEncoder;
use std::ffi::OsStr;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use tempfile::NamedTempFile;
use crate::error::Error;
pub trait FileSystem {
type HashedFileOut: HashedFileOut;
type HashedFileIn: HashedFileIn;
fn create_hashed_file(&self) -> Result<Self::HashedFileOut, Error>;
fn create_hashed_file_in(
&self,
path: impl AsRef<str>,
) -> Result<Self::HashedFileOut, Error>;
fn open_hashed_file(
&self,
path: impl AsRef<str>,
) -> Result<Self::HashedFileIn, Error>;
fn create_compressed_hashed_file(
&self,
) -> Result<CompressedHashedFileOut<Self::HashedFileOut>, Error> {
let file = self.create_hashed_file()?;
Ok(CompressedHashedFileOut::new(file))
}
fn create_compressed_hashed_file_in(
&self,
path: impl AsRef<str>,
) -> Result<CompressedHashedFileOut<Self::HashedFileOut>, Error> {
let file = self.create_hashed_file_in(path)?;
Ok(CompressedHashedFileOut::new(file))
}
fn open_compressed_hashed_file(
&self,
path: impl AsRef<str>,
) -> Result<CompressedHashedFileIn<Self::HashedFileIn>, Error> {
let file = self.open_hashed_file(path)?;
Ok(CompressedHashedFileIn::new(file))
}
}
pub trait HashedFileOut: Write {
fn persist(self, extension: impl AsRef<str>) -> Result<String, Error>;
}
pub trait HashedFileIn: Read {
fn verify(self) -> Result<(), Error>;
}
pub struct CompressedHashedFileOut<W>
where
W: std::io::Write,
{
encoder: ZlibEncoder<W>,
}
impl<W> CompressedHashedFileOut<W>
where
W: std::io::Write,
{
pub fn new(w: W) -> Self {
Self {
encoder: ZlibEncoder::new(w, Compression::default()),
}
}
}
impl<W> Write for CompressedHashedFileOut<W>
where
W: std::io::Write,
{
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.encoder.write(buf)
}
fn flush(&mut self) -> std::io::Result<()> {
self.encoder.flush()
}
}
impl<W> HashedFileOut for CompressedHashedFileOut<W>
where
W: HashedFileOut
{
fn persist(self, extension: impl AsRef<str>) -> Result<String, Error> {
self.encoder.finish()?.persist(extension)
}
}
pub struct CompressedHashedFileIn<R>
where
R: std::io::Read,
{
decoder: ZlibDecoder<R>,
}
impl<R> CompressedHashedFileIn<R>
where
R: std::io::Read,
{
pub fn new(r: R) -> Self {
Self {
decoder: ZlibDecoder::new(r),
}
}
}
impl<R> Read for CompressedHashedFileIn<R>
where
R: std::io::Read,
{
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.decoder.read(buf)
}
}
impl<R> HashedFileIn for CompressedHashedFileIn<R>
where
R: HashedFileIn,
{
fn verify(self) -> Result<(), Error> {
self.decoder.into_inner().verify()
}
}
pub struct LocalFileSystem {
base_path: PathBuf,
}
impl LocalFileSystem {
pub fn new(base_path: impl AsRef<Path>) -> Self {
Self {
base_path: base_path.as_ref().to_path_buf(),
}
}
}
impl FileSystem for LocalFileSystem {
type HashedFileOut = LocalHashedFileOut;
type HashedFileIn = LocalHashedFileIn;
fn create_hashed_file(&self) -> Result<Self::HashedFileOut, Error> {
LocalHashedFileOut::create(self.base_path.clone())
}
fn create_hashed_file_in(
&self,
path: impl AsRef<str>,
) -> Result<Self::HashedFileOut, Error> {
LocalHashedFileOut::create(self.base_path.join(path.as_ref()))
}
fn open_hashed_file(
&self,
path: impl AsRef<str>,
) -> Result<Self::HashedFileIn, Error> {
LocalHashedFileIn::open(self.base_path.join(path.as_ref()))
}
}
pub struct LocalHashedFileOut {
tempfile: NamedTempFile,
base_path: PathBuf,
context: ring::digest::Context,
}
impl LocalHashedFileOut {
fn create(base_path: PathBuf) -> Result<Self, Error> {
let tempfile = NamedTempFile::new()?;
Ok(LocalHashedFileOut {
tempfile,
base_path,
context: ring::digest::Context::new(&ring::digest::SHA256),
})
}
}
impl Write for LocalHashedFileOut {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.context.update(buf);
self.tempfile.write(buf)
}
fn flush(&mut self) -> std::io::Result<()> {
self.tempfile.flush()
}
}
impl HashedFileOut for LocalHashedFileOut {
fn persist(mut self, extension: impl AsRef<str>) -> Result<String, Error> {
self.flush()?;
if !self.base_path.exists() {
std::fs::create_dir_all(&self.base_path)?;
}
let hash = self.context.finish();
let hash = base64_engine.encode(&hash);
let path = self.base_path
.join(&hash)
.with_extension(extension.as_ref());
self.tempfile.persist(path)?;
Ok(hash)
}
}
pub struct LocalHashedFileIn {
file: std::fs::File,
path: PathBuf,
context: ring::digest::Context,
}
impl LocalHashedFileIn {
fn open(path: PathBuf) -> Result<Self, Error> {
let file = std::fs::File::open(&path)?;
Ok(LocalHashedFileIn {
file,
path,
context: ring::digest::Context::new(&ring::digest::SHA256),
})
}
}
impl Read for LocalHashedFileIn {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let n = self.file.read(buf)?;
self.context.update(&buf[..n]);
Ok(n)
}
}
impl HashedFileIn for LocalHashedFileIn {
fn verify(self) -> Result<(), Error> {
let hash = self.context.finish();
let hash = base64_engine.encode(&hash);
if hash.as_str() == self.path.file_stem().unwrap_or(OsStr::new("")) {
Ok(())
} else {
Err(Error::VerificationFailure(format!(
"Expected hash {:?}, but got {}",
self.path.file_stem(),
hash,
)))
}
}
}