From 645083913c7717361bf388eca2f65b96358314b8 Mon Sep 17 00:00:00 2001 From: Aiden McClelland Date: Thu, 8 Jan 2026 00:51:25 -0700 Subject: [PATCH] add start-cli flash-os --- core/src/context/rpc.rs | 2 +- core/src/lib.rs | 22 ++- core/src/os_install/gpt.rs | 297 +++++++++++++++++++++---------------- core/src/os_install/mbr.rs | 199 ++++++++++++++++--------- core/src/os_install/mod.rs | 224 ++++++++++++++++++++++------ 5 files changed, 499 insertions(+), 245 deletions(-) diff --git a/core/src/context/rpc.rs b/core/src/context/rpc.rs index c934459c3..a90cae869 100644 --- a/core/src/context/rpc.rs +++ b/core/src/context/rpc.rs @@ -279,7 +279,7 @@ impl RpcContext { .arg("100000") .invoke(ErrorKind::Filesystem) .await?; - tmp.unmount_and_delete().await?; + // tmp.unmount_and_delete().await?; } BlockDev::new(&sqfs) .mount(NVIDIA_OVERLAY_PATH, ReadOnly) diff --git a/core/src/lib.rs b/core/src/lib.rs index 347d46002..4be4448dd 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -8,7 +8,7 @@ pub use std::env::consts::ARCH; lazy_static::lazy_static! { pub static ref PLATFORM: String = { if let Ok(platform) = std::fs::read_to_string("/usr/lib/startos/PLATFORM.txt") { - platform + platform.trim().to_string() } else { ARCH.to_string() } @@ -18,6 +18,17 @@ lazy_static::lazy_static! { }; } +/// Map a platform string to its architecture +pub fn platform_to_arch(platform: &str) -> &str { + if let Some(arch) = platform.strip_suffix("-nonfree") { + return arch; + } + match platform { + "raspberrypi" | "rockchip64" => "aarch64", + _ => platform, + } +} + mod cap { #![allow(non_upper_case_globals)] @@ -246,6 +257,15 @@ pub fn main_api() -> ParentHandler { if &*PLATFORM != "raspberrypi" { api = api.subcommand("kiosk", kiosk::()); } + #[cfg(target_os = "linux")] + { + api = api.subcommand( + "flash-os", + from_fn_async(os_install::cli_install_os) + .no_display() + .with_about("Flash StartOS to a disk from a squashfs"), + ); + } api } diff --git a/core/src/os_install/gpt.rs b/core/src/os_install/gpt.rs index 9625209e3..3a5e9fb5b 100644 --- a/core/src/os_install/gpt.rs +++ b/core/src/os_install/gpt.rs @@ -1,148 +1,193 @@ -use std::path::Path; +use std::path::{Path, PathBuf}; use gpt::GptConfig; use gpt::disk::LogicalBlockSize; use crate::disk::OsPartitionInfo; -use crate::disk::util::DiskInfo; use crate::os_install::partition_for; use crate::prelude::*; -pub async fn partition(disk: &DiskInfo, overwrite: bool) -> Result { - let (efi, data_part) = { - let disk = disk.clone(); - tokio::task::spawn_blocking(move || { - let use_efi = Path::new("/sys/firmware/efi").exists(); - let mut device = Box::new( - std::fs::File::options() - .read(true) - .write(true) - .open(&disk.logicalname)?, - ); - let (mut gpt, guid_part) = if overwrite { - let mbr = gpt::mbr::ProtectiveMBR::with_lb_size( - u32::try_from((disk.capacity / 512) - 1).unwrap_or(0xFF_FF_FF_FF), - ); - mbr.overwrite_lba0(&mut device)?; - ( - GptConfig::new() - .writable(true) - .logical_block_size(LogicalBlockSize::Lb512) - .create_from_device(device, None)?, - None, - ) - } else { - let gpt = GptConfig::new() - .writable(true) +pub async fn partition( + disk_path: &Path, + capacity: u64, + protect: Option<&Path>, + use_efi: bool, +) -> Result { + // Guard: cannot protect the whole disk + if let Some(p) = protect { + if p == disk_path { + return Err(Error::new( + eyre!( + "Cannot protect the entire disk {}; must specify a partition", + disk_path.display() + ), + crate::ErrorKind::InvalidRequest, + )); + } + } + + let disk_path = disk_path.to_owned(); + let disk_path_clone = disk_path.clone(); + let protect = protect.map(|p| p.to_owned()); + let (efi, data_part) = tokio::task::spawn_blocking(move || { + let disk_path = disk_path_clone; + + let protected_partition_info: Option<(u64, u64, PathBuf)> = + if let Some(ref protect_path) = protect { + let existing_gpt = GptConfig::new() + .writable(false) .logical_block_size(LogicalBlockSize::Lb512) - .open_from_device(device)?; - let mut guid_part = None; - for (idx, part_info) in disk - .partitions + .open_from_device(Box::new( + std::fs::File::options().read(true).open(&disk_path)?, + ))?; + let info = existing_gpt + .partitions() .iter() - .enumerate() - .map(|(idx, x)| (idx + 1, x)) - { - if let Some(entry) = gpt.partitions().get(&(idx as u32)) { - if part_info.guid.is_some() { - if entry.first_lba < if use_efi { 33759266 } else { 33570850 } { - return Err(Error::new( - eyre!("Not enough space before StartOS data"), - crate::ErrorKind::InvalidRequest, - )); - } - guid_part = Some(entry.clone()); - break; - } - } + .find(|(num, _)| partition_for(&disk_path, **num) == *protect_path) + .map(|(_, p)| (p.first_lba, p.last_lba, protect_path.clone())); + if info.is_none() { + return Err(Error::new( + eyre!( + "Protected partition {} not found in GPT on {}", + protect_path.display(), + disk_path.display() + ), + crate::ErrorKind::NotFound, + )); } - (gpt, guid_part) - }; - - gpt.update_partitions(Default::default())?; - - let efi = if use_efi { - gpt.add_partition("efi", 100 * 1024 * 1024, gpt::partition_types::EFI, 0, None)?; - true + info } else { - gpt.add_partition( - "bios-grub", - 8 * 1024 * 1024, - gpt::partition_types::BIOS, - 0, - None, - )?; - false + None }; - gpt.add_partition( - "boot", - 1024 * 1024 * 1024, - gpt::partition_types::LINUX_FS, - 0, - None, - )?; - gpt.add_partition( - "root", - 15 * 1024 * 1024 * 1024, - match crate::ARCH { - "x86_64" => gpt::partition_types::LINUX_ROOT_X64, - "aarch64" => gpt::partition_types::LINUX_ROOT_ARM_64, - _ => gpt::partition_types::LINUX_FS, - }, - 0, - None, - )?; - let mut data_part = None; - if overwrite { - gpt.add_partition( - "data", - gpt.find_free_sectors() - .iter() - .map(|(_, size)| *size * u64::from(*gpt.logical_block_size())) - .max() - .ok_or_else(|| { - Error::new( - eyre!("No free space left on device"), - crate::ErrorKind::BlockDevice, - ) - })?, - gpt::partition_types::LINUX_LVM, - 0, - None, - )?; - data_part = gpt - .partitions() - .last_key_value() - .map(|(num, _)| partition_for(&disk.logicalname, *num)); - } else if let Some(guid_part) = guid_part { - let mut parts = gpt.partitions().clone(); - parts.insert( - gpt.find_next_partition_id().ok_or_else(|| { - Error::new(eyre!("Partition table is full"), ErrorKind::DiskManagement) - })?, - guid_part, - ); - gpt.update_partitions(parts)?; - data_part = gpt - .partitions() - .last_key_value() - .map(|(num, _)| partition_for(&disk.logicalname, *num)); + let mut device = Box::new( + std::fs::File::options() + .read(true) + .write(true) + .open(&disk_path)?, + ); + + let mbr = gpt::mbr::ProtectiveMBR::with_lb_size( + u32::try_from((capacity / 512) - 1).unwrap_or(0xFF_FF_FF_FF), + ); + mbr.overwrite_lba0(&mut device)?; + let mut gpt = GptConfig::new() + .writable(true) + .logical_block_size(LogicalBlockSize::Lb512) + .create_from_device(device, None)?; + + gpt.update_partitions(Default::default())?; + + // Calculate where the OS partitions will end + // EFI/BIOS: 100MB or 8MB, Boot: 1GB, Root: 15GB + let efi_size_sectors = if use_efi { + 100 * 1024 * 1024 / 512 + } else { + 8 * 1024 * 1024 / 512 + }; + let boot_size_sectors = 1024 * 1024 * 1024 / 512; + let root_size_sectors = 15 * 1024 * 1024 * 1024 / 512; + // GPT typically starts partitions at sector 2048 + let os_partitions_end_sector = + 2048 + efi_size_sectors + boot_size_sectors + root_size_sectors; + + // Check if protected partition would be overwritten + if let Some((first_lba, _, ref path)) = protected_partition_info { + if first_lba < os_partitions_end_sector { + return Err(Error::new( + eyre!( + concat!( + "Protected partition {} starts at sector {}", + " which would be overwritten by OS partitions ending at sector {}" + ), + path.display(), + first_lba, + os_partitions_end_sector + ), + crate::ErrorKind::DiskManagement, + )); } + } - gpt.write()?; + let efi = if use_efi { + gpt.add_partition("efi", 100 * 1024 * 1024, gpt::partition_types::EFI, 0, None)?; + true + } else { + gpt.add_partition( + "bios-grub", + 8 * 1024 * 1024, + gpt::partition_types::BIOS, + 0, + None, + )?; + false + }; + gpt.add_partition( + "boot", + 1024 * 1024 * 1024, + gpt::partition_types::LINUX_FS, + 0, + None, + )?; + gpt.add_partition( + "root", + 15 * 1024 * 1024 * 1024, + match crate::ARCH { + "x86_64" => gpt::partition_types::LINUX_ROOT_X64, + "aarch64" => gpt::partition_types::LINUX_ROOT_ARM_64, + _ => gpt::partition_types::LINUX_FS, + }, + 0, + None, + )?; - Ok((efi, data_part)) - }) - .await - .unwrap()? - }; + let data_part = if let Some((first_lba, last_lba, path)) = protected_partition_info { + // Re-create the data partition entry at the same location + let length_lba = last_lba - first_lba + 1; + let next_id = gpt.partitions().keys().max().map(|k| k + 1).unwrap_or(1); + gpt.add_partition_at( + "data", + next_id, + first_lba, + length_lba, + gpt::partition_types::LINUX_LVM, + 0, + )?; + Some(path) + } else { + gpt.add_partition( + "data", + gpt.find_free_sectors() + .iter() + .map(|(_, size)| *size * u64::from(*gpt.logical_block_size())) + .max() + .ok_or_else(|| { + Error::new( + eyre!("No free space left on device"), + crate::ErrorKind::BlockDevice, + ) + })?, + gpt::partition_types::LINUX_LVM, + 0, + None, + )?; + gpt.partitions() + .last_key_value() + .map(|(num, _)| partition_for(&disk_path, *num)) + }; + + gpt.write()?; + + Ok::<_, Error>((efi, data_part)) + }) + .await + .unwrap()?; Ok(OsPartitionInfo { - efi: efi.then(|| partition_for(&disk.logicalname, 1)), - bios: (!efi).then(|| partition_for(&disk.logicalname, 1)), - boot: partition_for(&disk.logicalname, 2), - root: partition_for(&disk.logicalname, 3), + efi: efi.then(|| partition_for(&disk_path, 1)), + bios: (!efi).then(|| partition_for(&disk_path, 1)), + boot: partition_for(&disk_path, 2), + root: partition_for(&disk_path, 3), data: data_part, }) } diff --git a/core/src/os_install/mbr.rs b/core/src/os_install/mbr.rs index ca9e272b4..ea68b4695 100644 --- a/core/src/os_install/mbr.rs +++ b/core/src/os_install/mbr.rs @@ -1,92 +1,149 @@ +use std::path::{Path, PathBuf}; + use color_eyre::eyre::eyre; use mbrman::{CHS, MBR, MBRPartitionEntry}; -use crate::Error; use crate::disk::OsPartitionInfo; -use crate::disk::util::DiskInfo; use crate::os_install::partition_for; +use crate::prelude::*; -pub async fn partition(disk: &DiskInfo, overwrite: bool) -> Result { - let data_part = { - let sectors = (disk.capacity / 512) as u32; - let disk = disk.clone(); - tokio::task::spawn_blocking(move || { - let mut file = std::fs::File::options() - .read(true) - .write(true) - .open(&disk.logicalname)?; - let (mut mbr, guid_part) = if overwrite { - (MBR::new_from(&mut file, 512, rand::random())?, None) - } else { - let mut mbr = MBR::read_from(&mut file, 512)?; - let mut guid_part = None; - for (idx, part_info) in disk - .partitions - .iter() - .enumerate() - .map(|(idx, x)| (idx + 1, x)) - { - if let Some(entry) = mbr.get_mut(idx) { - if part_info.guid.is_some() { - if entry.starting_lba < 33556480 { - return Err(Error::new( - eyre!("Not enough space before embassy data"), - crate::ErrorKind::InvalidRequest, - )); - } - guid_part = Some(std::mem::replace(entry, MBRPartitionEntry::empty())); +pub async fn partition( + disk_path: &Path, + capacity: u64, + protect: Option<&Path>, +) -> Result { + // Guard: cannot protect the whole disk + if let Some(p) = protect { + if p == disk_path { + return Err(Error::new( + eyre!( + "Cannot protect the entire disk {}; must specify a partition", + disk_path.display() + ), + crate::ErrorKind::InvalidRequest, + )); + } + } + + let disk_path = disk_path.to_owned(); + let disk_path_clone = disk_path.clone(); + let protect = protect.map(|p| p.to_owned()); + let sectors = (capacity / 512) as u32; + let data_part = tokio::task::spawn_blocking(move || { + let disk_path = disk_path_clone; + + // If protecting a partition, read its location from the existing MBR + let protected_partition_info: Option<(u32, u32, PathBuf)> = + if let Some(ref protect_path) = protect { + let mut file = std::fs::File::options().read(true).open(&disk_path)?; + let existing_mbr = MBR::read_from(&mut file, 512)?; + // Find the partition matching the protected path (check partitions 1-4) + let info = (1..=4u32) + .find(|&idx| partition_for(&disk_path, idx) == *protect_path) + .and_then(|idx| { + let entry = &existing_mbr[idx as usize]; + if entry.sectors > 0 { + Some((entry.starting_lba, entry.sectors, protect_path.clone())) + } else { + None } - *entry = MBRPartitionEntry::empty(); - } + }); + if info.is_none() { + return Err(Error::new( + eyre!( + "Protected partition {} not found in MBR on {}", + protect_path.display(), + disk_path.display() + ), + crate::ErrorKind::NotFound, + )); } - (mbr, guid_part) + info + } else { + None }; - mbr[1] = MBRPartitionEntry { - boot: 0x80, - first_chs: CHS::empty(), - sys: 0x0b, - last_chs: CHS::empty(), - starting_lba: 2048, - sectors: 2099200 - 2048, - }; - mbr[2] = MBRPartitionEntry { + // MBR partition layout: + // Partition 1 (boot): starts at 2048, ends at 2099200 (sectors: 2097152 = 1GB) + // Partition 2 (root): starts at 2099200, ends at 33556480 (sectors: 31457280 = 15GB) + // OS partitions end at sector 33556480 + let os_partitions_end_sector: u32 = 33556480; + + // Check if protected partition would be overwritten + if let Some((starting_lba, _, ref path)) = protected_partition_info { + if starting_lba < os_partitions_end_sector { + return Err(Error::new( + eyre!( + concat!( + "Protected partition {} starts at sector {}", + " which would be overwritten by OS partitions ending at sector {}" + ), + path.display(), + starting_lba, + os_partitions_end_sector + ), + crate::ErrorKind::DiskManagement, + )); + } + } + + let mut file = std::fs::File::options() + .read(true) + .write(true) + .open(&disk_path)?; + let mut mbr = MBR::new_from(&mut file, 512, rand::random())?; + + mbr[1] = MBRPartitionEntry { + boot: 0x80, + first_chs: CHS::empty(), + sys: 0x0b, + last_chs: CHS::empty(), + starting_lba: 2048, + sectors: 2099200 - 2048, + }; + mbr[2] = MBRPartitionEntry { + boot: 0, + first_chs: CHS::empty(), + sys: 0x83, + last_chs: CHS::empty(), + starting_lba: 2099200, + sectors: 33556480 - 2099200, + }; + + let data_part = if let Some((starting_lba, part_sectors, path)) = protected_partition_info { + // Re-create the data partition entry at the same location + mbr[3] = MBRPartitionEntry { boot: 0, first_chs: CHS::empty(), - sys: 0x83, + sys: 0x8e, last_chs: CHS::empty(), - starting_lba: 2099200, - sectors: 33556480 - 2099200, + starting_lba, + sectors: part_sectors, }; + Some(path) + } else { + mbr[3] = MBRPartitionEntry { + boot: 0, + first_chs: CHS::empty(), + sys: 0x8e, + last_chs: CHS::empty(), + starting_lba: 33556480, + sectors: sectors - 33556480, + }; + Some(partition_for(&disk_path, 3)) + }; + mbr.write_into(&mut file)?; - let mut data_part = true; - if overwrite { - mbr[3] = MBRPartitionEntry { - boot: 0, - first_chs: CHS::empty(), - sys: 0x8e, - last_chs: CHS::empty(), - starting_lba: 33556480, - sectors: sectors - 33556480, - } - } else if let Some(guid_part) = guid_part { - mbr[3] = guid_part; - } else { - data_part = false; - } - mbr.write_into(&mut file)?; - - Ok(data_part) - }) - .await - .unwrap()? - }; + Ok::<_, Error>(data_part) + }) + .await + .unwrap()?; Ok(OsPartitionInfo { efi: None, bios: None, - boot: partition_for(&disk.logicalname, 1), - root: partition_for(&disk.logicalname, 2), - data: data_part.then(|| partition_for(&disk.logicalname, 3)), + boot: partition_for(&disk_path, 1), + root: partition_for(&disk_path, 2), + data: data_part, }) } diff --git a/core/src/os_install/mod.rs b/core/src/os_install/mod.rs index 227343f1d..3f0fec743 100644 --- a/core/src/os_install/mod.rs +++ b/core/src/os_install/mod.rs @@ -6,8 +6,9 @@ use serde::{Deserialize, Serialize}; use tokio::process::Command; use ts_rs::TS; -use crate::context::SetupContext; +use crate::Error; use crate::context::config::ServerConfig; +use crate::context::{CliContext, SetupContext}; use crate::disk::OsPartitionInfo; use crate::disk::mount::filesystem::bind::Bind; use crate::disk::mount::filesystem::block_dev::BlockDev; @@ -15,18 +16,30 @@ use crate::disk::mount::filesystem::efivarfs::EfiVarFs; use crate::disk::mount::filesystem::overlayfs::OverlayFs; use crate::disk::mount::filesystem::{MountType, ReadWrite}; use crate::disk::mount::guard::{GenericMountGuard, MountGuard, TmpMountGuard}; -use crate::disk::util::{DiskInfo, PartitionTable}; +use crate::disk::util::PartitionTable; use crate::prelude::*; use crate::s9pk::merkle_archive::source::multi_cursor_file::MultiCursorFile; use crate::setup::SetupInfo; use crate::util::Invoke; use crate::util::io::{TmpDir, delete_file, open_file, write_file_atomic}; use crate::util::serde::IoFormat; -use crate::{ARCH, Error}; mod gpt; mod mbr; +/// Probe a squashfs image to determine its target architecture +async fn probe_squashfs_arch(squashfs_path: &Path) -> Result { + let output = String::from_utf8( + Command::new("unsquashfs") + .arg("-cat") + .arg(squashfs_path) + .arg("usr/lib/startos/PLATFORM.txt") + .invoke(ErrorKind::ParseSysInfo) + .await?, + )?; + Ok(crate::platform_to_arch(&output.trim()).into()) +} + pub fn partition_for(disk: impl AsRef, idx: u32) -> PathBuf { let disk_path = disk.as_ref(); let (root, leaf) = if let (Some(root), Some(leaf)) = ( @@ -44,18 +57,51 @@ pub fn partition_for(disk: impl AsRef, idx: u32) -> PathBuf { } } -async fn partition(disk: &mut DiskInfo, overwrite: bool) -> Result { - let partition_type = match (overwrite, disk.partition_table) { +async fn partition( + disk_path: &Path, + capacity: u64, + partition_table: Option, + protect: Option<&Path>, + use_efi: bool, +) -> Result { + let partition_type = match (protect.is_none(), partition_table) { (true, _) | (_, None) => PartitionTable::Gpt, (_, Some(t)) => t, }; - disk.partition_table = Some(partition_type); match partition_type { - PartitionTable::Gpt => gpt::partition(disk, overwrite).await, - PartitionTable::Mbr => mbr::partition(disk, overwrite).await, + PartitionTable::Gpt => gpt::partition(disk_path, capacity, protect, use_efi).await, + PartitionTable::Mbr => mbr::partition(disk_path, capacity, protect).await, } } +async fn get_block_device_size(path: impl AsRef) -> Result { + let path = path.as_ref(); + let device_name = path.file_name().and_then(|s| s.to_str()).ok_or_else(|| { + Error::new( + eyre!("Invalid block device path: {}", path.display()), + ErrorKind::BlockDevice, + ) + })?; + let size_path = Path::new("/sys/block").join(device_name).join("size"); + let sectors: u64 = tokio::fs::read_to_string(&size_path) + .await + .with_ctx(|_| { + ( + ErrorKind::BlockDevice, + format!("reading {}", size_path.display()), + ) + })? + .trim() + .parse() + .map_err(|e| { + Error::new( + eyre!("Failed to parse block device size: {}", e), + ErrorKind::BlockDevice, + ) + })?; + Ok(sectors * 512) +} + #[derive(Deserialize, Serialize, Parser, TS)] #[serde(rename_all = "camelCase")] #[command(rename_all = "kebab-case")] @@ -75,36 +121,25 @@ struct DataDrive { wipe: bool, } -pub async fn install_os( - ctx: SetupContext, - InstallOsParams { - os_drive, - data_drive, - }: InstallOsParams, -) -> Result { - let mut disks = crate::disk::util::list(&Default::default()).await?; - let disk = disks - .iter_mut() - .find(|d| &d.logicalname == &os_drive) - .ok_or_else(|| { - Error::new( - eyre!("Unknown disk {}", os_drive.display()), - crate::ErrorKind::DiskManagement, - ) - })?; +pub struct InstallOsResult { + pub part_info: OsPartitionInfo, + pub rootfs: TmpMountGuard, +} - let overwrite = if let Some(data_drive) = &data_drive { - data_drive.wipe - || ((disk.guid.is_none() || disk.logicalname != data_drive.logicalname) - && disk - .partitions - .iter() - .all(|p| p.guid.is_none() || p.logicalname != data_drive.logicalname)) - } else { - true - }; +pub async fn install_os_to( + squashfs_path: impl AsRef, + disk_path: impl AsRef, + capacity: u64, + partition_table: Option, + protect: Option>, + arch: &str, + use_efi: bool, +) -> Result { + let squashfs_path = squashfs_path.as_ref(); + let disk_path = disk_path.as_ref(); + let protect = protect.as_ref().map(|p| p.as_ref()); - let part_info = partition(disk, overwrite).await?; + let part_info = partition(disk_path, capacity, partition_table, protect, use_efi).await?; if let Some(efi) = &part_info.efi { Command::new("mkfs.vfat") @@ -128,7 +163,7 @@ pub async fn install_os( .invoke(crate::ErrorKind::DiskManagement) .await?; - if !overwrite { + if protect.is_some() { if let Ok(guard) = TmpMountGuard::mount(&BlockDev::new(part_info.root.clone()), MountType::ReadWrite).await { @@ -189,13 +224,13 @@ pub async fn install_os( tokio::fs::create_dir_all(&images_path).await?; let image_path = images_path .join(hex::encode( - &MultiCursorFile::from(open_file("/run/live/medium/live/filesystem.squashfs").await?) + &MultiCursorFile::from(open_file(squashfs_path).await?) .blake3_mmap() .await? .as_bytes()[..16], )) .with_extension("rootfs"); - tokio::fs::copy("/run/live/medium/live/filesystem.squashfs", &image_path).await?; + tokio::fs::copy(squashfs_path, &image_path).await?; // TODO: check hash of fs let unsquash_target = TmpDir::new().await?; let bootfs = MountGuard::mount( @@ -209,7 +244,7 @@ pub async fn install_os( .arg("-f") .arg("-d") .arg(&*unsquash_target) - .arg("/run/live/medium/live/filesystem.squashfs") + .arg(squashfs_path) .arg("boot") .invoke(crate::ErrorKind::Filesystem) .await?; @@ -230,8 +265,6 @@ pub async fn install_os( })?, ) .await?; - ctx.config - .mutate(|c| c.os_partitions = Some(part_info.clone())); let lower = TmpMountGuard::mount(&BlockDev::new(&image_path), MountType::ReadOnly).await?; let work = config_path.join("work"); @@ -313,13 +346,13 @@ pub async fn install_os( let mut install = Command::new("chroot"); install.arg(overlay.path()).arg("grub-install"); - if tokio::fs::metadata("/sys/firmware/efi").await.is_err() { - match ARCH { + if !use_efi { + match arch { "x86_64" => install.arg("--target=i386-pc"), _ => &mut install, }; } else { - match ARCH { + match arch { "x86_64" => install.arg("--target=x86_64-efi"), "aarch64" => install.arg("--target=arm64-efi"), "riscv64" => install.arg("--target=riscv64-efi"), @@ -327,7 +360,7 @@ pub async fn install_os( }; } install - .arg(&disk.logicalname) + .arg(disk_path) .invoke(crate::ErrorKind::Grub) .await?; @@ -352,6 +385,62 @@ pub async fn install_os( tokio::fs::remove_dir_all(&work).await?; lower.unmount().await?; + Ok(InstallOsResult { part_info, rootfs }) +} + +pub async fn install_os( + ctx: SetupContext, + InstallOsParams { + os_drive, + data_drive, + }: InstallOsParams, +) -> Result { + let mut disks = crate::disk::util::list(&Default::default()).await?; + let disk = disks + .iter_mut() + .find(|d| &d.logicalname == &os_drive) + .ok_or_else(|| { + Error::new( + eyre!("Unknown disk {}", os_drive.display()), + crate::ErrorKind::DiskManagement, + ) + })?; + + let protect: Option = data_drive.as_ref().and_then(|dd| { + if dd.wipe { + return None; + } + if disk.guid.as_ref().map_or(false, |g| { + g.starts_with("EMBASSY_") || g.starts_with("STARTOS_") + }) && disk.logicalname == dd.logicalname + { + return Some(disk.logicalname.clone()); + } + disk.partitions + .iter() + .find(|p| { + p.guid.as_ref().map_or(false, |g| { + g.starts_with("EMBASSY_") || g.starts_with("STARTOS_") + }) + }) + .map(|p| p.logicalname.clone()) + }); + + let use_efi = tokio::fs::metadata("/sys/firmware/efi").await.is_ok(); + let InstallOsResult { part_info, rootfs } = install_os_to( + "/run/live/medium/live/filesystem.squashfs", + &disk.logicalname, + disk.capacity, + disk.partition_table, + protect.as_ref(), + crate::ARCH, + use_efi, + ) + .await?; + + ctx.config + .mutate(|c| c.os_partitions = Some(part_info.clone())); + let mut setup_info = SetupInfo::default(); if let Some(data_drive) = data_drive { @@ -396,3 +485,46 @@ pub async fn install_os( Ok(setup_info) } + +#[derive(Deserialize, Serialize, Parser)] +#[serde(rename_all = "camelCase")] +#[command(rename_all = "kebab-case")] +pub struct CliInstallOsParams { + #[arg(help = "Path to the squashfs image to install")] + squashfs: PathBuf, + #[arg(help = "Target disk to install to (e.g., /dev/sda or /dev/loop0)")] + disk: PathBuf, + #[arg(long, help = "Use EFI boot (default: true for GPT disks)")] + efi: Option, +} + +pub async fn cli_install_os( + _ctx: CliContext, + CliInstallOsParams { + squashfs, + disk, + efi, + }: CliInstallOsParams, +) -> Result { + let capacity = get_block_device_size(&disk).await?; + let partition_table = crate::disk::util::get_partition_table(&disk).await?; + + let arch = probe_squashfs_arch(&squashfs).await?; + + let use_efi = efi.unwrap_or_else(|| !matches!(partition_table, Some(PartitionTable::Mbr))); + + let InstallOsResult { part_info, rootfs } = install_os_to( + &squashfs, + &disk, + capacity, + partition_table, + None::<&str>, + &*arch, + use_efi, + ) + .await?; + + rootfs.unmount().await?; + + Ok(part_info) +}