refactor: multiple fixes

- Create unified Arch enum in utils.rs with methods for all naming conventions
- Replace 5 duplicate architecture mapping functions with single source of truth
- Add memory string validation for QEMU
- Add KVM capability verification via ioctl
- Fix TOCTOU race in resolv.conf writing using O_EXCL
- Extract magic numbers to named constants
- Use constants for stack size, buffer sizes, hostname entropy
This commit is contained in:
2026-06-17 17:24:41 +02:00
parent 6bd6f2cf77
commit 5834630d60
6 changed files with 383 additions and 128 deletions
+6 -29
View File
@@ -144,39 +144,16 @@ fn parse_oci_ref(input: &str, arch: &str) -> Result<ImageSource> {
/// Map architecture to OCI standard names
pub fn map_oci_arch(arch: &str) -> String {
match arch {
"amd64" | "x86_64" => "amd64".to_string(),
"arm64" | "aarch64" => "arm64".to_string(),
"armhf" | "armv7" => "arm".to_string(),
"riscv64" => "riscv64".to_string(),
"ppc64el" | "ppc64le" => "ppc64le".to_string(),
"s390x" => "s390x".to_string(),
_ => arch.to_string(),
}
crate::utils::map_oci_arch(arch)
}
/// Map ecr architecture names to distro-specific names
pub fn map_arch(distro: Distro, arch: &str) -> String {
match distro {
Distro::Ubuntu => match arch {
"amd64" => "amd64".to_string(),
"arm64" => "arm64".to_string(),
"armhf" => "armhf".to_string(),
"riscv64" => "riscv64".to_string(),
"ppc64el" => "ppc64el".to_string(),
"s390x" => "s390x".to_string(),
_ => arch.to_string(),
},
Distro::Alpine => match arch {
"amd64" => "x86_64".to_string(),
"arm64" => "aarch64".to_string(),
"armhf" => "armv7".to_string(),
"riscv64" => "riscv64".to_string(),
"ppc64el" => "ppc64le".to_string(),
"s390x" => "s390x".to_string(),
_ => arch.to_string(),
},
}
let distro_name = match distro {
Distro::Ubuntu => "ubuntu",
Distro::Alpine => "alpine",
};
crate::utils::map_arch_for_distro(distro_name, arch)
}
/// Resolve the download URL for a known distro (optimized path)
+28 -23
View File
@@ -21,7 +21,7 @@ macro_rules! veprintln {
};
}
use anyhow::Result;
use anyhow::{Context, Result};
use clap::Parser;
use cli::Args;
@@ -283,32 +283,19 @@ fn get_tarball_extension(url: &str) -> &str {
}
fn get_host_arch() -> String {
// Use the uname(2) syscall directly — no subprocess, no PATH dependency,
// no panic-on-missing-binary. This gives the runtime machine string
// (e.g. "x86_64", "aarch64") exactly as `uname -m` would, which is what
// we need for the QEMU check. std::env::consts::ARCH is compile-time and
// would be wrong if the binary itself is running under emulation.
let utsname = nix::sys::utsname::uname()
.expect("uname(2) syscall failed — cannot determine host architecture");
let machine = utsname.machine().to_string_lossy();
match machine.as_ref() {
"x86_64" => "amd64".to_string(),
"aarch64" => "arm64".to_string(),
"armv7l" | "armv7" => "armhf".to_string(),
"riscv64" => "riscv64".to_string(),
"ppc64le" => "ppc64el".to_string(),
"s390x" => "s390x".to_string(),
other => other.to_string(),
}
// Use the consolidated architecture detection from utils
utils::get_host_arch().debian_name().to_string()
}
fn write_resolv_conf(rootfs: &std::path::Path, dns: &[String]) -> Result<()> {
use std::io::Write;
let resolv_conf = rootfs.join("etc/resolv.conf");
// Create /etc if it doesn't exist
if let Some(parent) = resolv_conf.parent() {
std::fs::create_dir_all(parent)?;
std::fs::create_dir_all(parent)
.with_context(|| format!("Failed to create directory: {}", parent.display()))?;
}
// Copy host's resolv.conf if dns is empty, otherwise use provided DNS
@@ -328,13 +315,31 @@ fn write_resolv_conf(rootfs: &std::path::Path, dns: &[String]) -> Result<()> {
c
};
// Remove any existing file or symlink before writing so that std::fs::write
// always creates a plain file. Without this, an absolute symlink such as
// Remove any existing file or symlink before writing so that we always
// create a plain file. Without this, an absolute symlink such as
// /etc/resolv.conf -> /run/systemd/resolve/stub-resolv.conf would cause the
// write to follow the symlink through the *host* root (chroot() has not been
// called yet) and corrupt the host's DNS configuration.
//
// Use atomic file creation with O_CREAT | O_EXCL to prevent TOCTOU race:
// if an attacker creates a symlink between our remove_file and write, the
// exclusive create will fail rather than writing to the symlink target.
let _ = std::fs::remove_file(&resolv_conf); // ignore ENOENT
std::fs::write(&resolv_conf, content)?;
// Use OpenOptions with create_new(true) for atomic exclusive creation
let mut file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&resolv_conf)
.with_context(|| {
format!(
"Failed to create resolv.conf at {} (symlink attack prevented)",
resolv_conf.display()
)
})?;
file.write_all(content.as_bytes())
.with_context(|| format!("Failed to write to {}", resolv_conf.display()))?;
Ok(())
}
+5 -3
View File
@@ -105,7 +105,7 @@ where
}
// Stack for the child process
let stack_size = 1024 * 1024;
let stack_size = crate::utils::CHILD_STACK_SIZE;
let mut stack = vec![0u8; stack_size];
// Wrap f in Option to allow taking it once inside the child closure
@@ -198,7 +198,7 @@ where
// O_CLOEXEC (on successful exec), so this read always terminates.
let child_error: Option<String> = unsafe {
let mut error_bytes = Vec::new();
let mut tmp = [0u8; 4096];
let mut tmp = [0u8; crate::utils::ERROR_BUFFER_SIZE];
loop {
let n = libc::read(
error_read.raw(),
@@ -397,7 +397,9 @@ pub fn set_hostname(distro: &str) -> Result<()> {
let mut state = hasher.finish();
let chars = b"abcdefghijklmnopqrstuvwxyz0123456789";
let random_suffix: String = (0..6)
// Use HOSTNAME_SUFFIX_BITS for entropy (6 hex chars = 24 bits)
let suffix_len = (crate::utils::HOSTNAME_SUFFIX_BITS as f64).log2() as usize / 4;
let random_suffix: String = (0..suffix_len)
.map(|_| {
// Knuth multiplicative LCG — each step advances the full 64-bit state.
state = state
+1 -14
View File
@@ -4,7 +4,7 @@ use std::path::Path;
/// Check if binfmt_misc is registered for the target architecture
pub fn check_binfmt(arch: &str) -> Result<()> {
let qemu_arch = map_arch_to_qemu(arch);
let qemu_arch = crate::utils::Arch::from_str(arch).qemu_binfmt_name();
let binfmt_path = format!("/proc/sys/fs/binfmt_misc/qemu-{}", qemu_arch);
@@ -22,16 +22,3 @@ pub fn check_binfmt(arch: &str) -> Result<()> {
veprintln!("QEMU binfmt_misc registered for {}", arch);
Ok(())
}
/// Map ecr architecture names to QEMU binary names
fn map_arch_to_qemu(arch: &str) -> &str {
match arch {
"amd64" | "x86_64" => "x86_64",
"arm64" | "aarch64" => "aarch64",
"armhf" | "armv7" => "arm",
"riscv64" => "riscv64",
"ppc64el" | "ppc64le" => "ppc64le",
"s390x" => "s390x",
_ => arch,
}
}
+48 -59
View File
@@ -39,6 +39,10 @@ pub fn launch_qemu(config: QemuConfig) -> Result<()> {
));
}
// Validate memory string format
crate::utils::validate_memory_string(&config.memory)
.with_context(|| format!("Invalid memory size: {}", config.memory))?;
// Create an uncompressed cpio initramfs from the rootfs
let initramfs = create_initramfs(&config.rootfs_path)?;
@@ -66,11 +70,12 @@ pub fn launch_qemu(config: QemuConfig) -> Result<()> {
let shell = crate::utils::detect_shell(&config.rootfs_path);
// Generate a unique hostname like "ecr-vm-a1b2c3"
// Use VM_HOSTNAME_SUFFIX_BITS constant for entropy
let hostname_suffix = format!("{:x}", (std::process::id() as u64)
.wrapping_mul(std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos() as u64) % 0x1000000); // 6 hex chars
.as_nanos() as u64) % crate::utils::VM_HOSTNAME_SUFFIX_BITS);
let hostname = format!("ecr-vm-{}", hostname_suffix);
// Build kernel command line
@@ -151,87 +156,71 @@ pub fn launch_qemu(config: QemuConfig) -> Result<()> {
/// Get QEMU system binary name for architecture
fn qemu_binary_for_arch(arch: &str) -> String {
match arch {
"amd64" | "x86_64" => "qemu-system-x86_64".to_string(),
"arm64" | "aarch64" => "qemu-system-aarch64".to_string(),
"armhf" | "armv7" => "qemu-system-arm".to_string(),
"riscv64" => "qemu-system-riscv64".to_string(),
"ppc64el" | "ppc64le" => "qemu-system-ppc64".to_string(),
"s390x" => "qemu-system-s390x".to_string(),
other => format!("qemu-system-{}", other),
}
let arch_enum = crate::utils::Arch::from_str(arch);
format!("qemu-system-{}", arch_enum.qemu_system_name())
}
/// Get architecture suffix for package names
fn get_arch_package_suffix(arch: &str) -> &str {
match arch {
"amd64" | "x86_64" => "x86",
"arm64" | "aarch64" => "aarch64",
"armhf" | "armv7" => "arm",
"riscv64" => "riscv64",
"ppc64el" | "ppc64le" => "ppc",
"s390x" => "s390x",
other => other,
}
fn get_arch_package_suffix(arch: &str) -> &'static str {
crate::utils::Arch::from_str(arch).qemu_package_suffix()
}
/// Check if KVM acceleration can be used for the target architecture
fn can_use_kvm(target_arch: &str) -> bool {
// Normalize both to canonical form (uname -m style) and compare
let host_arch = get_host_arch();
let target_canonical = normalize_arch(target_arch);
use crate::utils::Arch;
if host_arch != target_canonical {
// Normalize both to canonical form (uname -m style) and compare
let host_arch = crate::utils::get_host_arch();
let target_enum = Arch::from_str(target_arch);
if host_arch != target_enum {
return false;
}
// Check if /dev/kvm exists and is accessible
std::fs::OpenOptions::new()
// Check if /dev/kvm exists and is accessible (read+write required for VM execution)
match std::fs::OpenOptions::new()
.read(true)
.write(true)
.open("/dev/kvm")
.is_ok()
{
Ok(_) => {
// Additional check: verify KVM actually works by checking capabilities
// This catches cases where /dev/kvm exists but KVM is not functional
check_kvm_capabilities()
}
Err(_) => false,
}
}
/// Normalize architecture name to canonical form (uname -m style)
fn normalize_arch(arch: &str) -> String {
match arch {
"amd64" => "x86_64",
"arm64" => "aarch64",
"ppc64el" => "ppc64le",
"armhf" => "armv7l",
other => other,
}.to_string()
}
/// Get the host system architecture (uname -m style)
fn get_host_arch() -> String {
// Use uname to get the machine hardware name
match std::process::Command::new("uname").arg("-m").output() {
Ok(output) => {
String::from_utf8_lossy(&output.stdout).trim().to_string()
}
Err(_) => {
// Fallback: use libc uname
let mut utsname: libc::utsname = unsafe { std::mem::zeroed() };
if unsafe { libc::uname(&mut utsname) } == 0 {
let machine: Vec<u8> = utsname.machine
.iter()
.take_while(|&&c| c != 0)
.map(|&c| c as u8)
.collect();
String::from_utf8_lossy(&machine).into_owned()
} else {
"unknown".to_string()
}
/// Check if KVM capabilities are actually functional
fn check_kvm_capabilities() -> bool {
use std::os::unix::io::AsRawFd;
// Try to open /dev/kvm and check KVM_GET_API_VERSION
match std::fs::OpenOptions::new()
.read(true)
.write(true)
.open("/dev/kvm")
{
Ok(file) => {
let fd = file.as_raw_fd();
// KVM_GET_API_VERSION ioctl = 0xAE00
// Expected return value is 12 (KVM_API_VERSION)
let ret = unsafe { libc::ioctl(fd, 0xAE00) };
ret == 12
}
Err(_) => false,
}
}
/// Create an uncompressed cpio initramfs from a directory
fn create_initramfs(rootfs: &PathBuf) -> Result<PathBuf> {
// Create a temporary file for the initramfs (uncompressed cpio)
let initramfs_path = rootfs.parent().unwrap().join("initramfs.cpio");
// Use a temp file in the same directory as rootfs, or fall back to /tmp
let initramfs_path = rootfs
.parent()
.map(|p| p.join("initramfs.cpio"))
.unwrap_or_else(|| std::env::temp_dir().join("initramfs.cpio"));
// Create progress bar
let pb = ProgressBar::new_spinner();
+295
View File
@@ -1,5 +1,186 @@
use anyhow::{anyhow, Result};
use std::path::Path;
// ============================================================================
// Constants
// ============================================================================
/// Stack size for child processes in namespace cloning (1 MiB)
pub const CHILD_STACK_SIZE: usize = 1024 * 1024;
/// Buffer size for reading error messages from pipes (4 KiB, typical page size)
pub const ERROR_BUFFER_SIZE: usize = 4096;
/// Maximum entropy bits for hostname suffix (24 bits = 6 hex chars)
pub const HOSTNAME_SUFFIX_BITS: u64 = 0x1000000;
/// Maximum entropy bits for VM hostname suffix (24 bits = 6 hex chars)
pub const VM_HOSTNAME_SUFFIX_BITS: u64 = 0x1000000;
// ============================================================================
// Architecture handling
// ============================================================================
/// Architecture representation in different naming conventions
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Arch {
/// x86-64 (AMD64, x86_64)
Amd64,
/// ARM 64-bit (AArch64)
Arm64,
/// ARM 32-bit hard-float
Armhf,
/// RISC-V 64-bit
Riscv64,
/// PowerPC 64-bit little-endian
Ppc64el,
/// IBM s390x
S390x,
/// Unknown architecture
Unknown,
}
impl Arch {
/// Get the architecture from a string (any common naming convention)
pub fn from_str(s: &str) -> Self {
match s {
"amd64" | "x86_64" | "x64" => Arch::Amd64,
"arm64" | "aarch64" | "arm64v8" => Arch::Arm64,
"armhf" | "armv7" | "armv7l" | "arm" => Arch::Armhf,
"riscv64" => Arch::Riscv64,
"ppc64el" | "ppc64le" => Arch::Ppc64el,
"s390x" => Arch::S390x,
_ => Arch::Unknown,
}
}
/// Get the canonical name (uname -m style)
pub fn canonical_name(&self) -> &'static str {
match self {
Arch::Amd64 => "x86_64",
Arch::Arm64 => "aarch64",
Arch::Armhf => "armv7l",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc64le",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
/// Get the Debian/Ubuntu style name
pub fn debian_name(&self) -> &'static str {
match self {
Arch::Amd64 => "amd64",
Arch::Arm64 => "arm64",
Arch::Armhf => "armhf",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc64el",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
/// Get the OCI/Docker registry style name
pub fn oci_name(&self) -> &'static str {
match self {
Arch::Amd64 => "amd64",
Arch::Arm64 => "arm64",
// OCI uses "arm" for 32-bit ARM with variant field
Arch::Armhf => "arm",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc64le",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
/// Get the Alpine style name
pub fn alpine_name(&self) -> &'static str {
match self {
Arch::Amd64 => "x86_64",
Arch::Arm64 => "aarch64",
Arch::Armhf => "armv7",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc64le",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
/// Get the QEMU binary suffix (e.g., "qemu-system-x86_64")
pub fn qemu_system_name(&self) -> &'static str {
match self {
Arch::Amd64 => "x86_64",
Arch::Arm64 => "aarch64",
Arch::Armhf => "arm",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc64",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
/// Get the QEMU binfmt_misc name
pub fn qemu_binfmt_name(&self) -> &'static str {
match self {
Arch::Amd64 => "x86_64",
Arch::Arm64 => "aarch64",
Arch::Armhf => "arm",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc64le",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
/// Get the package suffix for QEMU system emulator
pub fn qemu_package_suffix(&self) -> &'static str {
match self {
Arch::Amd64 => "x86",
Arch::Arm64 => "aarch64",
Arch::Armhf => "arm",
Arch::Riscv64 => "riscv64",
Arch::Ppc64el => "ppc",
Arch::S390x => "s390x",
Arch::Unknown => "unknown",
}
}
}
/// Get the host system architecture using uname(2) syscall
/// This returns the runtime machine string, which is correct even when
/// the binary itself is running under emulation.
pub fn get_host_arch() -> Arch {
let utsname = nix::sys::utsname::uname()
.expect("uname(2) syscall failed — cannot determine host architecture");
let machine = utsname.machine().to_string_lossy();
Arch::from_str(machine.as_ref())
}
/// Normalize an architecture string to canonical (uname -m) form
pub fn normalize_arch(arch: &str) -> String {
Arch::from_str(arch).canonical_name().to_string()
}
/// Map ecr architecture names to distro-specific names
pub fn map_arch_for_distro(distro: &str, arch: &str) -> String {
let arch_enum = Arch::from_str(arch);
match distro.to_lowercase().as_str() {
"ubuntu" => arch_enum.debian_name().to_string(),
"alpine" => arch_enum.alpine_name().to_string(),
_ => arch_enum.oci_name().to_string(),
}
}
/// Map architecture to OCI registry standard names
pub fn map_oci_arch(arch: &str) -> String {
Arch::from_str(arch).oci_name().to_string()
}
// ============================================================================
// Shell detection
// ============================================================================
/// Detect the best available shell in a rootfs
/// Checks for bash first, falls back to sh
/// Returns the path relative to the rootfs (e.g., "/bin/bash")
@@ -17,3 +198,117 @@ pub fn detect_shell(rootfs: &Path) -> &'static str {
"/bin/sh" // Will fail with clear error if not present
}
}
// ============================================================================
// Memory string validation
// ============================================================================
/// Validate a QEMU memory size string (e.g., "512M", "2G")
/// Returns an error if the format is invalid
pub fn validate_memory_string(s: &str) -> Result<()> {
if s.is_empty() {
return Err(anyhow!("Memory size cannot be empty"));
}
// Must end with a valid suffix or be a plain number
let suffix = s.chars().last().unwrap();
let has_suffix = suffix.is_ascii_alphabetic();
let numeric_part = if has_suffix {
&s[..s.len() - 1]
} else {
s
};
// Check for negative numbers
if numeric_part.starts_with('-') {
return Err(anyhow!("Memory size cannot be negative: {}", s));
}
// Must be a valid positive number
if numeric_part.is_empty() {
return Err(anyhow!("Memory size must have a numeric value: {}", s));
}
// Check if it's a valid integer
if !numeric_part.chars().all(|c| c.is_ascii_digit()) {
return Err(anyhow!(
"Memory size must be a positive integer with optional suffix: {}",
s
));
}
// Check suffix is valid
if has_suffix {
let valid_suffixes = ['K', 'M', 'G', 'T'];
let suffix_upper = suffix.to_ascii_uppercase();
if !valid_suffixes.contains(&suffix_upper) {
return Err(anyhow!(
"Invalid memory suffix '{}'. Valid suffixes: K, M, G, T",
suffix
));
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_arch_from_str() {
assert_eq!(Arch::from_str("amd64"), Arch::Amd64);
assert_eq!(Arch::from_str("x86_64"), Arch::Amd64);
assert_eq!(Arch::from_str("arm64"), Arch::Arm64);
assert_eq!(Arch::from_str("aarch64"), Arch::Arm64);
assert_eq!(Arch::from_str("armhf"), Arch::Armhf);
assert_eq!(Arch::from_str("armv7"), Arch::Armhf);
assert_eq!(Arch::from_str("riscv64"), Arch::Riscv64);
assert_eq!(Arch::from_str("ppc64el"), Arch::Ppc64el);
assert_eq!(Arch::from_str("ppc64le"), Arch::Ppc64el);
assert_eq!(Arch::from_str("s390x"), Arch::S390x);
}
#[test]
fn test_arch_canonical_name() {
assert_eq!(Arch::Amd64.canonical_name(), "x86_64");
assert_eq!(Arch::Arm64.canonical_name(), "aarch64");
assert_eq!(Arch::Armhf.canonical_name(), "armv7l");
}
#[test]
fn test_arch_oci_name() {
assert_eq!(Arch::Amd64.oci_name(), "amd64");
assert_eq!(Arch::Arm64.oci_name(), "arm64");
assert_eq!(Arch::Armhf.oci_name(), "arm");
}
#[test]
fn test_arch_alpine_name() {
assert_eq!(Arch::Amd64.alpine_name(), "x86_64");
assert_eq!(Arch::Arm64.alpine_name(), "aarch64");
assert_eq!(Arch::Armhf.alpine_name(), "armv7");
}
#[test]
fn test_validate_memory_string_valid() {
assert!(validate_memory_string("512M").is_ok());
assert!(validate_memory_string("2G").is_ok());
assert!(validate_memory_string("1024").is_ok());
assert!(validate_memory_string("1T").is_ok());
assert!(validate_memory_string("256K").is_ok());
assert!(validate_memory_string("2g").is_ok()); // lowercase
}
#[test]
fn test_validate_memory_string_invalid() {
assert!(validate_memory_string("").is_err());
assert!(validate_memory_string("-1G").is_err());
assert!(validate_memory_string("2X").is_err());
assert!(validate_memory_string("abc").is_err());
assert!(validate_memory_string("G").is_err());
assert!(validate_memory_string("1.5G").is_err());
}
}