partitioning: Respect boundaries of disks for GPT LBA offsets

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

Author: ikeycode

crates/partitioning/src/planner.rs

@@ -46,6 +46,10 @@ pub enum Change {
 
 /// A disk partitioning planner.
 pub struct Planner {
+    /// First usable LBA position on disk in bytes
+    usable_start: u64,
+    /// Last usable LBA position on disk in bytes
+    usable_end: u64,
     /// The original block device state that we're planning changes for
     device: BlockDevice,
     /// Stack of changes that can be undone
@@ -192,30 +196,49 @@ impl Planner {
     /// Creates a new partitioning planner for the given disk.
     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();
+
         Self {
+            usable_start: 0,
+            usable_end: device.size(),
             device,
             changes: VecDeque::new(),
             original_regions,
         }
     }
 
+    /// Set the usable disk region offsets
+    pub fn with_start_offset(self, offset: u64) -> Self {
+        Self {
+            usable_start: offset,
+            ..self
+        }
+    }
+
+    /// Set the usable disk region offsets
+    pub fn with_end_offset(self, offset: u64) -> Self {
+        Self {
+            usable_end: offset,
+            ..self
+        }
+    }
+
     /// Get a human readable description of pending changes
     pub fn describe_changes(&self) -> String {
         if self.changes.is_empty() {
             return "No pending changes".to_string();
         }
 
-        let disk_size = self.device.size();
         let mut description = "Pending changes:\n".to_string();
 
         for (i, change) in self.changes.iter().enumerate() {
-            description.push_str(&format!("  {}: {}\n", i + 1, change.describe(disk_size)));
+            description.push_str(&format!("  {}: {}\n", i + 1, change.describe(self.usable_size())));
         }
 
         description
@@ -272,9 +295,9 @@ impl Planner {
         let aligned_end = align_down(end, PARTITION_ALIGNMENT);
         debug!("Aligned positions: {}..{}", aligned_start, aligned_end);
 
-        // Validate bounds
-        if aligned_end > self.device.size() {
-            warn!("Partition would exceed disk bounds");
+        // Validate bounds against usable disk region
+        if aligned_start < self.usable_start || aligned_end > self.usable_end {
+            warn!("Partition would be outside usable disk region");
             return Err(PlanError::RegionOutOfBounds {
                 start: aligned_start,
                 end: aligned_end,
@@ -321,8 +344,8 @@ impl Planner {
         if index >= self.original_regions.len() {
             warn!("Invalid partition index {}", index);
             return Err(PlanError::RegionOutOfBounds {
-                start: 0,
-                end: self.device.size(),
+                start: self.usable_start,
+                end: self.usable_size(),
             });
         }
 
@@ -362,6 +385,17 @@ impl Planner {
     pub fn original_device(&self) -> &BlockDevice {
         &self.device
     }
+
+    /// Get the size of the usable disk region in bytes
+    pub fn usable_size(&self) -> u64 {
+        self.usable_end - self.usable_start
+    }
+
+    /// Get the usable disk region offsets
+    pub fn offsets(&self) -> (u64, u64) {
+        (self.usable_start, self.usable_end)
+    }
+
     /// Plan to initialize a clean partition layout
     pub fn plan_initialize_disk(&mut self) -> Result<(), PlanError> {
         debug!("Planning to create new GPT partition table");

crates/partitioning/src/strategy.rs

@@ -85,9 +85,8 @@ impl Strategy {
 
     /// Find available free regions on the disk
     fn find_free_regions(&self, planner: &Planner) -> Vec<Region> {
-        let disk_size = planner.original_device().size();
         let mut regions = Vec::new();
-        let mut current = 0;
+        let (mut current, disk_size) = planner.offsets();
 
         // Sort existing partitions by start position
         let mut layout = planner.current_layout();
@@ -147,7 +146,8 @@ impl Strategy {
             AllocationStrategy::InitializeWholeDisk => {
                 // Clear existing partitions and start fresh
                 planner.plan_initialize_disk()?;
-                Region::new(0, planner.original_device().size())
+                let (start, end) = planner.offsets();
+                Region::new(start, end)
             }
             AllocationStrategy::LargestFree => {
                 let free_regions = self.find_free_regions(planner);

disk-test/src/main.rs

@@ -9,7 +9,7 @@ use disks::BlockDevice;
 use partitioning::{
     gpt::{disk::LogicalBlockSize, mbr::ProtectiveMBR, partition_types, GptConfig},
     loopback,
-    planner::Planner,
+    planner::{format_size, Planner},
     sparsefile,
     strategy::{AllocationStrategy, PartitionRequest, SizeRequirement, Strategy},
 };
@@ -68,10 +68,17 @@ where
     let mut gpt_disk = gpt_config.create(&path)?;
     gpt_disk.write_inplace()?;
 
+    eprintln!("GPT: {:?}", gpt_disk);
+
+    let first_usable = gpt_disk.header().first_usable * 512;
+    let last_usable = gpt_disk.header().last_usable * 512;
+
     // Connect the planner.
     let disk = disks::loopback::Device::from_device_path(path).unwrap();
     let block = BlockDevice::loopback_device(disk);
-    let mut planner = Planner::new(block);
+    let mut planner = Planner::new(block)
+        .with_start_offset(first_usable)
+        .with_end_offset(last_usable);
     let mut strategy = Strategy::new(AllocationStrategy::InitializeWholeDisk);
 
     // efi
@@ -111,9 +118,19 @@ where
     info!("Computed: {}", planner.describe_changes());
 
     // TODO: Track the types in the API and use them here
-    for partition in planner.current_layout() {
+    for (n, partition) in planner.current_layout().iter().enumerate() {
         info!("Adding partition: {:?}", &partition);
-        gpt_disk.add_partition("", partition.size(), partition_types::LINUX_FS, 0, None)?;
+        assert_ne!(0, partition.size());
+        let size = partitioning::planner::format_size(partition.size());
+        info!("Partition {} size: {}, at {}", n, size, format_size(partition.start));
+        gpt_disk.add_partition_at(
+            "",
+            n as u32 + 1,
+            partition.start / 512,
+            partition.size() / 512,
+            partition_types::LINUX_FS,
+            0,
+        )?;
     }
 
     let _ = gpt_disk.write()?;