bootm: Factor out common parts of image decompression code

Adjust the code so that the error reporting can all be done at the end,
and is the same for each decompression method. Try to detect when
decompression fails due to lack of space. Keep the behaviour of
resetting on failure even though there should be no memory corruption
now.

Signed-off-by: Simon Glass <sjg@chromium.org>

common/bootm.c

@@ -283,97 +283,103 @@ static void print_decomp_msg(int comp_type, int type, bool is_xip)
 		printf("   Uncompressing %s ... ", name);
 }
 
-#if defined(CONFIG_GZIP) || defined(CONFIG_GZIP) || defined(CONFIG_BZIP2) || \
-	defined(CONFIG_LZMA) || defined(CONFIG_LZO)
-static int handle_decomp_error(const char *algo, size_t size, size_t unc_len,
-			       int ret)
+/**
+ * handle_decomp_error() - display a decompression error
+ *
+ * This function tries to produce a useful message. In the case where the
+ * uncompressed size is the same as the available space, we can assume that
+ * the image is too large for the buffer.
+ *
+ * @comp_type:		Compression type being used (IH_COMP_...)
+ * @uncomp_size:	Number of bytes uncompressed
+ * @unc_len:		Amount of space available for decompression
+ * @ret:		Error code to report
+ * @return BOOTM_ERR_RESET, indicating that the board must be reset
+ */
+static int handle_decomp_error(int comp_type, size_t uncomp_size,
+			       size_t unc_len, int ret)
 {
-	if (size >= unc_len)
-		puts("Image too large: increase CONFIG_SYS_BOOTM_LEN\n");
+	const char *name = genimg_get_comp_name(comp_type);
+
+	if (uncomp_size >= unc_len)
+		printf("Image too large: increase CONFIG_SYS_BOOTM_LEN\n");
 	else
-		printf("%s: uncompress or overwrite error %d\n", algo, ret);
-	puts("Must RESET board to recover\n");
+		printf("%s: uncompress error %d\n", name, ret);
+
+	/*
+	 * The decompression routines are now safe, so will not write beyond
+	 * their bounds. Probably it is not necessary to reset, but maintain
+	 * the current behaviour for now.
+	 */
+	printf("Must RESET board to recover\n");
 #ifndef USE_HOSTCC
 	bootstage_error(BOOTSTAGE_ID_DECOMP_IMAGE);
 #endif
 
 	return BOOTM_ERR_RESET;
 }
-#endif
 
 int bootm_decomp_image(int comp, ulong load, ulong image_start, int type,
 		       void *load_buf, void *image_buf, ulong image_len,
 		       uint unc_len, ulong *load_end)
 {
+	int ret = 0;
+
 	*load_end = load;
 	print_decomp_msg(comp, type, load == image_start);
+
+	/*
+	 * Load the image to the right place, decompressing if needed. After
+	 * this, image_len will be set to the number of uncompressed bytes
+	 * loaded, ret will be non-zero on error.
+	 */
 	switch (comp) {
 	case IH_COMP_NONE:
-		if (load != image_start)
+		if (load == image_start)
+			break;
+		if (image_len <= unc_len)
 			memmove_wd(load_buf, image_buf, image_len, CHUNKSZ);
-		*load_end = load + image_len;
+		else
+			ret = 1;
 		break;
 #ifdef CONFIG_GZIP
 	case IH_COMP_GZIP: {
-		int ret;
-
 		ret = gunzip(load_buf, unc_len, image_buf, &image_len);
-		if (ret != 0) {
-			return handle_decomp_error("GUNZIP", image_len,
-						   unc_len, ret);
-		}
-
-		*load_end = load + image_len;
 		break;
 	}
 #endif /* CONFIG_GZIP */
 #ifdef CONFIG_BZIP2
 	case IH_COMP_BZIP2: {
-		size_t size = unc_len;
+		uint size = unc_len;
 
 		/*
 		 * If we've got less than 4 MB of malloc() space,
 		 * use slower decompression algorithm which requires
 		 * at most 2300 KB of memory.
 		 */
-		int i = BZ2_bzBuffToBuffDecompress(load_buf, &unc_len,
+		ret = BZ2_bzBuffToBuffDecompress(load_buf, &size,
 			image_buf, image_len,
 			CONFIG_SYS_MALLOC_LEN < (4096 * 1024), 0);
-		if (i != BZ_OK) {
-			return handle_decomp_error("BUNZIP2", size, unc_len,
-						   i);
-		}
-
-		*load_end = load + unc_len;
+		image_len = size;
 		break;
 	}
 #endif /* CONFIG_BZIP2 */
 #ifdef CONFIG_LZMA
 	case IH_COMP_LZMA: {
 		SizeT lzma_len = unc_len;
-		int ret;
 
 		ret = lzmaBuffToBuffDecompress(load_buf, &lzma_len,
 					       image_buf, image_len);
-		if (ret != SZ_OK) {
-			return handle_decomp_error("LZMA", lzma_len, unc_len,
-						   ret);
-		}
-		unc_len = lzma_len;
-		*load_end = load + unc_len;
+		image_len = lzma_len;
 		break;
 	}
 #endif /* CONFIG_LZMA */
 #ifdef CONFIG_LZO
 	case IH_COMP_LZO: {
 		size_t size = unc_len;
-		int ret;
 
 		ret = lzop_decompress(image_buf, image_len, load_buf, &size);
-		if (ret != LZO_E_OK)
-			return handle_decomp_error("LZO", size, unc_len, ret);
-
-		*load_end = load + size;
+		image_len = size;
 		break;
 	}
 #endif /* CONFIG_LZO */
@@ -382,6 +388,10 @@ int bootm_decomp_image(int comp, ulong load, ulong image_start, int type,
 		return BOOTM_ERR_UNIMPLEMENTED;
 	}
 
+	if (ret)
+		return handle_decomp_error(comp, image_len, unc_len, ret);
+	*load_end = load + image_len;
+
 	puts("OK\n");
 
 	return 0;