download: restore saved partitions while writing back first MiB

If we write back the image's GPT and then save partitions afterward, we
leave a window where a signal or system crash loses the saved partitions.
Do the merge and writeback in one step, within the image_copy callback.
This requires SavedPartitions.merge() to read the image's GPT from a
separate Reader.

src/blockdev.rs

@@ -23,7 +23,7 @@ use std::fs::{
     canonicalize, metadata, read_dir, read_to_string, remove_dir, symlink_metadata, File,
     OpenOptions,
 };
-use std::io::{Seek, SeekFrom};
+use std::io::{Read, Seek, SeekFrom};
 use std::num::{NonZeroU32, NonZeroU64};
 use std::os::linux::fs::MetadataExt;
 use std::os::raw::c_int;
@@ -629,17 +629,18 @@ impl SavedPartitions {
         Ok(())
     }
 
-    /// If any partitions are saved, merge them into the current GPT, which
-    /// must be valid.  Updating the kernel partition table is the caller's
-    /// responsibility.
-    pub fn merge(&self, disk: &mut File) -> Result<()> {
+    /// If any partitions are saved, merge them into the GPT from source,
+    /// which must be valid.  Updating the kernel partition table is the
+    /// caller's responsibility.
+    pub fn merge(&self, source: &mut (impl Read + Seek), disk: &mut File) -> Result<()> {
         if self.partitions.is_empty() {
             return Ok(());
         }
 
         // read GPT
         Self::verify_disk_sector_size(disk, self.sector_size)?;
-        let mut gpt = GPT::find_from(disk).chain_err(|| "couldn't read partition table")?;
+        let mut gpt =
+            GPT::find_from(source).chain_err(|| "couldn't read partition table from source")?;
         if gpt.sector_size != self.sector_size {
             // install will fail on an image that doesn't match the disk,
             // so this shouldn't happen
@@ -876,6 +877,12 @@ pub fn get_block_device_size(file: &File) -> Result<NonZeroU64> {
     }
 }
 
+/// Get the size of the GPT metadata at the start of the disk.
+pub fn get_gpt_size(file: &mut (impl Read + Seek)) -> Result<u64> {
+    let gpt = GPT::find_from(file).chain_err(|| "reading GPT")?;
+    Ok(gpt.header.first_usable_lba * gpt.sector_size)
+}
+
 pub fn udev_settle() -> Result<()> {
     // "udevadm settle" silently no-ops if the udev socket is missing, and
     // then lsblk can't find partition labels.  Catch this early.
@@ -1053,6 +1060,7 @@ mod tests {
             make_part(2, "EFI-SYSTEM", 384, 512),
             make_part(4, "root", 1024, 2200),
         ];
+        let merge_base_parts = vec![make_part(2, "unused", 500, 3500)];
 
         let index = |i| Some(NonZeroU32::new(i).unwrap());
         let label = |l| Label(glob::Pattern::new(l).unwrap());
@@ -1153,39 +1161,36 @@ mod tests {
             (
                 vec![Index(index(15), None)],
                 vec![],
-                vec![
-                    make_part(1, "boot", 1, 384),
-                    make_part(2, "EFI-SYSTEM", 384, 512),
-                    make_part(4, "root", 1024, 2200),
-                ],
+                merge_base_parts.clone(),
             ),
             // No filters
-            (
-                vec![],
-                vec![],
-                vec![
-                    make_part(1, "boot", 1, 384),
-                    make_part(2, "EFI-SYSTEM", 384, 512),
-                    make_part(4, "root", 1024, 2200),
-                ],
-            ),
+            (vec![], vec![], merge_base_parts.clone()),
         ];
 
         let mut base = make_disk(512, &base_parts);
+        let mut image = make_disk(512, &image_parts);
         for (testnum, (filter, expected_blank, expected_image)) in tests.iter().enumerate() {
             // try overwriting on blank disk
             let saved = SavedPartitions::new(&mut base, filter).unwrap();
             let mut disk = make_unformatted_disk();
             saved.overwrite(&mut disk).unwrap();
             assert!(disk_has_mbr(&mut disk), "test {}", testnum);
             let result = GPT::find_from(&mut disk).unwrap();
+            assert_eq!(
+                get_gpt_size(&mut disk).unwrap(),
+                512 * result.header.first_usable_lba
+            );
             assert_partitions_eq(expected_blank, &result, &format!("test {} blank", testnum));
 
-            // try merging with image disk
-            let mut disk = make_disk(512, &image_parts);
-            saved.merge(&mut disk).unwrap();
+            // try merging with image disk onto merge_base disk
+            let mut disk = make_disk(512, &merge_base_parts);
+            saved.merge(&mut image, &mut disk).unwrap();
             assert!(!disk_has_mbr(&mut disk), "test {}", testnum);
             let result = GPT::find_from(&mut disk).unwrap();
+            assert_eq!(
+                get_gpt_size(&mut disk).unwrap(),
+                512 * result.header.first_usable_lba
+            );
             assert_partitions_eq(expected_image, &result, &format!("test {} image", testnum));
             assert_eq!(
                 saved.get_offset().unwrap(),
@@ -1207,24 +1212,35 @@ mod tests {
         // test merging with unformatted initial disk
         let mut disk = make_unformatted_disk();
         let saved = SavedPartitions::new(&mut disk, &vec![label("z")]).unwrap();
-        let mut disk = make_disk(512, &image_parts);
-        saved.merge(&mut disk).unwrap();
+        let mut disk = make_disk(512, &merge_base_parts);
+        saved.merge(&mut image, &mut disk).unwrap();
         let result = GPT::find_from(&mut disk).unwrap();
-        assert_partitions_eq(&image_parts, &result, "unformatted disk");
+        assert_partitions_eq(&merge_base_parts, &result, "unformatted disk");
 
         // test overlapping partitions
         let saved = SavedPartitions::new(&mut base, &vec![Index(index(1), index(1))]).unwrap();
-        let mut disk = make_disk(512, &image_parts);
+        let mut disk = make_disk(512, &merge_base_parts);
         assert_eq!(
-            saved.merge(&mut disk).unwrap_err().to_string(),
+            saved.merge(&mut image, &mut disk).unwrap_err().to_string(),
             "disk partition 4 ('root') ends after start of saved partition 1 ('one')",
         );
 
         // test sector size mismatch
         let saved = SavedPartitions::new(&mut base, &vec![label("*i*")]).unwrap();
-        let mut disk = make_disk(4096, &image_parts);
+        let mut image_4096 = make_disk(4096, &image_parts);
+        assert_eq!(
+            get_gpt_size(&mut image_4096).unwrap(),
+            4096 * GPT::find_from(&mut image_4096)
+                .unwrap()
+                .header
+                .first_usable_lba
+        );
+        let mut disk = make_disk(4096, &merge_base_parts);
         assert_eq!(
-            saved.merge(&mut disk).unwrap_err().to_string(),
+            saved
+                .merge(&mut image_4096, &mut disk)
+                .unwrap_err()
+                .to_string(),
             "sector size 4096 of disk GPT doesn't match expected sector size 512"
         );
 

src/download.rs

@@ -16,14 +16,14 @@ use byte_unit::Byte;
 use error_chain::bail;
 use nix::unistd::isatty;
 use std::fs::{remove_file, File, OpenOptions};
-use std::io::{self, copy, stderr, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
+use std::io::{self, copy, stderr, BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write};
 use std::num::{NonZeroU32, NonZeroU64};
 use std::os::unix::io::AsRawFd;
 use std::path::{Path, PathBuf};
 use std::result;
 use std::time::{Duration, Instant};
 
-use crate::blockdev::{detect_formatted_sector_size, SavedPartitions};
+use crate::blockdev::{detect_formatted_sector_size, get_gpt_size, SavedPartitions};
 use crate::cmdline::*;
 use crate::errors::*;
 use crate::io::*;
@@ -295,10 +295,28 @@ pub fn image_copy_default(
     let dest = buf_dest.into_inner().map_err(|_| "flushing data to disk")?;
 
     // verify_reader has now checked the signature, so fill in the first MiB
-    dest.seek(SeekFrom::Start(0))
-        .chain_err(|| "seeking to start of disk")?;
-    dest.write_all(first_mb)
-        .chain_err(|| "writing to first MiB of disk")?;
+    let offset = match saved {
+        Some(saved) if saved.is_saved() => {
+            // write merged GPT
+            let mut cursor = Cursor::new(first_mb);
+            saved
+                .merge(&mut cursor, dest)
+                .chain_err(|| "writing updated GPT")?;
+
+            // copy all remaining bytes from first_mb (probably not
+            // important but can't hurt)
+            get_gpt_size(dest).chain_err(|| "getting GPT size")?
+        }
+        _ => {
+            // copy all of first_mb
+            0
+        }
+    };
+    // do the copy
+    dest.seek(SeekFrom::Start(offset))
+        .chain_err(|| format!("seeking disk to offset {}", offset))?;
+    dest.write_all(&first_mb[offset as usize..first_mb.len()])
+        .chain_err(|| "writing first MiB of disk")?;
 
     Ok(())
 }
@@ -428,6 +446,9 @@ impl<'a, R: Read> Drop for ProgressReader<'a, R> {
 mod tests {
     use super::*;
     use error_chain::ChainedError;
+    use gptman::{GPTPartitionEntry, GPT};
+    use std::io::{Seek, SeekFrom};
+    use uuid::Uuid;
 
     #[test]
     fn test_write_image_limit() {
@@ -464,4 +485,106 @@ mod tests {
         dest.read_exact(&mut buf).unwrap();
         assert_eq!(buf, precious.as_bytes());
     }
+
+    #[test]
+    fn test_image_copy_default_first_mb() {
+        let len: usize = 2 * 1024 * 1024;
+        let mb: usize = 1024 * 1024;
+
+        let mut data = vec![0u8; len];
+        for i in 0..data.len() {
+            data[i] = (i % 256) as u8;
+        }
+
+        // no saved partitions
+        let mut source = Cursor::new(&data);
+        let mut dest = tempfile::tempfile().unwrap();
+        // copy
+        source.seek(SeekFrom::Start(mb as u64)).unwrap();
+        image_copy_default(&data[0..mb], &mut source, &mut dest, Path::new("/z"), None).unwrap();
+        // compare
+        dest.seek(SeekFrom::Start(0)).unwrap();
+        let mut result = vec![0u8; len];
+        dest.read_exact(&mut result).unwrap();
+        assert_eq!(data, result);
+
+        // SavedPartitions but nothing saved
+        let mut source = Cursor::new(&data);
+        let mut dest = tempfile::tempfile().unwrap();
+        // gptman requires a fixed disk length
+        dest.set_len(len as u64).unwrap();
+        // create saved
+        let saved = SavedPartitions::new(&mut dest, &vec![]).unwrap();
+        assert!(!saved.is_saved());
+        // copy
+        source.seek(SeekFrom::Start(mb as u64)).unwrap();
+        image_copy_default(
+            &data[0..mb],
+            &mut source,
+            &mut dest,
+            Path::new("/z"),
+            Some(&saved),
+        )
+        .unwrap();
+        // compare
+        dest.seek(SeekFrom::Start(0)).unwrap();
+        let mut result = vec![0u8; len];
+        dest.read_exact(&mut result).unwrap();
+        assert_eq!(data, result);
+
+        // saved partition
+        let mut source = Cursor::new(data.clone());
+        let mut dest = tempfile::tempfile().unwrap();
+        // source must have a partition table
+        partition(&mut source, None);
+        let data_partitioned = source.into_inner();
+        let mut source = Cursor::new(&data_partitioned);
+        // gptman requires a fixed disk length
+        dest.set_len(2 * len as u64).unwrap();
+        // create partition to save
+        partition(&mut dest, Some(2));
+        // create saved
+        let saved = SavedPartitions::new(
+            &mut dest,
+            &vec![PartitionFilter::Label(glob::Pattern::new("bovik").unwrap())],
+        )
+        .unwrap();
+        assert!(saved.is_saved());
+        // copy
+        source.seek(SeekFrom::Start(mb as u64)).unwrap();
+        image_copy_default(
+            &data_partitioned[0..mb],
+            &mut source,
+            &mut dest,
+            Path::new("/z"),
+            Some(&saved),
+        )
+        .unwrap();
+        // compare
+        dest.seek(SeekFrom::Start(0)).unwrap();
+        let mut result = vec![0u8; len];
+        dest.read_exact(&mut result).unwrap();
+        assert_eq!(detect_formatted_sector_size(&result), NonZeroU32::new(512));
+        let gpt_size = get_gpt_size(&mut dest).unwrap();
+        assert!(gpt_size < 24576);
+        assert_eq!(
+            data_partitioned[gpt_size as usize..],
+            result[gpt_size as usize..]
+        );
+    }
+
+    fn partition(f: &mut (impl Read + Write + Seek), start_mb: Option<u64>) {
+        let mut gpt = GPT::new_from(f, 512, *Uuid::new_v4().as_bytes()).unwrap();
+        if let Some(start_mb) = start_mb {
+            gpt[1] = GPTPartitionEntry {
+                partition_type_guid: [1u8; 16],
+                unique_parition_guid: [1u8; 16],
+                starting_lba: start_mb * 2048,
+                ending_lba: (start_mb + 1) * 2048,
+                attribute_bits: 0,
+                partition_name: "bovik".into(),
+            };
+        }
+        gpt.write_into(f).unwrap();
+    }
 }

src/install.rs

@@ -172,11 +172,6 @@ fn write_disk(
             .map(|(offset, desc)| (offset, format!("collision with {}", desc))),
         Some(sector_size),
     )?;
-
-    // restore saved partitions, if any, and reread table
-    saved
-        .merge(dest)
-        .chain_err(|| format!("restoring saved partitions to {}", config.device))?;
     table.reread()?;
 
     // postprocess