From 6cbf7eda3cbe0f8cbaa84b4daaa86dfa2a696a77 Mon Sep 17 00:00:00 2001
From: Patrick Wildt <patrick@blueri.se>
Date: Wed, 10 May 2017 22:18:54 +0200
Subject: [PATCH] arm: mvebu: kwbimage: inline function to fix use-after-free

image_version_file()'s only use is to return the version number of the
specified image, and it's only called by kwbimage_generate().  This
version function mallocs "image_cfg" and reads the contents of the image
into that buffer.  Before return to its caller it frees the buffer.

After extracting the version, kwb_image_generate() tries to calculate
the header size by calling image_headersz_v1().  This function now
accesses "image_cfg", which has already been freed.

Since image_version_file() is only used by a single function, inline it
into kwbimage_generate() and only free the buffer after it is no longer
needed.  This also improves code readability since the code is mostly
equal to kwbimage_set_header().

Signed-off-by: Patrick Wildt <patrick@blueri.se>
Signed-off-by: Stefan Roese <sr@denx.de>
---
 tools/kwbimage.c | 93 +++++++++++++++++++++++++++++---------------------------
 1 file changed, 48 insertions(+), 45 deletions(-)

diff --git a/tools/kwbimage.c b/tools/kwbimage.c
index 8c0e730..edef560 100644
--- a/tools/kwbimage.c
+++ b/tools/kwbimage.c
@@ -1476,47 +1476,6 @@ static int image_get_version(void)
 	return e->version;
 }
 
-static int image_version_file(const char *input)
-{
-	FILE *fcfg;
-	int version;
-	int ret;
-
-	fcfg = fopen(input, "r");
-	if (!fcfg) {
-		fprintf(stderr, "Could not open input file %s\n", input);
-		return -1;
-	}
-
-	image_cfg = malloc(IMAGE_CFG_ELEMENT_MAX *
-			   sizeof(struct image_cfg_element));
-	if (!image_cfg) {
-		fprintf(stderr, "Cannot allocate memory\n");
-		fclose(fcfg);
-		return -1;
-	}
-
-	memset(image_cfg, 0,
-	       IMAGE_CFG_ELEMENT_MAX * sizeof(struct image_cfg_element));
-	rewind(fcfg);
-
-	ret = image_create_config_parse(fcfg);
-	fclose(fcfg);
-	if (ret) {
-		free(image_cfg);
-		return -1;
-	}
-
-	version = image_get_version();
-	/* Fallback to version 0 is no version is provided in the cfg file */
-	if (version == -1)
-		version = 0;
-
-	free(image_cfg);
-
-	return version;
-}
-
 static void kwbimage_set_header(void *ptr, struct stat *sbuf, int ifd,
 				struct image_tool_params *params)
 {
@@ -1657,18 +1616,62 @@ static int kwbimage_verify_header(unsigned char *ptr, int image_size,
 static int kwbimage_generate(struct image_tool_params *params,
 			     struct image_type_params *tparams)
 {
+	FILE *fcfg;
 	int alloc_len;
+	int version;
 	void *hdr;
-	int version = 0;
+	int ret;
 
-	version = image_version_file(params->imagename);
-	if (version == 0) {
+	fcfg = fopen(params->imagename, "r");
+	if (!fcfg) {
+		fprintf(stderr, "Could not open input file %s\n",
+			params->imagename);
+		exit(EXIT_FAILURE);
+	}
+
+	image_cfg = malloc(IMAGE_CFG_ELEMENT_MAX *
+			   sizeof(struct image_cfg_element));
+	if (!image_cfg) {
+		fprintf(stderr, "Cannot allocate memory\n");
+		fclose(fcfg);
+		exit(EXIT_FAILURE);
+	}
+
+	memset(image_cfg, 0,
+	       IMAGE_CFG_ELEMENT_MAX * sizeof(struct image_cfg_element));
+	rewind(fcfg);
+
+	ret = image_create_config_parse(fcfg);
+	fclose(fcfg);
+	if (ret) {
+		free(image_cfg);
+		exit(EXIT_FAILURE);
+	}
+
+	version = image_get_version();
+	switch (version) {
+		/*
+		 * Fallback to version 0 if no version is provided in the
+		 * cfg file
+		 */
+	case -1:
+	case 0:
 		alloc_len = sizeof(struct main_hdr_v0) +
 			sizeof(struct ext_hdr_v0);
-	} else {
+		break;
+
+	case 1:
 		alloc_len = image_headersz_v1(NULL);
+		break;
+
+	default:
+		fprintf(stderr, "Unsupported version %d\n", version);
+		free(image_cfg);
+		exit(EXIT_FAILURE);
 	}
 
+	free(image_cfg);
+
 	hdr = malloc(alloc_len);
 	if (!hdr) {
 		fprintf(stderr, "%s: malloc return failure: %s\n",