add GptDisk::add_partition_at (#99)

Implement a helper function add_partition_at, similar to add_partition, but can specify more things

src/lib.rs

@@ -117,6 +117,8 @@ pub enum GptError {
     OverflowPartitionCount,
     /// The partition count changes but you did not allow that
     PartitionCountWouldChange,
+    /// The id is already been used
+    PartitionIdAlreadyUsed,
 }
 
 impl From<io::Error> for GptError {
@@ -152,6 +154,7 @@ impl fmt::Display for GptError {
                 "partition would change but is not \
             allowed"
             }
+            PartitionIdAlreadyUsed => "partition id already used",
         };
         write!(fmt, "{desc}")
     }
@@ -537,6 +540,79 @@ where
 
         Err(GptError::NotEnoughSpace)
     }
+    /// create a new partition with a specific id
+    /// a specific name
+    /// a specific first_lba
+    /// a specific length_lba
+    /// a specific part_type
+    /// a specific flags
+    /// ## Panics
+    /// If length is empty panics
+    /// If id zero panics
+    pub fn add_partition_at(
+        &mut self,
+        name: &str,
+        id: u32,
+        first_lba: u64,
+        length_lba: u64,
+        part_type: partition_types::Type,
+        flags: u64,
+    ) -> Result<u32, GptError> {
+        assert!(length_lba > 0, "length must be greater than zero");
+        assert!(id > 0, "id must be greater than zero");
+        //check id
+        match self.partitions.get(&id) {
+            Some(p) if p.is_used() => return Err(GptError::PartitionIdAlreadyUsed),
+            // Allow unused ids , because we can allow to modify the part count
+            _ => {
+                // will override unused partition
+            }
+        }
+        //check partition segment
+        let free_sections = self.find_free_sectors();
+        for (starting_lba, length) in free_sections {
+            if first_lba >= starting_lba && length_lba <= length {
+                // part segment is legal
+                debug!(
+                    "starting_lba {}, length {}, id {}",
+                    first_lba, length_lba, id
+                );
+                debug!(
+                    "Adding partition id: {} {:?}.  first_lba: {} last_lba: {}",
+                    id,
+                    part_type,
+                    first_lba,
+                    first_lba + length_lba - 1_u64
+                );
+                // let's try to increase the num parts
+                // because partition_id 0 will never exist the num_parts is without + 1
+                let num_parts_changes = self.header().num_parts_would_change(id);
+                if num_parts_changes && !self.config.change_partition_count {
+                    return Err(GptError::PartitionCountWouldChange);
+                }
+                let part = partition::Partition {
+                    part_type_guid: part_type,
+                    part_guid: uuid::Uuid::new_v4(),
+                    first_lba,
+                    last_lba: first_lba + length_lba - 1_u64,
+                    flags,
+                    name: name.to_string(),
+                };
+                if let Some(p) = self.partitions.insert(id, part.clone()) {
+                    debug!("Replacing\n{}\nwith\n{}", p, part);
+                }
+                if num_parts_changes {
+                    // update headers
+                    self.init_headers()?;
+                }
+                return Ok(id);
+            }
+        }
+
+        //given segment is illegal
+        Err(GptError::NotEnoughSpace)
+    }
+
     /// Remove partition from this disk.
     pub fn remove_partition(&mut self, id: u32) -> Option<u32> {
         self.partitions.remove(&id).map(|_| {

tests/gpt.rs

@@ -122,7 +122,7 @@ fn test_gpt_disk_write_fidelity_with_device() {
 
 #[test]
 fn test_create_simple_on_device() {
-    const TOTAL_BYTES: usize = 1024 * 64;
+    const TOTAL_BYTES: usize = 1024 * 66;
     let mut mem_device = Box::new(std::io::Cursor::new(vec![0_u8; TOTAL_BYTES]));
 
     // Create a protective MBR at LBA0
@@ -143,6 +143,11 @@ fn test_create_simple_on_device() {
     gdisk
         .add_partition("test2", 1024 * 18, gpt::partition_types::LINUX_FS, 0, None)
         .unwrap();
+    let id = gdisk.find_next_partition_id().unwrap();
+    gdisk
+        .add_partition_at("test3", id, 94, 2, gpt::partition_types::BASIC, 0)
+        .unwrap();
+
     let mut mem_device = gdisk.write().unwrap();
     mem_device.seek(std::io::SeekFrom::Start(0)).unwrap();
     let mut final_bytes = vec![0_u8; TOTAL_BYTES];
@@ -161,7 +166,7 @@ fn test_only_valid_headers() {
     // write a valid disk
     let mut valid_disk = GptConfig::new()
         .writable(true)
-        .create_from_device(Cursor::new(vec![0; 1024 * 68]), None)
+        .create_from_device(Cursor::new(vec![0; 1024 * 70]), None)
         .unwrap();
 
     valid_disk
@@ -170,7 +175,10 @@ fn test_only_valid_headers() {
     valid_disk
         .add_partition("test2", 1024 * 18, gpt::partition_types::LINUX_FS, 0, None)
         .unwrap();
-
+    let id = valid_disk.find_next_partition_id().unwrap();
+    valid_disk
+        .add_partition_at("test3", id, 94, 2, gpt::partition_types::BASIC, 0)
+        .unwrap();
     // now write to memory
     let valid_disk = valid_disk.write().unwrap();
     let mut corrupt_disk = valid_disk.clone();
@@ -201,7 +209,7 @@ fn test_readonly_backup() {
     // write a valid disk
     let mut valid_disk = GptConfig::new()
         .writable(true)
-        .create_from_device(Cursor::new(vec![0; 1024 * 68]), None)
+        .create_from_device(Cursor::new(vec![0; 1024 * 70]), None)
         .unwrap();
 
     valid_disk
@@ -210,7 +218,10 @@ fn test_readonly_backup() {
     valid_disk
         .add_partition("test2", 1024 * 18, gpt::partition_types::LINUX_FS, 0, None)
         .unwrap();
-
+    let id = valid_disk.find_next_partition_id().unwrap();
+    valid_disk
+        .add_partition_at("test4", id, 94, 2, gpt::partition_types::BASIC, 0)
+        .unwrap();
     // now write to memory
     valid_disk.write_inplace().unwrap();
     let valid_disk_bytes = valid_disk.device_ref();