use anyhow::{Context, Result};
use std::path::Path;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum TmpfsFraction {
Half,
NinetyPercent,
}
impl TmpfsFraction {
fn ratio(self) -> (u64, u64) {
match self {
TmpfsFraction::Half => (1, 2),
TmpfsFraction::NinetyPercent => (9, 10),
}
}
pub(crate) fn for_kernel_version(version: Option<(u16, u16, u16)>) -> Self {
let Some((major, minor, patch)) = version else {
return TmpfsFraction::Half;
};
if (major, minor) >= (6, 18) {
return TmpfsFraction::NinetyPercent;
}
let floor = match (major, minor) {
(5, 4) => 301,
(5, 10) => 246,
(5, 15) => 195,
(6, 1) => 157,
(6, 6) => 113,
(6, 12) => 54,
(6, 17) => 4,
_ => return TmpfsFraction::Half,
};
if patch >= floor {
TmpfsFraction::NinetyPercent
} else {
TmpfsFraction::Half
}
}
}
pub(crate) struct MemoryBudget {
pub uncompressed_initramfs_bytes: u64,
pub compressed_initrd_bytes: u64,
pub kernel_init_size: u64,
pub init_coverage_instrumented: bool,
pub instrumented_reserve_bytes: u64,
pub tmpfs_fraction: TmpfsFraction,
}
pub(crate) fn read_kernel_init_size(kernel_path: &Path) -> Result<u64> {
use std::io::{Read, Seek, SeekFrom};
let mut f = std::fs::File::open(kernel_path)
.with_context(|| format!("open kernel for init_size: {}", kernel_path.display()))?;
#[cfg(target_arch = "x86_64")]
{
f.seek(SeekFrom::Start(0x260))
.context("seek to init_size in bzImage")?;
let mut buf = [0u8; 4];
f.read_exact(&mut buf)
.context("read init_size from bzImage")?;
Ok(u32::from_le_bytes(buf) as u64)
}
#[cfg(target_arch = "aarch64")]
{
let mut magic = [0u8; 2];
f.read_exact(&mut magic).context("read kernel magic")?;
if magic == [0x1f, 0x8b] {
f.seek(SeekFrom::Start(0))
.context("seek vmlinuz to start")?;
let mut decoder = flate2::read::GzDecoder::new(&mut f);
let mut header = [0u8; 24];
decoder
.read_exact(&mut header)
.context("decompress arm64 vmlinuz header for image_size")?;
return Ok(u64::from_le_bytes(header[16..24].try_into().unwrap()));
}
f.seek(SeekFrom::Start(16))
.context("seek to image_size in arm64 Image")?;
let mut buf = [0u8; 8];
f.read_exact(&mut buf)
.context("read image_size from arm64 Image")?;
Ok(u64::from_le_bytes(buf))
}
}
pub(crate) fn read_kernel_version(kernel_path: &Path) -> Option<(u16, u16, u16)> {
#[cfg(target_arch = "x86_64")]
{
use std::io::{Read, Seek, SeekFrom};
let mut f = std::fs::File::open(kernel_path).ok()?;
f.seek(SeekFrom::Start(0x202)).ok()?;
let mut magic = [0u8; 4];
f.read_exact(&mut magic).ok()?;
if &magic != b"HdrS" {
return None;
}
f.seek(SeekFrom::Start(0x20E)).ok()?;
let mut vbuf = [0u8; 2];
f.read_exact(&mut vbuf).ok()?;
let ver_ptr = u16::from_le_bytes(vbuf);
if ver_ptr == 0 {
return None;
}
f.seek(SeekFrom::Start(0x200u64 + ver_ptr as u64)).ok()?;
let mut window = [0u8; 256];
let n = f.read(&mut window).ok()?;
let bytes = &window[..n];
let end = bytes
.iter()
.position(|&b| b == 0 || b == b' ')
.unwrap_or(bytes.len());
let release = std::str::from_utf8(&bytes[..end]).ok()?;
parse_kernel_version(release)
}
#[cfg(target_arch = "aarch64")]
{
let _ = kernel_path;
None
}
}
fn parse_kernel_version(release: &str) -> Option<(u16, u16, u16)> {
let mut parts = release.split('.');
let major: u16 = parts.next()?.parse().ok()?;
let minor_raw = parts.next()?;
let minor_digits: String = minor_raw
.chars()
.take_while(|c| c.is_ascii_digit())
.collect();
if minor_digits.is_empty() {
return None;
}
let minor: u16 = minor_digits.parse().ok()?;
let patch: u16 = parts
.next()
.map(|p| {
p.chars()
.take_while(|c| c.is_ascii_digit())
.collect::<String>()
})
.and_then(|d| d.parse().ok())
.unwrap_or(0);
Some((major, minor, patch))
}
pub(crate) fn read_kernel_version_from_metadata_sidecar(
kernel_path: &Path,
) -> Option<(u16, u16, u16)> {
#[derive(serde::Deserialize)]
struct VersionProbe {
version: Option<String>,
}
let sidecar = kernel_path.parent()?.join("metadata.json");
let json = std::fs::read_to_string(sidecar).ok()?;
let probe: VersionProbe = serde_json::from_str(&json).ok()?;
parse_kernel_version(&probe.version?)
}
const WORKLOAD_MIB: u64 = 256;
pub(crate) fn initramfs_min_memory_mib(budget: &MemoryBudget) -> u32 {
let ceil_mib = |bytes: u64| -> u64 { bytes.saturating_add((1 << 20) - 1) >> 20 };
let init_size_mib = ceil_mib(budget.kernel_init_size);
let compressed_mib = ceil_mib(budget.compressed_initrd_bytes);
let uncompressed_mib = ceil_mib(budget.uncompressed_initramfs_bytes);
let (frac_num, frac_den) = budget.tmpfs_fraction.ratio();
let uncompressed_scaled = uncompressed_mib.saturating_mul(frac_den).div_ceil(frac_num);
let content_mib = uncompressed_scaled
.saturating_add(init_size_mib)
.saturating_add(compressed_mib);
let boot_mib = content_mib.saturating_mul(64).div_ceil(63);
let coverage_reserve_mib = if budget.init_coverage_instrumented {
ceil_mib(budget.instrumented_reserve_bytes)
} else {
0
};
let total_mib = boot_mib
.saturating_add(WORKLOAD_MIB)
.saturating_add(coverage_reserve_mib);
u32::try_from(total_mib).unwrap_or_else(|_| {
panic!(
"initramfs_min_memory_mib: computed floor {total_mib}MiB exceeds u32 \
(boot={boot_mib}MiB, workload={WORKLOAD_MIB}MiB, \
coverage_reserve={coverage_reserve_mib}MiB)"
)
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn workload_mib_is_256() {
assert_eq!(WORKLOAD_MIB, 256);
}
#[test]
fn initramfs_min_memory_mib_zeros_returns_workload_budget() {
let budget = MemoryBudget {
uncompressed_initramfs_bytes: 0,
compressed_initrd_bytes: 0,
kernel_init_size: 0,
init_coverage_instrumented: false,
instrumented_reserve_bytes: 0,
tmpfs_fraction: TmpfsFraction::Half,
};
assert_eq!(initramfs_min_memory_mib(&budget), WORKLOAD_MIB as u32);
}
#[test]
fn initramfs_min_memory_mib_known_input() {
let budget = MemoryBudget {
uncompressed_initramfs_bytes: 10 * (1 << 20),
compressed_initrd_bytes: 2 * (1 << 20),
kernel_init_size: 5 * (1 << 20),
init_coverage_instrumented: false,
instrumented_reserve_bytes: 0,
tmpfs_fraction: TmpfsFraction::Half,
};
assert_eq!(initramfs_min_memory_mib(&budget), 284);
}
#[test]
fn initramfs_min_memory_mib_subbyte_uncompressed_rounds_up() {
let budget = MemoryBudget {
uncompressed_initramfs_bytes: 1,
compressed_initrd_bytes: 0,
kernel_init_size: 0,
init_coverage_instrumented: false,
instrumented_reserve_bytes: 0,
tmpfs_fraction: TmpfsFraction::Half,
};
assert_eq!(initramfs_min_memory_mib(&budget), 259);
}
#[test]
fn initramfs_min_memory_mib_larger_input() {
let budget = MemoryBudget {
uncompressed_initramfs_bytes: 200 * (1 << 20),
compressed_initrd_bytes: 50 * (1 << 20),
kernel_init_size: 30 * (1 << 20),
init_coverage_instrumented: false,
instrumented_reserve_bytes: 0,
tmpfs_fraction: TmpfsFraction::Half,
};
assert_eq!(initramfs_min_memory_mib(&budget), 744);
}
#[test]
fn initramfs_min_memory_mib_instrumented_reserve_raises_floor() {
let base = MemoryBudget {
uncompressed_initramfs_bytes: 1200 * (1 << 20),
compressed_initrd_bytes: 300 * (1 << 20),
kernel_init_size: 30 * (1 << 20),
init_coverage_instrumented: false,
instrumented_reserve_bytes: 3500 * (1 << 20),
tmpfs_fraction: TmpfsFraction::Half,
};
let instrumented = MemoryBudget {
init_coverage_instrumented: true,
..base
};
let base_floor = initramfs_min_memory_mib(&base);
let instrumented_floor = initramfs_min_memory_mib(&instrumented);
assert_eq!(
base_floor, 3030,
"non-instrumented floor must NOT include the reserve \
(2774 boot + 256 workload)"
);
assert_eq!(
instrumented_floor, 6530,
"instrumented floor = 2774 boot + 256 workload + 3500 reserve"
);
assert!(
instrumented_floor > base_floor,
"instrumented reserve must raise the floor ({instrumented_floor} \
vs {base_floor})"
);
assert!(
instrumented_floor > 4096,
"instrumented floor must clear the old memory_deferred_min(4096) \
workaround (got {instrumented_floor})"
);
}
#[test]
fn tmpfs_fraction_gates_on_honoring_versions() {
use TmpfsFraction::{Half, NinetyPercent};
let frac = TmpfsFraction::for_kernel_version;
assert_eq!(frac(Some((6, 18, 0))), NinetyPercent);
assert_eq!(frac(Some((6, 19, 0))), NinetyPercent);
assert_eq!(frac(Some((7, 0, 5))), NinetyPercent);
assert_eq!(frac(Some((6, 16, 0))), Half);
assert_eq!(frac(Some((6, 17, 4))), NinetyPercent);
assert_eq!(frac(Some((6, 17, 3))), Half);
assert_eq!(frac(Some((6, 12, 54))), NinetyPercent);
assert_eq!(frac(Some((6, 12, 53))), Half);
assert_eq!(frac(Some((6, 6, 113))), NinetyPercent);
assert_eq!(frac(Some((6, 6, 112))), Half);
assert_eq!(frac(Some((6, 1, 157))), NinetyPercent);
assert_eq!(frac(Some((6, 1, 156))), Half);
assert_eq!(frac(Some((5, 15, 195))), NinetyPercent);
assert_eq!(frac(Some((5, 15, 194))), Half);
assert_eq!(frac(Some((5, 10, 246))), NinetyPercent);
assert_eq!(frac(Some((5, 10, 245))), Half);
assert_eq!(frac(Some((5, 4, 301))), NinetyPercent);
assert_eq!(frac(Some((5, 4, 300))), Half);
assert_eq!(frac(Some((6, 9, 999))), Half);
assert_eq!(frac(Some((6, 13, 999))), Half);
assert_eq!(frac(None), Half);
}
#[test]
fn parse_kernel_version_shapes() {
assert_eq!(parse_kernel_version("6.18.0-rc1"), Some((6, 18, 0)));
assert_eq!(
parse_kernel_version("7.1.0-rc7-gc80ba8d32ec3"),
Some((7, 1, 0))
);
assert_eq!(parse_kernel_version("6.12.54"), Some((6, 12, 54)));
assert_eq!(parse_kernel_version("6.6.113"), Some((6, 6, 113)));
assert_eq!(parse_kernel_version("6.18-rc1"), Some((6, 18, 0)));
assert_eq!(parse_kernel_version("6.6.x"), Some((6, 6, 0)));
assert_eq!(parse_kernel_version(""), None);
assert_eq!(parse_kernel_version("garbage"), None);
assert_eq!(parse_kernel_version("6"), None);
assert_eq!(parse_kernel_version("6."), None);
assert_eq!(parse_kernel_version("x.18"), None);
}
#[test]
fn ninety_percent_fraction_sizes_less_ram_than_half() {
let make = |frac: TmpfsFraction| MemoryBudget {
uncompressed_initramfs_bytes: 1200 * (1 << 20),
compressed_initrd_bytes: 300 * (1 << 20),
kernel_init_size: 30 * (1 << 20),
init_coverage_instrumented: false,
instrumented_reserve_bytes: 0,
tmpfs_fraction: frac,
};
let half = initramfs_min_memory_mib(&make(TmpfsFraction::Half));
let ninety = initramfs_min_memory_mib(&make(TmpfsFraction::NinetyPercent));
assert!(
ninety < half,
"90% tmpfs fraction must size less RAM than 50% \
(ninety={ninety}MiB, half={half}MiB)"
);
assert_eq!(half, 3030, "50% floor: 2774 boot + 256 workload");
assert_eq!(ninety, 1947, "90% floor: 1691 boot + 256 workload");
}
#[cfg(target_arch = "x86_64")]
#[test]
fn read_kernel_init_size_x86_64_reads_offset_0x260() {
use std::io::Write;
let mut f = tempfile::NamedTempFile::new().expect("tempfile");
let pad = vec![0u8; 0x260];
f.write_all(&pad).expect("write pad");
let init_size: u32 = 0x1234_5678;
f.write_all(&init_size.to_le_bytes())
.expect("write init_size");
f.flush().expect("flush");
let got = read_kernel_init_size(f.path()).expect("read init_size");
assert_eq!(got, init_size as u64);
}
#[cfg(target_arch = "x86_64")]
#[test]
fn read_kernel_init_size_x86_64_short_file_errors() {
use std::io::Write;
let mut f = tempfile::NamedTempFile::new().expect("tempfile");
let truncated = vec![0u8; 0x100];
f.write_all(&truncated).expect("write truncated");
f.flush().expect("flush");
let result = read_kernel_init_size(f.path());
assert!(result.is_err(), "truncated file must fail; got: {result:?}",);
}
#[cfg(target_arch = "aarch64")]
#[test]
fn read_kernel_init_size_aarch64_reads_offset_16() {
use std::io::Write;
let mut f = tempfile::NamedTempFile::new().expect("tempfile");
let prefix = [0u8; 16];
f.write_all(&prefix).expect("write prefix");
let image_size: u64 = 0x1234_5678_9abc_def0;
f.write_all(&image_size.to_le_bytes())
.expect("write image_size");
f.flush().expect("flush");
let got = read_kernel_init_size(f.path()).expect("read image_size");
assert_eq!(got, image_size);
}
#[cfg(target_arch = "aarch64")]
#[test]
fn read_kernel_init_size_aarch64_short_file_errors() {
use std::io::Write;
let mut f = tempfile::NamedTempFile::new().expect("tempfile");
let truncated = vec![0u8; 8];
f.write_all(&truncated).expect("write truncated");
f.flush().expect("flush");
let result = read_kernel_init_size(f.path());
assert!(result.is_err(), "truncated file must fail; got: {result:?}",);
}
#[cfg(target_arch = "x86_64")]
#[test]
fn read_kernel_version_x86_64_offset_chase_and_magic_gate() {
use std::io::{Seek, SeekFrom, Write};
let ver_ptr: u16 = 0x0100; let string_off = 0x200u64 + ver_ptr as u64;
let write_image = |magic: &[u8; 4]| {
let mut f = tempfile::NamedTempFile::new().expect("tempfile");
f.write_all(&vec![0u8; (string_off as usize) + 64])
.expect("pad");
f.seek(SeekFrom::Start(0x202)).expect("seek magic");
f.write_all(magic).expect("write magic");
f.seek(SeekFrom::Start(0x20E)).expect("seek ver_ptr");
f.write_all(&ver_ptr.to_le_bytes()).expect("write ver_ptr");
f.seek(SeekFrom::Start(string_off)).expect("seek string");
f.write_all(b"6.18.0-rc1 (builder@host)\0")
.expect("write string");
f.flush().expect("flush");
f
};
let good = write_image(b"HdrS");
assert_eq!(
read_kernel_version(good.path()),
Some((6, 18, 0)),
"valid bzImage offset-chase must parse (6, 18, 0)",
);
let bad = write_image(b"XXXX");
assert_eq!(
read_kernel_version(bad.path()),
None,
"wrong HdrS magic must return None (the safe-50% direction)",
);
}
#[test]
fn read_kernel_version_from_metadata_sidecar_parses_and_guards() {
use std::io::Write;
let dir = tempfile::tempdir().expect("tempdir");
let image = dir.path().join("Image");
let sidecar = dir.path().join("metadata.json");
let write_sidecar = |json: &str| {
let mut f = std::fs::File::create(&sidecar).expect("create sidecar");
f.write_all(json.as_bytes()).expect("write sidecar");
f.flush().expect("flush");
};
write_sidecar(r#"{"version":"6.18.2","arch":"aarch64"}"#);
assert_eq!(
read_kernel_version_from_metadata_sidecar(&image),
Some((6, 18, 2)),
"sidecar version must parse to (6, 18, 2)",
);
write_sidecar(r#"{"arch":"aarch64"}"#);
assert_eq!(
read_kernel_version_from_metadata_sidecar(&image),
None,
"absent version key must return None",
);
write_sidecar(r#"{"version":null}"#);
assert_eq!(read_kernel_version_from_metadata_sidecar(&image), None);
write_sidecar(r#"{"version":"not-a-version"}"#);
assert_eq!(read_kernel_version_from_metadata_sidecar(&image), None);
write_sidecar("{not json");
assert_eq!(read_kernel_version_from_metadata_sidecar(&image), None);
std::fs::remove_file(&sidecar).expect("remove sidecar");
assert_eq!(
read_kernel_version_from_metadata_sidecar(&image),
None,
"missing sidecar must return None (raw --kernel path)",
);
assert_eq!(
read_kernel_version_from_metadata_sidecar(std::path::Path::new("/")),
None,
"root path (no parent) must return None",
);
}
}