From 8667990506da8642344adb5da7f1b76c9d6cbfec Mon Sep 17 00:00:00 2001
From: Antoine Faure <antoine.faure@spidertracks.co.nz>
Date: Thu, 13 Jan 2022 11:41:14 +1300
Subject: [PATCH] Allow to completely erase a page when size is not aligned

---
 inc/stlink.h         |  2 +-
 src/common.c         | 13 ++++++++-----
 src/st-flash/flash.c |  2 +-
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/inc/stlink.h b/inc/stlink.h
index 852e95e..083b346 100644
--- a/inc/stlink.h
+++ b/inc/stlink.h
@@ -274,7 +274,7 @@ int stlink_trace_enable(stlink_t* sl, uint32_t frequency);
 int stlink_trace_disable(stlink_t* sl);
 int stlink_trace_read(stlink_t* sl, uint8_t* buf, size_t size);
 int stlink_erase_flash_mass(stlink_t* sl);
-int stlink_erase_flash_section(stlink_t *sl, stm32_addr_t base_addr, size_t size);
+int stlink_erase_flash_section(stlink_t *sl, stm32_addr_t base_addr, size_t size, bool align_size);
 int stlink_write_flash(stlink_t* sl, stm32_addr_t address, uint8_t* data, uint32_t length, uint8_t eraseonly);
 int stlink_parse_ihex(const char* path, uint8_t erased_pattern, uint8_t * * mem, size_t * size, uint32_t * begin);
 uint8_t stlink_get_erased_pattern(stlink_t *sl);
diff --git a/src/common.c b/src/common.c
index 4e276d2..f72cad5 100644
--- a/src/common.c
+++ b/src/common.c
@@ -2951,7 +2951,7 @@ int stlink_erase_flash_page(stlink_t *sl, stm32_addr_t flashaddr) {
   return check_flash_error(sl);
 }
 
-int stlink_erase_flash_section(stlink_t *sl, stm32_addr_t base_addr, size_t size) {
+int stlink_erase_flash_section(stlink_t *sl, stm32_addr_t base_addr, size_t size, bool align_size) {
   if (base_addr < sl->flash_base || (base_addr + size) > (sl->flash_base + sl->flash_size)) {
     ELOG("Invalid address or size\n");
     return (-1);
@@ -2971,8 +2971,8 @@ int stlink_erase_flash_section(stlink_t *sl, stm32_addr_t base_addr, size_t size
   do {
     size_t page_size = stlink_calculate_pagesize(sl, addr);
 
-    // Check if we are going further than the requested size
-    if ((addr + page_size) > (base_addr + size)) {
+    // Check if size is aligned with a page, unless we want to completely erase the last page
+    if ((addr + page_size) > (base_addr + size) && !align_size) {
       ELOG("Invalid size (not aligned with a page). Page size at address %#x is %#lx\n", addr, page_size);
       return -1;
     }
@@ -3000,7 +3000,7 @@ int stlink_erase_flash_mass(stlink_t *sl) {
   if (sl->flash_type == STLINK_FLASH_TYPE_L0 ||
       sl->flash_type == STLINK_FLASH_TYPE_WB) {
 
-    stlink_erase_flash_section(sl, sl->flash_base, sl->flash_size);
+    err = stlink_erase_flash_section(sl, sl->flash_base, sl->flash_size, false);
 
   } else {
     wait_flash_busy(sl);
@@ -3525,7 +3525,10 @@ int stlink_write_flash(stlink_t *sl, stm32_addr_t addr, uint8_t *base,
   stlink_core_id(sl);
 
   // Erase this section of the flash
-  stlink_erase_flash_section(sl, addr, len);
+  if (stlink_erase_flash_section(sl, addr, len, true) < 0) {
+    ELOG("Failed to erase the flash prior to writing\n");
+    return (-1);
+  }
 
   if (eraseonly) {
     return (0);
diff --git a/src/st-flash/flash.c b/src/st-flash/flash.c
index b1cdd02..6ac2ba4 100644
--- a/src/st-flash/flash.c
+++ b/src/st-flash/flash.c
@@ -169,7 +169,7 @@ int main(int ac, char** av) {
         }
     } else if (o.cmd == FLASH_CMD_ERASE) {
         if (o.size > 0 && o.addr > 0)
-          err = stlink_erase_flash_section(sl, o.addr, o.size);
+          err = stlink_erase_flash_section(sl, o.addr, o.size, false);
         else
           err = stlink_erase_flash_mass(sl);