From b63061a505b6abf58dc16fa50627564cb030f8ac Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 21 Feb 2019 08:23:33 -0800
Subject: [PATCH] sept: prepare for reboot immediately

---
 sept/sept-secondary/src/main.c  | 7 +++++--
 sept/sept-secondary/src/utils.c | 7 ++++++-
 sept/sept-secondary/src/utils.h | 1 +
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/sept/sept-secondary/src/main.c b/sept/sept-secondary/src/main.c
index 08500c300..211ab3987 100644
--- a/sept/sept-secondary/src/main.c
+++ b/sept/sept-secondary/src/main.c
@@ -57,6 +57,10 @@ static void exfiltrate_keys_and_reboot_if_needed(void) {
     uint8_t *dec_se_state = (uint8_t *)0x4003F000;
     
     if (!has_rebooted()) {
+        /* Prepare for a reboot before doing anything else. */
+        prepare_for_reboot_to_self();
+        set_has_rebooted(true);
+        
         /* Save the security engine context. */
         se_get_regs()->_0x4 = 0x0;
         se_set_in_context_save_mode(true);
@@ -67,8 +71,7 @@ static void exfiltrate_keys_and_reboot_if_needed(void) {
         for (size_t k = 0; k < 0x10; k++) {
             clear_aes_keyslot(k);
         }
-                
-        set_has_rebooted(true);
+        
         reboot_to_self();
     } else {
         /* Decrypt the security engine state. */
diff --git a/sept/sept-secondary/src/utils.c b/sept/sept-secondary/src/utils.c
index ac3d7b61f..4b95bf289 100644
--- a/sept/sept-secondary/src/utils.c
+++ b/sept/sept-secondary/src/utils.c
@@ -65,7 +65,7 @@ __attribute__((noreturn)) void pmc_reboot(uint32_t scratch0) {
     }
 }
 
-__attribute__((noreturn)) void reboot_to_self(void) {
+void prepare_for_reboot_to_self(void) {
     /* Patch SDRAM init to perform an SVC immediately after second write */
     APBDEV_PMC_SCRATCH45_0 = 0x2E38DFFF;
     APBDEV_PMC_SCRATCH46_0 = 0x6001DC28;
@@ -77,6 +77,11 @@ __attribute__((noreturn)) void reboot_to_self(void) {
     for (size_t i = 0; i < rebootstub_bin_size; i += sizeof(uint32_t)) {
         write32le((void *)0x4003F000, i, read32le(rebootstub_bin, i));
     }
+}
+
+__attribute__((noreturn)) void reboot_to_self(void) {
+    /* Prep IRAM for reboot. */
+    prepare_for_reboot_to_self();
     
     /* Trigger warm reboot. */
     pmc_reboot(1 << 0);
diff --git a/sept/sept-secondary/src/utils.h b/sept/sept-secondary/src/utils.h
index 58e53bffe..ee16c6e40 100644
--- a/sept/sept-secondary/src/utils.h
+++ b/sept/sept-secondary/src/utils.h
@@ -121,6 +121,7 @@ void hexdump(const void* data, size_t size, uintptr_t addrbase);
 
 __attribute__((noreturn)) void watchdog_reboot(void);
 __attribute__((noreturn)) void pmc_reboot(uint32_t scratch0);
+void prepare_for_reboot_to_self(void);
 __attribute__((noreturn)) void reboot_to_self(void);
 __attribute__((noreturn)) void wait_for_button_and_reboot(void);