use std::path::{Path, PathBuf};
use crate::cache::RootfsCache;
use crate::oci::OciRootfsBuilder;
use crate::vmm::TeeInstanceConfig;
use a3s_box_core::config::TeeConfig;
use a3s_box_core::error::{BoxError, Result};
use super::{BoxLayout, VmManager};
impl VmManager {
pub(crate) async fn prepare_layout(&self) -> Result<BoxLayout> {
let box_dir = self.home_dir.join("boxes").join(&self.box_id);
let socket_dir = self.socket_dir();
let logs_dir = box_dir.join("logs");
std::fs::create_dir_all(&socket_dir).map_err(|e| BoxError::BoxBootError {
message: format!("Failed to create socket directory: {}", e),
hint: None,
})?;
std::fs::create_dir_all(&logs_dir).map_err(|e| BoxError::BoxBootError {
message: format!("Failed to create logs directory: {}", e),
hint: None,
})?;
let workspace_path = if self.config.workspace.as_os_str().is_empty() {
box_dir.join("workspace")
} else {
PathBuf::from(&self.config.workspace)
};
if !workspace_path.exists() {
std::fs::create_dir_all(&workspace_path).map_err(|e| BoxError::BoxBootError {
message: format!("Failed to create workspace directory: {}", e),
hint: None,
})?;
}
let workspace_path = workspace_path
.canonicalize()
.map_err(|e| BoxError::BoxBootError {
message: format!(
"Failed to resolve workspace path {}: {}",
workspace_path.display(),
e
),
hint: None,
})?;
let prebuilt_rootfs = box_dir.join("rootfs");
let restore_marker = box_dir.join(".snapshot-rootfs");
let prebuilt_is_populated = restore_marker.exists()
&& std::fs::read_dir(&prebuilt_rootfs)
.map(|mut it| it.next().is_some())
.unwrap_or(false);
if prebuilt_is_populated {
tracing::info!(
rootfs = %prebuilt_rootfs.display(),
"Booting from pre-populated rootfs (snapshot restore)"
);
if let Ok(guest_init_path) = Self::find_guest_init() {
if let Err(e) = OciRootfsBuilder::new(&prebuilt_rootfs)
.with_guest_init(guest_init_path)
.install_guest_init_only()
{
tracing::warn!(error = %e, "Failed to refresh guest init on restored rootfs");
}
}
let tee_instance_config = self.generate_tee_config(&box_dir)?;
return Ok(BoxLayout {
rootfs_path: prebuilt_rootfs,
exec_socket_path: socket_dir.join("exec.sock"),
pty_socket_path: socket_dir.join("pty.sock"),
attest_socket_path: socket_dir.join("attest.sock"),
port_forward_socket_path: socket_dir.join("portfwd.sock"),
workspace_path,
console_output: Some(logs_dir.join("console.log")),
oci_config: None,
tee_instance_config,
});
}
let reference = &self.config.image;
#[cfg(unix)]
if super::is_restore_mode(&self.config) {
let cache_key = RootfsCache::compute_key(reference, &[], &[], &[]);
if let Some(cached_path) = self.try_rootfs_cache_path(&cache_key)? {
let rootfs_path = self.rootfs_provider.prepare(&box_dir, &cached_path)?;
let tee_instance_config = self.generate_tee_config(&box_dir)?;
return Ok(BoxLayout {
rootfs_path,
exec_socket_path: socket_dir.join("exec.sock"),
pty_socket_path: socket_dir.join("pty.sock"),
attest_socket_path: socket_dir.join("attest.sock"),
port_forward_socket_path: socket_dir.join("portfwd.sock"),
workspace_path,
console_output: Some(logs_dir.join("console.log")),
oci_config: None,
tee_instance_config,
});
}
}
let images_dir = self.home_dir.join("images");
let store = crate::oci::ImageStore::new(&images_dir, crate::DEFAULT_IMAGE_CACHE_SIZE)?;
let mut puller = crate::oci::ImagePuller::new(
std::sync::Arc::new(store),
crate::oci::RegistryAuth::from_env(),
);
if let Some(ref m) = self.prom {
puller = puller.set_metrics(m.clone());
}
if let Some(ref f) = self.pull_progress_fn {
puller = puller.with_progress_fn(f.clone());
}
tracing::info!(reference = %reference, "Pulling OCI image from registry");
let oci_image = puller.pull(reference).await?;
let image_path = oci_image.root_dir().to_path_buf();
let cache_key = RootfsCache::compute_key(reference, &[], &[], &[]);
let (rootfs_path, oci_config) =
if let Some(cached_path) = self.try_rootfs_cache_path(&cache_key)? {
tracing::info!(
cache_key = %&cache_key[..12],
reference = %reference,
provider = self.rootfs_provider.name(),
"Rootfs cache hit"
);
if let Some(ref prom) = self.prom {
prom.rootfs_cache_hits.inc();
}
let rootfs_path = self.rootfs_provider.prepare(&box_dir, &cached_path)?;
if let Ok(guest_init_path) = Self::find_guest_init() {
tracing::info!(
guest_init = %guest_init_path.display(),
"Refreshing guest init on cached rootfs"
);
OciRootfsBuilder::new(&rootfs_path)
.with_guest_init(guest_init_path)
.install_guest_init_only()?;
}
let builder = OciRootfsBuilder::new(&rootfs_path).with_image(&image_path);
(rootfs_path, Some(builder.image_config()?))
} else {
tracing::info!(
image = %image_path.display(),
"Building rootfs from pulled OCI image (cache miss)"
);
if let Some(ref prom) = self.prom {
prom.rootfs_cache_misses.inc();
}
let rootfs_path = box_dir.join("rootfs");
let mut builder = OciRootfsBuilder::new(&rootfs_path).with_image(&image_path);
if let Ok(guest_init_path) = Self::find_guest_init() {
tracing::info!(
guest_init = %guest_init_path.display(),
"Installing guest init"
);
builder = builder.with_guest_init(guest_init_path);
} else {
tracing::warn!(
"Guest init binary not found; container entrypoint will run as PID 1"
);
}
builder.build()?;
let config = builder.image_config()?;
self.store_rootfs_cache(&cache_key, &rootfs_path, reference);
(rootfs_path, Some(config))
};
let tee_instance_config = self.generate_tee_config(&box_dir)?;
Ok(BoxLayout {
rootfs_path,
exec_socket_path: socket_dir.join("exec.sock"),
pty_socket_path: socket_dir.join("pty.sock"),
attest_socket_path: socket_dir.join("attest.sock"),
port_forward_socket_path: socket_dir.join("portfwd.sock"),
workspace_path,
console_output: Some(logs_dir.join("console.log")),
oci_config,
tee_instance_config,
})
}
pub(crate) fn socket_dir(&self) -> PathBuf {
#[cfg(all(unix, target_os = "macos"))]
{
PathBuf::from("/private/tmp")
.join("a3s-box-sockets")
.join(&self.box_id)
}
#[cfg(all(unix, not(target_os = "macos")))]
{
PathBuf::from("/tmp")
.join("a3s-box-sockets")
.join(&self.box_id)
}
#[cfg(not(unix))]
{
self.home_dir
.join("boxes")
.join(&self.box_id)
.join("sockets")
}
}
#[cfg(test)]
pub(crate) fn try_rootfs_cache(
&self,
cache_key: &str,
target_path: &Path,
) -> Result<Option<PathBuf>> {
if !self.config.cache.enabled {
return Ok(None);
}
let cache_dir = self.resolve_cache_dir().join("rootfs");
let cache = match RootfsCache::new(&cache_dir) {
Ok(c) => c,
Err(e) => {
tracing::warn!(error = %e, "Failed to open rootfs cache, skipping");
return Ok(None);
}
};
match cache.get(cache_key)? {
Some(cached_path) => {
crate::cache::layer_cache::copy_dir_recursive(&cached_path, target_path)?;
Ok(Some(target_path.to_path_buf()))
}
None => Ok(None),
}
}
pub(crate) fn try_rootfs_cache_path(&self, cache_key: &str) -> Result<Option<PathBuf>> {
if !self.config.cache.enabled {
return Ok(None);
}
let cache_dir = self.resolve_cache_dir().join("rootfs");
let cache = match RootfsCache::new(&cache_dir) {
Ok(c) => c,
Err(e) => {
tracing::warn!(error = %e, "Failed to open rootfs cache, skipping");
return Ok(None);
}
};
cache.get(cache_key)
}
pub(crate) fn store_rootfs_cache(
&self,
cache_key: &str,
rootfs_path: &Path,
description: &str,
) {
if !self.config.cache.enabled {
return;
}
let cache_dir = self.resolve_cache_dir().join("rootfs");
let cache = match RootfsCache::new(&cache_dir) {
Ok(c) => c,
Err(e) => {
tracing::warn!(error = %e, "Failed to open rootfs cache for storing");
return;
}
};
match cache.put(cache_key, rootfs_path, description) {
Ok(_) => {
tracing::debug!(
cache_key = %&cache_key[..cache_key.len().min(12)],
description = %description,
"Stored rootfs in cache"
);
if let Err(e) = cache.prune(
self.config.cache.max_rootfs_entries,
self.config.cache.max_cache_bytes,
) {
tracing::warn!(error = %e, "Failed to prune rootfs cache");
}
}
Err(e) => {
tracing::warn!(error = %e, "Failed to store rootfs in cache");
}
}
}
pub(crate) fn resolve_cache_dir(&self) -> PathBuf {
self.config
.cache
.cache_dir
.clone()
.unwrap_or_else(|| self.home_dir.join("cache"))
}
#[cfg(unix)]
pub(crate) fn generate_tee_config(&self, box_dir: &Path) -> Result<Option<TeeInstanceConfig>> {
match &self.config.tee {
TeeConfig::None => Ok(None),
TeeConfig::SevSnp {
workload_id,
generation,
simulate,
} => {
if *simulate {
tracing::warn!("TEE simulation mode: skipping hardware check and TEE config");
return Ok(None);
}
crate::tee::require_sev_snp_support()?;
let config = serde_json::json!({
"workload_id": workload_id,
"cpus": self.config.resources.vcpus,
"ram_mib": self.config.resources.memory_mb,
"tee": "snp",
"tee_data": format!(r#"{{"gen":"{}"}}"#, generation.as_str()),
"attestation_url": ""
});
let config_path = box_dir.join("tee-config.json");
std::fs::write(&config_path, serde_json::to_string_pretty(&config)?).map_err(
|e| {
BoxError::TeeConfig(format!(
"Failed to write TEE config to {}: {}",
config_path.display(),
e
))
},
)?;
tracing::info!(
workload_id = %workload_id,
generation = %generation.as_str(),
config_path = %config_path.display(),
"Generated TEE configuration"
);
Ok(Some(TeeInstanceConfig {
config_path,
tee_type: "snp".to_string(),
}))
}
TeeConfig::Tdx {
workload_id,
simulate,
} => {
if *simulate {
tracing::warn!("TDX simulation mode: skipping hardware check and TEE config");
return Ok(None);
}
Err(BoxError::TeeConfig(format!(
"Intel TDX is not yet supported at runtime (workload_id='{}'). \
Use tee=sev-snp or tee=none.",
workload_id
)))
}
}
}
#[cfg(windows)]
pub(crate) fn generate_tee_config(&self, _box_dir: &Path) -> Result<Option<TeeInstanceConfig>> {
match &self.config.tee {
TeeConfig::None => Ok(None),
_ => Err(BoxError::TeeConfig(
"TEE configuration is not supported on Windows".to_string(),
)),
}
}
pub(crate) fn find_guest_init() -> Result<PathBuf> {
let mut candidates = Self::find_binary_candidates("a3s-box-guest-init");
candidates.sort_by_key(|path| {
let path_str = path.to_string_lossy();
if path_str.contains("-unknown-linux-musl") {
0
} else {
1
}
});
for path in candidates {
if Self::is_linux_elf(&path) {
return Ok(path);
}
tracing::debug!(
path = %path.display(),
"Skipping guest init (not a Linux ELF binary)"
);
}
Err(BoxError::BoxBootError {
message: "Linux guest init binary not found".to_string(),
hint: Some(
"Cross-compile the static guest init for your guest arch, e.g.: \
cargo build -p a3s-box-guest-init --release --target x86_64-unknown-linux-musl \
(or aarch64-unknown-linux-musl). A glibc-dynamic host build is rejected because \
it cannot run as PID 1 inside a minimal guest rootfs."
.to_string(),
),
})
}
fn find_binary_candidates(name: &str) -> Vec<PathBuf> {
let mut candidates = Vec::new();
if let Ok(exe_path) = std::env::current_exe() {
if let Some(exe_dir) = exe_path.parent() {
let path = exe_dir.join(name);
if path.exists() {
candidates.push(path);
}
if let Some(target_root) = exe_dir.parent() {
let cross_dirs = [
"aarch64-unknown-linux-musl/debug",
"aarch64-unknown-linux-musl/release",
"x86_64-unknown-linux-musl/debug",
"x86_64-unknown-linux-musl/release",
];
for dir in &cross_dirs {
let path = target_root.join(dir).join(name);
if path.exists() {
candidates.push(path);
}
}
}
}
}
let target_dirs = [
"target/aarch64-unknown-linux-musl/debug",
"target/aarch64-unknown-linux-musl/release",
"target/x86_64-unknown-linux-musl/debug",
"target/x86_64-unknown-linux-musl/release",
"target/debug",
"target/release",
];
for dir in &target_dirs {
let path = PathBuf::from(dir).join(name);
if path.exists() {
candidates.push(path);
}
}
let home_bin = a3s_box_core::dirs_home().join("bin").join(name);
if home_bin.exists() {
candidates.push(home_bin);
}
if let Ok(path_var) = std::env::var("PATH") {
for dir in std::env::split_paths(&path_var) {
let path = dir.join(name);
if path.exists() {
candidates.push(path);
}
}
}
candidates
}
fn is_linux_elf(path: &std::path::Path) -> bool {
let Ok(data) = std::fs::read(path) else {
return false;
};
if data.len() < 64 || data[0..4] != [0x7f, b'E', b'L', b'F'] {
return false;
}
if !matches!(data[7], 0x00 | 0x03) {
return false;
}
let is_elf64 = data[4] == 2;
let is_le = data[5] == 1;
if !is_elf64 || !is_le {
return true;
}
let u16_at = |off: usize| u16::from_le_bytes([data[off], data[off + 1]]);
let u64_at =
|off: usize| u64::from_le_bytes(data[off..off + 8].try_into().unwrap_or([0; 8]));
let e_phoff = u64_at(0x20) as usize; let e_phentsize = u16_at(0x36) as usize;
let e_phnum = u16_at(0x38) as usize;
if e_phoff == 0 || e_phentsize < 4 {
return true; }
const PT_INTERP: u32 = 3;
for i in 0..e_phnum {
let ph = e_phoff + i * e_phentsize;
if ph + 4 > data.len() {
break;
}
let p_type = u32::from_le_bytes(data[ph..ph + 4].try_into().unwrap_or([0; 4]));
if p_type == PT_INTERP {
return false;
}
}
true
}
}
#[cfg(test)]
mod tests {
use super::super::BoxState;
use super::*;
use crate::cache::RootfsCache;
use a3s_box_core::config::BoxConfig;
use std::sync::Arc;
use tempfile::TempDir;
use tokio::sync::RwLock;
fn make_vm_manager_with_home(home_dir: &Path) -> VmManager {
use a3s_box_core::event::EventEmitter;
let config = BoxConfig::default();
let emitter = EventEmitter::new(10);
VmManager {
config,
box_id: "test-box".to_string(),
state: Arc::new(RwLock::new(BoxState::Created)),
event_emitter: emitter,
provider: None,
handler: Arc::new(RwLock::new(None)),
#[cfg(unix)]
exec_client: None,
net_manager: None,
home_dir: home_dir.to_path_buf(),
anonymous_volumes: Vec::new(),
created_anonymous_volumes: Vec::new(),
image_config: None,
#[cfg(unix)]
tee: None,
rootfs_provider: crate::rootfs::default_provider(),
exec_socket_path: None,
pty_socket_path: None,
port_forward_socket_path: None,
prom: None,
shim_exit_code: None,
pull_progress_fn: None,
log_config: a3s_box_core::log::LogConfig::default(),
}
}
#[test]
fn test_resolve_cache_dir_default() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
let cache_dir = vm.resolve_cache_dir();
assert_eq!(cache_dir, tmp.path().join("cache"));
}
#[test]
fn test_resolve_cache_dir_custom() {
let tmp = TempDir::new().unwrap();
let mut vm = make_vm_manager_with_home(tmp.path());
vm.config.cache.cache_dir = Some(PathBuf::from("/custom/cache"));
let cache_dir = vm.resolve_cache_dir();
assert_eq!(cache_dir, PathBuf::from("/custom/cache"));
}
#[test]
fn test_try_rootfs_cache_disabled() {
let tmp = TempDir::new().unwrap();
let mut vm = make_vm_manager_with_home(tmp.path());
vm.config.cache.enabled = false;
let target = tmp.path().join("target");
let result = vm.try_rootfs_cache("some_key", &target).unwrap();
assert!(result.is_none());
}
#[test]
fn test_try_rootfs_cache_miss() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
let target = tmp.path().join("target");
let result = vm.try_rootfs_cache("nonexistent_key", &target).unwrap();
assert!(result.is_none());
}
#[test]
fn test_try_rootfs_cache_hit() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
let cache_dir = tmp.path().join("cache").join("rootfs");
let cache = RootfsCache::new(&cache_dir).unwrap();
let source = tmp.path().join("source_rootfs");
std::fs::create_dir_all(&source).unwrap();
std::fs::write(source.join("agent.bin"), "binary").unwrap();
cache.put("test_key", &source, "test").unwrap();
let target = tmp.path().join("target_rootfs");
let result = vm.try_rootfs_cache("test_key", &target).unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap(), target);
assert!(target.join("agent.bin").is_file());
assert_eq!(
std::fs::read_to_string(target.join("agent.bin")).unwrap(),
"binary"
);
}
#[test]
fn test_store_rootfs_cache_disabled() {
let tmp = TempDir::new().unwrap();
let mut vm = make_vm_manager_with_home(tmp.path());
vm.config.cache.enabled = false;
let source = tmp.path().join("rootfs");
std::fs::create_dir_all(&source).unwrap();
std::fs::write(source.join("f.txt"), "data").unwrap();
vm.store_rootfs_cache("key", &source, "test");
let cache_dir = tmp.path().join("cache").join("rootfs");
assert!(!cache_dir.exists());
}
#[test]
fn test_store_rootfs_cache_success() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
let source = tmp.path().join("rootfs");
std::fs::create_dir_all(&source).unwrap();
std::fs::write(source.join("agent.bin"), "binary").unwrap();
vm.store_rootfs_cache("store_key", &source, "test image");
let cache_dir = tmp.path().join("cache").join("rootfs");
let cache = RootfsCache::new(&cache_dir).unwrap();
let result = cache.get("store_key").unwrap();
assert!(result.is_some());
}
#[test]
fn test_store_rootfs_cache_prunes_on_store() {
let tmp = TempDir::new().unwrap();
let mut vm = make_vm_manager_with_home(tmp.path());
vm.config.cache.max_rootfs_entries = 2;
let source = tmp.path().join("rootfs");
std::fs::create_dir_all(&source).unwrap();
std::fs::write(source.join("f.txt"), "data").unwrap();
for i in 0..3 {
vm.store_rootfs_cache(&format!("key{}", i), &source, &format!("entry {}", i));
std::thread::sleep(std::time::Duration::from_millis(10));
}
let cache_dir = tmp.path().join("cache").join("rootfs");
let cache = RootfsCache::new(&cache_dir).unwrap();
assert!(cache.entry_count().unwrap() <= 2);
}
#[cfg(unix)]
#[tokio::test]
async fn test_exec_command_rejects_created_state() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
let result = vm.exec_command(vec!["echo".to_string()], 0).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("not yet booted"));
}
#[cfg(unix)]
#[tokio::test]
async fn test_exec_command_rejects_stopped_state() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
*vm.state.write().await = BoxState::Stopped;
let result = vm.exec_command(vec!["echo".to_string()], 0).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("stopped"));
}
#[cfg(unix)]
#[tokio::test]
async fn test_exec_command_no_client() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
*vm.state.write().await = BoxState::Ready;
let result = vm.exec_command(vec!["echo".to_string()], 0).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("not connected"));
}
#[cfg(unix)]
#[tokio::test]
async fn test_exec_request_rejects_empty_command() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
*vm.state.write().await = BoxState::Ready;
let request = a3s_box_core::exec::ExecRequest {
cmd: vec![],
timeout_ns: 0,
env: vec!["ENV=test".to_string()],
working_dir: Some("/app".to_string()),
rootfs: None,
stdin: None,
stdin_streaming: false,
user: None,
streaming: false,
};
let result = vm.exec_request(&request).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("non-empty command"));
}
#[cfg(unix)]
#[tokio::test]
async fn test_exec_request_no_client_preserves_request_fields() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
*vm.state.write().await = BoxState::Ready;
let request = a3s_box_core::exec::ExecRequest {
cmd: vec!["printenv".to_string()],
timeout_ns: 123,
env: vec!["ENV=test".to_string()],
working_dir: Some("/app".to_string()),
rootfs: Some("/run/a3s/cri/container-rootfs/sb/c/rootfs".to_string()),
stdin: Some(b"input".to_vec()),
stdin_streaming: false,
user: Some("1000:1000".to_string()),
streaming: false,
};
let result = vm.exec_request(&request).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("not connected"));
assert_eq!(request.env, vec!["ENV=test".to_string()]);
assert_eq!(request.working_dir, Some("/app".to_string()));
assert_eq!(request.stdin, Some(b"input".to_vec()));
assert_eq!(request.user, Some("1000:1000".to_string()));
}
#[test]
fn test_try_and_store_roundtrip() {
let tmp = TempDir::new().unwrap();
let vm = make_vm_manager_with_home(tmp.path());
let target1 = tmp.path().join("target1");
let result = vm.try_rootfs_cache("roundtrip_key", &target1).unwrap();
assert!(result.is_none());
let built_rootfs = tmp.path().join("built");
std::fs::create_dir_all(&built_rootfs).unwrap();
std::fs::write(built_rootfs.join("init"), "init_binary").unwrap();
std::fs::create_dir_all(built_rootfs.join("etc")).unwrap();
std::fs::write(built_rootfs.join("etc/config"), "config_data").unwrap();
vm.store_rootfs_cache("roundtrip_key", &built_rootfs, "roundtrip test");
let target2 = tmp.path().join("target2");
let result = vm.try_rootfs_cache("roundtrip_key", &target2).unwrap();
assert!(result.is_some());
assert!(target2.join("init").is_file());
assert_eq!(
std::fs::read_to_string(target2.join("init")).unwrap(),
"init_binary"
);
assert_eq!(
std::fs::read_to_string(target2.join("etc/config")).unwrap(),
"config_data"
);
}
}