From 227ca616ca54884e6dd25992cb191f80cfd9a132 Mon Sep 17 00:00:00 2001
From: M4x1m3 <M4x1me@protonmail.com>
Date: Mon, 28 Feb 2022 09:13:29 +0100
Subject: [PATCH] [bootloader] Boot other slot if a slot is invalid

---
 bootloader/boot.cpp | 30 +++++++++++++++++++++++-------
 bootloader/boot.h   |  4 ++--
 bootloader/main.cpp | 10 +++++++---
 bootloader/slot.cpp |  7 ++++++-
 4 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/bootloader/boot.cpp b/bootloader/boot.cpp
index 268ee125b..1c965d180 100644
--- a/bootloader/boot.cpp
+++ b/bootloader/boot.cpp
@@ -29,18 +29,34 @@ void Boot::setMode(BootMode mode) {
   Ion::ExamMode::IncrementExamMode(deltaMode);
 }
 
-[[ noreturn ]] void Boot::boot() {
+__attribute__((noreturn)) void Boot::boot() {
   assert(mode() != BootMode::Unknown);
 
-  if (mode() == BootMode::SlotA)
-    Slot::A().boot();
-  else if (mode() == BootMode::SlotB)
+  if (!Slot::A().kernelHeader()->isValid() && !Slot::B().kernelHeader()->isValid()) {
+    // Bootloader if both invalid
+    bootloader();
+  } else if (!Slot::A().kernelHeader()->isValid()) {
+    // If slot A is invalid and B valid, boot B
+    setMode(BootMode::SlotB);
     Slot::B().boot();
-  
-  for(;;);
+  } else if (!Slot::B().kernelHeader()->isValid()) {
+    // If slot B is invalid and A valid, boot A
+    setMode(BootMode::SlotA);
+    Slot::A().boot();
+  } else {
+    // Both valid, boot the selected one
+    if (mode() == BootMode::SlotA) {
+      Slot::A().boot();
+    } else if (mode() == BootMode::SlotB) {
+      Slot::B().boot();
+    }
+  }
+
+  // Achivement unlocked: How Did We Get Here?
+  bootloader();
 }
 
-[[ noreturn ]] void Boot::bootloader() {
+__attribute__ ((noreturn)) void Boot::bootloader() {
   Bootloader::Interface::draw();
   for(;;) {
     Ion::USB::enable();
diff --git a/bootloader/boot.h b/bootloader/boot.h
index 23967f87c..bb5df17c1 100644
--- a/bootloader/boot.h
+++ b/bootloader/boot.h
@@ -19,8 +19,8 @@ class Boot {
 public:
   static BootMode mode();
   static void setMode(BootMode mode);
-  [[ noreturn ]] static void boot();
-  [[ norteurn ]] static void bootloader();
+  __attribute__ ((noreturn)) static void boot();
+  __attribute__ ((noreturn)) static void bootloader();
 };
 
 }
diff --git a/bootloader/main.cpp b/bootloader/main.cpp
index 9f237bf0c..8e07dc8a1 100644
--- a/bootloader/main.cpp
+++ b/bootloader/main.cpp
@@ -1,19 +1,20 @@
 
 #include <ion.h>
 #include <assert.h>
-#include <ion/src/device/shared/drivers/board.h>
 
 #include <bootloader/boot.h>
 
-void ion_main(int argc, const char * const argv[]) {
+__attribute__ ((noreturn)) void ion_main(int argc, const char * const argv[]) {
   // Clear the screen
   Ion::Display::pushRectUniform(KDRect(0,0,320,240), KDColorBlack);
   // Initialize the backlight
   Ion::Backlight::init();
 
+  // Set the mode to slot A if undefined
   if (Bootloader::Boot::mode() == Bootloader::BootMode::Unknown)
     Bootloader::Boot::setMode(Bootloader::BootMode::SlotA);
 
+  // Handle rebooting to bootloader
   if (Bootloader::Boot::mode() == Bootloader::BootMode::SlotABootloader) {
     Bootloader::Boot::setMode(Bootloader::BootMode::SlotA);
     Bootloader::Boot::bootloader();
@@ -24,14 +25,17 @@ void ion_main(int argc, const char * const argv[]) {
 
   uint64_t scan = Ion::Keyboard::scan();
 
+  // Reset+4 => Launch bootloader
   if (scan == Ion::Keyboard::State(Ion::Keyboard::Key::Four)) {
     Bootloader::Boot::bootloader();
+  // Reset+1 => Launch slot A
   } else if (scan == Ion::Keyboard::State(Ion::Keyboard::Key::One)) {
     Bootloader::Boot::setMode(Bootloader::BootMode::SlotA);
+  // Reset+2 => Launch slot B
   } else if (scan == Ion::Keyboard::State(Ion::Keyboard::Key::Two)) {
     Bootloader::Boot::setMode(Bootloader::BootMode::SlotB);
   }
 
-  Ion::Device::Board::bootloaderMPU();
+  // Boot the firmware
   Bootloader::Boot::boot();
 }
diff --git a/bootloader/slot.cpp b/bootloader/slot.cpp
index 4430807a0..b23655aed 100644
--- a/bootloader/slot.cpp
+++ b/bootloader/slot.cpp
@@ -1,4 +1,5 @@
 #include <bootloader/slot.h>
+#include <ion/src/device/shared/drivers/board.h>
 
 extern "C" void jump_to_firmware(const uint32_t* stackPtr, const void(*startPtr)(void));
 
@@ -21,8 +22,12 @@ const UserlandHeader* Slot::userlandHeader() const {
 }
 
 [[ noreturn ]] void Slot::boot() const {
+  // Configure the MPU for the booted firmware
+  Ion::Device::Board::bootloaderMPU();
+
+  // Jump
   jump_to_firmware(kernelHeader()->stackPointer(), kernelHeader()->startPointer());
   for(;;);
 }
 
-}
\ No newline at end of file
+}