Refactor of main test code

- Some helper structs added to track state
- Add new trailer checking routine
- Add extra test for upgrade+revert with failure
- Misc improvements

Signed-off-by: Fabio Utzig <utzig@apache.org>
diff --git a/sim/src/main.rs b/sim/src/main.rs
index 086cc21..16400a3 100644
--- a/sim/src/main.rs
+++ b/sim/src/main.rs
@@ -110,6 +110,7 @@
 
         let align = args.flag_align.map(|x| x.0).unwrap_or(1);
 
+
         let device = match args.flag_device {
             None => panic!("Missing mandatory device argument"),
             Some(dev) => dev,
@@ -149,8 +150,6 @@
     }
 
     fn run_single(&mut self, device: DeviceName, align: u8) {
-        let mut failed = false;
-
         warn!("Running on device {} with alignment {}", device, align);
 
         let (mut flash, areadesc) = match device {
@@ -208,103 +207,57 @@
         assert_eq!(slot1_base, slot0_base + slot0_len);
         assert_eq!(scratch_base, slot1_base + slot1_len);
 
+        let offset_from_end = c::boot_magic_sz() + c::boot_max_align() * 2;
+
         // println!("Areas: {:#?}", areadesc.get_c());
 
         // Install the boot trailer signature, so that the code will start an upgrade.
         // TODO: This must be a multiple of flash alignment, add support for an image that is smaller,
         // and just gets padded.
-        let primary = install_image(&mut flash, slot0_base, 32784);
 
-        // Install an upgrade image.
-        let upgrade = install_image(&mut flash, slot1_base, 41928);
+        // Create original and upgrade images
+        let slot0 = SlotInfo {
+            base_off: slot0_base as usize,
+            trailer_off: slot1_base - offset_from_end,
+        };
+
+        let slot1 = SlotInfo {
+            base_off: slot1_base as usize,
+            trailer_off: scratch_base - offset_from_end,
+        };
+
+        let images = Images {
+            slot0: slot0,
+            slot1: slot1,
+            primary: install_image(&mut flash, slot0_base, 32784),
+            upgrade: install_image(&mut flash, slot1_base, 41928),
+        };
+
+        let mut failed = false;
 
         // Set an alignment, and position the magic value.
         c::set_sim_flash_align(align);
 
-        // Mark the upgrade as ready to install.  (This looks like it might be a bug in the code,
-        // however.)
-        mark_upgrade(&mut flash, scratch_base - c::boot_magic_sz() as usize);
+        mark_upgrade(&mut flash, &images.slot1);
 
-        let (fl2, total_count) = try_upgrade(&flash, &areadesc, None);
-        info!("First boot, count={}", total_count);
-        if !verify_image(&fl2, slot0_base, &upgrade) {
-            error!("Image mismatch after first boot");
-            // This isn't really recoverable, and more tests aren't likely to reveal much.
-            self.failures += 1;
-            return;
-        }
-
-        let mut bad = 0;
-        // Let's try an image halfway through.
-        for i in 1 .. total_count {
-            info!("Try interruption at {}", i);
-            let (fl3, count) = try_upgrade(&flash, &areadesc, Some(i));
-            info!("Second boot, count={}", count);
-            if !verify_image(&fl3, slot0_base, &upgrade) {
-                warn!("FAIL at step {} of {}", i, total_count);
-                bad += 1;
-            }
-            if Caps::SwapUpgrade.present() {
-                if !verify_image(&fl3, slot1_base, &primary) {
-                    warn!("Slot 1 FAIL at step {} of {}", i, total_count);
-                    bad += 1;
-                }
-            }
-        }
-        error!("{} out of {} failed {:.2}%",
-               bad, total_count,
-               bad as f32 * 100.0 / total_count as f32);
-        if bad > 0 {
-            failed = true;
-        }
-
-        let (fl4, total_counts) = try_random_fails(&flash, &areadesc, total_count, 5);
-        info!("Random interruptions at reset points={:?}", total_counts);
-        let slot0_ok = verify_image(&fl4, slot0_base, &upgrade);
-        let slot1_ok = if Caps::SwapUpgrade.present() {
-            verify_image(&fl4, slot1_base, &primary)
-        } else {
-            true
+        // upgrades without fails, counts number of flash operations
+        let total_count = match run_basic_upgrade(&flash, &areadesc, &images) {
+            Ok(v)  => v,
+            Err(_) => {
+                self.failures += 1;
+                return;
+            },
         };
-        if !slot0_ok /* || !slot1_ok */ {
-            error!("Image mismatch after random interrupts: slot0={} slot1={}",
-                   if slot0_ok { "ok" } else { "fail" },
-                   if slot1_ok { "ok" } else { "fail" });
-            self.failures += 1;
-            return;
-        }
 
-        if Caps::SwapUpgrade.present() {
-            for count in 2 .. 5 {
-                info!("Try revert: {}", count);
-                let fl2 = try_revert(&flash, &areadesc, count);
-                if !verify_image(&fl2, slot0_base, &primary) {
-                    warn!("Revert failure on count {}", count);
-                    failed = true;
-                }
-            }
-        }
+        failed |= run_basic_revert(&flash, &areadesc, &images);
+        failed |= run_revert_with_fails(&flash, &areadesc, &images, total_count);
+        failed |= run_perm_with_fails(&flash, &areadesc, &images, total_count);
+        failed |= run_perm_with_random_fails(&flash, &areadesc, &images,
+                                                    total_count, 5);
+        failed |= run_norevert(&flash, &areadesc, &images);
 
-        info!("Try norevert");
-        let fl2 = try_norevert(&flash, &areadesc);
-        if !verify_image(&fl2, slot0_base, &upgrade) {
-            warn!("No revert failed");
-            failed = true;
-        }
+        //show_flash(&flash);
 
-        /*
-        // show_flash(&flash);
-
-        println!("First boot for upgrade");
-        // c::set_flash_counter(570);
-        c::boot_go(&mut flash, &areadesc);
-        // println!("{} flash ops", c::get_flash_counter());
-
-        verify_image(&flash, slot0_base, &upgrade);
-
-        println!("\n------------------\nSecond boot");
-        c::boot_go(&mut flash, &areadesc);
-        */
         if failed {
             self.failures += 1;
         } else {
@@ -313,20 +266,190 @@
     }
 }
 
-/// Test a boot, optionally stopping after 'n' flash options.  Returns a count of the number of
-/// flash operations done total.
-fn try_upgrade(flash: &Flash, areadesc: &AreaDesc, stop: Option<i32>) -> (Flash, i32) {
+/// A simple upgrade without forced failures.
+///
+/// Returns the number of flash operations which can later be used to
+/// inject failures at chosen steps.
+fn run_basic_upgrade(flash: &Flash, areadesc: &AreaDesc, images: &Images)
+                     -> Result<i32, ()> {
+    let (fl, total_count) = try_upgrade(&flash, &areadesc, &images, None);
+    info!("Total flash operation count={}", total_count);
+
+    if !verify_image(&fl, images.slot0.base_off, &images.upgrade) {
+        warn!("Image mismatch after first boot");
+        Err(())
+    } else {
+        Ok(total_count)
+    }
+}
+
+fn run_basic_revert(flash: &Flash, areadesc: &AreaDesc, images: &Images) -> bool {
+    let mut fails = 0;
+
+    if Caps::SwapUpgrade.present() {
+        for count in 2 .. 5 {
+            info!("Try revert: {}", count);
+            let fl = try_revert(&flash, &areadesc, count);
+            if !verify_image(&fl, images.slot0.base_off, &images.primary) {
+                warn!("Revert failure on count {}", count);
+                fails += 1;
+            }
+        }
+    }
+
+    fails > 0
+}
+
+fn run_perm_with_fails(flash: &Flash, areadesc: &AreaDesc, images: &Images,
+                       total_flash_ops: i32) -> bool {
+    let mut fails = 0;
+
+    // Let's try an image halfway through.
+    for i in 1 .. total_flash_ops {
+        info!("Try interruption at {}", i);
+        let (fl, count) = try_upgrade(&flash, &areadesc, &images, Some(i));
+        info!("Second boot, count={}", count);
+        if !verify_image(&fl, images.slot0.base_off, &images.upgrade) {
+            warn!("FAIL at step {} of {}", i, total_flash_ops);
+            fails += 1;
+        }
+
+        if !verify_trailer(&fl, images.slot0.trailer_off, MAGIC_VALID, IMAGE_OK,
+                           COPY_DONE) {
+            warn!("Mismatched trailer for Slot 0");
+            fails += 1;
+        }
+
+        if !verify_trailer(&fl, images.slot1.trailer_off, MAGIC_UNSET, UNSET,
+                           UNSET) {
+            warn!("Mismatched trailer for Slot 1");
+            fails += 1;
+        }
+
+        if Caps::SwapUpgrade.present() {
+            if !verify_image(&fl, images.slot1.base_off, &images.primary) {
+                warn!("Slot 1 FAIL at step {} of {}", i, total_flash_ops);
+                fails += 1;
+            }
+        }
+    }
+
+    info!("{} out of {} failed {:.2}%", fails, total_flash_ops,
+           fails as f32 * 100.0 / total_flash_ops as f32);
+
+    fails > 0
+}
+
+fn run_perm_with_random_fails(flash: &Flash, areadesc: &AreaDesc,
+                              images: &Images, total_flash_ops: i32,
+                              total_fails: usize) -> bool {
+    let mut fails = 0;
+    let (fl, total_counts) = try_random_fails(&flash, &areadesc, &images,
+                                              total_flash_ops, total_fails);
+    info!("Random interruptions at reset points={:?}", total_counts);
+
+    let slot0_ok = verify_image(&fl, images.slot0.base_off, &images.upgrade);
+    let slot1_ok = if Caps::SwapUpgrade.present() {
+        verify_image(&fl, images.slot1.base_off, &images.primary)
+    } else {
+        true
+    };
+    if !slot0_ok || !slot1_ok {
+        error!("Image mismatch after random interrupts: slot0={} slot1={}",
+               if slot0_ok { "ok" } else { "fail" },
+               if slot1_ok { "ok" } else { "fail" });
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot0.trailer_off, MAGIC_VALID, IMAGE_OK,
+                       COPY_DONE) {
+        error!("Mismatched trailer for Slot 0");
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot1.trailer_off, MAGIC_UNSET, UNSET,
+                       UNSET) {
+        error!("Mismatched trailer for Slot 1");
+        fails += 1;
+    }
+
+    fails > 0
+}
+
+fn run_revert_with_fails(flash: &Flash, areadesc: &AreaDesc, images: &Images,
+                         total_count: i32) -> bool {
+    let mut fails = 0;
+
+    if Caps::SwapUpgrade.present() {
+        for i in 1 .. (total_count - 1) {
+            info!("Try interruption at {}", i);
+            if try_revert_with_fail_at(&flash, &areadesc, &images, i) {
+                fails += 1;
+            }
+        }
+    }
+
+    fails > 0
+}
+
+fn run_norevert(flash: &Flash, areadesc: &AreaDesc, images: &Images) -> bool {
+    let mut fl = flash.clone();
+    let mut fails = 0;
+
+    info!("Try norevert");
+    c::set_flash_counter(0);
+
+    // First do a normal upgrade...
+    if c::boot_go(&mut fl, &areadesc) != 0 {
+        warn!("Failed first boot");
+        fails += 1;
+    }
+
+    if !verify_image(&fl, images.slot0.base_off, &images.upgrade) {
+        warn!("Slot 0 image verification FAIL");
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot0.trailer_off, MAGIC_VALID, UNSET,
+                       COPY_DONE) {
+        warn!("Mismatched trailer for Slot 0");
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot1.trailer_off, MAGIC_UNSET, UNSET,
+                       UNSET) {
+        warn!("Mismatched trailer for Slot 1");
+        fails += 1;
+    }
+
+    // Marks image in slot0 as permanent, no revert should happen...
+    mark_permanent_upgrade(&mut fl, &images.slot0);
+
+    if c::boot_go(&mut fl, &areadesc) != 0 {
+        warn!("Failed second boot");
+        fails += 1;
+    }
+
+    if !verify_trailer(&fl, images.slot0.trailer_off, MAGIC_VALID, IMAGE_OK,
+                       COPY_DONE) {
+        warn!("Mismatched trailer for Slot 0");
+        fails += 1;
+    }
+    if !verify_image(&fl, images.slot0.base_off, &images.upgrade) {
+        warn!("Failed image verification");
+        fails += 1;
+    }
+
+    fails > 0
+}
+
+/// Test a boot, optionally stopping after 'n' flash options.  Returns a count
+/// of the number of flash operations done total.
+fn try_upgrade(flash: &Flash, areadesc: &AreaDesc, images: &Images,
+               stop: Option<i32>) -> (Flash, i32) {
     // Clone the flash to have a new copy.
     let mut fl = flash.clone();
 
-    // mark permanent upgrade
-    let ok = [1u8, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
-    let (base, _) = areadesc.find(FlashId::ImageScratch);
-    let align = c::get_sim_flash_align() as usize;
-    fl.write(base - c::boot_magic_sz() - c::boot_max_align(), &ok[..align]).unwrap();
+    mark_permanent_upgrade(&mut fl, &images.slot1);
 
     c::set_flash_counter(stop.unwrap_or(0));
-    let (first_interrupted, cnt1) = match c::boot_go(&mut fl, &areadesc) {
+    let (first_interrupted, count) = match c::boot_go(&mut fl, &areadesc) {
         -0x13579 => (true, stop.unwrap()),
         0 => (false, -c::get_flash_counter()),
         x => panic!("Unknown return: {}", x),
@@ -342,9 +465,7 @@
         }
     }
 
-    let cnt2 = cnt1 - c::get_flash_counter();
-
-    (fl, cnt2)
+    (fl, count - c::get_flash_counter())
 }
 
 fn try_revert(flash: &Flash, areadesc: &AreaDesc, count: usize) -> Flash {
@@ -359,30 +480,85 @@
     fl
 }
 
-fn try_norevert(flash: &Flash, areadesc: &AreaDesc) -> Flash {
+fn try_revert_with_fail_at(flash: &Flash, areadesc: &AreaDesc, images: &Images,
+                           stop: i32) -> bool {
     let mut fl = flash.clone();
-    c::set_flash_counter(0);
-    let align = c::get_sim_flash_align() as usize;
+    let mut x: i32;
+    let mut fails = 0;
 
-    assert_eq!(c::boot_go(&mut fl, &areadesc), 0);
-    // Write boot_ok
-    let ok = [1u8, 0, 0, 0, 0, 0, 0, 0];
-    let (slot0_base, slot0_len) = areadesc.find(FlashId::Image0);
-    fl.write(slot0_base + slot0_len - c::boot_magic_sz() - c::boot_max_align(),
-             &ok[..align]).unwrap();
-    assert_eq!(c::boot_go(&mut fl, &areadesc), 0);
-    fl
+    c::set_flash_counter(stop);
+    x = c::boot_go(&mut fl, &areadesc);
+    if x != -0x13579 {
+        warn!("Should have stopped at interruption point");
+        fails += 1;
+    }
+
+    if !verify_trailer(&fl, images.slot0.trailer_off, None, None, UNSET) {
+        warn!("copy_done should be unset");
+        fails += 1;
+    }
+
+    c::set_flash_counter(0);
+    x = c::boot_go(&mut fl, &areadesc);
+    if x != 0 {
+        warn!("Should have finished upgrade");
+        fails += 1;
+    }
+
+    if !verify_image(&fl, images.slot0.base_off, &images.upgrade) {
+        warn!("Image in slot 0 before revert is invalid at stop={}", stop);
+        fails += 1;
+    }
+    if !verify_image(&fl, images.slot1.base_off, &images.primary) {
+        warn!("Image in slot 1 before revert is invalid at stop={}", stop);
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot0.trailer_off, MAGIC_VALID, UNSET,
+                       COPY_DONE) {
+        warn!("Mismatched trailer for Slot 0 before revert");
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot1.trailer_off, MAGIC_UNSET, UNSET,
+                       UNSET) {
+        warn!("Mismatched trailer for Slot 1 before revert");
+        fails += 1;
+    }
+
+    // Do Revert
+    c::set_flash_counter(0);
+    x = c::boot_go(&mut fl, &areadesc);
+    if x != 0 {
+        warn!("Should have finished a revert");
+        fails += 1;
+    }
+
+    if !verify_image(&fl, images.slot0.base_off, &images.primary) {
+        warn!("Image in slot 0 after revert is invalid at stop={}", stop);
+        fails += 1;
+    }
+    if !verify_image(&fl, images.slot1.base_off, &images.upgrade) {
+        warn!("Image in slot 1 after revert is invalid at stop={}", stop);
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot0.trailer_off, MAGIC_VALID, IMAGE_OK,
+                       COPY_DONE) {
+        warn!("Mismatched trailer for Slot 1 after revert");
+        fails += 1;
+    }
+    if !verify_trailer(&fl, images.slot1.trailer_off, MAGIC_UNSET, UNSET,
+                       UNSET) {
+        warn!("Mismatched trailer for Slot 1 after revert");
+        fails += 1;
+    }
+
+    fails > 0
 }
 
-fn try_random_fails(flash: &Flash, areadesc: &AreaDesc, total_ops: i32, 
-                    count: usize) -> (Flash, Vec<i32>) {
+fn try_random_fails(flash: &Flash, areadesc: &AreaDesc, images: &Images,
+                    total_ops: i32,  count: usize) -> (Flash, Vec<i32>) {
     let mut fl = flash.clone();
 
-    // mark permanent upgrade
-    let ok = [1u8, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
-    let (base, _) = areadesc.find(FlashId::ImageScratch);
-    let align = c::get_sim_flash_align() as usize;
-    fl.write(base - c::boot_magic_sz() - c::boot_max_align(), &ok[..align]).unwrap();
+    mark_permanent_upgrade(&mut fl, &images.slot1);
 
     let mut rng = rand::thread_rng();
     let mut resets = vec![0i32; count];
@@ -414,7 +590,7 @@
 fn show_flash(flash: &Flash) {
     println!("---- Flash configuration ----");
     for sector in flash.sector_iter() {
-        println!("    {:2}: 0x{:08x}, 0x{:08x}",
+        println!("    {:3}: 0x{:08x}, 0x{:08x}",
                  sector.num, sector.base, sector.size);
     }
     println!("");
@@ -484,6 +660,53 @@
     }
 }
 
+fn verify_trailer(flash: &Flash, offset: usize,
+                  magic: Option<&[u8]>, image_ok: Option<u8>,
+                  copy_done: Option<u8>) -> bool {
+    let mut copy = vec![0u8; c::boot_magic_sz() + c::boot_max_align() * 2];
+    let mut failed = false;
+
+    flash.read(offset, &mut copy).unwrap();
+
+    failed |= match magic {
+        Some(v) => {
+            if &copy[16..] != v  {
+                warn!("\"magic\" mismatch at {:#x}", offset);
+                true
+            } else {
+                false
+            }
+        },
+        None => false,
+    };
+
+    failed |= match image_ok {
+        Some(v) => {
+            if copy[8] != v {
+                warn!("\"image_ok\" mismatch at {:#x}", offset);
+                true
+            } else {
+                false
+            }
+        },
+        None => false,
+    };
+
+    failed |= match copy_done {
+        Some(v) => {
+            if copy[0] != v {
+                warn!("\"copy_done\" mismatch at {:#x}", offset);
+                true
+            } else {
+                false
+            }
+        },
+        None => false,
+    };
+
+    !failed
+}
+
 /// The image header
 #[repr(C)]
 pub struct ImageHeader {
@@ -509,13 +732,41 @@
     build_num: u32,
 }
 
+struct SlotInfo {
+    base_off: usize,
+    trailer_off: usize,
+}
+
+struct Images {
+    slot0: SlotInfo,
+    slot1: SlotInfo,
+    primary: Vec<u8>,
+    upgrade: Vec<u8>,
+}
+
+const MAGIC_VALID: Option<&[u8]> = Some(&[0x77, 0xc2, 0x95, 0xf3,
+                                          0x60, 0xd2, 0xef, 0x7f,
+                                          0x35, 0x52, 0x50, 0x0f,
+                                          0x2c, 0xb6, 0x79, 0x80]);
+const MAGIC_UNSET: Option<&[u8]> = Some(&[0xff; 16]);
+
+const COPY_DONE: Option<u8> = Some(1);
+const IMAGE_OK: Option<u8> = Some(1);
+const UNSET: Option<u8> = Some(0xff);
+
 /// Write out the magic so that the loader tries doing an upgrade.
-fn mark_upgrade(flash: &mut Flash, offset: usize) {
-    let magic = vec![0x77, 0xc2, 0x95, 0xf3,
-                     0x60, 0xd2, 0xef, 0x7f,
-                     0x35, 0x52, 0x50, 0x0f,
-                     0x2c, 0xb6, 0x79, 0x80];
-    flash.write(offset, &magic).unwrap();
+fn mark_upgrade(flash: &mut Flash, slot: &SlotInfo) {
+    let offset = slot.trailer_off + c::boot_max_align() * 2;
+    flash.write(offset, MAGIC_VALID.unwrap()).unwrap();
+}
+
+/// Writes the image_ok flag which, guess what, tells the bootloader
+/// the this image is ok (not a test, and no revert is to be performed).
+fn mark_permanent_upgrade(flash: &mut Flash, slot: &SlotInfo) {
+    let ok = [1u8, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
+    let align = c::get_sim_flash_align() as usize;
+    let off = slot.trailer_off + c::boot_max_align();
+    flash.write(off, &ok[..align]).unwrap();
 }
 
 // Drop some pseudo-random gibberish onto the data.