mkwim 0.1.0

// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (c) 2026 Jarkko Sakkinen

//! Output setup: GPT partitioning and FAT32 formatting for either a block device
//! or a raw disk image.
//!
//! Block-device output is mounted through the kernel. Raw images are populated
//! directly through [`fatfs`], so creating an image does not require root or a
//! loop device.

use std::fs::{self, File, OpenOptions};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::os::unix::{ffi::OsStrExt, fs::FileTypeExt};
use std::path::{Path, PathBuf};

use anyhow::{Context, Result, anyhow};
use fatfs::{FatType, FileSystem, FormatVolumeOptions, FsOptions, format_volume};
use gptman::linux::{get_sector_size, reread_partition_table};
use gptman::{GPT, GPTPartitionEntry};
use nix::mount::{MsFlags, mount, umount};

/// EFI System Partition type GUID: `C12A7328-F81F-11D2-BA4B-00A0C93EC93B`,
/// stored in the mixed-endian wire format GPT expects.
const ESP_TYPE_GUID: [u8; 16] = [
    0x28, 0x73, 0x2A, 0xC1, 0x1F, 0xF8, 0xD2, 0x11, 0xBA, 0x4B, 0x00, 0xA0, 0xC9, 0x3E, 0xC9, 0x3B,
];

const SECTOR_SIZE: u64 = 512;
const PARTITION_ALIGNMENT_BYTES: u64 = 1024 * 1024;
const GPT_BACKUP_SECTORS: u64 = 34;
const IMAGE_HEADROOM: u64 = 128 * 1024 * 1024;
const IMAGE_SIZE_ALIGNMENT: u64 = 1024 * 1024;

/// The kind of destination represented by an output path.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OutputKind {
    /// An existing whole-disk block device.
    BlockDevice,
    /// A regular file, including an output path that does not yet exist.
    Image,
}

pub(crate) type FatPartition = Partition;
pub(crate) type FatDir<'a> = fatfs::Dir<'a, FatPartition>;
type ImageFileSystem = FileSystem<FatPartition>;
type FatFile<'a> = fatfs::File<'a, FatPartition>;

/// Determine whether `output` names a block device or a raw image file.
pub fn output_kind(output: &Path) -> Result<OutputKind> {
    match fs::metadata(output) {
        Ok(metadata) if metadata.file_type().is_block_device() => Ok(OutputKind::BlockDevice),
        Ok(metadata) if metadata.is_file() => Ok(OutputKind::Image),
        Ok(_) => Err(anyhow!(
            "output {} must be a whole-disk block device or regular image file",
            output.display()
        )),
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(OutputKind::Image),
        Err(error) => Err(error).with_context(|| format!("inspect output {}", output.display())),
    }
}

/// A raw disk image whose EFI System Partition is exposed as a FAT filesystem.
pub struct Image {
    filesystem: ImageFileSystem,
}

impl Image {
    /// Return the root directory of the image's EFI System Partition.
    pub(crate) fn root_dir(&self) -> FatDir<'_> {
        self.filesystem.root_dir()
    }

    /// Flush filesystem metadata before the image is closed.
    pub fn finish(self) -> Result<()> {
        self.filesystem.unmount().context("flush image filesystem")
    }
}

/// Adapter allowing ISO data to be streamed through the `fatfs` writer.
pub(crate) struct ImageFileWriter<'a> {
    file: FatFile<'a>,
}

impl Write for ImageFileWriter<'_> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.file.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.file.flush()
    }
}

pub(crate) fn create_image_file<'a>(
    dir: &FatDir<'a>,
    path: &str,
) -> io::Result<ImageFileWriter<'a>> {
    dir.create_file(path).map(|file| ImageFileWriter { file })
}

/// A seekable view of one partition within a raw disk image.
pub(crate) struct Partition {
    file: File,
    start: u64,
    len: u64,
    position: u64,
}

impl Partition {
    fn new(file: File, start: u64, len: u64) -> Self {
        Self {
            file,
            start,
            len,
            position: 0,
        }
    }

    fn seek_file(&mut self) -> io::Result<()> {
        let position = self
            .start
            .checked_add(self.position)
            .ok_or_else(|| io::Error::other("partition offset overflow"))?;
        self.file.seek(SeekFrom::Start(position))?;
        Ok(())
    }
}

impl Read for Partition {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.position == self.len {
            return Ok(0);
        }

        let buffer_len = u64::try_from(buf.len())
            .map_err(|_| io::Error::other("read buffer is larger than u64"))?;
        let read_len = usize::try_from((self.len - self.position).min(buffer_len))
            .map_err(|_| io::Error::other("read length does not fit usize"))?;
        self.seek_file()?;
        let count = self.file.read(&mut buf[..read_len])?;
        self.position = self
            .position
            .checked_add(
                u64::try_from(count)
                    .map_err(|_| io::Error::other("read count does not fit u64"))?,
            )
            .ok_or_else(|| io::Error::other("partition position overflow"))?;
        Ok(count)
    }
}

impl Write for Partition {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let buffer_len = u64::try_from(buf.len())
            .map_err(|_| io::Error::other("write buffer is larger than u64"))?;
        let end = self
            .position
            .checked_add(buffer_len)
            .ok_or_else(|| io::Error::other("partition position overflow"))?;
        if end > self.len {
            return Err(io::Error::new(
                io::ErrorKind::WriteZero,
                "write exceeds image partition",
            ));
        }

        self.seek_file()?;
        let count = self.file.write(buf)?;
        self.position = self
            .position
            .checked_add(
                u64::try_from(count)
                    .map_err(|_| io::Error::other("write count does not fit u64"))?,
            )
            .ok_or_else(|| io::Error::other("partition position overflow"))?;
        Ok(count)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.file.flush()
    }
}

impl Seek for Partition {
    fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
        let position = match from {
            SeekFrom::Start(position) => i128::from(position),
            SeekFrom::Current(offset) => i128::from(self.position) + i128::from(offset),
            SeekFrom::End(offset) => i128::from(self.len) + i128::from(offset),
        };
        if !(0..=i128::from(self.len)).contains(&position) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "seek is outside image partition",
            ));
        }

        self.position = u64::try_from(position)
            .map_err(|_| io::Error::other("seek position does not fit u64"))?;
        Ok(self.position)
    }
}

struct PartitionRange {
    start: u64,
    len: u64,
}

/// Create and format a raw disk image sized for `contents_size` bytes of files.
pub fn create_image(path: &Path, contents_size: u64) -> Result<Image> {
    let image_size = image_size(contents_size)?;
    let mut file = OpenOptions::new()
        .create(true)
        .read(true)
        .truncate(true)
        .write(true)
        .open(path)
        .with_context(|| format!("create image {}", path.display()))?;
    file.set_len(image_size)
        .with_context(|| format!("size image {}", path.display()))?;

    let range = write_partition_table(&mut file, path, SECTOR_SIZE)?;
    file.sync_all()
        .with_context(|| format!("sync image {}", path.display()))?;

    let mut partition = Partition::new(file, range.start, range.len);
    format_image_partition(&mut partition, path)?;
    partition.seek(SeekFrom::Start(0))?;
    let filesystem = FileSystem::new(partition, FsOptions::new())
        .with_context(|| format!("open FAT32 filesystem in image {}", path.display()))?;

    Ok(Image { filesystem })
}

/// Partition a whole-disk block device with a single EFI System Partition that
/// fills the disk, then ask the kernel to re-read the partition table.
/// Returns the device path of partition 1 (e.g. `/dev/sdb1`).
pub fn partition_drive(drive: &Path) -> Result<PathBuf> {
    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .open(drive)
        .with_context(|| format!("open {}", drive.display()))?;
    let sector_size = get_sector_size(&mut file)
        .with_context(|| format!("get sector size of {}", drive.display()))?;

    write_partition_table(&mut file, drive, sector_size)?;
    reread_partition_table(&mut file)
        .with_context(|| format!("reread partition table on {}", drive.display()))?;
    file.sync_all()
        .with_context(|| format!("sync {}", drive.display()))?;

    Ok(partition_device_path(drive, 1))
}

/// Write a new GPT with one EFI System Partition and return its byte range.
fn write_partition_table(
    file: &mut File,
    target: &Path,
    sector_size: u64,
) -> Result<PartitionRange> {
    let total_bytes = file
        .seek(SeekFrom::End(0))
        .with_context(|| format!("determine size of {}", target.display()))?;
    let total_sectors = total_bytes / sector_size;
    let start = PARTITION_ALIGNMENT_BYTES.div_ceil(sector_size);
    let end = total_sectors
        .checked_sub(GPT_BACKUP_SECTORS)
        .ok_or_else(|| {
            anyhow!(
                "{} is too small for an EFI System Partition",
                target.display()
            )
        })?;
    if end <= start {
        return Err(anyhow!(
            "{} is too small for an EFI System Partition",
            target.display()
        ));
    }

    let disk_guid = crate::util::random_bytes::<16>()?;
    file.seek(SeekFrom::Start(0))?;
    let mut gpt = GPT::new_from(file, sector_size, disk_guid)?;
    let part_guid = crate::util::random_bytes::<16>()?;
    gpt[1] = GPTPartitionEntry {
        partition_type_guid: ESP_TYPE_GUID,
        unique_partition_guid: part_guid,
        starting_lba: start,
        ending_lba: end,
        attribute_bits: 0,
        partition_name: "EFI".into(),
    };

    GPT::write_protective_mbr_into(file, sector_size)?;
    gpt.write_into(file)?;

    let sectors = end
        .checked_sub(start)
        .and_then(|count| count.checked_add(1))
        .ok_or_else(|| anyhow!("partition size overflow for {}", target.display()))?;
    let start = start
        .checked_mul(sector_size)
        .ok_or_else(|| anyhow!("partition offset overflow for {}", target.display()))?;
    let len = sectors
        .checked_mul(sector_size)
        .ok_or_else(|| anyhow!("partition size overflow for {}", target.display()))?;
    Ok(PartitionRange { start, len })
}

fn image_size(contents_size: u64) -> Result<u64> {
    let size = contents_size
        .checked_add(contents_size / 20)
        .and_then(|size| size.checked_add(IMAGE_HEADROOM))
        .and_then(|size| size.checked_add(PARTITION_ALIGNMENT_BYTES))
        .and_then(|size| size.checked_add(GPT_BACKUP_SECTORS * SECTOR_SIZE))
        .ok_or_else(|| anyhow!("image size overflow"))?;
    let remainder = size % IMAGE_SIZE_ALIGNMENT;
    if remainder == 0 {
        Ok(size)
    } else {
        size.checked_add(IMAGE_SIZE_ALIGNMENT - remainder)
            .ok_or_else(|| anyhow!("image size overflow"))
    }
}

/// Derive the kernel-visible device node for partition `index` of `drive`.
///
/// SATA/SCSI names (`/dev/sdb`) append the index directly (`/dev/sdb1`);
/// `NVMe`, `eMMC`, and loop names require a `p` separator (`/dev/nvme0n1p1`).
fn partition_device_path(drive: &Path, index: u32) -> PathBuf {
    let needs_p = drive.file_name().is_some_and(|name| {
        ["nvme", "mmcblk", "loop"]
            .iter()
            .any(|prefix| name.as_bytes().starts_with(prefix.as_bytes()))
    });
    let mut partition = drive.as_os_str().to_owned();
    if needs_p {
        partition.push("p");
    }
    partition.push(index.to_string());
    PathBuf::from(partition)
}

/// Write a fresh FAT32 filesystem to `part` (a partition device path).
pub fn format_fat32(part: &Path) -> Result<()> {
    let mut device = OpenOptions::new()
        .read(true)
        .write(true)
        .open(part)
        .with_context(|| format!("open {}", part.display()))?;
    format_volume(&mut device, fat32_options()?)
        .map_err(|error| anyhow!("format FAT32 on {}: {error}", part.display()))?;
    device
        .sync_all()
        .with_context(|| format!("sync {}", part.display()))?;
    Ok(())
}

fn format_image_partition(partition: &mut Partition, path: &Path) -> Result<()> {
    format_volume(&mut *partition, fat32_options()?)
        .with_context(|| format!("format FAT32 in image {}", path.display()))?;
    partition
        .flush()
        .with_context(|| format!("flush FAT32 image {}", path.display()))?;
    Ok(())
}

fn fat32_options() -> Result<FormatVolumeOptions> {
    let mut label = [b' '; 11];
    label[..3].copy_from_slice(b"EFI");
    let volume_id = u32::from_le_bytes(crate::util::random_bytes::<4>()?);
    Ok(FormatVolumeOptions::new()
        .fat_type(FatType::Fat32)
        .volume_id(volume_id)
        .volume_label(label))
}

/// RAII guard that unmounts (and removes the mountpoint) when dropped.
pub struct MountGuard {
    dir: tempfile::TempDir,
}

impl MountGuard {
    pub fn point(&self) -> &Path {
        self.dir.path()
    }
}

impl Drop for MountGuard {
    fn drop(&mut self) {
        if let Err(error) = umount(self.dir.path()) {
            tracing::warn!("umount {}: {error}", self.dir.path().display());
        }
        // self.dir drops here, removing the now-empty mountpoint.
    }
}

/// Mount `part` (a vfat partition) at a fresh temporary directory.
pub fn mount_partition(part: &Path) -> Result<MountGuard> {
    let dir = tempfile::Builder::new()
        .prefix("mkwim-")
        .tempdir()
        .context("create mountpoint")?;

    mount(
        Some(part),
        dir.path(),
        Some("vfat"),
        MsFlags::empty(),
        None::<&str>,
    )
    .with_context(|| format!("mount {} at {}", part.display(), dir.path().display()))?;

    Ok(MountGuard { dir })
}