From 7f0bfa5c30aaac0cc258d67b280e2897b75fcbeb Mon Sep 17 00:00:00 2001
From: Raphael Coeffic <raphaelcoeffic@users.noreply.github.com>
Date: Tue, 5 Nov 2024 13:20:40 +0100
Subject: [PATCH] feat: enable building base image with qemu-user-static (#3)

Co-authored-by: raphaelcoeffic <1050031+raphaelcoeffic@users.noreply.github.com>
---
 .github/workflows/rust.yml | 17 +++++++++++++----
 image_builder/src/lib.rs   | 28 ++++++++++++++++++++--------
 src/bin/build-img.rs       | 26 ++++++++++++++++++++++++--
 3 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index f2b47d5..f2dac5e 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -12,13 +12,17 @@ env:
 
 jobs:
   build:
-    name: Build - ${{ matrix.platform.os-name }}
+    name: Build - ${{ matrix.platform.arch }}
     strategy:
       matrix:
         platform:
-          - os-name: Linux-x86_64
-            runs-on: ubuntu-latest
+          - arch: x86_64
             target: x86_64-unknown-linux-musl
+            runs-on: ubuntu-latest
+
+          - arch: aarch64
+            target: aarch64-unknown-linux-musl
+            runs-on: ubuntu-latest
 
     runs-on: ${{ matrix.platform.runs-on }}
     steps:
@@ -26,8 +30,13 @@ jobs:
       uses: actions/checkout@v4
 
     - name: Build base image
+      env:
+        ARCH: ${{ matrix.platform.arch }}
       run: |
-        cargo run --bin build-img
+        if [ "$ARCH" != "$(uname -m)" ]; then
+          sudo apt-get install qemu-user-static
+        fi
+        cargo run --bin build-img -- -a $ARCH
         cargo clean
 
     - name: Upload artifacts - base image
diff --git a/image_builder/src/lib.rs b/image_builder/src/lib.rs
index a60b693..13c5b98 100644
--- a/image_builder/src/lib.rs
+++ b/image_builder/src/lib.rs
@@ -92,7 +92,11 @@ impl BaseImageBuilder {
     }
 
     pub fn build_base(&self) -> Result<()> {
-        install_nix(&self.nix_dir)?;
+        self.build_base_with_arch(ARCH)
+    }
+
+    pub fn build_base_with_arch(&self, arch: &str) -> Result<()> {
+        install_nix(arch, &self.nix_dir)?;
         log::info!("building base image");
 
         let tmp = tempdir()?;
@@ -208,10 +212,17 @@ impl BaseImageBuilder {
             tar_cmd.args(["-I", "xz -T0"]);
         }
 
+        // prefer native tar over emulated one
+        let current_path = std::env::var_os("PATH")
+            .filter(|path| !path.is_empty())
+            .map(|path| path.to_string_lossy().into_owned() + ":")
+            .unwrap_or_default();
+        let path_env = format!("{current_path}/nix/.base/bin");
+
         tar_cmd
             .args([".bin", ".base", "etc", "var/nix"])
             .args(nix_set.union(&base_set).map(|p| "store/".to_owned() + p))
-            .env("PATH", "/nix/.base/bin");
+            .env("PATH", path_env);
 
         let tar_status = tar_cmd.status()?;
 
@@ -227,9 +238,9 @@ impl BaseImageBuilder {
     }
 }
 
-fn nix_installer_url(version: &str) -> String {
+fn nix_installer_url(version: &str, arch: &str) -> String {
     const NIX_BASE_URL: &str = "https://releases.nixos.org/nix";
-    format!("{NIX_BASE_URL}/nix-{version}/nix-{version}-{ARCH}-linux.tar.xz")
+    format!("{NIX_BASE_URL}/nix-{version}/nix-{version}-{arch}-linux.tar.xz")
 }
 
 fn write_nix_paths(nix_dir: &Path) -> Result<(), io::Error> {
@@ -317,6 +328,7 @@ pub fn progress_bar(len: u64) -> ProgressBar {
 /// Download and install Nix.
 fn download_and_install_nix(
     version: &str,
+    arch: &str,
     url: &str,
     dest: &Path,
 ) -> Result<()> {
@@ -331,7 +343,7 @@ fn download_and_install_nix(
     fs::create_dir_all(&store_dir)?;
 
     // unpack files
-    let tar_prefix = format!("nix-{version}-{ARCH}-linux");
+    let tar_prefix = format!("nix-{version}-{arch}-linux");
     for file in ar.entries()? {
         let mut f = file?;
         let fpath = f.path()?;
@@ -370,7 +382,7 @@ fn find_nix(store_dir: &Path, version: &str) -> Result<PathBuf> {
 }
 
 /// Install Nix into `dest`
-fn install_nix<P>(dest: P) -> Result<()>
+fn install_nix<P>(arch: &str, dest: P) -> Result<()>
 where
     P: AsRef<Path>,
 {
@@ -385,8 +397,8 @@ where
     let nix_store = dest.join("store");
     if !nix_store.exists() {
         log::info!("installing Nix in {}", dest.display());
-        let nix_url = nix_installer_url(NIX_VERSION);
-        download_and_install_nix(NIX_VERSION, &nix_url, dest)?;
+        let nix_url = nix_installer_url(NIX_VERSION, arch);
+        download_and_install_nix(NIX_VERSION, arch, &nix_url, dest)?;
     }
 
     let nix_bin = dest.join(".bin");
diff --git a/src/bin/build-img.rs b/src/bin/build-img.rs
index 497782f..d6445d8 100644
--- a/src/bin/build-img.rs
+++ b/src/bin/build-img.rs
@@ -3,7 +3,7 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 use clap::Parser;
 
 use image_builder::*;
@@ -20,6 +20,10 @@ struct Args {
     #[arg(short, long, env)]
     flake_dir: Option<PathBuf>,
 
+    /// Architecture
+    #[arg(short, long, env)]
+    arch: Option<String>,
+
     /// Compress base image
     #[arg(short, long, env)]
     uncompressed: bool,
@@ -71,6 +75,14 @@ impl Drop for BaseDir {
     }
 }
 
+fn is_native_arch(arch: &str) -> bool {
+    arch == std::env::consts::ARCH
+}
+
+fn is_qemu_supported_arch(arch: &str) -> bool {
+    fs::exists(format!("/proc/sys/fs/binfmt_misc/qemu-{arch}")).is_ok_and(|x| x)
+}
+
 fn init_logging() {
     env_logger::Builder::new()
         .filter_level(log::LevelFilter::Info)
@@ -92,5 +104,15 @@ fn main() -> Result<()> {
         base_builder.flake_dir(flake_dir);
     }
 
-    base_builder.build_base()
+    if let Some(arch) = args.arch {
+        if !is_native_arch(&arch) && !is_qemu_supported_arch(&arch) {
+            bail!(
+                "{arch} does not seem to be supported. \
+                 Try installing 'qemu-user-static' to enable support."
+            )
+        }
+        base_builder.build_base_with_arch(&arch)
+    } else {
+        base_builder.build_base()
+    }
 }