From 72eaf8d97886fade64c5049f7ead33eb83600bf7 Mon Sep 17 00:00:00 2001
From: David Garske <david@wolfssl.com>
Date: Fri, 23 Aug 2024 15:44:30 -0700
Subject: [PATCH] Fixes for QSPI driver.

---
 hal/zynq.c       | 67 +++++++++++++++++++++++++-----------------------
 hal/zynq.h       |  6 +++--
 src/update_ram.c |  4 +--
 3 files changed, 41 insertions(+), 36 deletions(-)

diff --git a/hal/zynq.c b/hal/zynq.c
index 702237608..368e05c73 100644
--- a/hal/zynq.c
+++ b/hal/zynq.c
@@ -948,46 +948,49 @@ int RAMFUNCTION ext_flash_write(uintptr_t address, const uint8_t *data, int len)
 {
     int ret = 0;
     uint8_t cmd[8]; /* size multiple of uint32_t */
-    uint32_t xferSz, page, pages, idx = 0;
+    uint32_t xferSz, page, pages, idx;
     uintptr_t addr;
 
     /* write by page */
     pages = ((len + (FLASH_PAGE_SIZE-1)) / FLASH_PAGE_SIZE);
     for (page = 0; page < pages; page++) {
         ret = qspi_write_enable(&mDev);
-        if (ret == GQSPI_CODE_SUCCESS) {
-            xferSz = len;
-            if (xferSz > FLASH_PAGE_SIZE)
-                xferSz = FLASH_PAGE_SIZE;
-
-            addr = address + (page * FLASH_PAGE_SIZE);
-            if (mDev.stripe) {
-                /* For dual parallel the address divide by 2 */
-                addr /= 2;
-            }
+        if (ret != GQSPI_CODE_SUCCESS) {
+            break;
+        }
+        xferSz = len;
+        if (xferSz > FLASH_PAGE_SIZE)
+            xferSz = FLASH_PAGE_SIZE;
+
+        addr = address + (page * FLASH_PAGE_SIZE);
+        if (mDev.stripe) {
+            /* For dual parallel the address divide by 2 */
+            addr /= 2;
+        }
 
-            /* ------ Write Flash (page at a time) ------ */
-            memset(cmd, 0, sizeof(cmd));
-            cmd[idx++] = PAGE_PROG_CMD;
-        #if GQPI_USE_4BYTE_ADDR == 1
-            cmd[idx++] = ((addr >> 24) & 0xFF);
-        #endif
-            cmd[idx++] = ((addr >> 16) & 0xFF);
-            cmd[idx++] = ((addr >> 8)  & 0xFF);
-            cmd[idx++] = ((addr >> 0)  & 0xFF);
-            ret = qspi_transfer(&mDev, cmd, idx,
-                (const uint8_t*)(data + (page * FLASH_PAGE_SIZE)),
-                xferSz, NULL, 0, 0, GQSPI_GEN_FIFO_MODE_SPI);
-            wolfBoot_printf("Flash Page %d Write: Ret %d\n", page, ret);
-            if (ret != GQSPI_CODE_SUCCESS)
-                break;
+        /* ------ Write Flash (page at a time) ------ */
+        memset(cmd, 0, sizeof(cmd));
+        idx = 0;
+        cmd[idx++] = PAGE_PROG_CMD;
+    #if GQPI_USE_4BYTE_ADDR == 1
+        cmd[idx++] = ((addr >> 24) & 0xFF);
+    #endif
+        cmd[idx++] = ((addr >> 16) & 0xFF);
+        cmd[idx++] = ((addr >> 8)  & 0xFF);
+        cmd[idx++] = ((addr >> 0)  & 0xFF);
+        ret = qspi_transfer(&mDev, cmd, idx,
+            (const uint8_t*)(data + (page * FLASH_PAGE_SIZE)),
+            xferSz, NULL, 0, 0, GQSPI_GEN_FIFO_MODE_SPI);
+        wolfBoot_printf("Flash Page %d Write: Ret %d\n", page, ret);
+        if (ret != GQSPI_CODE_SUCCESS)
+            break;
 
-            ret = qspi_wait_ready(&mDev); /* Wait for not busy */
-            if (ret != GQSPI_CODE_SUCCESS) {
-                break;
-            }
-            qspi_write_disable(&mDev);
+        ret = qspi_wait_ready(&mDev); /* Wait for not busy */
+        if (ret != GQSPI_CODE_SUCCESS) {
+            break;
         }
+        qspi_write_disable(&mDev);
+        len -= xferSz;
     }
 
     return ret;
@@ -1101,7 +1104,7 @@ static int test_ext_flash(QspiDev_t* dev)
 {
     int ret;
     uint32_t i;
-    uint8_t pageData[FLASH_PAGE_SIZE];
+    uint8_t pageData[FLASH_PAGE_SIZE*4];
 
 #ifndef TEST_FLASH_READONLY
     /* Erase sector */
diff --git a/hal/zynq.h b/hal/zynq.h
index 6e0a7b29f..1e59949f7 100644
--- a/hal/zynq.h
+++ b/hal/zynq.h
@@ -195,9 +195,11 @@
 #endif
 #ifndef FLASH_PAGE_SIZE
     #ifdef ZCU102
-        #define FLASH_PAGE_SIZE 256 /* MT25QU512ABB */
+        /* MT25QU512ABB */
+        #define FLASH_PAGE_SIZE 256
     #else
-        #define FLASH_PAGE_SIZE 512 /* MT25QU01GBBB */
+        /* MT25QU01GBBB */
+        #define FLASH_PAGE_SIZE 512
     #endif
 #endif
 #define FLASH_NUM_SECTORS      (FLASH_DEVICE_SIZE/WOLFBOOT_SECTOR_SIZE)
diff --git a/src/update_ram.c b/src/update_ram.c
index edd8591fc..9856d6ede 100644
--- a/src/update_ram.c
+++ b/src/update_ram.c
@@ -39,8 +39,8 @@ extern int wolfBoot_get_dts_size(void *dts_addr);
 extern uint32_t kernel_load_addr;
 extern uint32_t dts_load_addr;
 
-#if (defined(EXT_FLASH) && defined(NO_XIP)) || \
-    (defined(EXT_ENCRYPTED) && defined(MMU)) && \
+#if ((defined(EXT_FLASH) && defined(NO_XIP)) || \
+    (defined(EXT_ENCRYPTED) && defined(MMU))) && \
     !defined(WOLFBOOT_NO_RAMBOOT)
     /* Load firmware to RAM on boot (single flash read) */
     #undef  WOLFBOOT_USE_RAMBOOT