From f21eb7442bae5ee93930927dd478583041b756e3 Mon Sep 17 00:00:00 2001
From: Andrzej Puzdrowski <andrzej.puzdrowski@nordicsemi.no>
Date: Mon, 7 Oct 2024 17:41:24 +0200
Subject: [PATCH] susbys/dfu/img_util: refined ERASE PROGRESSIVELY
 implementation

Moved MCUboot trailer's status erase ahead any write operation.

Signed-off-by: Andrzej Puzdrowski <andrzej.puzdrowski@nordicsemi.no>
Signed-off-by: Dominik Ermel <dominik.ermel@nordicsemi.no>
---
 subsys/dfu/img_util/flash_img.c | 78 +++++++++++++++++++++------------
 1 file changed, 51 insertions(+), 27 deletions(-)

diff --git a/subsys/dfu/img_util/flash_img.c b/subsys/dfu/img_util/flash_img.c
index 9f23be50118d8c2..8e0db2fb4052729 100644
--- a/subsys/dfu/img_util/flash_img.c
+++ b/subsys/dfu/img_util/flash_img.c
@@ -43,43 +43,67 @@ BUILD_ASSERT((CONFIG_IMG_BLOCK_BUF_SIZE % FLASH_WRITE_BLOCK_SIZE == 0),
 	     "FLASH_WRITE_BLOCK_SIZE");
 #endif
 
+static int scramble_mcuboot_trailer(struct flash_img_context *ctx)
+{
+	int rc = 0;
+
+#ifdef CONFIG_IMG_ERASE_PROGRESSIVELY
+	if (stream_flash_bytes_written(&ctx->stream) == 0) {
+		off_t toff = boot_get_trailer_status_offset(ctx->flash_area->fa_size);
+		off_t offset;
+		size_t size;
+		const struct flash_parameters *fparams =
+				flash_get_parameters(flash_area_get_device(ctx->flash_area));
+
+		if (flash_params_get_erase_cap(fparams) & FLASH_ERASE_C_EXPLICIT) {
+			/* On devices with explicit erase we are aligning to page
+			 * layout.
+			 */
+			struct flash_pages_info info;
+
+			rc = flash_get_page_info_by_offs(flash_area_get_device(ctx->flash_area), toff,
+							 &info);
+			if (rc != 0) {
+				return rc;
+			}
+			offset = info.start_offset;
+			size = info.size;
+
+		} else {
+			/* On devices with no erase, we are aligning to write block
+			 * size.
+			 */
+			offset = (toff + fparams->write_block_size - 1) &
+				 ~(fparams->write_block_size - 1);
+			/* No alignment correction neede here, offset is corrected already */
+			size = ctx->flash_area->fa_size - offset;
+		}
+
+		rc = flash_area_flatten(ctx->flash_area, offset, size);
+	}
+#endif
+
+	return rc;
+}
+
+
 int flash_img_buffered_write(struct flash_img_context *ctx, const uint8_t *data,
 			     size_t len, bool flush)
 {
 	int rc;
 
-	rc = stream_flash_buffered_write(&ctx->stream, data, len, flush);
-	if (!flush) {
+	/* If there is need to erase trailer, that should happen before any
+	 * write is done to partition.
+	 */
+	rc = scramble_mcuboot_trailer(ctx);
+	if (rc != 0) {
 		return rc;
 	}
 
-#ifdef CONFIG_IMG_ERASE_PROGRESSIVELY
-	ssize_t status_offset = boot_get_trailer_status_offset(
-		ctx->flash_area->fa_size);
-
-#ifdef CONFIG_STREAM_FLASH_ERASE
-	const struct flash_parameters *fparams = flash_get_parameters(flash_area_get_device(ctx->flash_area));
-
-	if ((flash_params_get_erase_cap(fparams) & FLASH_ERASE_C_EXPLICIT)) {
-		/* use pistine-page-erase procedure for a device which needs it */
-		rc = stream_flash_erase_page(&ctx->stream,
-					ctx->flash_area->fa_off +
-					status_offset);
-	} else
-#endif
-	{
-		if (status_offset > stream_flash_bytes_written(&ctx->stream)) {
-			rc = flash_area_flatten(ctx->flash_area, status_offset,
-									ctx->flash_area->fa_off - status_offset);
-		} else {
-			rc = 0;
-		}
-	}
-
-	if (rc) {
+	rc = stream_flash_buffered_write(&ctx->stream, data, len, flush);
+	if (!flush) {
 		return rc;
 	}
-#endif
 
 	flash_area_close(ctx->flash_area);
 	ctx->flash_area = NULL;