fdtoverlay: Rework output allocation

At present the fdtoverlay tool allocates space for its output based on a
calculation of the worse case size based on the size of the input blobs.

Except.. that certain edge cases with very long target paths can actually
exceed that "worst case" calculation.

This reworks the code to instead dynamically reallocate the output buffer
if we run out of space at any point.

Signed-off-by: David Gibson <david@gibson.dropbear.id.au>

1 files changed, 58 insertions, 16 deletions

diff --git a/fdtoverlay.c b/fdtoverlay.c
index 40c249e..8a1d11b 100644
--- a/fdtoverlay.c
+++ b/fdtoverlay.c
@@ -18,6 +18,8 @@
 
 #include "util.h"
 
+#define BUF_INCREMENT	65536
+
 /* Usage related data. */
 static const char usage_synopsis[] =
 	"apply a number of overlays to a base blob\n"
@@ -40,34 +42,81 @@ static const char * const usage_opts_help[] = {
 
 int verbose = 0;
 
+static void *apply_one(char *base, const char *overlay, size_t *buf_len,
+		       const char *name)
+{
+	char *tmp = NULL;
+	char *tmpo;
+	int ret;
+
+	/*
+	 * We take a copies first, because a a failed apply can trash
+	 * both the base blob and the overlay
+	 */
+	tmpo = xmalloc(fdt_totalsize(overlay));
+
+	do {
+		tmp = xrealloc(tmp, *buf_len);
+		ret = fdt_open_into(base, tmp, *buf_len);
+		if (ret) {
+			fprintf(stderr,
+				"\nFailed to make temporary copy: %s\n",
+				fdt_strerror(ret));
+			goto fail;
+		}
+
+		memcpy(tmpo, overlay, fdt_totalsize(overlay));
+
+		ret = fdt_overlay_apply(tmp, tmpo);
+		if (ret == -FDT_ERR_NOSPACE) {
+			*buf_len += BUF_INCREMENT;
+		}
+	} while (ret == -FDT_ERR_NOSPACE);
+
+	if (ret) {
+		fprintf(stderr, "\nFailed to apply '%s': %s\n",
+			name, fdt_strerror(ret));
+		goto fail;
+	}
+
+	free(base);
+	free(tmpo);
+	return tmp;
+
+fail:
+	free(tmpo);
+	if (tmp)
+		free(tmp);
+
+	return NULL;
+}
 static int do_fdtoverlay(const char *input_filename,
 			 const char *output_filename,
 			 int argc, char *argv[])
 {
 	char *blob = NULL;
 	char **ovblob = NULL;
-	size_t blob_len, total_len;
+	size_t buf_len;
 	int i, ret = -1;
 
-	blob = utilfdt_read(input_filename, &blob_len);
+	blob = utilfdt_read(input_filename, &buf_len);
 	if (!blob) {
 		fprintf(stderr, "\nFailed to read '%s'\n", input_filename);
 		goto out_err;
 	}
-	if (fdt_totalsize(blob) > blob_len) {
+	if (fdt_totalsize(blob) > buf_len) {
 		fprintf(stderr,
  "\nBase blob is incomplete (%lu / %" PRIu32 " bytes read)\n",
-			(unsigned long)blob_len, fdt_totalsize(blob));
+			(unsigned long)buf_len, fdt_totalsize(blob));
 		goto out_err;
 	}
 	ret = 0;
 
 	/* allocate blob pointer array */
-	ovblob = malloc(sizeof(*ovblob) * argc);
+	ovblob = xmalloc(sizeof(*ovblob) * argc);
 	memset(ovblob, 0, sizeof(*ovblob) * argc);
 
 	/* read and keep track of the overlay blobs */
-	total_len = 0;
 	for (i = 0; i < argc; i++) {
 		size_t ov_len;
 		ovblob[i] = utilfdt_read(argv[i], &ov_len);
@@ -82,22 +131,15 @@ static int do_fdtoverlay(const char *input_filename,
 				fdt_totalsize(ovblob[i]));
 			goto out_err;
 		}
-		total_len += ov_len;
 	}
 
-	/* grow the blob to worst case */
-	blob_len = fdt_totalsize(blob) + total_len;
-	blob = xrealloc(blob, blob_len);
-	fdt_open_into(blob, blob, blob_len);
+	buf_len = fdt_totalsize(blob);
 
 	/* apply the overlays in sequence */
 	for (i = 0; i < argc; i++) {
-		ret = fdt_overlay_apply(blob, ovblob[i]);
-		if (ret) {
-			fprintf(stderr, "\nFailed to apply '%s': %s\n",
-				argv[i], fdt_strerror(ret));
+		blob = apply_one(blob, ovblob[i], &buf_len, argv[i]);
+		if (!blob)
 			goto out_err;
-		}
 	}
 
 	fdt_pack(blob);