use avml::{Error, ONE_MB, Result, image};
use clap::{Parser, ValueEnum};
use snap::read::FrameDecoder;
use std::{
fs::{File, metadata},
io::{Read, Seek, SeekFrom, Write, copy},
path::{Path, PathBuf},
};
fn convert(src: &Path, dst: &Path, compress: bool) -> Result<()> {
let src_len = metadata(src)
.map_err(|source| image::Error::Io {
context: "unable to read source size",
source,
})?
.len();
let mut image = image::Image::<File, File>::new(1, src, dst)?;
convert_image(&mut image, src_len, compress)
}
fn convert_image<R, W>(image: &mut image::Image<R, W>, src_len: u64, compress: bool) -> Result<()>
where
R: Read + Seek,
W: Write,
{
image.version = if compress { 2 } else { 1 };
loop {
let current = image
.src
.stream_position()
.map_err(|source| image::Error::Io {
context: "unable to get current offset into the memory source",
source,
})?;
if current >= src_len {
break;
}
image.convert_block()?;
}
Ok(())
}
fn convert_to_raw_image<R, W>(image: &mut image::Image<R, W>, src_len: u64) -> Result<()>
where
R: Read + Seek,
W: Write + Seek,
{
image.version = 1;
loop {
let current = image
.src
.stream_position()
.map_err(|source| image::Error::Io {
context: "unable to get the current offset into the memory source",
source,
})?;
if current >= src_len {
break;
}
let current_dst = image
.dst
.stream_position()
.map_err(|source| image::Error::Io {
context: "unable to get the current offset into the destination stream",
source,
})?;
let header = image.read_header()?;
let mut zeros = vec![0; ONE_MB];
let mut unmapped = usize::try_from(header.range.start.saturating_sub(current_dst))
.map_err(image::Error::IntConversion)?;
while unmapped > ONE_MB {
image
.dst
.write_all(&zeros)
.map_err(|source| image::Error::Io {
context: "unable to write padding bytes",
source,
})?;
unmapped = unmapped.saturating_sub(ONE_MB);
}
if unmapped > 0 {
zeros.resize(unmapped, 0);
image
.dst
.write_all(&zeros)
.map_err(|source| image::Error::Io {
context: "unable to write padding bytes",
source,
})?;
}
let size = header.size()?;
match header.version {
1 => {
let mut handle =
(&mut image.src).take(size.try_into().map_err(image::Error::IntConversion)?);
copy(&mut handle, &mut image.dst).map_err(|source| image::Error::Io {
context: "unable to copy image data",
source,
})?;
}
2 => {
let mut decoder = FrameDecoder::new(&mut image.src)
.take(size.try_into().map_err(image::Error::IntConversion)?);
copy(&mut decoder, &mut image.dst).map_err(|source| image::Error::Io {
context: "unable to copy image data",
source,
})?;
image
.src
.seek(SeekFrom::Current(8))
.map_err(|source| image::Error::Io {
context: "unable to seek past the compressed size",
source,
})?;
}
_ => return Err(image::Error::UnimplementedVersion.into()),
}
}
Ok(())
}
fn convert_to_raw(src: &Path, dst: &Path) -> Result<()> {
let src_len = metadata(src)
.map_err(|source| image::Error::Io {
context: "unable to get source file size",
source,
})?
.len();
let mut image = image::Image::<File, File>::new(1, src, dst)?;
convert_to_raw_image(&mut image, src_len)
}
fn encode_raw_image<R, W>(
image: &mut image::Image<R, W>,
raw_len: u64,
compress: bool,
) -> Result<()>
where
R: Read + Seek,
W: Write,
{
image.version = if compress { 2 } else { 1 };
let mut start = 0_u64;
while start < raw_len {
let end = raw_len.min(start.saturating_add(image::MAX_BLOCK_SIZE));
image.copy_block(start..end)?;
start = end;
}
Ok(())
}
fn convert_from_raw(src: &Path, dst: &Path, compress: bool) -> Result<()> {
let src_len = metadata(src)
.map_err(|source| image::Error::Io {
context: "unable to read source size",
source,
})?
.len();
let mut image = image::Image::<File, File>::new(1, src, dst)?;
encode_raw_image(&mut image, src_len, compress)
}
#[derive(Parser)]
#[command(version)]
struct Config {
#[arg(long, value_enum, default_value_t = Format::LimeCompressed)]
source_format: Format,
#[arg(long, value_enum, default_value_t = Format::Lime)]
format: Format,
src: PathBuf,
dst: PathBuf,
}
#[derive(ValueEnum, Clone)]
enum Format {
Raw,
Lime,
#[value(rename_all = "snake_case")]
LimeCompressed,
}
fn main() -> Result<()> {
let config = Config::parse();
match (config.source_format, config.format) {
(Format::Lime | Format::LimeCompressed, Format::Raw) => {
convert_to_raw(&config.src, &config.dst)
}
(Format::Lime, Format::LimeCompressed) => convert(&config.src, &config.dst, true),
(Format::LimeCompressed, Format::Lime) => convert(&config.src, &config.dst, false),
(Format::Raw, Format::Lime) => convert_from_raw(&config.src, &config.dst, false),
(Format::Raw, Format::LimeCompressed) => convert_from_raw(&config.src, &config.dst, true),
(Format::Lime, Format::Lime)
| (Format::LimeCompressed, Format::LimeCompressed)
| (Format::Raw, Format::Raw) => Err(Error::NoConversionRequired),
}
}
#[cfg(test)]
mod tests {
use crate::{convert_image, convert_to_raw_image, encode_raw_image};
use avml::{Result, image};
use rand::{Rng as _, SeedableRng as _, rngs::SmallRng};
use std::io::Cursor;
fn memory_image(src: &[u8]) -> image::Image<Cursor<&[u8]>, Cursor<Vec<u8>>> {
image::Image {
version: 1,
align_src: false,
src: Cursor::new(src),
dst: Cursor::new(Vec::new()),
}
}
fn block_size() -> Result<usize> {
usize::try_from(image::MAX_BLOCK_SIZE)
.map_err(image::Error::IntConversion)
.map_err(Into::into)
}
fn random_bytes(rng: &mut SmallRng, len: usize) -> Vec<u8> {
let mut bytes = vec![0; len];
rng.fill_bytes(&mut bytes);
bytes
}
fn build_sparse_raw() -> Result<Vec<u8>> {
let mut rng = SmallRng::seed_from_u64(0);
let block_size = block_size()?;
let partial_head_len = block_size / 4;
let partial_tail_len = block_size.saturating_sub(partial_head_len);
let chunks = vec![
vec![0; block_size],
random_bytes(&mut rng, block_size),
vec![0; block_size],
random_bytes(&mut rng, partial_head_len),
vec![0; block_size],
random_bytes(&mut rng, partial_tail_len),
random_bytes(&mut rng, block_size),
vec![0; block_size],
random_bytes(&mut rng, block_size),
];
Ok(chunks.concat())
}
fn encode_raw(raw: &[u8], version: u32) -> Result<Vec<u8>> {
let mut image = memory_image(raw);
let total = u64::try_from(raw.len()).map_err(image::Error::IntConversion)?;
encode_raw_image(&mut image, total, version == 2)?;
Ok(image.dst.into_inner())
}
fn convert_encoded(encoded: &[u8], compress: bool) -> Result<Vec<u8>> {
let encoded_len = u64::try_from(encoded.len()).map_err(image::Error::IntConversion)?;
let mut image = memory_image(encoded);
convert_image(&mut image, encoded_len, compress)?;
Ok(image.dst.into_inner())
}
fn decode_to_raw(encoded: &[u8]) -> Result<Vec<u8>> {
let encoded_len = u64::try_from(encoded.len()).map_err(image::Error::IntConversion)?;
let mut image = memory_image(encoded);
convert_to_raw_image(&mut image, encoded_len)?;
Ok(image.dst.into_inner())
}
fn header_version(encoded: &[u8]) -> Result<u32> {
Ok(image::Header::read(Cursor::new(encoded))?.version)
}
#[test]
fn convert_sparse_raw_between_lime_and_compressed_formats() -> Result<()> {
let raw = build_sparse_raw()?;
let lime = encode_raw(&raw, 1)?;
assert_eq!(header_version(&lime)?, 1);
assert_eq!(decode_to_raw(&lime)?, raw);
let compressed = convert_encoded(&lime, true)?;
assert_eq!(header_version(&compressed)?, 2);
let lime_roundtrip = convert_encoded(&compressed, false)?;
assert_eq!(header_version(&lime_roundtrip)?, 1);
assert_eq!(lime_roundtrip, lime);
assert_eq!(decode_to_raw(&compressed)?, raw);
assert_eq!(decode_to_raw(&lime_roundtrip)?, raw);
Ok(())
}
#[test]
fn trailing_zero_block_is_dropped_from_raw_roundtrip() -> Result<()> {
let mut raw = build_sparse_raw()?;
let expected_raw = raw.clone();
let block_size = block_size()?;
raw.extend(vec![0; block_size]);
let lime = encode_raw(&raw, 1)?;
assert_eq!(decode_to_raw(&lime)?, expected_raw);
let compressed = convert_encoded(&lime, true)?;
let lime_roundtrip = convert_encoded(&compressed, false)?;
assert_eq!(lime_roundtrip, lime);
assert_eq!(decode_to_raw(&compressed)?, expected_raw);
assert_eq!(decode_to_raw(&lime_roundtrip)?, expected_raw);
Ok(())
}
}