From d3eb6ecb9a9191a516ec83038c05f219e1012664 Mon Sep 17 00:00:00 2001
From: Zoltan Szabadka <szabadka@google.com>
Date: Fri, 2 Oct 2015 16:19:30 +0200
Subject: [PATCH] Add BROTLI_FREE macro, check malloc status after histogram
 allocations.

---
 dec/decode.c  | 10 +++++++---
 dec/huffman.c |  4 +---
 dec/port.h    |  5 +++++
 dec/state.c   | 36 +++++++++---------------------------
 4 files changed, 22 insertions(+), 33 deletions(-)

diff --git a/dec/decode.c b/dec/decode.c
index e9529a3e8..7c17f2e48 100644
--- a/dec/decode.c
+++ b/dec/decode.c
@@ -948,7 +948,7 @@ int BROTLI_NOINLINE BrotliAllocateRingBuffer(BrotliState* s,
   s->ringbuffer = (uint8_t*)malloc((size_t)(s->ringbuffer_size +
                                          kRingBufferWriteAheadSlack +
                                          kBrotliMaxDictionaryWordLength));
-  if (!s->ringbuffer) {
+  if (s->ringbuffer == 0) {
     return 0;
   }
   s->ringbuffer_end = s->ringbuffer + s->ringbuffer_size;
@@ -1069,8 +1069,7 @@ BrotliResult BrotliDecompressStreaming(BrotliInput input, BrotliOutput output,
         /* Allocate memory for both block_type_trees and block_len_trees. */
         s->block_type_trees = (HuffmanCode*)malloc(
             6 * BROTLI_HUFFMAN_MAX_TABLE_SIZE * sizeof(HuffmanCode));
-
-        if (s->block_type_trees == NULL) {
+        if (s->block_type_trees == 0) {
           result = BROTLI_FAILURE();
           break;
         }
@@ -1249,6 +1248,11 @@ BrotliResult BrotliDecompressStreaming(BrotliInput input, BrotliOutput output,
               s->num_block_types[1]);
           BrotliHuffmanTreeGroupInit(
               &s->distance_hgroup, num_distance_codes, s->num_dist_htrees);
+          if (s->literal_hgroup.codes == 0 ||
+              s->insert_copy_hgroup.codes == 0 ||
+              s->distance_hgroup.codes == 0) {
+            return BROTLI_FAILURE();
+          }
         }
         i = 0;
         s->state = BROTLI_STATE_TREE_GROUP;
diff --git a/dec/huffman.c b/dec/huffman.c
index 2407670f3..476c8d736 100644
--- a/dec/huffman.c
+++ b/dec/huffman.c
@@ -316,9 +316,7 @@ void BrotliHuffmanTreeGroupInit(HuffmanTreeGroup* group, int alphabet_size,
 }
 
 void BrotliHuffmanTreeGroupRelease(HuffmanTreeGroup* group) {
-  if (group->codes) {
-    free(group->codes);
-  }
+  BROTLI_FREE(group->codes);
 }
 
 #if defined(__cplusplus) || defined(c_plusplus)
diff --git a/dec/port.h b/dec/port.h
index af60dfc59..b9783e400 100644
--- a/dec/port.h
+++ b/dec/port.h
@@ -161,4 +161,9 @@ static BROTLI_INLINE unsigned BrotliRBit(unsigned input) {
 #endif  /* armv7 */
 #endif  /* gcc || clang */
 
+#define BROTLI_FREE(X) { \
+  free(X); \
+  X = NULL; \
+}
+
 #endif  /* BROTLI_DEC_PORT_H_ */
diff --git a/dec/state.c b/dec/state.c
index 71c748b52..150630197 100644
--- a/dec/state.c
+++ b/dec/state.c
@@ -103,18 +103,9 @@ void BrotliStateMetablockBegin(BrotliState* s) {
 }
 
 void BrotliStateCleanupAfterMetablock(BrotliState* s) {
-  if (s->context_modes != 0) {
-    free(s->context_modes);
-    s->context_modes = NULL;
-  }
-  if (s->context_map != 0) {
-    free(s->context_map);
-    s->context_map = NULL;
-  }
-  if (s->dist_context_map != 0) {
-    free(s->dist_context_map);
-    s->dist_context_map = NULL;
-  }
+  BROTLI_FREE(s->context_modes);
+  BROTLI_FREE(s->context_map);
+  BROTLI_FREE(s->dist_context_map);
 
   BrotliHuffmanTreeGroupRelease(&s->literal_hgroup);
   BrotliHuffmanTreeGroupRelease(&s->insert_copy_hgroup);
@@ -128,25 +119,16 @@ void BrotliStateCleanupAfterMetablock(BrotliState* s) {
 }
 
 void BrotliStateCleanup(BrotliState* s) {
-  if (s->context_modes != 0) {
-    free(s->context_modes);
-  }
-  if (s->context_map != 0) {
-    free(s->context_map);
-  }
-  if (s->dist_context_map != 0) {
-    free(s->dist_context_map);
-  }
+  BROTLI_FREE(s->context_modes);
+  BROTLI_FREE(s->context_map);
+  BROTLI_FREE(s->dist_context_map);
+
   BrotliHuffmanTreeGroupRelease(&s->literal_hgroup);
   BrotliHuffmanTreeGroupRelease(&s->insert_copy_hgroup);
   BrotliHuffmanTreeGroupRelease(&s->distance_hgroup);
 
-  if (s->ringbuffer != 0) {
-    free(s->ringbuffer);
-  }
-  if (s->block_type_trees != 0) {
-    free(s->block_type_trees);
-  }
+  BROTLI_FREE(s->ringbuffer);
+  BROTLI_FREE(s->block_type_trees);
 }
 
 #if defined(__cplusplus) || defined(c_plusplus)