From ecb365734ad408025b661be3644ca59559ae8f80 Mon Sep 17 00:00:00 2001
From: Eric Biggers <ebiggers3@gmail.com>
Date: Sat, 1 Sep 2012 21:22:47 -0500
Subject: [PATCH] Fix in write_wim_resource(), and add test cases

---
 src/resource.c    | 37 +++++++++++++++++++++++--------------
 tests/test-imagex | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/src/resource.c b/src/resource.c
index dee8bd15..3baa8f82 100644
--- a/src/resource.c
+++ b/src/resource.c
@@ -854,18 +854,27 @@ static int write_wim_resource(struct lookup_table_entry *lte,
 		offset += to_read;
 	} while (bytes_remaining);
 
-	/* If writing a compressed resource and not doing a raw copy, write the
-	 * chunk table, and finish_wim_resource_chunk_tab() will provide the
-	 * compressed size of the resource we wrote.  Otherwise, the compressed
-	 * size of the written resource is the same as the compressed size of
-	 * the existing resource. */
-	if (out_ctype != WIM_COMPRESSION_TYPE_NONE && !raw) {
-		ret = finish_wim_resource_chunk_tab(chunk_tab, out_fp,
-						    &new_compressed_size);
-		if (ret != 0)
-			goto out_fclose;
-	} else {
+	/* Raw copy:  The new compressed size is the same as the old compressed
+	 * size
+	 * 
+	 * Using WIM_COMPRESSION_TYPE_NONE:  The new compressed size is the
+	 * original size
+	 *
+	 * Using a different compression type:  Call
+	 * finish_wim_resource_chunk_tab() and it will provide the new
+	 * compressed size.
+	 */
+	if (raw) {
 		new_compressed_size = old_compressed_size;
+	} else {
+		if (out_ctype == WIM_COMPRESSION_TYPE_NONE)
+			new_compressed_size = original_size;
+		else {
+			ret = finish_wim_resource_chunk_tab(chunk_tab, out_fp,
+							    &new_compressed_size);
+			if (ret != 0)
+				goto out_fclose;
+		}
 	}
 
 	/* Verify SHA1 message digest of the resource, unless we are doing a raw
@@ -889,8 +898,8 @@ static int write_wim_resource(struct lookup_table_entry *lte,
 		}
 	}
 
-	if (new_compressed_size >= original_size &&
-	    out_ctype != WIM_COMPRESSION_TYPE_NONE && !raw)
+	if (!raw && new_compressed_size >= original_size &&
+	    out_ctype != WIM_COMPRESSION_TYPE_NONE)
 	{
 		/* Oops!  We compressed the resource to larger than the original
 		 * size.  Write the resource uncompressed instead. */
@@ -915,7 +924,7 @@ static int write_wim_resource(struct lookup_table_entry *lte,
 		}
 		goto out_fclose;
 	}
-	wimlib_assert(new_compressed_size <= original_size);
+	wimlib_assert(new_compressed_size <= original_size || raw);
 	if (out_res_entry) {
 		out_res_entry->size          = new_compressed_size;
 		out_res_entry->original_size = original_size;
diff --git a/tests/test-imagex b/tests/test-imagex
index a0b88a96..36e02534 100755
--- a/tests/test-imagex
+++ b/tests/test-imagex
@@ -4,6 +4,7 @@
 # tests every aspect of wimlib comprehensively.
 
 set -e
+srcdir=${srcdir:-.}
 srcdir=`realpath $srcdir`
 cd tests
 
@@ -611,6 +612,46 @@ if imagex export dir.wim all new.wim --boot; then
 	error "Successfully exported multiple images with --boot but with no bootable images"
 fi
 
+# Test exporting an image to another WIM, then applying it.
+# We try with 5 different combinations of compression types to make sure we go
+# through all paths in the resource-handling code.
+for i in `seq 1 3`; do
+	case $i in 
+	1)
+		cflag1="--compress=none";
+		cflag2="--compress=none";
+		;;
+	2)
+		cflag1="--compress=xpress";
+		cflag2="--compress=xpress";
+		;;
+	3)
+		cflag1="--compress=xpress"
+		cflag2="--compress=lzx"
+		;;
+	4)
+		cflag1="--compress=none"
+		cflag2="--compress=xpress"
+		;;
+	5)
+		cflag1="--compress=xpress"
+		cflag2="--compress=none"
+		;;
+	esac
+	echo "Testing exporting then applying an image (\"$cflag1\" => \"$cflag2\")"
+	rm -rf dir.wim new.wim tmp tmp2
+	imagex capture dir dir.wim $cflag1
+	imagex capture dir2 dir2.wim $cflag2
+	imagex export dir.wim dir dir2.wim
+	imagex apply dir.wim dir tmp
+	if ! imagex apply dir2.wim dir tmp2; then
+		error "Failed to apply image that was exported to a WIM"
+	fi
+	if ! diff -r tmp tmp2; then
+		error "Image that was exported to a WIM was not applied correctly"
+	fi
+done
+
 echo "**********************************************************"
 echo "                 Basic imagex tests passed                "
 echo "**********************************************************"
-- 
2.43.0