use std::fs::{self, File};
use std::io::{BufWriter, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use crate::format::{AssetEntry, AssetInventory, LayerEntry};
use crate::{PackError, Result};
fn digest_to_filename(digest: &str) -> Result<String> {
let short = digest.strip_prefix("sha256:").unwrap_or(digest);
if short.len() < 12 {
return Err(PackError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"digest too short for filename: '{}' ({} chars, need 12)",
digest,
short.len()
),
)));
}
Ok(format!("{}.tar", &short[..12]))
}
pub const ZSTD_LEVEL: i32 = 3;
fn find_existing_template(filename: &str) -> Option<PathBuf> {
if let Some(home) = dirs::home_dir() {
let path = home.join(".smolvm").join(filename);
if path.exists() {
return Some(path);
}
}
if let Ok(exe) = std::env::current_exe() {
if let Some(dir) = exe.parent() {
let path = dir.join(filename);
if path.exists() {
return Some(path);
}
}
}
None
}
pub struct AssetCollector {
staging_dir: PathBuf,
inventory: AssetInventory,
}
impl AssetCollector {
pub fn new(staging_dir: PathBuf) -> Result<Self> {
fs::create_dir_all(&staging_dir)?;
fs::create_dir_all(staging_dir.join("layers"))?;
Ok(Self {
staging_dir,
inventory: AssetInventory {
libraries: Vec::new(),
agent_rootfs: AssetEntry {
path: "agent-rootfs.tar".to_string(),
size: 0,
},
layers: Vec::new(),
storage_template: None,
overlay_template: None,
overlay_logical_size: None,
},
})
}
pub fn staging_dir(&self) -> &Path {
&self.staging_dir
}
pub fn collect_libraries(&mut self, lib_dir: &Path) -> Result<()> {
fs::create_dir_all(self.staging_dir.join("lib"))?;
let lib_names = if cfg!(target_os = "macos") {
vec!["libkrun.dylib", "libkrunfw.5.dylib"]
} else {
vec!["libkrun.so", "libkrunfw.so.5"]
};
for name in lib_names {
let src = lib_dir.join(name);
if !src.exists() {
return Err(PackError::AssetNotFound(format!(
"library not found: {}",
src.display()
)));
}
let dst = self.staging_dir.join("lib").join(name);
fs::copy(&src, &dst)?;
let metadata = fs::metadata(&dst)?;
self.inventory.libraries.push(AssetEntry {
path: format!("lib/{}", name),
size: metadata.len(),
});
}
#[cfg(target_os = "macos")]
{
let gpu_libs = [
"libvirglrenderer.1.dylib",
"libMoltenVK.dylib",
"libepoxy.0.dylib",
];
for name in &gpu_libs {
let src = lib_dir.join(name);
if src.exists() {
let dst = self.staging_dir.join("lib").join(name);
fs::copy(&src, &dst)?;
let metadata = fs::metadata(&dst)?;
self.inventory.libraries.push(AssetEntry {
path: format!("lib/{}", name),
size: metadata.len(),
});
}
}
}
#[cfg(target_os = "linux")]
{
let gpu_libs = ["libvirglrenderer.so.1", "libepoxy.so.0"];
for name in &gpu_libs {
let src = lib_dir.join(name);
if src.exists() {
let dst = self.staging_dir.join("lib").join(name);
fs::copy(&src, &dst)?;
let metadata = fs::metadata(&dst)?;
self.inventory.libraries.push(AssetEntry {
path: format!("lib/{}", name),
size: metadata.len(),
});
}
}
let server_src = lib_dir.join("virgl_render_server");
if server_src.exists() {
let server_dst = self.staging_dir.join("lib").join("virgl_render_server");
fs::copy(&server_src, &server_dst)?;
use std::os::unix::fs::PermissionsExt;
fs::set_permissions(&server_dst, fs::Permissions::from_mode(0o755))?;
let metadata = fs::metadata(&server_dst)?;
self.inventory.libraries.push(AssetEntry {
path: "lib/virgl_render_server".to_string(),
size: metadata.len(),
});
}
}
Ok(())
}
pub fn collect_agent_rootfs(&mut self, rootfs_dir: &Path) -> Result<()> {
if !rootfs_dir.exists() {
return Err(PackError::AssetNotFound(format!(
"agent rootfs not found: {}",
rootfs_dir.display()
)));
}
let tar_path = self.staging_dir.join("agent-rootfs.tar");
let tar_file = File::create(&tar_path)?;
let mut tar_builder = tar::Builder::new(BufWriter::new(tar_file));
tar_builder.follow_symlinks(false);
tar_builder
.append_dir_all(".", rootfs_dir)
.map_err(|e| PackError::Tar(e.to_string()))?;
tar_builder
.finish()
.map_err(|e| PackError::Tar(e.to_string()))?;
let metadata = fs::metadata(&tar_path)?;
self.inventory.agent_rootfs = AssetEntry {
path: "agent-rootfs.tar".to_string(),
size: metadata.len(),
};
Ok(())
}
pub fn add_layer(&mut self, digest: &str, layer_data: &[u8]) -> Result<()> {
let filename = digest_to_filename(digest)?;
let path = format!("layers/{}", filename);
let dst = self.staging_dir.join(&path);
fs::write(&dst, layer_data)?;
self.inventory.layers.push(LayerEntry {
digest: digest.to_string(),
path,
size: layer_data.len() as u64,
});
Ok(())
}
pub fn layer_staging_path(&self, digest: &str) -> PathBuf {
let filename = digest_to_filename(digest)
.expect("layer digest must be sha256:<hex> with at least 12 hex chars");
self.staging_dir.join(format!("layers/{}", filename))
}
pub fn register_layer(&mut self, digest: &str) -> Result<()> {
let filename = digest_to_filename(digest)?;
let path = format!("layers/{}", filename);
let dst = self.staging_dir.join(&path);
let metadata = fs::metadata(&dst)?;
self.inventory.layers.push(LayerEntry {
digest: digest.to_string(),
path,
size: metadata.len(),
});
Ok(())
}
pub fn add_layer_from_file(&mut self, digest: &str, layer_path: &Path) -> Result<()> {
let filename = digest_to_filename(digest)?;
let path = format!("layers/{}", filename);
let dst = self.staging_dir.join(&path);
fs::copy(layer_path, &dst)?;
let metadata = fs::metadata(&dst)?;
self.inventory.layers.push(LayerEntry {
digest: digest.to_string(),
path,
size: metadata.len(),
});
Ok(())
}
pub fn create_storage_template(&mut self) -> Result<()> {
use std::io::{Seek, SeekFrom, Write};
use std::process::Command;
const TEMPLATE_SIZE: u64 = 512 * 1024 * 1024; const TEMPLATE_NAME: &str = "storage.ext4";
let template_path = self.staging_dir.join(TEMPLATE_NAME);
if let Some(existing) = find_existing_template("storage-template.ext4") {
fs::copy(&existing, &template_path)?;
let metadata = fs::metadata(&template_path)?;
self.inventory.storage_template = Some(AssetEntry {
path: TEMPLATE_NAME.to_string(),
size: metadata.len(),
});
return Ok(());
}
let mut file = File::create(&template_path)?;
file.seek(SeekFrom::Start(TEMPLATE_SIZE - 1))?;
file.write_all(&[0])?;
file.sync_all()?;
drop(file);
let mkfs_paths = [
"/opt/homebrew/opt/e2fsprogs/sbin/mkfs.ext4",
"/usr/local/opt/e2fsprogs/sbin/mkfs.ext4",
"/opt/homebrew/sbin/mkfs.ext4",
"/usr/local/sbin/mkfs.ext4",
"/sbin/mkfs.ext4",
"/usr/sbin/mkfs.ext4",
"mkfs.ext4",
];
let mkfs_path = mkfs_paths
.iter()
.find(|p| {
if p.contains('/') {
std::path::Path::new(p).exists()
} else {
Command::new(p).arg("--version").output().is_ok()
}
})
.ok_or_else(|| {
PackError::AssetNotFound(
"mkfs.ext4 not found. Install e2fsprogs or place a pre-formatted \
storage-template.ext4 in ~/.smolvm/"
.into(),
)
})?;
#[cfg(unix)]
unsafe {
libc::signal(libc::SIGCHLD, libc::SIG_DFL);
}
let mut child = Command::new(mkfs_path)
.args([
"-F", "-q", "-m", "0", "-L", "smolvm", ])
.arg(&template_path)
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
.map_err(|e| PackError::AssetNotFound(format!("failed to spawn mkfs.ext4: {}", e)))?;
let status = child.wait().map_err(|e| {
PackError::AssetNotFound(format!("failed to wait for mkfs.ext4: {}", e))
})?;
if !status.success() {
return Err(PackError::AssetNotFound(
"mkfs.ext4 failed to format storage template".into(),
));
}
let metadata = fs::metadata(&template_path)?;
self.inventory.storage_template = Some(AssetEntry {
path: TEMPLATE_NAME.to_string(),
size: metadata.len(),
});
Ok(())
}
pub fn add_overlay_template(&mut self, path: &Path) -> Result<()> {
if !path.exists() {
return Err(PackError::AssetNotFound(format!(
"overlay disk not found: {}",
path.display()
)));
}
const OVERLAY_NAME: &str = "overlay.raw";
let dst = self.staging_dir.join(OVERLAY_NAME);
let (logical_size, truncated_size) = sparse_copy_overlay(path, &dst)?;
self.inventory.overlay_template = Some(AssetEntry {
path: OVERLAY_NAME.to_string(),
size: truncated_size,
});
if logical_size > truncated_size {
self.inventory.overlay_logical_size = Some(logical_size);
}
Ok(())
}
pub fn inventory(&self) -> &AssetInventory {
&self.inventory
}
pub fn into_inventory(self) -> AssetInventory {
self.inventory
}
pub fn compress(&self, output: &Path, exclude_libs: bool) -> Result<u64> {
let output_file = File::create(output)?;
let encoder = zstd::stream::Encoder::new(output_file, ZSTD_LEVEL)
.map_err(|e| PackError::Compression(e.to_string()))?;
let mut tar_builder = tar::Builder::new(encoder);
let mut entries: Vec<_> = fs::read_dir(&self.staging_dir)?
.filter_map(|e| e.ok())
.collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let name = entry.file_name();
if exclude_libs && name == "lib" {
continue; }
let path = entry.path();
if path.is_dir() {
tar_builder
.append_dir_all(name.to_string_lossy().as_ref(), &path)
.map_err(|e| PackError::Tar(e.to_string()))?;
} else {
tar_builder
.append_path_with_name(&path, name.to_string_lossy().as_ref())
.map_err(|e| PackError::Tar(e.to_string()))?;
}
}
let encoder = tar_builder
.into_inner()
.map_err(|e| PackError::Tar(e.to_string()))?;
encoder
.finish()
.map_err(|e| PackError::Compression(e.to_string()))?;
let metadata = fs::metadata(output)?;
Ok(metadata.len())
}
}
fn sparse_copy_overlay(src: &Path, dst: &Path) -> std::io::Result<(u64, u64)> {
let mut src_file = File::open(src)?;
let logical_size = src_file.metadata()?.len();
let truncated_size = find_last_data_byte(&mut src_file, logical_size)?;
let mut dst_file = File::create(dst)?;
dst_file.set_len(truncated_size)?;
if truncated_size == 0 {
return Ok((logical_size, 0));
}
src_file.seek(SeekFrom::Start(0))?;
let mut buf = vec![0u8; 512 * 1024];
let mut offset: u64 = 0;
while offset < truncated_size {
let to_read = (truncated_size - offset).min(buf.len() as u64) as usize;
let n = src_file.read(&mut buf[..to_read])?;
if n == 0 {
break;
}
let chunk = &buf[..n];
if chunk.iter().any(|&b| b != 0) {
dst_file.seek(SeekFrom::Start(offset))?;
dst_file.write_all(chunk)?;
}
offset += n as u64;
}
Ok((logical_size, truncated_size))
}
fn find_last_data_byte(file: &mut File, logical_size: u64) -> std::io::Result<u64> {
if logical_size == 0 {
return Ok(0);
}
const CHUNK: u64 = 1024 * 1024; let mut buf = vec![0u8; CHUNK as usize];
let mut pos = logical_size;
while pos > 0 {
let chunk_start = pos.saturating_sub(CHUNK);
let chunk_size = (pos - chunk_start) as usize;
file.seek(SeekFrom::Start(chunk_start))?;
let n = file.read(&mut buf[..chunk_size])?;
if n == 0 {
break;
}
for i in (0..n).rev() {
if buf[i] != 0 {
return Ok(chunk_start + i as u64 + 1);
}
}
pos = chunk_start;
}
Ok(0) }
pub fn decompress_assets(compressed: &[u8], output_dir: &Path) -> Result<()> {
fs::create_dir_all(output_dir)?;
let decoder = zstd::stream::Decoder::new(compressed)
.map_err(|e| PackError::Compression(e.to_string()))?;
let mut archive = tar::Archive::new(decoder);
archive
.unpack(output_dir)
.map_err(|e| PackError::Tar(e.to_string()))?;
Ok(())
}
pub fn decompress_assets_from_file(compressed_path: &Path, output_dir: &Path) -> Result<()> {
fs::create_dir_all(output_dir)?;
let file = File::open(compressed_path)?;
let decoder =
zstd::stream::Decoder::new(file).map_err(|e| PackError::Compression(e.to_string()))?;
let mut archive = tar::Archive::new(decoder);
archive
.unpack(output_dir)
.map_err(|e| PackError::Tar(e.to_string()))?;
Ok(())
}
pub fn crc32(data: &[u8]) -> u32 {
crc32fast::hash(data)
}
pub fn crc32_file(path: &Path) -> Result<u32> {
let mut file = File::open(path)?;
let mut hasher = crc32fast::Hasher::new();
let mut buf = [0u8; 64 * 1024];
loop {
let n = file.read(&mut buf)?;
if n == 0 {
break;
}
hasher.update(&buf[..n]);
}
Ok(hasher.finalize())
}
pub fn crc32_file_range(path: &Path, offset: u64, size: u64) -> Result<u32> {
use std::io::{Seek, SeekFrom};
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(offset))?;
let mut hasher = crc32fast::Hasher::new();
let mut remaining = size;
let mut buf = [0u8; 64 * 1024];
while remaining > 0 {
let to_read = remaining.min(buf.len() as u64) as usize;
let n = file.read(&mut buf[..to_read])?;
if n == 0 {
break;
}
hasher.update(&buf[..n]);
remaining -= n as u64;
}
Ok(hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_find_last_data_byte_all_zero() {
let temp = tempfile::NamedTempFile::new().unwrap();
fs::write(temp.path(), vec![0u8; 4096]).unwrap();
let mut file = File::open(temp.path()).unwrap();
assert_eq!(find_last_data_byte(&mut file, 4096).unwrap(), 0);
}
#[test]
fn test_find_last_data_byte_trailing_zeros() {
let temp = tempfile::NamedTempFile::new().unwrap();
let mut data = vec![0u8; 4100];
data[99] = 0xAB; fs::write(temp.path(), &data).unwrap();
let mut file = File::open(temp.path()).unwrap();
assert_eq!(find_last_data_byte(&mut file, 4100).unwrap(), 100);
}
#[test]
fn test_find_last_data_byte_nonzero_at_end() {
let temp = tempfile::NamedTempFile::new().unwrap();
let mut data = vec![0u8; 1024];
data[1023] = 1;
fs::write(temp.path(), &data).unwrap();
let mut file = File::open(temp.path()).unwrap();
assert_eq!(find_last_data_byte(&mut file, 1024).unwrap(), 1024);
}
#[test]
fn test_sparse_copy_overlay_all_zero() {
let temp_dir = tempfile::tempdir().unwrap();
let src = temp_dir.path().join("src.raw");
let dst = temp_dir.path().join("dst.raw");
fs::write(&src, vec![0u8; 8192]).unwrap();
let (logical, truncated) = sparse_copy_overlay(&src, &dst).unwrap();
assert_eq!(logical, 8192);
assert_eq!(truncated, 0);
assert_eq!(fs::metadata(&dst).unwrap().len(), 0);
}
#[test]
fn test_sparse_copy_overlay_trailing_zeros_and_interior_holes() {
let temp_dir = tempfile::tempdir().unwrap();
let src = temp_dir.path().join("src.raw");
let dst = temp_dir.path().join("dst.raw");
let mut data = vec![0u8; 8192];
data[0] = 0x01; data[511] = 0xFF; fs::write(&src, &data).unwrap();
let (logical, truncated) = sparse_copy_overlay(&src, &dst).unwrap();
assert_eq!(logical, 8192);
assert_eq!(truncated, 512);
let dst_data = fs::read(&dst).unwrap();
assert_eq!(dst_data[0], 0x01);
assert_eq!(dst_data[511], 0xFF);
assert_eq!(dst_data[256], 0x00); }
#[test]
fn test_crc32_basic() {
let data = b"hello world";
let checksum = crc32(data);
assert_eq!(checksum, 0x0D4A_1185); }
#[test]
fn test_crc32_empty() {
let data = b"";
let checksum = crc32(data);
assert_eq!(checksum, 0); }
#[test]
fn test_asset_collector_staging() {
let temp_dir = tempfile::tempdir().unwrap();
let staging = temp_dir.path().join("staging");
let _collector = AssetCollector::new(staging.clone()).unwrap();
assert!(!staging.join("lib").exists());
assert!(staging.join("layers").exists());
}
#[test]
fn test_compression_roundtrip() {
let temp_dir = tempfile::tempdir().unwrap();
let staging = temp_dir.path().join("staging");
let output = temp_dir.path().join("output");
fs::create_dir_all(&staging).unwrap();
let test_file = staging.join("test.txt");
fs::write(&test_file, b"hello world").unwrap();
let collector = AssetCollector::new(staging).unwrap();
let compressed = temp_dir.path().join("assets.tar.zst");
collector.compress(&compressed, false).unwrap();
decompress_assets_from_file(&compressed, &output).unwrap();
let restored = output.join("test.txt");
assert!(restored.exists());
assert_eq!(fs::read_to_string(&restored).unwrap(), "hello world");
}
}