Add support for allocating the LZMA2 dictionary dynamically.

Previously the dictionary was preallocated at initialization
time, which is useful since the decoder cannot then later run
out of memory, but in several cases it is just an annoying
limitation.

It is now possible to enable only the needed operation mode(s)
at build time, which saves a few bytes in code size if only
one or two modes are actually needed. Bigger savings would be
possible especially in single-call mode, but I'll think about
that later.

This commit changes the API by adding the mode argument to
xz_dec_init(). A new return value (XZ_MEM_ERROR) was also added,
but it is used only in the new XZ_DYNALLOC operation mode.

4 files changed, 95 insertions, 33 deletions

diff --git a/linux/lib/decompress_unxz.c b/linux/lib/decompress_unxz.c
index 880f301..af09213 100644
--- a/linux/lib/decompress_unxz.c
+++ b/linux/lib/decompress_unxz.c
@@ -278,9 +278,9 @@ STATIC int XZ_FUNC unxz(/*const*/ unsigned char *in, int in_size,
 	xz_crc32_init();
 
 	if (in != NULL && out != NULL)
-		s = xz_dec_init(0);
+		s = xz_dec_init(XZ_SINGLE, 0);
 	else
-		s = xz_dec_init(DICT_MAX);
+		s = xz_dec_init(XZ_PREALLOC /*FIXME*/, DICT_MAX);
 
 	if (s == NULL)
 		goto error_alloc_state;
diff --git a/linux/lib/xz/xz_dec_lzma2.c b/linux/lib/xz/xz_dec_lzma2.c
index 3710390..0b8ddb6 100644
--- a/linux/lib/xz/xz_dec_lzma2.c
+++ b/linux/lib/xz/xz_dec_lzma2.c
@@ -34,7 +34,8 @@
  *
  * In multi-call mode, also these are true:
  *    end == size
- *    size <= allocated
+ *    size <= size_max
+ *    allocated <= size
  *
  * Most of these variables are size_t to support single-call mode,
  * in which the dictionary variables address the actual output
@@ -74,11 +75,20 @@ struct dictionary {
 	uint32_t size;
 
 	/*
-	 * Amount of memory allocated for the dictionary. A special
-	 * value of zero indicates that we are in single-call mode,
-	 * where the output buffer works as the dictionary.
+	 * Maximum allowed dictionary size in multi-call mode.
+	 * This is ignored in single-call mode.
+	 */
+	uint32_t size_max;
+
+	/*
+	 * Amount of memory currently allocated for the dictionary.
+	 * This is used only with XZ_DYNALLOC. (With XZ_PREALLOC,
+	 * size_max is always the same as the allocated size.)
 	 */
 	uint32_t allocated;
+
+	/* Operation mode */
+	enum xz_mode mode;
 };
 
 /* Range decoder */
@@ -269,7 +279,7 @@ struct xz_dec_lzma2 {
  */
 static void XZ_FUNC dict_reset(struct dictionary *dict, struct xz_buf *b)
 {
-	if (dict->allocated == 0) {
+	if (DEC_IS_SINGLE(dict->mode)) {
 		dict->buf = b->out + b->out_pos;
 		dict->end = b->out_size - b->out_pos;
 	}
@@ -379,7 +389,7 @@ static void XZ_FUNC dict_uncompressed(
 		if (dict->full < dict->pos)
 			dict->full = dict->pos;
 
-		if (dict->allocated != 0) {
+		if (DEC_IS_MULTI(dict->mode)) {
 			if (dict->pos == dict->end)
 				dict->pos = 0;
 
@@ -404,7 +414,7 @@ static uint32_t XZ_FUNC dict_flush(struct dictionary *dict, struct xz_buf *b)
 {
 	size_t copy_size = dict->pos - dict->start;
 
-	if (dict->allocated != 0) {
+	if (DEC_IS_MULTI(dict->mode)) {
 		if (dict->pos == dict->end)
 			dict->pos = 0;
 
@@ -1088,28 +1098,27 @@ XZ_EXTERN enum xz_ret XZ_FUNC xz_dec_lzma2_run(
 	return XZ_OK;
 }
 
-XZ_EXTERN struct xz_dec_lzma2 * XZ_FUNC xz_dec_lzma2_create(uint32_t dict_max)
+XZ_EXTERN struct xz_dec_lzma2 * XZ_FUNC xz_dec_lzma2_create(
+		enum xz_mode mode, uint32_t dict_max)
 {
-	struct xz_dec_lzma2 *s;
-
-	/* Maximum supported dictionary by this implementation is 3 GiB. */
-	if (dict_max > ((uint32_t)3 << 30))
-		return NULL;
-
-	s = kmalloc(sizeof(*s), GFP_KERNEL);
+	struct xz_dec_lzma2 *s = kmalloc(sizeof(*s), GFP_KERNEL);
 	if (s == NULL)
 		return NULL;
 
-	if (dict_max > 0) {
+	s->dict.mode = mode;
+	s->dict.size_max = dict_max;
+
+	if (DEC_IS_PREALLOC(mode)) {
 		s->dict.buf = vmalloc(dict_max);
 		if (s->dict.buf == NULL) {
 			kfree(s);
 			return NULL;
 		}
+	} else if (DEC_IS_DYNALLOC(mode)) {
+		s->dict.buf = NULL;
+		s->dict.allocated = 0;
 	}
 
-	s->dict.allocated = dict_max;
-
 	return s;
 }
 
@@ -1123,10 +1132,23 @@ XZ_EXTERN enum xz_ret XZ_FUNC xz_dec_lzma2_reset(
 	s->dict.size = 2 + (props & 1);
 	s->dict.size <<= (props >> 1) + 11;
 
-	if (s->dict.allocated > 0 && s->dict.allocated < s->dict.size)
-		return XZ_MEMLIMIT_ERROR;
+	if (DEC_IS_MULTI(s->dict.mode)) {
+		if (s->dict.size > s->dict.size_max)
+			return XZ_MEMLIMIT_ERROR;
 
-	s->dict.end = s->dict.size;
+		s->dict.end = s->dict.size;
+
+		if (DEC_IS_DYNALLOC(s->dict.mode)) {
+			if (s->dict.allocated < s->dict.size) {
+				vfree(s->dict.buf);
+				s->dict.buf = vmalloc(s->dict.size);
+				if (s->dict.buf == NULL) {
+					s->dict.allocated = 0;
+					return XZ_MEM_ERROR;
+				}
+			}
+		}
+	}
 
 	s->lzma.len = 0;
 
@@ -1140,7 +1162,7 @@ XZ_EXTERN enum xz_ret XZ_FUNC xz_dec_lzma2_reset(
 
 XZ_EXTERN void XZ_FUNC xz_dec_lzma2_end(struct xz_dec_lzma2 *s)
 {
-	if (s->dict.allocated > 0)
+	if (DEC_IS_MULTI(s->dict.mode))
 		vfree(s->dict.buf);
 
 	kfree(s);
diff --git a/linux/lib/xz/xz_dec_stream.c b/linux/lib/xz/xz_dec_stream.c
index 21db283..bdcbf1b 100644
--- a/linux/lib/xz/xz_dec_stream.c
+++ b/linux/lib/xz/xz_dec_stream.c
@@ -48,8 +48,8 @@ struct xz_dec {
 	/* Type of the integrity check calculated from uncompressed data */
 	enum xz_check check_type;
 
-	/* True if we are operating in single-call mode. */
-	bool single_call;
+	/* Operation mode */
+	enum xz_mode mode;
 
 	/*
 	 * True if the next call to xz_dec_run() is allowed to return
@@ -737,14 +737,14 @@ XZ_EXTERN enum xz_ret XZ_FUNC xz_dec_run(struct xz_dec *s, struct xz_buf *b)
 	size_t out_start;
 	enum xz_ret ret;
 
-	if (s->single_call)
+	if (DEC_IS_SINGLE(s->mode))
 		xz_dec_reset(s);
 
 	in_start = b->in_pos;
 	out_start = b->out_pos;
 	ret = dec_main(s, b);
 
-	if (s->single_call) {
+	if (DEC_IS_SINGLE(s->mode)) {
 		if (ret == XZ_OK)
 			ret = b->in_pos == b->in_size
 					? XZ_DATA_ERROR : XZ_BUF_ERROR;
@@ -767,21 +767,22 @@ XZ_EXTERN enum xz_ret XZ_FUNC xz_dec_run(struct xz_dec *s, struct xz_buf *b)
 	return ret;
 }
 
-XZ_EXTERN struct xz_dec * XZ_FUNC xz_dec_init(uint32_t dict_max)
+XZ_EXTERN struct xz_dec * XZ_FUNC xz_dec_init(
+		enum xz_mode mode, uint32_t dict_max)
 {
 	struct xz_dec *s = kmalloc(sizeof(*s), GFP_KERNEL);
 	if (s == NULL)
 		return NULL;
 
-	s->single_call = dict_max == 0;
+	s->mode = mode;
 
 #ifdef XZ_DEC_BCJ
-	s->bcj = xz_dec_bcj_create(s->single_call);
+	s->bcj = xz_dec_bcj_create(DEC_IS_SINGLE(mode));
 	if (s->bcj == NULL)
 		goto error_bcj;
 #endif
 
-	s->lzma2 = xz_dec_lzma2_create(dict_max);
+	s->lzma2 = xz_dec_lzma2_create(mode, dict_max);
 	if (s->lzma2 == NULL)
 		goto error_lzma2;
 
diff --git a/linux/lib/xz/xz_private.h b/linux/lib/xz/xz_private.h
index f4e0b40..145649a 100644
--- a/linux/lib/xz/xz_private.h
+++ b/linux/lib/xz/xz_private.h
@@ -53,6 +53,45 @@
 #	include "xz_config.h"
 #endif
 
+/* If no specific decoding mode is requested, enable support for all modes. */
+#if !defined(XZ_DEC_SINGLE) && !defined(XZ_DEC_PREALLOC) \
+		&& !defined(XZ_DEC_DYNALLOC)
+#	define XZ_DEC_SINGLE
+#	define XZ_DEC_PREALLOC
+#	define XZ_DEC_DYNALLOC
+#endif
+
+/*
+ * The DEC_IS_foo(mode) macros are used in "if" statements. If only some
+ * of the supported modes are enabled, these macros will evaluate to true or
+ * false at compile time and thus allow the compiler to omit unneeded code.
+ */
+#ifdef XZ_DEC_SINGLE
+#	define DEC_IS_SINGLE(mode) ((mode) == XZ_SINGLE)
+#else
+#	define DEC_IS_SINGLE(mode) (false)
+#endif
+
+#ifdef XZ_DEC_PREALLOC
+#	define DEC_IS_PREALLOC(mode) ((mode) == XZ_PREALLOC)
+#else
+#	define DEC_IS_PREALLOC(mode) (false)
+#endif
+
+#ifdef XZ_DEC_DYNALLOC
+#	define DEC_IS_DYNALLOC(mode) ((mode) == XZ_DYNALLOC)
+#else
+#	define DEC_IS_DYNALLOC(mode) (false)
+#endif
+
+#if !defined(XZ_DEC_SINGLE)
+#	define DEC_IS_MULTI(mode) (true)
+#elif defined(XZ_DEC_PREALLOC) || defined(XZ_DEC_DYNALLOC)
+#	define DEC_IS_MULTI(mode) ((mode) != XZ_SINGLE)
+#else
+#	define DEC_IS_MULTI(mode) (false)
+#endif
+
 /*
  * If any of the BCJ filter decoders are wanted, define XZ_DEC_BCJ.
  * XZ_DEC_BCJ is used to enable generic support for BCJ decoders.
@@ -71,7 +110,7 @@
  * before calling xz_dec_lzma2_run().
  */
 XZ_EXTERN struct xz_dec_lzma2 * XZ_FUNC xz_dec_lzma2_create(
-		uint32_t dict_max);
+		enum xz_mode mode, uint32_t dict_max);
 
 /*
  * Decode the LZMA2 properties (one byte) and reset the decoder. Return