From ccae1e3ec8f81e92bbc10e648f0d111466b5f659 Mon Sep 17 00:00:00 2001
From: Eric Callahan <arksine.code@gmail.com>
Date: Sat, 23 Apr 2022 06:25:26 -0400
Subject: [PATCH] spi_flash: validate mcu configuration

Prior to attempting a flash, verify that the MCU is successfully
configured.

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
---
 scripts/spi_flash/spi_flash.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/scripts/spi_flash/spi_flash.py b/scripts/spi_flash/spi_flash.py
index c16d76d5..e6aae54d 100644
--- a/scripts/spi_flash/spi_flash.py
+++ b/scripts/spi_flash/spi_flash.py
@@ -119,6 +119,9 @@ FINALIZE_CFG_CMD = "finalize_config crc=%d"
 class SPIFlashError(Exception):
     pass
 
+class MCUConfigError(SPIFlashError):
+    pass
+
 class SPIDirect:
     def __init__(self, ser):
         self.oid = SPI_OID
@@ -862,8 +865,7 @@ class MCUConnection:
         self._serial.disconnect()
         self.connected = False
 
-    def check_need_restart(self):
-        output("Checking Current MCU Configuration...")
+    def get_mcu_config(self):
         # Iterate through backwards compatible response strings
         for response in GET_CFG_RESPONSES:
             try:
@@ -875,7 +877,11 @@ class MCUConnection:
                 if response == GET_CFG_RESPONSES[-1]:
                     raise err
                 output("Trying fallback...")
-        params = get_cfg_cmd.send()
+        return get_cfg_cmd.send()
+
+    def check_need_restart(self):
+        output("Checking Current MCU Configuration...")
+        params = self.get_mcu_config()
         output_line("Done")
         if params['is_config'] or params['is_shutdown']:
             output_line("MCU needs restart: is_config=%d, is_shutdown=%d"
@@ -926,9 +932,12 @@ class MCUConnection:
         self._serial.send(bus_cmd)
         config_crc = zlib.crc32('\n'.join(cfg_cmds).encode()) & 0xffffffff
         self._serial.send(FINALIZE_CFG_CMD % (config_crc,))
+        config = self.get_mcu_config()
+        if not config["is_config"] or config["is_shutdown"]:
+            raise MCUConfigError("Failed to configure MCU")
+        printfunc("Initializing SD Card and Mounting file system...")
         self.fatfs = FatFS(self._serial)
         self.reactor.pause(self.reactor.monotonic() + .5)
-        printfunc("Initializing SD Card and Mounting file system...")
         try:
             self.fatfs.mount(printfunc)
         except OSError:
@@ -1098,7 +1107,14 @@ class SPIFlash:
         if not self.mcu_conn.connected:
             self.mcu_conn.connect()
         self.old_dictionary = self.mcu_conn.raw_dictionary
-        self.mcu_conn.configure_mcu(printfunc=output_line)
+        try:
+            self.mcu_conn.configure_mcu(printfunc=output_line)
+        except MCUConfigError:
+            output_line("MCU configuration failed, attempting restart")
+            self.need_upload = True
+            self.mcu_conn.reset()
+            self.task_complete = True
+            return
         self.firmware_checksum = self.mcu_conn.sdcard_upload()
         self.mcu_conn.reset()
         self.task_complete = True