From 899592ac341ed9f10ee554424f964f6880fc8c48 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Wed, 13 Mar 2024 18:03:19 +1100
Subject: [PATCH] stm32/boards/LEGO_HUB_NO6: Move robust logic to mboot.

Signed-off-by: Damien George <damien@micropython.org>
---
 ports/stm32/boards/LEGO_HUB_NO6/appupdate.py | 12 ------------
 ports/stm32/boards/LEGO_HUB_NO6/board_init.c |  9 ++++++++-
 2 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/ports/stm32/boards/LEGO_HUB_NO6/appupdate.py b/ports/stm32/boards/LEGO_HUB_NO6/appupdate.py
index 02398e81243f..9cfbc06e5659 100644
--- a/ports/stm32/boards/LEGO_HUB_NO6/appupdate.py
+++ b/ports/stm32/boards/LEGO_HUB_NO6/appupdate.py
@@ -7,9 +7,6 @@
 _IOCTL_BLOCK_COUNT = const(4)
 _IOCTL_BLOCK_SIZE = const(5)
 
-_SPIFLASH_UPDATE_KEY_ADDR = const(1020 * 1024)
-_SPIFLASH_UPDATE_KEY_VALUE = const(0x12345678)
-
 _FILESYSTEM_ADDR = const(0x8000_0000 + 1024 * 1024)
 
 # Roundabout way to get actual filesystem size from config.
@@ -26,9 +23,6 @@ def update_app(filename):
     if not elems:
         return
 
-    # Create the update key.
-    key = struct.pack("<I", _SPIFLASH_UPDATE_KEY_VALUE)
-
     # Create a SPI flash object.
     spi = machine.SoftSPI(
         sck=machine.Pin.board.FLASH_SCK,
@@ -41,11 +35,5 @@ def update_app(filename):
     # We can't use pyb.Flash() because we need to write to the "reserved" 1M area.
     flash = spiflash.SPIFlash(spi, cs)
 
-    # Write the update key and elements to the SPI flash.
-    flash.erase_block(_SPIFLASH_UPDATE_KEY_ADDR)
-    flash.write(_SPIFLASH_UPDATE_KEY_ADDR, key + elems)
-
     # Enter mboot with a request to do a filesystem-load update.
-    # If there is a power failure during the update (eg battery removed) then
-    # mboot will read the SPI flash update key and elements and retry.
     machine.bootloader(elems)
diff --git a/ports/stm32/boards/LEGO_HUB_NO6/board_init.c b/ports/stm32/boards/LEGO_HUB_NO6/board_init.c
index db4c07914af9..5e4e7791b50d 100644
--- a/ports/stm32/boards/LEGO_HUB_NO6/board_init.c
+++ b/ports/stm32/boards/LEGO_HUB_NO6/board_init.c
@@ -194,7 +194,14 @@ int board_mboot_get_reset_mode(uint32_t *initial_r0) {
 }
 
 void board_mboot_state_change(int state, uint32_t arg) {
-    if (state == MBOOT_STATE_FSLOAD_END) {
+    if (state == MBOOT_STATE_FSLOAD_START) {
+        // The FS-load update is about to start.  Program the update key and FS-load elements
+        // into the flash so they can be retrieved if there is a power failure during the update.
+        mp_spiflash_erase_block(MBOOT_SPIFLASH_SPIFLASH, SPIFLASH_UPDATE_KEY_ADDR);
+        uint32_t key = SPIFLASH_UPDATE_KEY_VALUE;
+        mp_spiflash_write(MBOOT_SPIFLASH_SPIFLASH, SPIFLASH_UPDATE_KEY_ADDR, 4, (const uint8_t *)&key);
+        mp_spiflash_write(MBOOT_SPIFLASH_SPIFLASH, SPIFLASH_UPDATE_KEY_ADDR + 4, ELEM_DATA_SIZE, ELEM_DATA_START);
+    } else if (state == MBOOT_STATE_FSLOAD_END) {
         // The FS-load update completed (either with success or failure), so erase the
         // update key and write the result of the FS-load operation into flash.
         mp_spiflash_erase_block(MBOOT_SPIFLASH_SPIFLASH, SPIFLASH_UPDATE_KEY_ADDR);