download: defer validity using saved partitions rather than zeroes

If we're installing to disk (rather than downloading an artifact to a
file), instead of writing zeroes to the first MiB, write out the saved
partitions.  That way we won't lose them if we die on a signal or if the
machine crashes.

src/blockdev.rs

@@ -1173,9 +1173,15 @@ mod tests {
 
         let mut base = make_disk(512, &base_parts);
         for (testnum, (filter, expected_blank, expected_image)) in tests.iter().enumerate() {
-            // perform writes in the same order as install()
-            // first, try merging with image disk
+            // 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_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();
             assert!(!disk_has_mbr(&mut disk), "test {}", testnum);
@@ -1196,13 +1202,6 @@ mod tests {
                 "test {}",
                 testnum
             );
-
-            // then, try overwriting on blank disk
-            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_partitions_eq(expected_blank, &result, &format!("test {} blank", testnum));
         }
 
         // test merging with unformatted initial disk

src/download.rs

@@ -23,7 +23,7 @@ use std::path::{Path, PathBuf};
 use std::result;
 use std::time::{Duration, Instant};
 
-use crate::blockdev::detect_formatted_sector_size;
+use crate::blockdev::{detect_formatted_sector_size, SavedPartitions};
 use crate::cmdline::*;
 use crate::errors::*;
 use crate::io::*;
@@ -151,6 +151,7 @@ fn write_image_and_sig(
         decompress,
         None,
         None,
+        None,
     )?;
 
     // write signature, if relevant
@@ -170,17 +171,19 @@ fn write_image_and_sig(
 }
 
 /// Copy the image to disk and verify its signature.
+#[allow(clippy::too_many_arguments)]
 pub fn write_image<F>(
     source: &mut ImageSource,
     dest: &mut File,
     dest_path: &Path,
     image_copy: F,
     decompress: bool,
+    saved: Option<&SavedPartitions>,
     byte_limit: Option<(u64, String)>, // limit and explanation
     expected_sector_size: Option<NonZeroU32>,
 ) -> Result<()>
 where
-    F: FnOnce(&[u8], &mut dyn Read, &mut File, &Path) -> Result<()>,
+    F: FnOnce(&[u8], &mut dyn Read, &mut File, &Path, Option<&SavedPartitions>) -> Result<()>,
 {
     // wrap source for GPG verification
     let mut verify_reader: Box<dyn Read> = {
@@ -238,7 +241,7 @@ where
     }
 
     // call the callback to copy the image
-    image_copy(&first_mb, &mut limit_reader, dest, dest_path)?;
+    image_copy(&first_mb, &mut limit_reader, dest, dest_path, saved)?;
 
     // finish I/O before closing the progress bar
     dest.sync_all().chain_err(|| "syncing data to disk")?;
@@ -251,12 +254,25 @@ pub fn image_copy_default(
     source: &mut dyn Read,
     dest: &mut File,
     _dest_path: &Path,
+    saved: Option<&SavedPartitions>,
 ) -> Result<()> {
-    // Cache the first MiB and write zeroes to dest instead.  This ensures
-    // that the disk image can't be used accidentally before its GPG signature
-    // is verified.
-    dest.write_all(&[0u8; 1024 * 1024])
-        .chain_err(|| "clearing first MiB of disk")?;
+    // Don't write the first MiB yet.  This ensures that the disk image
+    // can't be used accidentally before its GPG signature is verified.  If
+    // this is a real disk, write the saved partitions (so they don't get
+    // lost if we crash), and otherwise write zeroes.
+    match saved {
+        Some(saved) => {
+            saved
+                .overwrite(dest)
+                .chain_err(|| "overwriting disk partition table")?;
+            dest.seek(SeekFrom::Start(1024 * 1024))
+                .chain_err(|| "seeking disk")?;
+        }
+        None => dest
+            .write_all(&[0u8; 1024 * 1024])
+            .chain_err(|| "clearing first MiB of disk")?,
+    };
+    dest.sync_all().chain_err(|| "syncing data to disk")?;
 
     // do the rest of the copy
     // This physically writes any runs of zeroes, rather than sparsifying,
@@ -434,6 +450,7 @@ mod tests {
             &dest_path,
             image_copy_default,
             false,
+            None,
             Some((offset, explanation.into())),
             None,
         )

src/install.rs

@@ -166,6 +166,7 @@ fn write_disk(
         Path::new(&config.device),
         image_copy,
         true,
+        Some(&saved),
         saved
             .get_offset()?
             .map(|(offset, desc)| (offset, format!("collision with {}", desc))),

src/s390x/dasd.rs

@@ -21,7 +21,7 @@ use std::os::unix::io::AsRawFd;
 use std::path::Path;
 use std::process::{Command, Stdio};
 
-use crate::blockdev::{get_sector_size, udev_settle};
+use crate::blockdev::{get_sector_size, udev_settle, SavedPartitions};
 use crate::cmdline::*;
 use crate::errors::*;
 use crate::io::{copy_exactly_n, BUFFER_SIZE};
@@ -54,6 +54,7 @@ pub fn image_copy_s390x(
     source: &mut dyn Read,
     dest_file: &mut File,
     dest_path: &Path,
+    _saved: Option<&SavedPartitions>,
 ) -> Result<()> {
     let (ranges, partitions) = partition_ranges(first_mb, dest_file)?;
     make_partitions(