2 files changed, 50 insertions, 42 deletions

diff --git a/common/cmd_vbexport_test.c b/common/cmd_vbexport_test.c
index 3813ce1810..fe4a3434c7 100644
--- a/common/cmd_vbexport_test.c
+++ b/common/cmd_vbexport_test.c
@@ -436,6 +436,8 @@ static int show_images_and_delay(BmpBlockHeader *bmph, int index)
 	int i;
 	ScreenLayout *screen;
 	ImageInfo *info;
+	void *rawimg;
+        uint32_t inoutsize;
 
 	screen = (ScreenLayout *)(bmph + 1);
 	screen += index;
@@ -445,18 +447,37 @@ static int show_images_and_delay(BmpBlockHeader *bmph, int index)
 	     i++) {
 		info = (ImageInfo *)((uint8_t *)bmph +
 				screen->images[i].image_info_offset);
+		inoutsize = info->original_size;
+
+                if (COMPRESS_NONE == info->compression) {
+			rawimg = NULL;
+		} else {
+			rawimg = VbExMalloc(inoutsize);
+			if (VbExDecompress(info + 1, info->compressed_size,
+					   info->compression,
+					   rawimg, &inoutsize)) {
+				goto bad;
+			}
+		}
+
 		if (VbExDisplayImage(screen->images[i].x,
 				     screen->images[i].y,
-				     info,
-				     info + 1)) {
-			VbExDebug("Failed to display image, screen=%lu, "
-				  "image=%d!\n", index, i);
-			return 1;
+				     rawimg ? rawimg : info + 1,
+				     inoutsize)) {
+			goto bad;
 		}
+		if (rawimg)
+			VbExFree(rawimg);
 	}
 
 	VbExSleepMs(1000);
 	return 0;
+
+bad:
+	if (rawimg)
+		VbExFree(rawimg);
+	VbExDebug("Failed to display image, screen=%lu, image=%d!\n", index, i);
+	return 1;
 }
 
 static int do_vbexport_test_display(
diff --git a/lib/vbexport/display.c b/lib/vbexport/display.c
index 80f45128d9..c177f3bacf 100644
--- a/lib/vbexport/display.c
+++ b/lib/vbexport/display.c
@@ -127,56 +127,43 @@ VbError_t VbExDisplayScreen(uint32_t screen_type)
 	return VBERROR_SUCCESS;
 }
 
-static uint8_t *uncompress_lzma(uint8_t *in_addr, SizeT in_size,
-                                SizeT out_size)
+VbError_t VbExDecompress(void *inbuf, uint32_t in_size,
+                         uint32_t compression_type,
+                         void *outbuf, uint32_t *out_size)
 {
-	uint8_t *out_addr = VbExMalloc(out_size);
-	SizeT lzma_len = out_size;
+	SizeT input_size = in_size;
+	SizeT output_size = *out_size;
 	int ret;
 
-	ret = lzmaBuffToBuffDecompress(out_addr, &lzma_len, in_addr, in_size);
-	if (ret != SZ_OK) {
-		VbExFree(out_addr);
-		out_addr = NULL;
-	}
-	return out_addr;
-}
-
-VbError_t VbExDisplayImage(uint32_t x, uint32_t y, const ImageInfo *info,
-                           const void *buffer)
-{
-	int ret;
-	uint8_t *raw_data;
-
-	switch (info->compression) {
+	switch (compression_type) {
 	case COMPRESS_NONE:
-		raw_data = (uint8_t *)buffer;
-		break;
+		memcpy(outbuf, inbuf, in_size);
+                *out_size = in_size;
+		return VBERROR_SUCCESS;
 
 	case COMPRESS_LZMA1:
-		raw_data = uncompress_lzma((uint8_t *)buffer,
-				(SizeT)info->compressed_size,
-				(SizeT)info->original_size);
-		if (!raw_data) {
-			VBDEBUG("LZMA decompress failed.\n");
-			return 1;
+		ret = lzmaBuffToBuffDecompress(outbuf, &output_size,
+					       inbuf, input_size);
+		if (ret != SZ_OK) {
+			return ret;
 		}
-		break;
-
-	default:
-		VBDEBUG("Unsupported compression format: %08x\n",
-				info->compression);
-		return 1;
+                *out_size = output_size;
+		return VBERROR_SUCCESS;
 	}
 
-	ret = display_callbacks_.dc_display_bitmap((ulong)raw_data, x, y);
+	VBDEBUG("Unsupported compression format: %08x\n", compression_type);
+	return VBERROR_INVALID_PARAMETER;
+}
 
-	if (info->compression == COMPRESS_LZMA1)
-		VbExFree(raw_data);
 
+VbError_t VbExDisplayImage(uint32_t x, uint32_t y,
+                           void *buffer, uint32_t buffersize)
+{
+	int ret;
+	ret = display_callbacks_.dc_display_bitmap((ulong)buffer, x, y);
 	if (ret) {
 		VBDEBUG("LCD display error.\n");
-		return 1;
+		return VBERROR_UNKNOWN;
 	}
 
 	return VBERROR_SUCCESS;