From d3fe0a06e8799ef95716aaada0a95505aebab48c Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Sun, 28 Apr 2024 15:35:38 +1000
Subject: [PATCH] stm32/flash: Fix writing final words to flash on H5 and H7
 MCUs.

The calculations `num_word32 / 4` and `num_word32 / 8` were rounding down
the number of words to program to flash, and therefore possibly truncating
the data (eg mboot could miss writing the final few words of the firmware).

That's fixed in this commit by adding extra logic to program any remaining
words.  And the logic for H5 and H7 is combined.

Signed-off-by: Damien George <damien@micropython.org>
---
 ports/stm32/flash.c | 47 ++++++++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 18 deletions(-)

diff --git a/ports/stm32/flash.c b/ports/stm32/flash.c
index 10c5f296a2..7668b53a43 100644
--- a/ports/stm32/flash.c
+++ b/ports/stm32/flash.c
@@ -24,6 +24,7 @@
  * THE SOFTWARE.
  */
 
+#include <string.h>
 #include "py/mpconfig.h"
 #include "py/misc.h"
 #include "py/mphal.h"
@@ -450,28 +451,38 @@ int flash_write(uint32_t flash_dest, const uint32_t *src, uint32_t num_word32) {
         #endif
     }
 
-    #elif defined(STM32H5)
+    #elif defined(STM32H5) || defined(STM32H7)
 
-    // program the flash 128 bits (4 words) at a time
-    for (int i = 0; i < num_word32 / 4; i++) {
-        status = HAL_FLASH_Program(FLASH_TYPEPROGRAM_QUADWORD, flash_dest, (uint64_t)(uint32_t)src);
-        if (status != HAL_OK) {
-            break;
-        }
-        flash_dest += 16;
-        src += 4;
-    }
+    #if defined(STM32H5)
+    static const unsigned int WORD32_PER_FLASHWORD = 4;
+    static const unsigned int PROGRAM_TYPE = FLASH_TYPEPROGRAM_QUADWORD;
+    #else
+    static const unsigned int WORD32_PER_FLASHWORD = FLASH_NB_32BITWORD_IN_FLASHWORD;
+    static const unsigned int PROGRAM_TYPE = FLASH_TYPEPROGRAM_FLASHWORD;
+    #endif
 
-    #elif defined(STM32H7)
+    if (flash_dest % (WORD32_PER_FLASHWORD * sizeof(uint32_t))) {
+        // The flash_dest address is not aligned correctly.
+        status = HAL_ERROR;
+    } else {
+        // Program the flash WORD32_PER_FLASHWORD words at a time.
+        for (int i = 0; i < num_word32 / WORD32_PER_FLASHWORD; i++) {
+            status = HAL_FLASH_Program(PROGRAM_TYPE, flash_dest, (uint32_t)src);
+            if (status != HAL_OK) {
+                break;
+            }
+            flash_dest += WORD32_PER_FLASHWORD * sizeof(uint32_t);
+            src += WORD32_PER_FLASHWORD;
+        }
 
-    // program the flash 256 bits at a time
-    for (int i = 0; i < num_word32 / 8; i++) {
-        status = HAL_FLASH_Program(FLASH_TYPEPROGRAM_FLASHWORD, flash_dest, (uint64_t)(uint32_t)src);
-        if (status != HAL_OK) {
-            break;
+        // Program any remaining words from src.
+        // Additional bytes beyond that are programmed to 0xff.
+        if (num_word32 % WORD32_PER_FLASHWORD) {
+            uint32_t buf[WORD32_PER_FLASHWORD];
+            memset(buf, 0xff, sizeof(buf));
+            memcpy(buf, src, (num_word32 % WORD32_PER_FLASHWORD) * sizeof(uint32_t));
+            status = HAL_FLASH_Program(PROGRAM_TYPE, flash_dest, (uint32_t)buf);
         }
-        flash_dest += 32;
-        src += 8;
     }
 
     #else