partitioning/blkpg: Explicitly control blkpg rescan lifetime

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

Author: ikeycode

crates/partitioning/src/blkpg.rs

@@ -130,26 +130,17 @@ where
     Ok(())
 }
 
-/// Updates kernel partition representations to match the GPT table
+/// Removes all kernel partitions for the specified block device
 ///
 /// # Arguments
 /// * `path` - Path to the block device
 ///
 /// # Returns
 /// `Result<(), Error>` indicating success or partition operation failure
-pub fn sync_gpt_partitions<P: AsRef<Path>>(path: P) -> Result<(), Error> {
-    info!("Initiating GPT partition synchronization for {:?}", path.as_ref());
+pub fn remove_kernel_partitions<P: AsRef<Path>>(path: P) -> Result<(), Error> {
+    debug!("Beginning partition cleanup process for {:?}", path.as_ref());
     let file = File::open(&path)?;
 
-    // Read GPT table
-    debug!("Reading GPT partition table");
-    let gpt = gpt::GptConfig::new().writable(false).open(&path)?;
-    let partitions = gpt.partitions();
-    let block_size = 512;
-    info!("Located {} partitions (block size: {})", partitions.len(), block_size);
-
-    debug!("Beginning partition cleanup process");
-
     // Find the disk for enumeration purposes
     let base_name = path
         .as_ref()
@@ -164,6 +155,28 @@ pub fn sync_gpt_partitions<P: AsRef<Path>>(path: P) -> Result<(), Error> {
         let _ = delete_partition(file.as_raw_fd(), partition.number as i32);
     }
 
+    info!("Successfully removed all kernel partitions");
+    Ok(())
+}
+
+/// Creates kernel partitions based on the current GPT table
+///
+/// # Arguments
+/// * `path` - Path to the block device
+///
+/// # Returns
+/// `Result<(), Error>` indicating success or partition operation failure
+pub fn create_kernel_partitions<P: AsRef<Path>>(path: P) -> Result<(), Error> {
+    info!("Creating kernel partitions from GPT for {:?}", path.as_ref());
+    let file = File::open(&path)?;
+
+    // Read GPT table
+    debug!("Reading GPT partition table");
+    let gpt = gpt::GptConfig::new().writable(false).open(&path)?;
+    let partitions = gpt.partitions();
+    let block_size = 512;
+    info!("Located {} partitions (block size: {})", partitions.len(), block_size);
+
     // Add partitions from GPT
     debug!("Beginning partition creation from GPT table");
     for (i, partition) in partitions.iter() {
@@ -175,6 +188,23 @@ pub fn sync_gpt_partitions<P: AsRef<Path>>(path: P) -> Result<(), Error> {
         )?;
     }
 
+    info!("GPT partition creation completed successfully");
+    Ok(())
+}
+
+/// Updates kernel partition representations to match the GPT table
+///
+/// # Arguments
+/// * `path` - Path to the block device
+///
+/// # Returns
+/// `Result<(), Error>` indicating success or partition operation failure
+pub fn sync_gpt_partitions<P: AsRef<Path>>(path: P) -> Result<(), Error> {
+    info!("Initiating GPT partition synchronization for {:?}", path.as_ref());
+
+    remove_kernel_partitions(&path)?;
+    create_kernel_partitions(&path)?;
+
     info!("GPT partition synchronization completed successfully");
     Ok(())
 }

crates/partitioning/src/writer.rs

@@ -3,20 +3,28 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
-use std::{fs, io::{Write, Seek}};
+use std::{
+    fs,
+    io::{Seek, Write},
+};
 
 use disks::BlockDevice;
 use gpt::{mbr, partition_types, GptConfig};
 use thiserror::Error;
 
 use crate::{
+    blkpg,
     planner::{Change, Planner},
     GptAttributes,
 };
 
 /// Errors that can occur when writing changes to disk
 #[derive(Debug, Error)]
 pub enum WriteError {
+    // A blkpg error
+    #[error("error syncing partitions: {0}")]
+    Blkpg(#[from] blkpg::Error),
+
     /// A partition ID was used multiple times
     #[error("Duplicate partition ID: {0}")]
     DuplicatePartitionId(u32),
@@ -98,6 +106,11 @@ impl<'a> DiskWriter<'a> {
     /// - Creating or opening the GPT table
     /// - Applying each change in sequence
     fn apply_changes(&self, device: &mut fs::File, writable: bool) -> Result<(), WriteError> {
+        // Remove known partitions pre wipe
+        if writable {
+            blkpg::remove_kernel_partitions(self.device.device())?;
+        }
+
         let mut gpt_table = if self.planner.wipe_disk() {
             if writable {
                 // Zero out the first MiB to clear any old partition tables and boot sectors
@@ -182,6 +195,7 @@ impl<'a> DiskWriter<'a> {
         if writable {
             let original = gpt_table.write()?;
             original.sync_all()?;
+            blkpg::create_kernel_partitions(self.device.device())?;
         }
 
         Ok(())