From 78365edbee5c6fdcf56fd13cab6406487358133d Mon Sep 17 00:00:00 2001
From: nightwalker-87 <15526941+Nightwalker-87@users.noreply.github.com>
Date: Sun, 17 Sep 2023 00:33:25 +0200
Subject: [PATCH] Fixes for STM32H7 & STM32G0B1 devices

- Fixed flash lock for STM32H7 dual bank devices
- Fixed flash erase issue on STM32G0B1 (Closes #1321)
---
 src/stlink-lib/calculate.c    |  2 +-
 src/stlink-lib/common_flash.c | 28 ++++++++++++----------------
 2 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/src/stlink-lib/calculate.c b/src/stlink-lib/calculate.c
index a7344e0b0..1230875fd 100644
--- a/src/stlink-lib/calculate.c
+++ b/src/stlink-lib/calculate.c
@@ -66,7 +66,7 @@ uint32_t calculate_L4_page(stlink_t *sl, uint32_t flashaddr) {
   if (sl->chip_id == STM32_CHIPID_L4 ||
       sl->chip_id == STM32_CHIPID_L496x_L4A6x ||
       sl->chip_id == STM32_CHIPID_L4Rx) {
-    // this chip use dual banked flash
+    // these chips use dual bank flash
     if (flashopt & (uint32_t)(1lu << FLASH_L4_OPTR_DUALBANK)) {
       uint32_t banksize = sl->flash_size / 2;
 
diff --git a/src/stlink-lib/common_flash.c b/src/stlink-lib/common_flash.c
index 3038b53e7..62b3e260c 100644
--- a/src/stlink-lib/common_flash.c
+++ b/src/stlink-lib/common_flash.c
@@ -96,6 +96,11 @@ void lock_flash(stlink_t *sl) {
     cr_lock_shift = FLASH_Gx_CR_LOCK;
   } else if (sl->flash_type == STM32_FLASH_TYPE_H7) {
     cr_reg = FLASH_H7_CR1;
+    if (sl->chip_flags & CHIP_F_HAS_DUAL_BANK) {
+      cr2_reg = FLASH_H7_CR2;
+    }
+    cr_lock_shift = FLASH_H7_CR_LOCK;
+    cr_mask = ~(1u << FLASH_H7_CR_SER);
   } else if (sl->flash_type == STM32_FLASH_TYPE_L0_L1) {
     cr_reg = get_stm32l0_flash_base(sl) + FLASH_PECR_OFF;
     cr_lock_shift = FLASH_L0_PELOCK;
@@ -108,11 +113,6 @@ void lock_flash(stlink_t *sl) {
   } else if (sl->flash_type == STM32_FLASH_TYPE_WB_WL) {
     cr_reg = FLASH_WB_CR;
     cr_lock_shift = FLASH_WB_CR_LOCK;
-    if (sl->chip_flags & CHIP_F_HAS_DUAL_BANK) {
-      cr2_reg = FLASH_H7_CR2;
-    }
-    cr_lock_shift = FLASH_H7_CR_LOCK;
-    cr_mask = ~(1u << FLASH_H7_CR_SER);
   } else {
     ELOG("unsupported flash method, abort\n");
     return;
@@ -965,11 +965,7 @@ int32_t stlink_erase_flash_page(stlink_t *sl, stm32_addr_t flashaddr) {
     unlock_flash_if(sl);
 
     // select the page to erase
-    if ((sl->chip_id == STM32_CHIPID_L4) ||
-        (sl->chip_id == STM32_CHIPID_L43x_L44x) ||
-        (sl->chip_id == STM32_CHIPID_L45x_L46x) ||
-        (sl->chip_id == STM32_CHIPID_L496x_L4A6x) ||
-        (sl->chip_id == STM32_CHIPID_L4Rx)) {
+    if (sl->flash_type == STM32_FLASH_TYPE_L4) {
       // calculate the actual bank+page from the address
       uint32_t page = calculate_L4_page(sl, flashaddr);
 
@@ -1071,16 +1067,16 @@ int32_t stlink_erase_flash_page(stlink_t *sl, stm32_addr_t flashaddr) {
     if (sl->flash_type == STM32_FLASH_TYPE_G0) {
       uint32_t flash_page = ((flashaddr - STM32_FLASH_BASE) / sl->flash_pgsz);
       stlink_read_debug32(sl, FLASH_Gx_CR, &val);
-      // sec 3.7.5 - PNB[5:0] is offset by 3. PER is 0x2.
-      val &= ~(0x3F << 3);
-      val |= ((flash_page & 0x3F) << 3) | (1 << FLASH_CR_PER);
+      // sec 3.7.5 - PNB[9:0] is offset by 3. PER is 0x2.
+      val &= ~(0x3FF << 3);
+      val |= ((flash_page & 0x3FF) << 3) | (1 << FLASH_CR_PER);
       stlink_write_debug32(sl, FLASH_Gx_CR, val);
     } else if (sl->flash_type == STM32_FLASH_TYPE_G4) {
       uint32_t flash_page = ((flashaddr - STM32_FLASH_BASE) / sl->flash_pgsz);
       stlink_read_debug32(sl, FLASH_Gx_CR, &val);
-      // sec 3.7.5 - PNB[6:0] is offset by 3. PER is 0x2.
-      val &= ~(0x7F << 3);
-      val |= ((flash_page & 0x7F) << 3) | (1 << FLASH_CR_PER);
+      // sec 3.7.5 - PNB[9:0] is offset by 3. PER is 0x2.
+      val &= ~(0x7FF << 3);
+      val |= ((flash_page & 0x7FF) << 3) | (1 << FLASH_CR_PER);
       stlink_write_debug32(sl, FLASH_Gx_CR, val);
     } else if (sl->flash_type == STM32_FLASH_TYPE_L5_U5_H5) {
       uint32_t flash_page;