From 72c1bc090cb0e6503bc9fb80155e27e55b62f551 Mon Sep 17 00:00:00 2001
From: Michael Olbrich <m.olbrich@pengutronix.de>
Date: Fri, 21 Jun 2024 15:59:36 +0200
Subject: [PATCH] image-hd: split out partition image handling and fix the
 file_size

The last partition may not be the one at the end of the image, so it
could happen, that file_size is overwritten with a smaller value.

It's usually not notiable because the value is only used for the initial
image size. The file can grow larger when it is actually filled.

Signed-off-by: Michael Olbrich <m.olbrich@pengutronix.de>
---
 image-hd.c | 68 +++++++++++++++++++++++++++++++++---------------------
 1 file changed, 42 insertions(+), 26 deletions(-)

diff --git a/image-hd.c b/image-hd.c
index 4db8d6e..5658c50 100644
--- a/image-hd.c
+++ b/image-hd.c
@@ -916,6 +916,40 @@ static int setup_part_autoresize(struct image *image, struct partition *part, bo
 	return 0;
 }
 
+static int setup_part_image(struct image *image, struct partition *part)
+{
+	struct hdimage *hd = image->handler_priv;
+	struct image *child;
+
+	if (!part->image)
+		return 0;
+
+	child = image_get(part->image);
+	if (!child) {
+		image_error(image, "could not find %s\n",
+				part->image);
+		return -EINVAL;
+	}
+	if (!part->size) {
+		if (part->in_partition_table)
+			part->size = roundup(child->size, part->align);
+		else
+			part->size = child->size;
+	}
+	if (child->size > part->size) {
+		image_error(image, "part %s size (%lld) too small for %s (%lld)\n",
+				part->name, part->size, child->file, child->size);
+		return -EINVAL;
+	}
+	if (part->offset + child->size > hd->file_size) {
+		size_t file_size = part->offset + child->size;
+		if (file_size > hd->file_size)
+			hd->file_size = file_size;
+	}
+
+	return 0;
+}
+
 static int hdimage_setup(struct image *image, cfg_t *cfg)
 {
 	struct partition *part;
@@ -1106,25 +1140,10 @@ static int hdimage_setup(struct image *image, cfg_t *cfg)
 		if (ret < 0)
 			return ret;
 
-		if (part->image) {
-			struct image *child = image_get(part->image);
-			if (!child) {
-				image_error(image, "could not find %s\n",
-						part->image);
-				return -EINVAL;
-			}
-			if (!part->size) {
-				if (part->in_partition_table)
-					part->size = roundup(child->size, part->align);
-				else
-					part->size = child->size;
-			}
-			if (child->size > part->size) {
-				image_error(image, "part %s size (%lld) too small for %s (%lld)\n",
-						part->name, part->size, child->file, child->size);
-				return -EINVAL;
-			}
-		}
+		ret = setup_part_image(image, part);
+		if (ret < 0)
+			return ret;
+
 		/* the size of the extended partition will be filled in later */
 		if (!part->size && part != hd->extended_partition) {
 			image_error(image, "part %s size must not be zero\n",
@@ -1149,14 +1168,11 @@ static int hdimage_setup(struct image *image, cfg_t *cfg)
 		if (part->offset + part->size > now)
 			now = part->offset + part->size;
 
-		if (part->image) {
-			struct image *child = image_get(part->image);
-			if (part->offset + child->size > hd->file_size) {
-				hd->file_size = part->offset + child->size;
-			}
+		if (part->logical) {
+			size_t file_size = part->offset - hd->align + 512;
+			if (file_size > hd->file_size)
+				hd->file_size = file_size;
 		}
-		else if (part->logical)
-			hd->file_size = part->offset - hd->align + 512;
 
 		if (part->logical) {
 			hd->extended_partition->size = now - hd->extended_partition->offset;