From 2cbcb0c9e4824bb7fa5c103be95c936f660d37ec Mon Sep 17 00:00:00 2001
From: Moinak Ghosh <moinakg@gmail.com>
Date: Sat, 4 Aug 2012 17:55:20 +0530
Subject: [PATCH] Fix buffer sizing for LZ4. Fix exit condition checks in LZ4
 decompression wrapper.

---
 lz4_compress.c | 31 +++++++++++++++++++++++--------
 main.c         | 28 ++++++++++++++++++++++++++--
 pcompress.h    |  1 +
 3 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/lz4_compress.c b/lz4_compress.c
index 60d5cd2..c219e43 100644
--- a/lz4_compress.c
+++ b/lz4_compress.c
@@ -32,6 +32,8 @@
 #include <lz4hc.h>
 #include <allocator.h>
 
+#define	LZ4_MAX_CHUNK	2147450621L
+
 struct lz4_params {
 	int level;
 };
@@ -41,14 +43,23 @@ lz4_stats(int show)
 {
 }
 
+int
+lz4_buf_extra(ssize_t buflen)
+{
+	if (buflen > LZ4_MAX_CHUNK)
+		buflen = LZ4_MAX_CHUNK;
+	return (LZ4_compressBound(buflen) - buflen + sizeof(int));
+}
+
 int
 lz4_init(void **data, int *level, ssize_t chunksize)
 {
 	struct lz4_params *lzdat;
 	int lev;
 
-	if (chunksize > INT_MAX) {
-		fprintf(stderr, "Chunk size too big for LZ4.\n");
+	if (chunksize > LZ4_MAX_CHUNK) {
+		fprintf(stderr, "Max allowed chunk size for LZ4 is: %d \n",
+		    LZ4_MAX_CHUNK);
 		return (1);
 	}
 	lzdat = slab_alloc(NULL, sizeof (struct lz4_params));
@@ -85,12 +96,13 @@ lz4_compress(void *src, size_t srclen, void *dst, size_t *dstlen,
 
 	if (lzdat->level == 1) {
 		rv = LZ4_compress(src, dst, _srclen);
+
 	} else if (lzdat->level == 2) {
 		rv = LZ4_compress(src, dst, _srclen);
-		if (rv == 0) {
+		if (rv == 0 || rv > *dstlen) {
 			return (-1);
 		}
-		dst2 = slab_alloc(NULL, rv + sizeof (int));
+		dst2 = slab_alloc(NULL, rv + sizeof (int) + LZ4_compressBound(rv));
 		*((int *)dst2) = htonl(rv);
 		rv = LZ4_compressHC(dst, dst2 + sizeof (int), rv);
 		if (rv != 0) {
@@ -120,20 +132,23 @@ lz4_decompress(void *src, size_t srclen, void *dst, size_t *dstlen,
 
 	if (lzdat->level == 1 || lzdat->level == 3) {
 		rv = LZ4_uncompress(src, dst, _dstlen);
+		if (rv != srclen) {
+			return (-1);
+		}
 
 	} else if (lzdat->level == 2) {
 		int sz1;
 
 		sz1 = ntohl(*((int *)src));
 		rv = LZ4_uncompress(src + sizeof (int), dst, sz1);
-		if (rv != sz1) {
+		if (rv != srclen - sizeof (int)) {
 			return (-1);
 		}
 		memcpy(src, dst, sz1);
 		rv = LZ4_uncompress(src, dst, _dstlen);
-	}
-	if (rv != srclen) {
-		return (-1);
+		if (rv != sz1) {
+			return (-1);
+		}
 	}
 	return (0);
 }
diff --git a/main.c b/main.c
index 259ab01..2b51b0c 100644
--- a/main.c
+++ b/main.c
@@ -398,8 +398,20 @@ start_decompress(const char *filename, const char *to_filename)
 		goto uncomp_done;
 	}
 
-	compressed_chunksize = chunksize + (chunksize >> 6) + sizeof (uint64_t)
-	    + sizeof (chunksize);
+	compressed_chunksize = chunksize + sizeof (chunksize) +
+	    sizeof (uint64_t) + sizeof (chunksize) + zlib_buf_extra(chunksize);
+
+	/*
+	 * Adjust for LZ4 overflow size if LZ4 was selected.
+	 * TODO: A more generic way to handle this for various algos. I'm too
+	 * lazy to do this at present.
+	 */
+	if (strncmp(algo, "lz4", 3) == 0) {
+		if (chunksize + lz4_buf_extra(chunksize) > compressed_chunksize) {
+			compressed_chunksize += (chunksize + lz4_buf_extra(chunksize) -
+			    compressed_chunksize);
+		}
+	}
 
 	if (flags & FLAG_DEDUP) {
 		enable_rabin_scan = 1;
@@ -804,6 +816,18 @@ start_compress(const char *filename, uint64_t chunksize, int level)
 	compressed_chunksize = chunksize + sizeof (chunksize) +
 	    sizeof (uint64_t) + sizeof (chunksize) + zlib_buf_extra(chunksize);
 
+	/*
+	 * Adjust for LZ4 overflow size if LZ4 was selected.
+	 * TODO: A more generic way to handle this for various algos. I'm too
+	 * lazy to do this at present.
+	 */
+	if (strncmp(algo, "lz4", 3) == 0) {
+		if (chunksize + lz4_buf_extra(chunksize) > compressed_chunksize) {
+			compressed_chunksize += (chunksize + lz4_buf_extra(chunksize) -
+			    compressed_chunksize);
+		}
+	}
+
 	flags = 0;
 	if (enable_rabin_scan) {
 		flags |= FLAG_DEDUP;
diff --git a/pcompress.h b/pcompress.h
index a0b3f48..011f09d 100644
--- a/pcompress.h
+++ b/pcompress.h
@@ -56,6 +56,7 @@ extern uint64_t lzma_crc64(const uint8_t *buf, size_t size, uint64_t crc);
 extern uint64_t lzma_crc64_8bchk(const uint8_t *buf, size_t size,
 	uint64_t crc, uint64_t *cnt);
 extern uint32_t zlib_buf_extra(ssize_t buflen);
+extern int lz4_buf_extra(ssize_t buflen);
 
 extern int zlib_compress(void *src, size_t srclen, void *dst,
 	size_t *destlen, int level, uchar_t chdr, void *data);