[PATCH 2/3] Fix zlib/lzma decompression.

[Yousong Zhou]

Let {zlib,lzma}_decompress_file() return NULL if anything wrong happened
to allow the other method to have a chance to run.

Signed-off-by: Yousong Zhou <yszhou4tech at gmail.com>
---
 kexec/lzma.c |   33 ++++++++++++++++++++++-----------
 kexec/zlib.c |   57 +++++++++++++++++++++++++++++++++++----------------------
 2 files changed, 57 insertions(+), 33 deletions(-)

diff --git a/kexec/lzma.c b/kexec/lzma.c
index 939aeb3..5bfccb7 100644
--- a/kexec/lzma.c
+++ b/kexec/lzma.c
@@ -162,13 +162,16 @@ char *lzma_decompress_file(const char *filename, off_t *r_size)
 	off_t size, allocated;
 	ssize_t result;
 
-	if (!filename) {
-		*r_size = 0;
-		return 0;
-	}
+	dbgprintf("Try LZMA decompression.\n");
+
+	*r_size = 0;
+	if (!filename)
+		return NULL;
+
 	fp = lzopen(filename, "rb");
 	if (fp == 0) {
-		die("Cannot open `%s'\n", filename);
+		dbgprintf("Cannot open `%s'\n", filename);
+		return NULL;
 	}
 	size = 0;
 	allocated = 65536;
@@ -183,17 +186,25 @@ char *lzma_decompress_file(const char *filename, off_t *r_size)
 			if ((errno == EINTR) || (errno == EAGAIN))
 				continue;
 
-			die ("read on %s of %ld bytes failed\n",
-				filename, (allocated - size) + 0UL);
+			dbgprintf("%s: read on %s of %ld bytes failed\n",
+				__func__, filename, (allocated - size) + 0UL);
+			break;
 		}
 		size += result;
-	} while(result > 0);
-	result = lzclose(fp);
-	if (result != LZMA_OK) {
-		die ("Close of %s failed\n", filename);
+	} while (result > 0);
+
+	if (lzclose(fp) != LZMA_OK) {
+		dbgprintf("%s: Close of %s failed\n", __func__, filename);
+		goto fail;
 	}
+	if (result < 0)
+		goto fail;
+
 	*r_size =  size;
 	return buf;
+fail:
+	free(buf);
+	return NULL;
 }
 #else
 char *lzma_decompress_file(const char *UNUSED(filename), off_t *UNUSED(r_size))
diff --git a/kexec/zlib.c b/kexec/zlib.c
index d44df12..7170ac3 100644
--- a/kexec/zlib.c
+++ b/kexec/zlib.c
@@ -15,29 +15,39 @@
 #include <ctype.h>
 #include <zlib.h>
 
+static void _gzerror(gzFile fp, int *errnum, const char **errmsg)
+{
+	*errmsg = gzerror(fp, errnum);
+	if (*errnum == Z_ERRNO) {
+		*errmsg = strerror(*errnum);
+	}
+}
+
 char *zlib_decompress_file(const char *filename, off_t *r_size)
 {
 	gzFile fp;
 	int errnum;
 	const char *msg;
 	char *buf;
-	off_t size, allocated;
+	off_t size = 0, allocated;
 	ssize_t result;
 
+	dbgprintf("Try gzip decompression.\n");
+
+	*r_size = 0;
 	if (!filename) {
-		*r_size = 0;
-		return 0;
+		return NULL;
 	}
 	fp = gzopen(filename, "rb");
 	if (fp == 0) {
-		msg = gzerror(fp, &errnum);
-		if (errnum == Z_ERRNO) {
-			msg = strerror(errno);
-		}
-		fprintf(stderr, "Cannot open `%s': %s\n", filename, msg);
+		_gzerror(fp, &errnum, &msg);
+		dbgprintf("Cannot open `%s': %s\n", filename, msg);
+		return NULL;
+	}
+	if (gzdirect(fp)) {
+		/* It's not in gzip format */
 		return NULL;
 	}
-	size = 0;
 	allocated = 65536;
 	buf = xmalloc(allocated);
 	do {
@@ -49,25 +59,28 @@ char *zlib_decompress_file(const char *filename, off_t *r_size)
 		if (result < 0) {
 			if ((errno == EINTR) || (errno == EAGAIN))
 				continue;
-
-			msg = gzerror(fp, &errnum);
-			if (errnum == Z_ERRNO) {
-				msg = strerror(errno);
-			}
-			die ("read on %s of %ld bytes failed: %s\n",
-				filename, (allocated - size) + 0UL, msg);
+			_gzerror(fp, &errnum, &msg);
+			dbgprintf("Read on %s of %ld bytes failed: %s\n",
+					filename, (allocated - size) + 0UL, msg);
+			size = 0;
+			goto fail;
 		}
 		size += result;
 	} while(result > 0);
+
+fail:
 	result = gzclose(fp);
 	if (result != Z_OK) {
-		msg = gzerror(fp, &errnum);
-		if (errnum == Z_ERRNO) {
-			msg = strerror(errno);
-		}
-		die ("Close of %s failed: %s\n", filename, msg);
+		_gzerror(fp, &errnum, &msg);
+		dbgprintf(" Close of %s failed: %s\n", filename, msg);
+	}
+
+	if (size > 0) {
+		*r_size = size;
+	} else {
+		free(buf);
+		buf = NULL;
 	}
-	*r_size =  size;
 	return buf;
 }
 #else
-- 
1.7.10.4