use std::io::{self, Read};
use std::path::Path;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Compression {
None,
Gzip,
Zstd,
}
pub fn detect_compression(bytes: &[u8]) -> Compression {
if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
return Compression::Gzip;
}
if bytes.len() >= 4
&& bytes[0] == 0x28
&& bytes[1] == 0xb5
&& bytes[2] == 0x2f
&& bytes[3] == 0xfd
{
return Compression::Zstd;
}
Compression::None
}
pub fn detect_compression_from_extension(path: &Path) -> Compression {
match path.extension().and_then(|e| e.to_str()) {
Some("gz") => Compression::Gzip,
Some("zst") => Compression::Zstd,
_ => Compression::None,
}
}
const MMAP_THRESHOLD: u64 = 64 * 1024;
pub fn read_file_contents(path: &Path) -> Result<FileContents, Box<dyn std::error::Error>> {
let file = std::fs::File::open(path)?;
let metadata = file.metadata()?;
let mut magic = [0u8; 4];
let bytes_read = {
let mut f = &file;
f.read(&mut magic)?
};
let compression = if bytes_read > 0 {
detect_compression(&magic[..bytes_read])
} else {
Compression::None
};
match compression {
Compression::Gzip => {
let file = std::fs::File::open(path)?;
let mut decoder = flate2::read::GzDecoder::new(file);
let mut contents = String::new();
decoder.read_to_string(&mut contents)?;
Ok(FileContents::Owned(contents))
}
Compression::Zstd => {
#[cfg(feature = "zstd")]
{
let file = std::fs::File::open(path)?;
let mut decoder = zstd::stream::read::Decoder::new(file)?;
let mut contents = String::new();
decoder.read_to_string(&mut contents)?;
Ok(FileContents::Owned(contents))
}
#[cfg(not(feature = "zstd"))]
{
Err(io::Error::new(
io::ErrorKind::Unsupported,
"zstd-compressed input detected; rebuild readcon-core with --features zstd",
)
.into())
}
}
Compression::None => {
if metadata.len() < MMAP_THRESHOLD {
let contents = std::fs::read_to_string(path)?;
Ok(FileContents::Owned(contents))
} else {
let file = std::fs::File::open(path)?;
let mmap = unsafe { memmap2::Mmap::map(&file)? };
Ok(FileContents::Mapped(mmap))
}
}
}
}
pub enum FileContents {
Owned(String),
Mapped(memmap2::Mmap),
}
impl FileContents {
pub fn as_str(&self) -> Result<&str, std::str::Utf8Error> {
match self {
FileContents::Owned(s) => Ok(s.as_str()),
FileContents::Mapped(m) => std::str::from_utf8(m),
}
}
}
pub fn gzip_writer(path: &Path) -> io::Result<flate2::write::GzEncoder<std::fs::File>> {
let file = std::fs::File::create(path)?;
Ok(flate2::write::GzEncoder::new(
file,
flate2::Compression::default(),
))
}
#[cfg(feature = "zstd")]
pub fn zstd_writer(
path: &Path,
) -> io::Result<zstd::stream::write::AutoFinishEncoder<'static, std::fs::File>> {
let file = std::fs::File::create(path)?;
let encoder = zstd::stream::write::Encoder::new(file, 3)?;
Ok(encoder.auto_finish())
}