Infineon: Add cyw20829 platform, shared slot feature, json memory map, psoc6 xip

Based in 1.8.0 release of MCUBoot library

This commit adds CYW20829 Infineon platform support with following capabilities:
1. Overwrite and swap upgrade mode support
2. Multi-image with up to 4 images
3. Hardware security counter is supported for CYW20829 platform

Add XIP support for PSOC6 platform - place BOOT slot in external memory and execute it in place using SMIF in XIP mode

and some new features for Infineon devices.

1. Shared upgrade slot feature - use one shared area for upgrade slots of multiple images
2. Memory map defined using JSON file - define memory regions for bootloader and user app in conventional way using JSON file
diff --git a/sim/src/caps.rs b/sim/src/caps.rs
index c626ee6..3fbf4c3 100644
--- a/sim/src/caps.rs
+++ b/sim/src/caps.rs
@@ -1,6 +1,6 @@
 // Copyright (c) 2017-2019 Linaro LTD
 // Copyright (c) 2019 JUUL Labs
-// Copyright (c) 2019 Arm Limited
+// Copyright (c) 2019-2021 Arm Limited
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -25,7 +25,10 @@
     DowngradePrevention  = (1 << 12),
     EncX25519            = (1 << 13),
     Bootstrap            = (1 << 14),
-    SwapUsingStatus      = (1 << 15),
+    Aes256               = (1 << 15),
+    RamLoad              = (1 << 16),
+    DirectXip            = (1 << 17),
+    SwapUsingStatus      = (1 << 18),
 }
 
 impl Caps {
@@ -39,6 +42,12 @@
     pub fn get_num_images() -> usize {
         (unsafe { bootutil_get_num_images() }) as usize
     }
+
+    /// Query if this configuration performs some kind of upgrade by writing to flash.
+    pub fn modifies_flash() -> bool {
+        // All other configurations perform upgrades by writing to flash.
+        !(Self::RamLoad.present() || Self::DirectXip.present())
+    }
 }
 
 extern "C" {
diff --git a/sim/src/image.rs b/sim/src/image.rs
index 6bd14c5..29ddedd 100644
--- a/sim/src/image.rs
+++ b/sim/src/image.rs
@@ -1,6 +1,6 @@
 // Copyright (c) 2019 Linaro LTD
 // Copyright (c) 2019-2020 JUUL Labs
-// Copyright (c) 2019 Arm Limited
+// Copyright (c) 2019-2021 Arm Limited
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -19,22 +19,26 @@
     rngs::SmallRng,
 };
 use std::{
-    collections::HashSet,
+    collections::{BTreeMap, HashSet},
     io::{Cursor, Write},
     mem,
     slice,
 };
-use aes_ctr::{
+use aes::{
+    Aes128,
     Aes128Ctr,
-    stream_cipher::{
-        generic_array::GenericArray,
-        NewStreamCipher,
-        SyncStreamCipher,
-    },
+    Aes256,
+    Aes256Ctr,
+    NewBlockCipher,
 };
+use cipher::{
+    FromBlockCipher,
+    generic_array::GenericArray,
+    StreamCipher,
+    };
 
 use simflash::{Flash, SimFlash, SimMultiFlash};
-use mcuboot_sys::{c, AreaDesc, FlashId};
+use mcuboot_sys::{c, AreaDesc, FlashId, RamBlock};
 use crate::{
     ALL_DEVICES,
     DeviceName,
@@ -50,6 +54,11 @@
     UpgradeInfo,
 };
 use crate::tlv::{ManifestGen, TlvGen, TlvFlags};
+use typenum::{U32, U16};
+
+/// For testing, use a non-zero offset for the ram-load, to make sure the offset is getting used
+/// properly, but the value is not really that important.
+const RAM_LOAD_ADDR: u32 = 1024;
 
 /// A builder for Images.  This describes a single run of the simulator,
 /// capturing the configuration of a particular set of devices, including
@@ -59,6 +68,7 @@
     flash: SimMultiFlash,
     areadesc: AreaDesc,
     slots: Vec<[SlotInfo; 2]>,
+    ram: RamData,
 }
 
 /// Images represents the state of a simulation for a given set of images.
@@ -69,6 +79,7 @@
     areadesc: AreaDesc,
     images: Vec<OneImage>,
     total_count: Option<i32>,
+    ram: RamData,
 }
 
 /// When doing multi-image, there is an instance of this information for
@@ -83,10 +94,34 @@
 /// is just the unencrypted payload.  For encrypted images, we store both
 /// the encrypted and the plaintext.
 struct ImageData {
+    size: usize,
     plain: Vec<u8>,
     cipher: Option<Vec<u8>>,
 }
 
+/// For the RamLoad test cases, we need a contiguous area of RAM to load these images into.  For
+/// multi-image builds, these may not correspond with the offsets.  This has to be computed early,
+/// before images are built, because each image contains the offset where the image is to be loaded
+/// in the header, which is contained within the signature.
+#[derive(Clone, Debug)]
+struct RamData {
+    places: BTreeMap<SlotKey, SlotPlace>,
+    total: u32,
+}
+
+/// Every slot is indexed by this key.
+#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
+struct SlotKey {
+    dev_id: u8,
+    base_off: usize,
+}
+
+#[derive(Clone, Debug)]
+struct SlotPlace {
+    offset: u32,
+    size: u32,
+}
+
 impl ImagesBuilder {
     /// Construct a new image builder for the given device.  Returns
     /// Some(builder) if is possible to test this configuration, or None if
@@ -148,10 +183,13 @@
             slots.push([primary, secondary]);
         }
 
+        let ram = RamData::new(&slots);
+
         Ok(ImagesBuilder {
             flash,
             areadesc,
             slots,
+            ram,
         })
     }
 
@@ -174,16 +212,17 @@
     pub fn make_no_upgrade_image(self, deps: &DepTest) -> Images {
         let num_images = self.num_images();
         let mut flash = self.flash;
+        let ram = self.ram.clone();  // TODO: This is wasteful.
         let images = self.slots.into_iter().enumerate().map(|(image_num, slots)| {
             let dep: Box<dyn Depender> = if num_images > 1 {
                 Box::new(PairDep::new(num_images, image_num, deps))
             } else {
                 Box::new(BoringDep::new(image_num, deps))
             };
-            let primaries = install_image(&mut flash, &slots[0], 42784, &*dep, false);
+            let primaries = install_image(&mut flash, &slots[0], 42784, &ram, &*dep, false);
             let upgrades = match deps.depends[image_num] {
                 DepType::NoUpgrade => install_no_image(),
-                _ => install_image(&mut flash, &slots[1], 46928, &*dep, false)
+                _ => install_image(&mut flash, &slots[1], 46928, &ram, &*dep, false)
             };
             OneImage {
                 slots,
@@ -196,6 +235,7 @@
             areadesc: self.areadesc,
             images,
             total_count: None,
+            ram: self.ram,
         }
     }
 
@@ -205,6 +245,11 @@
             mark_upgrade(&mut images.flash, &image.slots[1]);
         }
 
+        // The count is meaningless if no flash operations are performed.
+        if !Caps::modifies_flash() {
+            return images;
+        }
+
         // upgrades without fails, counts number of flash operations
         let total_count = match images.run_basic_upgrade(permanent) {
             Some(v)  => v,
@@ -222,10 +267,11 @@
 
     pub fn make_bad_secondary_slot_image(self) -> Images {
         let mut bad_flash = self.flash;
+        let ram = self.ram.clone(); // TODO: Avoid this clone.
         let images = self.slots.into_iter().enumerate().map(|(image_num, slots)| {
             let dep = BoringDep::new(image_num, &NO_DEPS);
-            let primaries = install_image(&mut bad_flash, &slots[0], 32784, &dep, false);
-            let upgrades = install_image(&mut bad_flash, &slots[1], 41928, &dep, true);
+            let primaries = install_image(&mut bad_flash, &slots[0], 32784, &ram, &dep, false);
+            let upgrades = install_image(&mut bad_flash, &slots[1], 41928, &ram, &dep, true);
             OneImage {
                 slots,
                 primaries,
@@ -236,14 +282,16 @@
             areadesc: self.areadesc,
             images,
             total_count: None,
+            ram: self.ram,
         }
     }
 
     pub fn make_erased_secondary_image(self) -> Images {
         let mut flash = self.flash;
+        let ram = self.ram.clone(); // TODO: Avoid this clone.
         let images = self.slots.into_iter().enumerate().map(|(image_num, slots)| {
             let dep = BoringDep::new(image_num, &NO_DEPS);
-            let primaries = install_image(&mut flash, &slots[0], 32784, &dep, false);
+            let primaries = install_image(&mut flash, &slots[0], 32784, &ram, &dep, false);
             let upgrades = install_no_image();
             OneImage {
                 slots,
@@ -255,15 +303,17 @@
             areadesc: self.areadesc,
             images,
             total_count: None,
+            ram: self.ram,
         }
     }
 
     pub fn make_bootstrap_image(self) -> Images {
         let mut flash = self.flash;
+        let ram = self.ram.clone(); // TODO: Avoid this clone.
         let images = self.slots.into_iter().enumerate().map(|(image_num, slots)| {
             let dep = BoringDep::new(image_num, &NO_DEPS);
             let primaries = install_no_image();
-            let upgrades = install_image(&mut flash, &slots[1], 32784, &dep, false);
+            let upgrades = install_image(&mut flash, &slots[1], 32784, &ram, &dep, false);
             OneImage {
                 slots,
                 primaries,
@@ -274,15 +324,14 @@
             areadesc: self.areadesc,
             images,
             total_count: None,
+            ram: self.ram,
         }
     }
 
     /// Build the Flash and area descriptor for a given device.
     pub fn make_device(device: DeviceName, align: usize, erased_val: u8) -> (SimMultiFlash, AreaDesc, &'static [Caps]) {
-        info!(" +++ Make new device...");
         match device {
             DeviceName::Stm32f4 => {
-                info!("DeviceName::Stm32f4");
                 // STM style flash.  Large sectors, with a large scratch area.
                 let dev = SimFlash::new(vec![16 * 1024, 16 * 1024, 16 * 1024, 16 * 1024,
                                         64 * 1024,
@@ -300,7 +349,6 @@
                 (flash, areadesc, &[Caps::SwapUsingMove, Caps::SwapUsingStatus])
             }
             DeviceName::K64f => {
-                info!("DeviceName::K64f");
                 // NXP style flash.  Small sectors, one small sector for scratch.
                 let dev = SimFlash::new(vec![4096; 128], align as usize, erased_val);
 
@@ -316,7 +364,6 @@
                 (flash, areadesc, &[Caps::SwapUsingStatus])
             }
             DeviceName::K64fBig => {
-                info!("DeviceName::K64fBig");
                 // Simulating an STM style flash on top of an NXP style flash.  Underlying flash device
                 // uses small sectors, but we tell the bootloader they are large.
                 let dev = SimFlash::new(vec![4096; 128], align as usize, erased_val);
@@ -333,7 +380,6 @@
                 (flash, areadesc, &[Caps::SwapUsingMove, Caps::SwapUsingStatus])
             }
             DeviceName::Nrf52840 => {
-                info!("DeviceName::Nrf52840");
                 // Simulating the flash on the nrf52840 with partitions set up so that the scratch size
                 // does not divide into the image size.
                 let dev = SimFlash::new(vec![4096; 128], align as usize, erased_val);
@@ -350,7 +396,6 @@
                 (flash, areadesc, &[Caps::SwapUsingStatus])
             }
             DeviceName::Nrf52840UnequalSlots => {
-                info!("DeviceName::Nrf52840UnequalSlots");
                 let dev = SimFlash::new(vec![4096; 128], align as usize, erased_val);
 
                 let dev_id = 0;
@@ -364,7 +409,6 @@
                 (flash, areadesc, &[Caps::SwapUsingScratch, Caps::OverwriteUpgrade, Caps::SwapUsingStatus])
             }
             DeviceName::Nrf52840SpiFlash => {
-                info!("DeviceName::Nrf52840SpiFlash");
                 // Simulate nrf52840 with external SPI flash. The external SPI flash
                 // has a larger sector size so for now store scratch on that flash.
                 let dev0 = SimFlash::new(vec![4096; 128], align as usize, erased_val);
@@ -384,7 +428,6 @@
                 (flash, areadesc, &[Caps::SwapUsingMove, Caps::SwapUsingStatus])
             }
             DeviceName::K64fMulti => {
-                info!("DeviceName::K64fMulti");
                 // NXP style flash, but larger, to support multiple images.
                 let dev = SimFlash::new(vec![4096; 256], align as usize, erased_val);
 
@@ -401,32 +444,21 @@
                 flash.insert(dev_id, dev);
                 (flash, areadesc, &[Caps::SwapUsingStatus])
             }
-            DeviceName::PSoC6Multi => {
-                info!("DeviceName::PSoC6Multi");
-                // NXP style flash, but larger, to support multiple images.
+            DeviceName::PSoC6 => {
+            // PSoC style flash of 512K, single-image case
+                let dev = SimFlash::new(vec![512; 1024], align as usize, erased_val);
 
+                let dev_id = 0;
                 let mut areadesc = AreaDesc::new();
+                areadesc.add_flash_sectors(dev_id, &dev);
+                areadesc.add_image(0x018000, 0x010000, FlashId::Image0, dev_id);
+                areadesc.add_image(0x028000, 0x010000, FlashId::Image1, dev_id);
+                areadesc.add_image(0x039800, 0x001000, FlashId::ImageScratch, dev_id);
+                areadesc.add_image(0x038000, 0x001800, FlashId::ImageSwapStatus, dev_id);
+
                 let mut flash = SimMultiFlash::new();
-
-                // let dev0 = SimFlash::new(vec![4096; 256], align as usize, 0);
-                let mut dev0 = SimFlash::new(vec![512; 1024], align as usize, 0);
-                dev0.set_verify_writes(false);
-                dev0.set_erase_by_sector(true);
-
-                areadesc.add_flash_sectors(0, &dev0);
-                areadesc.add_image(0x020000, 0x020000, FlashId::Image0, 0);
-                areadesc.add_image(0x040000, 0x020000, FlashId::Image1, 0);
-                areadesc.add_image(0x060000, 0x008000, FlashId::ImageScratch, 0);
-                flash.insert(0, dev0);
-
-                // let dev1 = SimFlash::new(vec![4096; 256], align as usize, erased_val);
-                // areadesc.add_flash_sectors(1, &dev1);
-                // areadesc.add_image(0x080000, 0x020000, FlashId::Image2, 0);
-                // areadesc.add_image(0x0a0000, 0x020000, FlashId::Image3, 1);
-                // flash.insert(1, dev1);
-
-                // (flash, areadesc, &[])
-                (flash, areadesc, &[Caps::SwapUsingScratch, Caps::SwapUsingMove])
+                flash.insert(dev_id, dev);
+                (flash, areadesc, &[Caps::SwapUsingMove])
             }
         }
     }
@@ -461,8 +493,7 @@
         if Caps::Bootstrap.present() {
             info!("Try bootstraping image in the primary");
 
-            let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-            if result != 0 {
+            if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
                 warn!("Failed first boot");
                 fails += 1;
             }
@@ -490,17 +521,21 @@
     /// Test a simple upgrade, with dependencies given, and verify that the
     /// image does as is described in the test.
     pub fn run_check_deps(&self, deps: &DepTest) -> bool {
+        if !Caps::modifies_flash() {
+            return false;
+        }
+
         let (flash, _) = self.try_upgrade(None, true);
 
         self.verify_dep_images(&flash, deps)
     }
 
     fn is_swap_upgrade(&self) -> bool {
-        Caps::SwapUsingScratch.present() || Caps::SwapUsingMove.present() || Caps::SwapUsingStatus.present()
+        Caps::SwapUsingScratch.present() || Caps::SwapUsingMove.present()
     }
 
     pub fn run_basic_revert(&self) -> bool {
-        if Caps::OverwriteUpgrade.present() {
+        if Caps::OverwriteUpgrade.present() || !Caps::modifies_flash() {
             return false;
         }
 
@@ -522,6 +557,10 @@
     }
 
     pub fn run_perm_with_fails(&self) -> bool {
+        if !Caps::modifies_flash() {
+            return false;
+        }
+
         let mut fails = 0;
         let total_flash_ops = self.total_count.unwrap();
 
@@ -563,6 +602,10 @@
     }
 
     pub fn run_perm_with_random_fails(&self, total_fails: usize) -> bool {
+        if !Caps::modifies_flash() {
+            return false;
+        }
+
         let mut fails = 0;
         let total_flash_ops = self.total_count.unwrap();
         let (flash, total_counts) = self.try_random_fails(total_flash_ops, total_fails);
@@ -601,7 +644,7 @@
     }
 
     pub fn run_revert_with_fails(&self) -> bool {
-        if Caps::OverwriteUpgrade.present() {
+        if Caps::OverwriteUpgrade.present() || !Caps::modifies_flash() {
             return false;
         }
 
@@ -621,7 +664,7 @@
     }
 
     pub fn run_norevert(&self) -> bool {
-        if Caps::OverwriteUpgrade.present() {
+        if Caps::OverwriteUpgrade.present() || !Caps::modifies_flash() {
             return false;
         }
 
@@ -631,8 +674,7 @@
         info!("Try norevert");
 
         // First do a normal upgrade...
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed first boot");
             fails += 1;
         }
@@ -665,8 +707,7 @@
             fails += 1;
         }
 
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed second boot");
             fails += 1;
         }
@@ -701,8 +742,7 @@
         info!("Try no downgrade");
 
         // First, do a normal upgrade.
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed first boot");
             fails += 1;
         }
@@ -723,6 +763,11 @@
     // image_ok set while there is no image on the secondary slot, so no revert
     // should ever happen...
     pub fn run_norevert_newimage(&self) -> bool {
+        if !Caps::modifies_flash() {
+            info!("Skipping run_norevert_newimage, as configuration doesn't modify flash");
+            return false;
+        }
+
         let mut flash = self.flash.clone();
         let mut fails = 0;
 
@@ -739,8 +784,7 @@
         }
 
         // Run the bootloader...
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed first boot");
             fails += 1;
         }
@@ -777,6 +821,12 @@
 
         info!("Try upgrade image with bad signature");
 
+        // Only perform this test if an upgrade is expected to happen.
+        if !Caps::modifies_flash() {
+            info!("Skipping upgrade image with bad signature");
+            return false;
+        }
+
         self.mark_upgrades(&mut flash, 0);
         self.mark_permanent_upgrades(&mut flash, 0);
         self.mark_upgrades(&mut flash, 1);
@@ -788,8 +838,7 @@
         }
 
         // Run the bootloader...
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed first boot");
             fails += 1;
         }
@@ -815,6 +864,10 @@
     // Should detect there is a leftover trailer in an otherwise erased
     // secondary slot and erase its trailer.
     pub fn run_secondary_leftover_trailer(&self) -> bool {
+        if !Caps::modifies_flash() {
+            return false;
+        }
+
         let mut flash = self.flash.clone();
         let mut fails = 0;
 
@@ -825,8 +878,7 @@
         self.mark_upgrades(&mut flash, 1);
 
         // Run the bootloader...
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed first boot");
             fails += 1;
         }
@@ -861,7 +913,7 @@
     /// allowing for fails in the status area. This should run to the end
     /// and warn that write fails were detected...
     pub fn run_with_status_fails_complete(&self) -> bool {
-        if !Caps::ValidatePrimarySlot.present() {
+        if !Caps::ValidatePrimarySlot.present() || !Caps::modifies_flash() {
             return false;
         }
 
@@ -873,15 +925,15 @@
         self.mark_permanent_upgrades(&mut flash, 1);
         self.mark_bad_status_with_rate(&mut flash, 0, 1.0);
 
-        let (result, asserts) = c::boot_go(&mut flash, &self.areadesc, None, true);
-        if result != 0 {
+        let result = c::boot_go(&mut flash, &self.areadesc, None, true);
+        if !result.success() {
             warn!("Failed!");
             fails += 1;
         }
 
         // Failed writes to the marked "bad" region don't assert anymore.
         // Any detected assert() is happening in another part of the code.
-        if asserts != 0 {
+        if result.asserts() != 0 {
             warn!("At least one assert() was called");
             fails += 1;
         }
@@ -899,8 +951,7 @@
 
         info!("validate primary slot enabled; \
                re-run of boot_go should just work");
-        let (result, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if result != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Failed!");
             fails += 1;
         }
@@ -916,7 +967,7 @@
     /// allowing for fails in the status area. This should run to the end
     /// and warn that write fails were detected...
     pub fn run_with_status_fails_with_reset(&self) -> bool {
-        if Caps::OverwriteUpgrade.present() {
+        if Caps::OverwriteUpgrade.present() || !Caps::modifies_flash() {
             false
         } else if Caps::ValidatePrimarySlot.present() {
 
@@ -932,7 +983,7 @@
             self.mark_bad_status_with_rate(&mut flash, 0, 0.5);
 
             // Should not fail, writing to bad regions does not assert
-            let (_, asserts) = c::boot_go(&mut flash, &self.areadesc, Some(&mut count), true);
+            let asserts = c::boot_go(&mut flash, &self.areadesc, Some(&mut count), true).asserts();
             if asserts != 0 {
                 warn!("At least one assert() was called");
                 fails += 1;
@@ -941,7 +992,7 @@
             self.reset_bad_status(&mut flash, 0);
 
             info!("Resuming an interrupted swap operation");
-            let (_, asserts) = c::boot_go(&mut flash, &self.areadesc, None, true);
+            let asserts = c::boot_go(&mut flash, &self.areadesc, None, true).asserts();
 
             // This might throw no asserts, for large sector devices, where
             // a single failure writing is indistinguishable from no failure,
@@ -968,7 +1019,7 @@
             self.mark_bad_status_with_rate(&mut flash, 0, 1.0);
 
             // This is expected to fail while writing to bad regions...
-            let (_, asserts) = c::boot_go(&mut flash, &self.areadesc, None, true);
+            let asserts = c::boot_go(&mut flash, &self.areadesc, None, true).asserts();
             if asserts == 0 {
                 warn!("No assert() detected");
                 fails += 1;
@@ -978,6 +1029,78 @@
         }
     }
 
+    /// Test the direct XIP configuration.  With this mode, flash images are never moved, and the
+    /// bootloader merely selects which partition is the proper one to boot.
+    pub fn run_direct_xip(&self) -> bool {
+        if !Caps::DirectXip.present() {
+            return false;
+        }
+
+        // Clone the flash so we can tell if unchanged.
+        let mut flash = self.flash.clone();
+
+        let result = c::boot_go(&mut flash, &self.areadesc, None, true);
+
+        // Ensure the boot was successful.
+        let resp = if let Some(resp) = result.resp() {
+            resp
+        } else {
+            panic!("Boot didn't return a valid result");
+        };
+
+        // This configuration should always try booting from the first upgrade slot.
+        if let Some((offset, _, dev_id)) = self.areadesc.find(FlashId::Image1) {
+            assert_eq!(offset, resp.image_off as usize);
+            assert_eq!(dev_id, resp.flash_dev_id);
+        } else {
+            panic!("Unable to find upgrade image");
+        }
+        false
+    }
+
+    /// Test the ram-loading.
+    pub fn run_ram_load(&self) -> bool {
+        if !Caps::RamLoad.present() {
+            return false;
+        }
+
+        // Clone the flash so we can tell if unchanged.
+        let mut flash = self.flash.clone();
+
+        // Setup ram based on the ram configuration we determined earlier for the images.
+        let ram = RamBlock::new(self.ram.total - RAM_LOAD_ADDR, RAM_LOAD_ADDR);
+
+        // println!("Ram: {:#?}", self.ram);
+
+        // Verify that the images area loaded into this.
+        let result = ram.invoke(|| c::boot_go(&mut flash, &self.areadesc, None, true));
+        if !result.success() {
+            error!("Failed to execute ram-load");
+            return true;
+        }
+
+        // Verify each image.
+        for image in &self.images {
+            let place = self.ram.lookup(&image.slots[0]);
+            let ram_image = ram.borrow_part(place.offset as usize - RAM_LOAD_ADDR as usize,
+                place.size as usize);
+            let src_sz = image.upgrades.size();
+            if src_sz > ram_image.len() {
+                error!("Image ended up too large, nonsensical");
+                return true;
+            }
+            let src_image = &image.upgrades.plain[0..src_sz];
+            let ram_image = &ram_image[0..src_sz];
+            if ram_image != src_image {
+                error!("Image not loaded correctly");
+                return true;
+            }
+
+        }
+
+        return false;
+    }
+
     /// Adds a new flash area that fails statistically
     fn mark_bad_status_with_rate(&self, flash: &mut SimMultiFlash, slot: usize,
                                  rate: f32) {
@@ -1028,18 +1151,18 @@
         let mut counter = stop.unwrap_or(0);
 
         let (first_interrupted, count) = match c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false) {
-            (-0x13579, _) => (true, stop.unwrap()),
-            (0, _) => (false, -counter),
-            (x, _) => panic!("Unknown return: {}", x),
+            x if x.interrupted() => (true, stop.unwrap()),
+            x if x.success() => (false, -counter),
+            x => panic!("Unknown return: {:?}", x),
         };
 
         counter = 0;
         if first_interrupted {
             // fl.dump();
             match c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false) {
-                (-0x13579, _) => panic!("Shouldn't stop again"),
-                (0, _) => (),
-                (x, _) => panic!("Unknown return: {}", x),
+                x if x.interrupted() => panic!("Shouldn't stop again"),
+                x if x.success() => (),
+                x => panic!("Unknown return: {:?}", x),
             }
         }
 
@@ -1052,7 +1175,7 @@
         // fl.write_file("image0.bin").unwrap();
         for i in 0 .. count {
             info!("Running boot pass {}", i + 1);
-            assert_eq!(c::boot_go(&mut flash, &self.areadesc, None, false), (0, 0));
+            assert!(c::boot_go(&mut flash, &self.areadesc, None, false).success_no_asserts());
         }
         flash
     }
@@ -1062,8 +1185,7 @@
         let mut fails = 0;
 
         let mut counter = stop;
-        let (x, _) = c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false);
-        if x != -0x13579 {
+        if !c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false).interrupted() {
             warn!("Should have stopped test at interruption point");
             fails += 1;
         }
@@ -1075,8 +1197,7 @@
             fails += 1;
         }
 
-        let (x, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if x != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Should have finished test upgrade");
             fails += 1;
         }
@@ -1104,14 +1225,12 @@
 
         // Do Revert
         let mut counter = stop;
-        let (x, _) = c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false);
-        if x != -0x13579 {
+        if !c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false).interrupted() {
             warn!("Should have stopped revert at interruption point");
             fails += 1;
         }
 
-        let (x, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if x != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Should have finished revert upgrade");
             fails += 1;
         }
@@ -1138,8 +1257,7 @@
             fails += 1;
         }
 
-        let (x, _) = c::boot_go(&mut flash, &self.areadesc, None, false);
-        if x != 0 {
+        if !c::boot_go(&mut flash, &self.areadesc, None, false).success() {
             warn!("Should have finished 3rd boot");
             fails += 1;
         }
@@ -1166,20 +1284,20 @@
         let mut resets = vec![0i32; count];
         let mut remaining_ops = total_ops;
         for reset in &mut resets {
-            let reset_counter = rng.gen_range(1, remaining_ops / 2);
+            let reset_counter = rng.gen_range(1 ..= remaining_ops / 2);
             let mut counter = reset_counter;
             match c::boot_go(&mut flash, &self.areadesc, Some(&mut counter), false) {
-                (0, _) | (-0x13579, _) => (),
-                (x, _) => panic!("Unknown return: {}", x),
+                x if x.interrupted() => (),
+                x => panic!("Unknown return: {:?}", x),
             }
             remaining_ops -= reset_counter;
             *reset = reset_counter;
         }
 
         match c::boot_go(&mut flash, &self.areadesc, None, false) {
-            (-0x13579, _) => panic!("Should not be have been interrupted!"),
-            (0, _) => (),
-            (x, _) => panic!("Unknown return: {}", x),
+            x if x.interrupted() => panic!("Should not be have been interrupted!"),
+            x if x.success() => (),
+            x => panic!("Unknown return: {:?}", x),
         }
 
         (flash, resets)
@@ -1266,6 +1384,40 @@
     }
 }
 
+impl RamData {
+    // TODO: This is not correct. The second slot of each image should be at the same address as
+    // the primary.
+    fn new(slots: &[[SlotInfo; 2]]) -> RamData {
+        let mut addr = RAM_LOAD_ADDR;
+        let mut places = BTreeMap::new();
+        // println!("Setup:-------------");
+        for imgs in slots {
+            for si in imgs {
+                // println!("Setup: si: {:?}", si);
+                let offset = addr;
+                let size = si.len as u32;
+                places.insert(SlotKey {
+                    dev_id: si.dev_id,
+                    base_off: si.base_off,
+                }, SlotPlace { offset, size });
+                // println!("  load: offset: {}, size: {}", offset, size);
+            }
+            addr += imgs[0].len as u32;
+        }
+        RamData {
+            places,
+            total: addr,
+        }
+    }
+
+    /// Lookup the ram data associated with a given flash partition.  We just panic if not present,
+    /// because all slots used should be in the map.
+    fn lookup(&self, slot: &SlotInfo) -> &SlotPlace {
+        self.places.get(&SlotKey{dev_id: slot.dev_id, base_off: slot.base_off})
+            .expect("RamData should contain all slots")
+    }
+}
+
 /// Show the flash layout.
 #[allow(dead_code)]
 fn show_flash(flash: &dyn Flash) {
@@ -1280,6 +1432,7 @@
 /// Install a "program" into the given image.  This fakes the image header, or at least all of the
 /// fields used by the given code.  Returns a copy of the image that was written.
 fn install_image(flash: &mut SimMultiFlash, slot: &SlotInfo, len: usize,
+                 ram: &RamData,
                  deps: &dyn Depender, bad_sig: bool) -> ImageData {
     let offset = slot.base_off;
     let slot_len = slot.len;
@@ -1294,10 +1447,17 @@
 
     const HDR_SIZE: usize = 32;
 
+    let place = ram.lookup(&slot);
+    let load_addr = if Caps::RamLoad.present() {
+        place.offset
+    } else {
+        0
+    };
+
     // Generate a boot header.  Note that the size doesn't include the header.
     let header = ImageHeader {
         magic: tlv.get_magic(),
-        load_addr: 0,
+        load_addr,
         hdr_size: HDR_SIZE as u16,
         protect_tlv_size: tlv.protect_size(),
         img_size: len as u32,
@@ -1329,17 +1489,27 @@
     tlv.add_bytes(&b_img);
 
     // Generate encrypted images
-    let flag = TlvFlags::ENCRYPTED as u32;
-    let is_encrypted = (tlv.get_flags() & flag) == flag;
+    let flag = TlvFlags::ENCRYPTED_AES128 as u32 | TlvFlags::ENCRYPTED_AES256 as u32;
+    let is_encrypted = (tlv.get_flags() & flag) != 0;
     let mut b_encimg = vec![];
     if is_encrypted {
+        let flag = TlvFlags::ENCRYPTED_AES256 as u32;
+        let aes256 = (tlv.get_flags() & flag) == flag;
         tlv.generate_enc_key();
         let enc_key = tlv.get_enc_key();
-        let key = GenericArray::from_slice(enc_key.as_slice());
         let nonce = GenericArray::from_slice(&[0; 16]);
-        let mut cipher = Aes128Ctr::new(&key, &nonce);
         b_encimg = b_img.clone();
-        cipher.apply_keystream(&mut b_encimg);
+        if aes256 {
+            let key: &GenericArray<u8, U32> = GenericArray::from_slice(enc_key.as_slice());
+            let block = Aes256::new(&key);
+            let mut cipher = Aes256Ctr::from_block_cipher(block, &nonce);
+            cipher.apply_keystream(&mut b_encimg);
+        } else {
+            let key: &GenericArray<u8, U16> = GenericArray::from_slice(enc_key.as_slice());
+            let block = Aes128::new(&key);
+            let mut cipher = Aes128Ctr::from_block_cipher(block, &nonce);
+            cipher.apply_keystream(&mut b_encimg);
+        }
     }
 
     // Build the TLV itself.
@@ -1357,6 +1527,7 @@
 
     // Pad the buffer to a multiple of the flash alignment.
     let align = dev.align();
+    let image_sz = buf.len();
     while buf.len() % align != 0 {
         buf.push(dev.erased_val());
     }
@@ -1399,6 +1570,7 @@
         dev.read(offset, &mut copy).unwrap();
 
         ImageData {
+            size: image_sz,
             plain: copy,
             cipher: enc_copy,
         }
@@ -1425,6 +1597,7 @@
         }
 
         ImageData {
+            size: image_sz,
             plain: copy,
             cipher: enc_copy,
         }
@@ -1434,6 +1607,7 @@
 /// Install no image.  This is used when no upgrade happens.
 fn install_no_image() -> ImageData {
     ImageData {
+        size: 0,
         plain: vec![],
         cipher: None,
     }
@@ -1443,32 +1617,36 @@
     if Caps::EcdsaP224.present() {
         panic!("Ecdsa P224 not supported in Simulator");
     }
+    let mut aes_key_size = 128;
+    if Caps::Aes256.present() {
+        aes_key_size = 256;
+    }
 
     if Caps::EncKw.present() {
         if Caps::RSA2048.present() {
-            TlvGen::new_rsa_kw()
+            TlvGen::new_rsa_kw(aes_key_size)
         } else if Caps::EcdsaP256.present() {
-            TlvGen::new_ecdsa_kw()
+            TlvGen::new_ecdsa_kw(aes_key_size)
         } else {
-            TlvGen::new_enc_kw()
+            TlvGen::new_enc_kw(aes_key_size)
         }
     } else if Caps::EncRsa.present() {
         if Caps::RSA2048.present() {
-            TlvGen::new_sig_enc_rsa()
+            TlvGen::new_sig_enc_rsa(aes_key_size)
         } else {
-            TlvGen::new_enc_rsa()
+            TlvGen::new_enc_rsa(aes_key_size)
         }
     } else if Caps::EncEc256.present() {
         if Caps::EcdsaP256.present() {
-            TlvGen::new_ecdsa_ecies_p256()
+            TlvGen::new_ecdsa_ecies_p256(aes_key_size)
         } else {
-            TlvGen::new_ecies_p256()
+            TlvGen::new_ecies_p256(aes_key_size)
         }
     } else if Caps::EncX25519.present() {
         if Caps::Ed25519.present() {
-            TlvGen::new_ed25519_ecies_x25519()
+            TlvGen::new_ed25519_ecies_x25519(aes_key_size)
         } else {
-            TlvGen::new_ecies_x25519()
+            TlvGen::new_ecies_x25519(aes_key_size)
         }
     } else {
         // The non-encrypted configuration.
@@ -1499,6 +1677,10 @@
             _ => panic!("Invalid slot requested"),
         }
     }
+
+    fn size(&self) -> usize {
+        self.size
+    }
 }
 
 /// Verify that given image is present in the flash at the given offset.
@@ -1691,6 +1873,7 @@
 const BOOT_FLAG_UNSET: Option<u8> = Some(3);
 
 /// Write out the magic so that the loader tries doing an upgrade.
+#[cfg(not(feature = "swap-status"))]
 pub fn mark_upgrade(flash: &mut SimMultiFlash, slot: &SlotInfo) {
     let dev = flash.get_mut(&slot.dev_id).unwrap();
     let align = dev.align();
@@ -1707,8 +1890,24 @@
     }
 }
 
+/// Write out the magic so that the loader tries doing an upgrade.
+#[cfg(feature = "swap-status")]
+pub fn mark_upgrade(flash: &mut SimMultiFlash, slot: &SlotInfo) {
+    let dev = flash.get_mut(&slot.dev_id).unwrap();
+    let align = dev.align();
+    let offset = slot.trailer_off + c::boot_max_align() * 4;
+    let mask = align - 1;
+    let sector_off = offset & !mask;
+    let mut buf = vec![dev.erased_val(); align];
+    dev.read(sector_off, &mut buf).unwrap();
+    buf[(offset & mask)..].copy_from_slice(MAGIC);
+    dev.erase(sector_off, align).unwrap();
+    dev.write(sector_off, &buf).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).
+#[cfg(not(feature = "swap-status"))]
 fn mark_permanent_upgrade(flash: &mut SimMultiFlash, slot: &SlotInfo) {
     // Overwrite mode always is permanent, and only the magic is used in
     // the trailer.  To avoid problems with large write sizes, don't try to
@@ -1725,9 +1924,32 @@
     dev.write(off, &ok[..align]).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).
+#[cfg(feature = "swap-status")]
+fn mark_permanent_upgrade(flash: &mut SimMultiFlash, slot: &SlotInfo) {
+    // Overwrite mode always is permanent, and only the magic is used in
+    // the trailer.  To avoid problems with large write sizes, don't try to
+    // set anything in this case.
+    if Caps::OverwriteUpgrade.present() {
+        return;
+    }
+    
+    let dev = flash.get_mut(&slot.dev_id).unwrap();
+    let align:usize = dev.align();
+    let mask:usize = align - 1;
+    let ok_off:usize = slot.trailer_off + c::boot_max_align() * 3;
+    let sector_off:usize = ok_off & !mask;
+    let mut buf = vec![dev.erased_val(); align];
+    dev.read(sector_off, &mut buf).unwrap();
+    buf[ok_off & mask] = 1u8;
+    dev.erase(sector_off, align).unwrap();
+    dev.write(sector_off, &buf[..align]).unwrap();
+}
+
 // Drop some pseudo-random gibberish onto the data.
 fn splat(data: &mut [u8], seed: usize) {
-    let mut seed_block = [0u8; 16];
+    let mut seed_block = [0u8; 32];
     let mut buf = Cursor::new(&mut seed_block[..]);
     buf.write_u32::<LittleEndian>(0x135782ea).unwrap();
     buf.write_u32::<LittleEndian>(0x92184728).unwrap();
@@ -1754,11 +1976,18 @@
 }
 
 #[cfg(not(feature = "large-write"))]
+#[cfg(not(feature = "swap-status"))]
 fn test_alignments() -> &'static [usize] {
     &[1, 2, 4, 8]
 }
 
 #[cfg(feature = "large-write")]
+#[cfg(not(feature = "swap-status"))]
 fn test_alignments() -> &'static [usize] {
     &[1, 2, 4, 8, 128, 512]
 }
+
+#[cfg(feature = "swap-status")]
+fn test_alignments() -> &'static [usize] {
+    &[512]
+}
diff --git a/sim/src/lib.rs b/sim/src/lib.rs
index 4203c0e..69f13f1 100644
--- a/sim/src/lib.rs
+++ b/sim/src/lib.rs
@@ -63,8 +63,10 @@
 
 #[derive(Copy, Clone, Debug, Deserialize)]
 pub enum DeviceName {
-    Stm32f4, K64f, K64fBig, K64fMulti, Nrf52840, Nrf52840SpiFlash,
-    Nrf52840UnequalSlots, PSoC6Multi,
+    Stm32f4,
+    K64f, K64fBig, K64fMulti,
+    Nrf52840, Nrf52840SpiFlash, Nrf52840UnequalSlots,
+    PSoC6
 }
 
 pub static ALL_DEVICES: &[DeviceName] = &[
@@ -75,7 +77,7 @@
     DeviceName::Nrf52840,
     DeviceName::Nrf52840SpiFlash,
     DeviceName::Nrf52840UnequalSlots,
-    DeviceName::PSoC6Multi,
+    DeviceName::PSoC6
 ];
 
 impl fmt::Display for DeviceName {
@@ -88,7 +90,7 @@
             DeviceName::Nrf52840 => "nrf52840",
             DeviceName::Nrf52840SpiFlash => "Nrf52840SpiFlash",
             DeviceName::Nrf52840UnequalSlots => "Nrf52840UnequalSlots",
-            DeviceName::PSoC6Multi => "psoc6multi",
+            DeviceName::PSoC6 => "PSoC6"
         };
         f.write_str(name)
     }
diff --git a/sim/src/tlv.rs b/sim/src/tlv.rs
index b6c3c96..6680b4f 100644
--- a/sim/src/tlv.rs
+++ b/sim/src/tlv.rs
@@ -1,5 +1,6 @@
 // Copyright (c) 2017-2020 Linaro LTD
 // Copyright (c) 2017-2020 JUUL Labs
+// Copyright (c) 2021 Arm Limited
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -16,6 +17,8 @@
 use byteorder::{
     LittleEndian, WriteBytesExt,
 };
+use cipher::FromBlockCipher;
+use crate::caps::Caps;
 use crate::image::ImageVersion;
 use log::info;
 use ring::{digest, rand, agreement, hkdf, hmac};
@@ -27,15 +30,19 @@
     ECDSA_P256_SHA256_ASN1_SIGNING,
     Ed25519KeyPair,
 };
-use aes_ctr::{
+use aes::{
+    Aes128,
     Aes128Ctr,
-    stream_cipher::{
-        generic_array::GenericArray,
-        NewStreamCipher,
-        SyncStreamCipher,
-    },
+    Aes256,
+    Aes256Ctr,
+    NewBlockCipher
+};
+use cipher::{
+    generic_array::GenericArray,
+    StreamCipher,
 };
 use mcuboot_sys::c;
+use typenum::{U16, U32};
 
 #[repr(u16)]
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -49,7 +56,7 @@
     RSA3072 = 0x23,
     ED25519 = 0x24,
     ENCRSA2048 = 0x30,
-    ENCKW128 = 0x31,
+    ENCKW = 0x31,
     ENCEC256 = 0x32,
     ENCX25519 = 0x33,
     DEPENDENCY = 0x40,
@@ -59,8 +66,9 @@
 pub enum TlvFlags {
     PIC = 0x01,
     NON_BOOTABLE = 0x02,
-    ENCRYPTED = 0x04,
+    ENCRYPTED_AES128 = 0x04,
     RAM_LOAD = 0x20,
+    ENCRYPTED_AES256 = 0x08,
 }
 
 /// A generator for manifests.  The format of the manifest can be either a
@@ -115,8 +123,6 @@
     version: ImageVersion,
 }
 
-const AES_KEY_LEN: usize = 16;
-
 impl TlvGen {
     /// Construct a new tlv generator that will only contain a hash of the data.
     #[allow(dead_code)]
@@ -160,81 +166,126 @@
     }
 
     #[allow(dead_code)]
-    pub fn new_enc_rsa() -> TlvGen {
+    pub fn new_enc_rsa(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
+            flags: flag,
             kinds: vec![TlvKinds::SHA256, TlvKinds::ENCRSA2048],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_sig_enc_rsa() -> TlvGen {
+    pub fn new_sig_enc_rsa(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
+            flags: flag,
             kinds: vec![TlvKinds::SHA256, TlvKinds::RSA2048, TlvKinds::ENCRSA2048],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_enc_kw() -> TlvGen {
+    pub fn new_enc_kw(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
-            kinds: vec![TlvKinds::SHA256, TlvKinds::ENCKW128],
+            flags: flag,
+            kinds: vec![TlvKinds::SHA256, TlvKinds::ENCKW],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_rsa_kw() -> TlvGen {
+    pub fn new_rsa_kw(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
-            kinds: vec![TlvKinds::SHA256, TlvKinds::RSA2048, TlvKinds::ENCKW128],
+            flags: flag,
+            kinds: vec![TlvKinds::SHA256, TlvKinds::RSA2048, TlvKinds::ENCKW],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_ecdsa_kw() -> TlvGen {
+    pub fn new_ecdsa_kw(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
-            kinds: vec![TlvKinds::SHA256, TlvKinds::ECDSA256, TlvKinds::ENCKW128],
+            flags: flag,
+            kinds: vec![TlvKinds::SHA256, TlvKinds::ECDSA256, TlvKinds::ENCKW],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_ecies_p256() -> TlvGen {
+    pub fn new_ecies_p256(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
+            flags: flag,
             kinds: vec![TlvKinds::SHA256, TlvKinds::ENCEC256],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_ecdsa_ecies_p256() -> TlvGen {
+    pub fn new_ecdsa_ecies_p256(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
+            flags: flag,
             kinds: vec![TlvKinds::SHA256, TlvKinds::ECDSA256, TlvKinds::ENCEC256],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_ecies_x25519() -> TlvGen {
+    pub fn new_ecies_x25519(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
+            flags: flag,
             kinds: vec![TlvKinds::SHA256, TlvKinds::ENCX25519],
             ..Default::default()
         }
     }
 
     #[allow(dead_code)]
-    pub fn new_ed25519_ecies_x25519() -> TlvGen {
+    pub fn new_ed25519_ecies_x25519(aes_key_size: u32) -> TlvGen {
+        let flag = if aes_key_size == 256 {
+            TlvFlags::ENCRYPTED_AES256 as u32
+        } else {
+            TlvFlags::ENCRYPTED_AES128 as u32
+        };
         TlvGen {
-            flags: TlvFlags::ENCRYPTED as u32,
+            flags: flag,
             kinds: vec![TlvKinds::SHA256, TlvKinds::ED25519, TlvKinds::ENCX25519],
             ..Default::default()
         }
@@ -248,7 +299,12 @@
 
     /// Retrieve the header flags for this configuration.  This can be called at any time.
     fn get_flags(&self) -> u32 {
-        self.flags
+        // For the RamLoad case, add in the flag for this feature.
+        if Caps::RamLoad.present() {
+            self.flags | (TlvFlags::RAM_LOAD as u32)
+        } else {
+            self.flags
+        }
     }
 
     /// Add bytes to the covered hash.
@@ -479,20 +535,28 @@
             result.extend_from_slice(&encbuf);
         }
 
-        if self.kinds.contains(&TlvKinds::ENCKW128) {
-            let key_bytes = base64::decode(
-                include_str!("../../enc-aes128kw.b64").trim()).unwrap();
-
+        if self.kinds.contains(&TlvKinds::ENCKW) {
+            let flag = TlvFlags::ENCRYPTED_AES256 as u32;
+            let aes256 = (self.get_flags() & flag) == flag;
+            let key_bytes = if aes256 {
+                base64::decode(
+                    include_str!("../../enc-aes256kw.b64").trim()).unwrap()
+            } else {
+                base64::decode(
+                    include_str!("../../enc-aes128kw.b64").trim()).unwrap()
+            };
             let cipherkey = self.get_enc_key();
             let cipherkey = cipherkey.as_slice();
-            let encbuf = match c::kw_encrypt(&key_bytes, cipherkey) {
+            let keylen = if aes256 { 32 } else { 16 };
+            let encbuf = match c::kw_encrypt(&key_bytes, cipherkey, keylen) {
                 Ok(v) => v,
                 Err(_) => panic!("Failed to encrypt secret key"),
             };
 
-            assert!(encbuf.len() == 24);
-            result.write_u16::<LittleEndian>(TlvKinds::ENCKW128 as u16).unwrap();
-            result.write_u16::<LittleEndian>(24).unwrap();
+            let size = if aes256 { 40 } else { 24 };
+            assert!(encbuf.len() == size);
+            result.write_u16::<LittleEndian>(TlvKinds::ENCKW as u16).unwrap();
+            result.write_u16::<LittleEndian>(size as u16).unwrap();
             result.extend_from_slice(&encbuf);
         }
 
@@ -503,7 +567,6 @@
                 pem::parse(include_bytes!("../../enc-x25519-pub.pem").as_ref()).unwrap()
             };
             assert_eq!(key_bytes.tag, "PUBLIC KEY");
-
             let rng = rand::SystemRandom::new();
             let alg = if self.kinds.contains(&TlvKinds::ENCEC256) {
                 &agreement::ECDH_P256
@@ -535,15 +598,19 @@
                 }
             }
 
+            let flag = TlvFlags::ENCRYPTED_AES256 as u32;
+            let aes256 = (self.get_flags() & flag) == flag;
+
             let derived_key = match agreement::agree_ephemeral(
                 pk, &peer_pubk, ring::error::Unspecified, |shared| {
                     let salt = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]);
                     let prk = salt.extract(&shared);
-                    let okm = match prk.expand(&[b"MCUBoot_ECIES_v1"], OkmLen(48)) {
+                    let okm_len = if aes256 { 64 } else { 48 };
+                    let okm = match prk.expand(&[b"MCUBoot_ECIES_v1"], OkmLen(okm_len)) {
                         Ok(okm) => okm,
                         Err(_) => panic!("Failed building HKDF OKM"),
                     };
-                    let mut buf = [0u8; 48];
+                    let mut buf = if aes256 { vec![0u8; 64] } else { vec![0u8; 48] };
                     match okm.fill(&mut buf) {
                         Ok(_) => Ok(buf),
                         Err(_) => panic!("Failed generating HKDF output"),
@@ -554,13 +621,22 @@
                 Err(_) => panic!("Failed building HKDF"),
             };
 
-            let key = GenericArray::from_slice(&derived_key[..16]);
             let nonce = GenericArray::from_slice(&[0; 16]);
-            let mut cipher = Aes128Ctr::new(&key, &nonce);
             let mut cipherkey = self.get_enc_key();
-            cipher.apply_keystream(&mut cipherkey);
+            if aes256 {
+                let key: &GenericArray<u8, U32> = GenericArray::from_slice(&derived_key[..32]);
+                let block = Aes256::new(&key);
+                let mut cipher = Aes256Ctr::from_block_cipher(block, &nonce);
+                cipher.apply_keystream(&mut cipherkey);
+            } else {
+                let key: &GenericArray<u8, U16> = GenericArray::from_slice(&derived_key[..16]);
+                let block = Aes128::new(&key);
+                let mut cipher = Aes128Ctr::from_block_cipher(block, &nonce);
+                cipher.apply_keystream(&mut cipherkey);
+            }
 
-            let key = hmac::Key::new(hmac::HMAC_SHA256, &derived_key[16..]);
+            let size = if aes256 { 32 } else { 16 };
+            let key = hmac::Key::new(hmac::HMAC_SHA256, &derived_key[size..]);
             let tag = hmac::sign(&key, &cipherkey);
 
             let mut buf = vec![];
@@ -569,13 +645,15 @@
             buf.append(&mut cipherkey);
 
             if self.kinds.contains(&TlvKinds::ENCEC256) {
-                assert!(buf.len() == 113);
+                let size = if aes256 { 129 } else { 113 };
+                assert!(buf.len() == size);
                 result.write_u16::<LittleEndian>(TlvKinds::ENCEC256 as u16).unwrap();
-                result.write_u16::<LittleEndian>(113).unwrap();
+                result.write_u16::<LittleEndian>(size as u16).unwrap();
             } else {
-                assert!(buf.len() == 80);
+                let size = if aes256 { 96 } else { 80 };
+                assert!(buf.len() == size);
                 result.write_u16::<LittleEndian>(TlvKinds::ENCX25519 as u16).unwrap();
-                result.write_u16::<LittleEndian>(80).unwrap();
+                result.write_u16::<LittleEndian>(size as u16).unwrap();
             }
             result.extend_from_slice(&buf);
         }
@@ -590,7 +668,13 @@
 
     fn generate_enc_key(&mut self) {
         let rng = rand::SystemRandom::new();
-        let mut buf = vec![0u8; AES_KEY_LEN];
+        let flag = TlvFlags::ENCRYPTED_AES256 as u32;
+        let aes256 = (self.get_flags() & flag) == flag;
+        let mut buf = if aes256 {
+            vec![0u8; 32]
+        } else {
+            vec![0u8; 16]
+        };
         if rng.fill(&mut buf).is_err() {
             panic!("Error generating encrypted key");
         }
@@ -599,7 +683,7 @@
     }
 
     fn get_enc_key(&self) -> Vec<u8> {
-        if self.enc_key.len() != AES_KEY_LEN {
+        if self.enc_key.len() != 32 && self.enc_key.len() != 16 {
             panic!("No random key was generated");
         }
         self.enc_key.clone()