gadget: add --pid1 mode

apps/smoo-gadget-cli/src/main.rs

@@ -14,17 +14,21 @@ use smoo_proto::{Ident, OpCode, Request, Response, SMOO_STATUS_REQUEST, SMOO_STA
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     convert::Infallible,
+    ffi::{CString, OsStr},
     fs::File,
     io,
     net::SocketAddr,
     os::fd::{FromRawFd, IntoRawFd, OwnedFd},
+    os::unix::fs::FileTypeExt,
+    os::unix::process::CommandExt,
     path::{Path, PathBuf},
     sync::{
         atomic::{AtomicU64, Ordering},
         Arc,
     },
     time::{Duration, Instant},
 };
+use std::io::Write;
 use tokio::{
     io::AsyncReadExt,
     signal,
@@ -92,6 +96,12 @@ struct Args {
     /// Expose Prometheus metrics on this TCP port (0 disables).
     #[arg(long, default_value_t = 0)]
     metrics_port: u16,
+    /// Run as the initramfs PID1 wrapper (auto-enabled when argv0 == /init).
+    #[arg(long)]
+    pid1: bool,
+    /// Internal flag for the forked gadget child.
+    #[arg(long, hide = true)]
+    pid1_child: bool,
 }
 
 #[derive(Clone, Copy, Debug, ValueEnum)]
@@ -120,7 +130,14 @@ async fn main() -> Result<()> {
         .with(tracing_subscriber::fmt::layer())
         .init();
 
-    let args = Args::parse();
+    let mut args = Args::parse();
+    let argv0 = std::env::args().next().unwrap_or_default();
+    let auto_pid1 = argv0 == "/init";
+    if (args.pid1 || auto_pid1) && !args.pid1_child {
+        args.pid1 = true;
+        run_pid1().context("pid1 initramfs flow")?;
+        return Ok(());
+    }
     let metrics_shutdown = CancellationToken::new();
     let metrics_task = spawn_metrics_listener(args.metrics_port, metrics_shutdown.clone())?;
     let mut ublk = SmooUblk::new().context("init ublk")?;
@@ -226,6 +243,284 @@ async fn main() -> Result<()> {
     result
 }
 
+fn run_pid1() -> Result<()> {
+    ensure!(
+        unsafe { libc::getpid() } == 1,
+        "pid1 mode requires PID 1"
+    );
+
+    klog("pid1: starting smoo initramfs flow");
+    mount_fs(Some("proc"), "/proc", Some("proc"), 0, None).ok();
+    mount_fs(Some("sysfs"), "/sys", Some("sysfs"), 0, None).ok();
+    mount_fs(Some("devtmpfs"), "/dev", Some("devtmpfs"), 0, None).ok();
+    mount_fs(Some("tmpfs"), "/run", Some("tmpfs"), 0, None).ok();
+
+    let default_modules = [
+        "configfs",
+        "ublk",
+        "ublk_drv",
+        "overlay",
+        "erofs",
+        "libcomposite",
+        "usb_f_fs",
+    ];
+    let modules = load_modules_from_dir("/etc/modules-load.d")
+        .filter(|list| !list.is_empty())
+        .unwrap_or_else(|| default_modules.iter().map(|s| s.to_string()).collect());
+    for module in modules {
+        modprobe(&module);
+    }
+
+    std::fs::create_dir_all("/sys/kernel/config").ok();
+    mount_fs(Some("configfs"), "/sys/kernel/config", Some("configfs"), 0, None).ok();
+
+    let udc_wait_secs = 15;
+    if !wait_for_udc(Duration::from_secs(udc_wait_secs))? {
+        return Err(anyhow!("UDC not ready after {udc_wait_secs}s"));
+    }
+
+    let mut child = spawn_gadget_child().context("spawn gadget child")?;
+    let ublk_dev = "/dev/ublkb0";
+    let wait_secs = 30;
+    if !wait_for_block_device(ublk_dev, Duration::from_secs(wait_secs), &mut child)? {
+        return Err(anyhow!("timed out waiting for {ublk_dev}"));
+    }
+    klog(&format!("pid1: found ublk device {ublk_dev}"));
+
+    std::fs::create_dir_all("/lower").ok();
+    std::fs::create_dir_all("/upper").ok();
+    std::fs::create_dir_all("/newroot").ok();
+
+    mount_fs(
+        Some(ublk_dev),
+        "/lower",
+        Some("erofs"),
+        libc::MS_RDONLY as libc::c_ulong,
+        None,
+    )
+    .context("mount erofs lower")?;
+    mount_fs(Some("tmpfs"), "/upper", Some("tmpfs"), 0, None).context("mount tmpfs upper")?;
+    std::fs::create_dir_all("/upper/upper").ok();
+    std::fs::create_dir_all("/upper/work").ok();
+    if !filesystem_available("overlay")? {
+        return Err(anyhow!("overlayfs not available in kernel"));
+    }
+    mount_fs(
+        Some("overlay"),
+        "/newroot",
+        Some("overlay"),
+        0,
+        Some("lowerdir=/lower,upperdir=/upper/upper,workdir=/upper/work"),
+    )
+    .context("mount overlay root")?;
+
+    // Avoid EINVAL from pivot_root on shared mount trees.
+    mount_fs(None, "/", None, libc::MS_PRIVATE | libc::MS_REC, None)
+        .context("make / private")?;
+
+    if matches!(cmdline_value("smoo.break").as_deref(), Some("1")) {
+        debug_shell("smoo.break requested")?;
+    }
+
+    std::fs::create_dir_all("/newroot/proc").ok();
+    std::fs::create_dir_all("/newroot/sys").ok();
+    std::fs::create_dir_all("/newroot/dev").ok();
+    std::fs::create_dir_all("/newroot/run").ok();
+
+    move_mount("/proc", "/newroot/proc").ok();
+    move_mount("/sys", "/newroot/sys").ok();
+    move_mount("/dev", "/newroot/dev").ok();
+    move_mount("/run", "/newroot/run").ok();
+
+    std::env::set_current_dir("/newroot").ok();
+    mount_fs(Some("/newroot"), "/", None, libc::MS_MOVE as libc::c_ulong, None)
+        .context("move newroot to /")?;
+    chroot_to(".").context("chroot to new root")?;
+    std::env::set_current_dir("/").ok();
+
+    klog("pid1: exec /sbin/init");
+    let err = std::process::Command::new("/sbin/init").exec();
+    Err(anyhow!("exec /sbin/init failed: {err}"))
+}
+
+fn klog(msg: &str) {
+    let line = format!("<6>[smoo-pid1] {msg}\n");
+    if let Ok(mut file) = File::options().write(true).open("/dev/kmsg") {
+        let _ = file.write_all(line.as_bytes());
+    }
+    if let Ok(mut file) = File::options().write(true).open("/dev/console") {
+        let _ = file.write_all(format!("[smoo-pid1] {msg}\n").as_bytes());
+    }
+}
+
+fn cmdline_value(key: &str) -> Option<String> {
+    let data = std::fs::read_to_string("/proc/cmdline").ok()?;
+    for token in data.split_whitespace() {
+        if let Some(value) = token.strip_prefix(&format!("{key}=")) {
+            return Some(value.to_string());
+        }
+    }
+    None
+}
+
+fn modprobe(module: &str) {
+    let status = std::process::Command::new("/sbin/modprobe")
+        .arg(module)
+        .status();
+    if let Ok(status) = status {
+        if !status.success() {
+            klog(&format!("modprobe {module} exited {status}"));
+        }
+    } else {
+        klog(&format!("modprobe {module} failed to exec"));
+    }
+}
+
+fn load_modules_from_dir(path: &str) -> Option<Vec<String>> {
+    let mut modules = Vec::new();
+    let entries = std::fs::read_dir(path).ok()?;
+    for entry in entries.filter_map(Result::ok) {
+        let name = entry.file_name();
+        if name.to_string_lossy().ends_with(".conf") {
+            if let Ok(contents) = std::fs::read_to_string(entry.path()) {
+                for line in contents.lines() {
+                    let trimmed = line.trim();
+                    if trimmed.is_empty() || trimmed.starts_with('#') {
+                        continue;
+                    }
+                    modules.push(trimmed.to_string());
+                }
+            }
+        }
+    }
+    Some(modules)
+}
+
+fn debug_shell(reason: &str) -> Result<()> {
+    klog(&format!("pid1: dropping to shell ({reason})"));
+    for dev in ["/dev/ttyMSM0", "/dev/console"] {
+        if let Ok(meta) = std::fs::metadata(dev) {
+            if meta.file_type().is_char_device() {
+                let file = std::fs::OpenOptions::new()
+                    .read(true)
+                    .write(true)
+                    .open(dev)
+                    .with_context(|| format!("open {dev}"))?;
+                let _ = unsafe { libc::setsid() };
+                let err = std::process::Command::new("/bin/sh")
+                    .arg("-i")
+                    .stdin(file.try_clone()?)
+                    .stdout(file.try_clone()?)
+                    .stderr(file)
+                    .exec();
+                return Err(anyhow!("exec /bin/sh failed: {err}"));
+            }
+        }
+    }
+    Err(anyhow!("no console device available for debug shell"))
+}
+
+fn wait_for_udc(timeout: Duration) -> Result<bool> {
+    let start = Instant::now();
+    loop {
+        if Path::new("/sys/class/udc").exists() {
+            if let Ok(entries) = std::fs::read_dir("/sys/class/udc") {
+                if entries.filter_map(Result::ok).next().is_some() {
+                    return Ok(true);
+                }
+            }
+        }
+        if start.elapsed() >= timeout {
+            return Ok(false);
+        }
+        std::thread::sleep(Duration::from_secs(1));
+    }
+}
+
+fn wait_for_block_device(path: &str, timeout: Duration, child: &mut std::process::Child) -> Result<bool> {
+    let start = Instant::now();
+    loop {
+        if let Ok(meta) = std::fs::metadata(path) {
+            if meta.file_type().is_block_device() {
+                return Ok(true);
+            }
+        }
+        if let Ok(Some(status)) = child.try_wait() {
+            return Err(anyhow!("gadget child exited: {status}"));
+        }
+        if start.elapsed() >= timeout {
+            return Ok(false);
+        }
+        std::thread::sleep(Duration::from_secs(1));
+    }
+}
+
+fn spawn_gadget_child() -> Result<std::process::Child> {
+    let exe = std::env::current_exe().context("locate self")?;
+    let mut child_args: Vec<_> = std::env::args_os().collect();
+    child_args.retain(|arg| arg != OsStr::new("--pid1") && arg != OsStr::new("--pid1-child"));
+    child_args.push(OsStr::new("--pid1-child").to_os_string());
+    let mut cmd = std::process::Command::new(exe);
+    cmd.args(child_args.iter().skip(1));
+    cmd.stdin(std::process::Stdio::null());
+    cmd.spawn().context("spawn gadget process")
+}
+
+fn filesystem_available(name: &str) -> Result<bool> {
+    let data = std::fs::read_to_string("/proc/filesystems").context("read /proc/filesystems")?;
+    Ok(data.lines().any(|line| line.split_whitespace().last() == Some(name)))
+}
+
+fn move_mount(src: &str, dst: &str) -> Result<()> {
+    mount_fs(
+        Some(src),
+        dst,
+        None,
+        libc::MS_MOVE as libc::c_ulong,
+        None,
+    )
+    .with_context(|| format!("move mount {src} -> {dst}"))
+}
+
+fn chroot_to(path: &str) -> Result<()> {
+    let path = CString::new(path)?;
+    let res = unsafe { libc::chroot(path.as_ptr()) };
+    if res != 0 {
+        return Err(io::Error::last_os_error()).context("chroot syscall failed");
+    }
+    Ok(())
+}
+
+fn mount_fs(
+    source: Option<&str>,
+    target: &str,
+    fstype: Option<&str>,
+    flags: libc::c_ulong,
+    data: Option<&str>,
+) -> Result<()> {
+    let target = CString::new(target)?;
+    let source = source.map(CString::new).transpose()?;
+    let fstype = fstype.map(CString::new).transpose()?;
+    let data = data.map(CString::new).transpose()?;
+    let data_ptr = data
+        .as_ref()
+        .map(|s| s.as_ptr() as *const libc::c_void)
+        .unwrap_or(std::ptr::null());
+    let res = unsafe {
+        libc::mount(
+            source.as_ref().map(|s| s.as_ptr()).unwrap_or(std::ptr::null()),
+            target.as_ptr(),
+            fstype.as_ref().map(|s| s.as_ptr()).unwrap_or(std::ptr::null()),
+            flags,
+            data_ptr,
+        )
+    };
+    if res != 0 {
+        return Err(io::Error::last_os_error()).context("mount failed");
+    }
+    Ok(())
+}
+
 fn spawn_metrics_listener(
     port: u16,
     shutdown: CancellationToken,