add efi-size-check

Signed-off-by: Nianyu Shen <[email protected]>

Earthfile

@@ -59,6 +59,10 @@ ARG CMDLINE="stylus.registration"
 ARG BRANDING="Palette eXtended Kubernetes Edge"
 ARG ETCD_VERSION="v3.5.13"
 
+# EFI size check
+ARG EFI_MAX_SIZE=2048
+ARG EFI_IMG_SIZE=2200
+
 # internal variables
 ARG GOLANG_VERSION=1.22
 
@@ -240,7 +244,7 @@ uki-iso:
     SAVE ARTIFACT /build/* AS LOCAL ./build/
 
 uki-provider-image:
-    FROM --platform=linux/${ARCH} gcr.io/spectro-images-public/ubuntu-systemd:22.04
+    FROM --platform=linux/${ARCH} +ubuntu-systemd
     WORKDIR /
     COPY +luet/luet /usr/bin/luet
     COPY +kairos-agent/kairos-agent /usr/bin/kairos-agent
@@ -433,7 +437,7 @@ download-sbctl:
     SAVE ARTIFACT /usr/bin/sbctl
 
 uki-byok:
-    FROM --platform=linux/${ARCH} ubuntu:latest
+    FROM +ubuntu-systemd
 
     RUN apt-get update && apt-get install -y efitools curl
     COPY +download-sbctl/sbctl /usr/bin/sbctl
@@ -795,6 +799,48 @@ internal-slink:
 
     SAVE ARTIFACT slink
 
+rust-deps:
+    FROM rust:1.78-bookworm
+    RUN apt-get update -qq
+    RUN apt-get install --no-install-recommends -qq autoconf autotools-dev libtool-bin clang cmake bsdmainutils
+
+build-efi-size-check:
+    FROM +rust-deps
+
+    WORKDIR /build
+    COPY --keep-ts efi-size-check efi-size-check
+
+    WORKDIR /build/efi-size-check
+    RUN cargo build --target x86_64-unknown-uefi
+
+    SAVE ARTIFACT target/x86_64-unknown-uefi/debug/efi-size-check.efi
+
+iso-efi-size-check:
+    FROM +ubuntu-systemd
+
+    RUN apt-get update
+    RUN apt-get install -y mtools xorriso
+
+    WORKDIR /build
+
+    COPY +build-efi-size-check/efi-size-check.efi /build/efi-size-check.efi
+    RUN mkdir -p esp
+    RUN dd if=/dev/urandom of=esp/ABC bs=1M count=$EFI_MAX_SIZE
+    RUN dd if=/dev/zero of=fat.img bs=1M count=$EFI_IMG_SIZE
+    RUN mformat -i fat.img -F ::
+    RUN mmd -i fat.img ::/EFI
+    RUN mmd -i fat.img ::/EFI/BOOT
+    RUN mcopy -i fat.img efi-size-check.efi ::/EFI/BOOT/BOOTX64.EFI
+    RUN mcopy -i fat.img esp/ABC ::
+    RUN mkdir -p iso
+    RUN cp fat.img iso
+    RUN xorriso -as mkisofs -e fat.img -no-emul-boot -o efi-size-check.iso iso
+
+    SAVE ARTIFACT efi-size-check.iso AS LOCAL ./build/
+
+ubuntu-systemd:
+    FROM $SPECTRO_PUB_REPO/ubuntu-systemd:22.04
+
 OS_RELEASE:
     COMMAND
     ARG OS_ID=${OS_DISTRIBUTION}

efi-size-check/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "efi-size-check"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+log = "0.4.21"
+uefi = { git = "https://github.com/nianyush/uefi-rs.git", branch = "unchunked" }
+uefi-services = { git = "https://github.com/nianyush/uefi-rs.git", branch = "unchunked" }

efi-size-check/rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+targets = ["aarch64-unknown-uefi", "i686-unknown-uefi", "x86_64-unknown-uefi"]

efi-size-check/src/main.rs

@@ -0,0 +1,128 @@
+#![no_main]
+#![no_std]
+
+extern crate alloc;
+extern crate uefi;
+extern crate uefi_services;
+
+use alloc::format;
+use alloc::string::String;
+use alloc::vec::Vec;
+use log::info;
+use uefi::prelude::*;
+use uefi::proto::media::file::{File, FileAttribute, FileMode};
+use uefi::table::boot::{AllocateType, MemoryType};
+use uefi::CStr16;
+
+const KB: usize = 1024;
+const MB: usize = KB * 1024;
+const GB: usize = MB * 1024;
+
+const MAX_SIZE: usize = 2 * GB;
+const INCREMENT: usize = 100 * MB;
+
+#[entry]
+fn main(image: Handle, mut st: SystemTable<Boot>) -> Status {
+    match uefi_services::init(&mut st) {
+        Ok(_) => {}
+        Err(err) => {
+            panic!("Failed to initialize uefi services: {:?}", err);
+        }
+    }
+    info!("Starting EFI size checker...");
+
+    let bs = st.boot_services();
+
+    let mut sfs = bs.get_image_file_system(image).unwrap();
+    let mut root = sfs.open_volume().unwrap();
+
+    let mut buf = [0; 4];
+    let file_name: &CStr16 = CStr16::from_str_with_buf("ABC", &mut buf).unwrap(); // Ensure this file exists on your ESP with a large size for testing
+    let mut buffer_size = MB * 100; // Start with 100 MiB
+    info!("start to do file read check...");
+
+    loop {
+        if buffer_size > MAX_SIZE {
+            info!("Max buffer size reached: {}", format_bytes(MAX_SIZE));
+            break;
+        }
+        let mut file = match root.open(file_name, FileMode::Read, FileAttribute::empty()) {
+            Ok(f) => f.into_regular_file().unwrap(),
+            Err(_) => {
+                panic!("Failed to open file: {:?}", file_name);
+            }
+        };
+
+        info!("Reading {} bytes into buffer", format_bytes(buffer_size));
+
+        let mut buffer = Vec::with_capacity(buffer_size);
+        unsafe {
+            buffer.set_len(buffer_size);
+        } // Unsafe due to uninitialized memory
+
+        match file.read_unchunked(&mut buffer) {
+            Ok(_) => {
+                info!(
+                    "Successfully read {} into buffer",
+                    format_bytes(buffer_size)
+                );
+                buffer_size += INCREMENT;
+            }
+            Err(e) => {
+                panic!("Failed to read into a {} byte buffer: {:?}", buffer_size, e);
+            }
+        }
+    }
+
+    mem_test(bs);
+
+    bs.stall(100_000_000);
+    Status::SUCCESS
+}
+
+fn mem_test(bs: &BootServices) {
+    let mut buffer_size = MB * 100; // Start with 100 MiB
+    info!("start to do memory allocation check...");
+
+    loop {
+        if buffer_size > MAX_SIZE {
+            info!("Max buffer size reached: {}", format_bytes(MAX_SIZE));
+            break;
+        }
+        match bs.allocate_pages(
+            AllocateType::AnyPages,
+            MemoryType::ACPI_RECLAIM,
+            buffer_size,
+        ) {
+            Ok(addr) => {
+                info!("Successfully allocated {} mem", buffer_size);
+                unsafe {
+                    match bs.free_pages(addr, buffer_size) {
+                        Ok(_) => {
+                            info!("Successfully freed {} mem", buffer_size)
+                        }
+                        Err(err) => {
+                            panic!("Failed to free {} mem: {:?}", buffer_size, err);
+                        }
+                    }
+                }
+            }
+            Err(err) => {
+                panic!("Failed to allocate {} mem: {:?}", buffer_size, err);
+            }
+        }
+        buffer_size += INCREMENT;
+    }
+}
+
+fn format_bytes(size: usize) -> String {
+    if size >= GB {
+        format!("{:.2} GB", size as f64 / GB as f64)
+    } else if size >= MB {
+        format!("{:.2} MB", size as f64 / MB as f64)
+    } else if size >= KB {
+        format!("{:.2} KB", size as f64 / KB as f64)
+    } else {
+        format!("{} bytes", size)
+    }
+}