Add BROTLI_FREE macro, check malloc status after histogram allocations.

google · Oct 2, 2015 · d3eb6ec · d3eb6ec
1 parent 9d9b95e
commit d3eb6ec
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 33 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)