partitioning: Add basic writer that tracks partition ids

Signed-off-by: Ikey Doherty <[email protected]>

Author: ikeycode

crates/partitioning/src/lib.rs

@@ -10,3 +10,5 @@ pub use gpt;
 
 pub mod planner;
 pub mod strategy;
+
+pub mod writer;

crates/partitioning/src/planner.rs

@@ -39,9 +39,9 @@ pub enum PlanError {
 #[derive(Debug, Clone)]
 pub enum Change {
     /// Add a new partition
-    AddPartition { start: u64, end: u64 },
+    AddPartition { start: u64, end: u64, partition_id: u32 },
     /// Delete an existing partition
-    DeletePartition { original_index: usize },
+    DeletePartition { original_index: usize, partition_id: u32 },
 }
 
 /// A disk partitioning planner.
@@ -55,6 +55,12 @@ pub struct Planner {
     changes: VecDeque<Change>,
     /// Original partition layout for reference
     original_regions: Vec<Region>,
+    /// Track original partition IDs
+    original_partition_ids: Vec<u32>,
+    /// Next available partition ID for new partitions
+    next_partition_id: u32,
+
+    wipe_disk: bool,
 }
 
 /// A contiguous region of disk space defined by absolute start and end positions
@@ -189,16 +195,24 @@ impl Change {
     /// Get a human readable description of this change
     pub fn describe(&self, disk_size: u64) -> String {
         match self {
-            Change::AddPartition { start, end } => {
+            Change::AddPartition {
+                start,
+                end,
+                partition_id,
+            } => {
                 format!(
-                    "Add new partition: {} ({} at {})",
+                    "Add new partition #{}: {} ({} at {})",
+                    partition_id,
                     format_size(end - start),
                     Region::new(*start, *end).describe(disk_size),
                     format_position(*start, disk_size)
                 )
             }
-            Change::DeletePartition { original_index } => {
-                format!("Delete partition #{}", original_index + 1)
+            Change::DeletePartition {
+                original_index,
+                partition_id,
+            } => {
+                format!("Delete partition #{} (index {})", partition_id, original_index + 1)
             }
         }
     }
@@ -209,18 +223,25 @@ impl Planner {
     pub fn new(device: &BlockDevice) -> Self {
         debug!("Creating new partition planner for device of size {}", device.size());
 
-        // Extract original regions from device
-        let original_regions = device
-            .partitions()
-            .iter()
-            .map(|p| Region::new(p.start, p.end))
-            .collect();
+        // Extract original regions and partition IDs from device
+        let mut original_regions = Vec::new();
+        let mut original_partition_ids = Vec::new();
+        let mut max_id = 0u32;
+
+        for part in device.partitions() {
+            original_regions.push(Region::new(part.start, part.end));
+            original_partition_ids.push(part.number);
+            max_id = max_id.max(part.number);
+        }
 
         Self {
             usable_start: 0,
             usable_end: device.size(),
             changes: VecDeque::new(),
             original_regions,
+            original_partition_ids,
+            next_partition_id: max_id + 1,
+            wipe_disk: false,
         }
     }
 
@@ -262,7 +283,11 @@ impl Planner {
 
         // First pass: collect indices to delete
         for change in &self.changes {
-            if let Change::DeletePartition { original_index } = change {
+            if let Change::DeletePartition {
+                original_index,
+                partition_id: _,
+            } = change
+            {
                 deleted_indices.push(*original_index);
             }
         }
@@ -276,8 +301,13 @@ impl Planner {
 
         // Second pass: add new partitions
         for change in &self.changes {
-            if let Change::AddPartition { start, end } = change {
-                debug!("Adding partition {}..{}", start, end);
+            if let Change::AddPartition {
+                start,
+                end,
+                partition_id,
+            } = change
+            {
+                debug!("Adding partition {}..{} (ID: {})", start, end, partition_id);
                 layout.push(Region {
                     start: *start,
                     end: *end,
@@ -358,10 +388,12 @@ impl Planner {
             }
         }
 
-        debug!("Adding new partition to change queue");
+        let partition_id = self.allocate_partition_id();
+        debug!("Adding new partition with ID {} to change queue", partition_id);
         self.changes.push_back(Change::AddPartition {
             start: aligned_start,
             end: aligned_end,
+            partition_id,
         });
         Ok(())
     }
@@ -378,9 +410,18 @@ impl Planner {
             });
         }
 
-        debug!("Adding partition deletion to change queue");
-        self.changes
-            .push_back(Change::DeletePartition { original_index: index });
+        let partition_id = self
+            .get_original_partition_id(index)
+            .ok_or(PlanError::RegionOutOfBounds {
+                start: self.usable_start,
+                end: self.usable_size(),
+            })?;
+
+        debug!("Adding deletion of partition ID {} to change queue", partition_id);
+        self.changes.push_back(Change::DeletePartition {
+            original_index: index,
+            partition_id,
+        });
         Ok(())
     }
 
@@ -425,8 +466,24 @@ impl Planner {
         debug!("Planning to create new GPT partition table");
         self.changes.clear(); // Clear any existing changes
         self.original_regions.clear(); // Clear original partitions
+        self.wipe_disk = true;
         Ok(())
     }
+
+    pub fn wipe_disk(&self) -> bool {
+        self.wipe_disk
+    }
+    /// Get the next available partition ID and increment the counter
+    pub fn allocate_partition_id(&mut self) -> u32 {
+        let id = self.next_partition_id;
+        self.next_partition_id += 1;
+        id
+    }
+
+    /// Get the original partition ID for a given index
+    pub fn get_original_partition_id(&self, index: usize) -> Option<u32> {
+        self.original_partition_ids.get(index).copied()
+    }
 }
 
 #[cfg(test)]

crates/partitioning/src/writer.rs

@@ -0,0 +1,158 @@
+// SPDX-FileCopyrightText: Copyright © 2025 Serpent OS Developers
+// SPDX-FileCopyrightText: Copyright © 2025 AerynOS Developers
+//
+// SPDX-License-Identifier: MPL-2.0
+
+use std::fs;
+
+use disks::BlockDevice;
+use gpt::{mbr, partition_types, GptConfig};
+use thiserror::Error;
+
+use crate::planner::{Change, Planner};
+
+/// Errors that can occur when writing changes to disk
+#[derive(Debug, Error)]
+pub enum WriteError {
+    /// Device size has changed since the plan was created
+    #[error("Device size changed since planning")]
+    DeviceSizeChanged,
+
+    /// A partition ID was used multiple times
+    #[error("Duplicate partition ID: {0}")]
+    DuplicatePartitionId(u32),
+
+    /// Error from GPT library
+    #[error("GPT error: {0}")]
+    Gpt(#[from] gpt::GptError),
+
+    /// Error from MBR handling
+    #[error("GPT error: {0}")]
+    Mbr(#[from] gpt::mbr::MBRError),
+
+    /// Underlying I/O error
+    #[error("I/O error: {0}")]
+    IoError(#[from] std::io::Error),
+}
+
+/// A writer that applies the layouts from the Planner to the disk.
+pub struct DiskWriter<'a> {
+    /// The block device to write to
+    pub device: &'a BlockDevice,
+    /// The planner containing the changes to apply
+    pub planner: &'a Planner,
+}
+
+impl<'a> DiskWriter<'a> {
+    /// Create a new DiskWriter.
+    pub fn new(device: &'a BlockDevice, planner: &'a Planner) -> Self {
+        Self { device, planner }
+    }
+
+    /// Simulate changes without writing to disk
+    pub fn simulate(&self) -> Result<(), WriteError> {
+        let mut device = fs::OpenOptions::new()
+            .read(true)
+            .write(false)
+            .open(self.device.device())?;
+        self.validate_changes(&device)?;
+        self.apply_changes(&mut device, false)?;
+        Ok(())
+    }
+
+    /// Actually write changes to disk
+    pub fn write(&self) -> Result<(), WriteError> {
+        let mut device = fs::OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(self.device.device())?;
+
+        self.validate_changes(&device)?;
+        self.apply_changes(&mut device, true)?;
+        Ok(())
+    }
+
+    /// Validate all planned changes before applying them by checking:
+    /// - Device size matches the planned size
+    /// - No duplicate partition IDs exist
+    fn validate_changes(&self, device: &fs::File) -> Result<(), WriteError> {
+        // Verify device size matches what we planned for
+        let metadata = device.metadata()?;
+        if metadata.len() != self.device.size() {
+            return Err(WriteError::DeviceSizeChanged);
+        }
+
+        // Verify partition IDs don't conflict
+        let mut used_ids = std::collections::HashSet::new();
+        for change in self.planner.changes() {
+            match change {
+                Change::AddPartition { partition_id, .. } => {
+                    if !used_ids.insert(*partition_id) {
+                        return Err(WriteError::DuplicatePartitionId(*partition_id));
+                    }
+                }
+                Change::DeletePartition { partition_id, .. } => {
+                    used_ids.remove(partition_id);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Apply the changes to disk by:
+    /// - Creating or opening the GPT table
+    /// - Applying each change in sequence
+    fn apply_changes(&self, device: &mut fs::File, writable: bool) -> Result<(), WriteError> {
+        let mut gpt_table = if self.planner.wipe_disk() {
+            if writable {
+                let mbr = mbr::ProtectiveMBR::with_lb_size(
+                    u32::try_from((self.device.size() / 512) - 1).unwrap_or(0xFF_FF_FF_FF),
+                );
+                mbr.overwrite_lba0(device)?;
+            }
+
+            GptConfig::default()
+                .writable(writable)
+                .logical_block_size(gpt::disk::LogicalBlockSize::Lb512)
+                .create_from_device(device, None)?
+        } else {
+            GptConfig::default().writable(writable).open_from_device(device)?
+        };
+
+        let _layout = self.planner.current_layout();
+        let changes = self.planner.changes();
+
+        for change in changes {
+            match change {
+                Change::DeletePartition {
+                    partition_id,
+                    original_index,
+                } => {
+                    if let Some(id) = gpt_table.remove_partition(*partition_id) {
+                        println!(
+                            "Deleted partition {} (index {}): {:?}",
+                            partition_id, original_index, id
+                        );
+                    }
+                }
+                Change::AddPartition {
+                    start,
+                    end,
+                    partition_id,
+                } => {
+                    let start_lba = *start / 512;
+                    let size_lba = (*end - *start) / 512;
+                    let part_type = partition_types::BASIC;
+
+                    let id = gpt_table.add_partition_at("", *partition_id, start_lba, size_lba, part_type, 0)?;
+                    println!("Added partition {}: {:?}", partition_id, id);
+                }
+            }
+        }
+
+        eprintln!("GPT is now: {gpt_table:?}");
+
+        Ok(())
+    }
+}