Do uncompressed size validation in raw encoder. This way

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

diff --git a/src/liblzma/common/raw_encoder.c b/src/liblzma/common/raw_encoder.c
index cd1f34dbef7a3371791fad329e2364809a5df01c..fb12862b2f1e7887a54e473a528270c9d9ff6730 100644 (file)
--- 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;
 }