Do uncompressed size validation in raw encoder. This way

it gets done for not only raw encoder, but also Block
and LZMA_Alone encoders.

1 files changed, 73 insertions, 17 deletions

diff --git a/src/liblzma/common/raw_encoder.c b/src/liblzma/common/raw_encoder.c
index cd1f34db..fb12862b 100644
--- a/src/liblzma/common/raw_encoder.c
+++ b/src/liblzma/common/raw_encoder.c
@@ -25,6 +25,12 @@
 #include "lzma_encoder.h"
 
 
+struct lzma_coder_s {
+	lzma_next_coder next;
+	lzma_vli uncompressed_size;
+};
+
+
 static lzma_init_function
 get_function(lzma_vli id)
 {
@@ -84,22 +90,78 @@ get_function(lzma_vli id)
 }
 
 
-extern lzma_ret
-lzma_raw_encoder_init(lzma_next_coder *next, lzma_allocator *allocator,
+static lzma_ret
+raw_encode(lzma_coder *coder, lzma_allocator *allocator,
+		const uint8_t *restrict in, size_t *restrict in_pos,
+		size_t in_size, uint8_t *restrict out,
+		size_t *restrict out_pos, size_t out_size, lzma_action action)
+{
+	// Check that our amount of input stays in proper limits.
+	if (coder->uncompressed_size != LZMA_VLI_VALUE_UNKNOWN) {
+		if (action == LZMA_FINISH) {
+			if (coder->uncompressed_size != in_size - *in_pos)
+				return LZMA_PROG_ERROR;
+		} else {
+			if (coder->uncompressed_size < in_size - *in_pos)
+				return LZMA_PROG_ERROR;
+		}
+	}
+
+	const size_t in_start = *in_pos;
+
+	const lzma_ret ret = coder->next.code(coder->next.coder, allocator,
+			in, in_pos, in_size, out, out_pos, out_size, action);
+
+	if (coder->uncompressed_size != LZMA_VLI_VALUE_UNKNOWN)
+		coder->uncompressed_size -= *in_pos - in_start;
+
+	return ret;
+}
+
+
+static void
+raw_encoder_end(lzma_coder *coder, lzma_allocator *allocator)
+{
+	lzma_next_coder_end(&coder->next, allocator);
+	lzma_free(coder, allocator);
+	return;
+}
+
+
+static lzma_ret
+raw_encoder_init(lzma_next_coder *next, lzma_allocator *allocator,
 		const lzma_options_filter *options,
 		lzma_vli uncompressed_size, bool allow_implicit)
 {
+	if (next->coder == NULL) {
+		next->coder = lzma_alloc(sizeof(lzma_coder), allocator);
+		if (next->coder == NULL)
+			return LZMA_MEM_ERROR;
+
+		next->code = &raw_encode;
+		next->end = &raw_encoder_end;
+
+		next->coder->next = LZMA_NEXT_CODER_INIT;
+	}
+
+	next->coder->uncompressed_size = uncompressed_size;
+
 	// lzma_raw_coder_init() accesses get_function() via function pointer,
 	// because this way linker doesn't statically link both encoder and
 	// decoder functions if user needs only encoder or decoder.
-	const lzma_ret ret = lzma_raw_coder_init(next, allocator,
-			options, uncompressed_size, &get_function,
-			allow_implicit, true);
+	return lzma_raw_coder_init(&next->coder->next, allocator,
+			options, uncompressed_size,
+			&get_function, allow_implicit, true);
+}
 
-	if (ret != LZMA_OK)
-		lzma_next_coder_end(next, allocator);
 
-	return ret;
+extern lzma_ret
+lzma_raw_encoder_init(lzma_next_coder *next, lzma_allocator *allocator,
+		const lzma_options_filter *options,
+		lzma_vli uncompressed_size, bool allow_implicit)
+{
+	lzma_next_coder_init(raw_encoder_init, next, allocator,
+			options, uncompressed_size, allow_implicit);
 }
 
 
@@ -107,18 +169,12 @@ extern LZMA_API lzma_ret
 lzma_raw_encoder(lzma_stream *strm, const lzma_options_filter *options,
 		lzma_vli uncompressed_size, lzma_bool allow_implicit)
 {
-	return_if_error(lzma_strm_init(strm));
+	lzma_next_strm_init(strm, raw_encoder_init,
+			options, uncompressed_size, allow_implicit);
 
 	strm->internal->supported_actions[LZMA_RUN] = true;
 	strm->internal->supported_actions[LZMA_SYNC_FLUSH] = true;
 	strm->internal->supported_actions[LZMA_FINISH] = true;
 
-	const lzma_ret ret = lzma_raw_coder_init(&strm->internal->next,
-			strm->allocator, options, uncompressed_size,
-			&get_function, allow_implicit, true);
-
-	if (ret != LZMA_OK)
-		lzma_end(strm);
-
-	return ret;
+	return LZMA_OK;
 }