+ Fixed bug in processing Huffmann compression

2 files changed, 9 insertions, 5 deletions

diff --git a/src/huffman/huff.cpp b/src/huffman/huff.cpp
index 877a294..9de5acb 100644
--- a/src/huffman/huff.cpp
+++ b/src/huffman/huff.cpp
@@ -517,7 +517,7 @@ unsigned int THuffmannTree::FixupItemPosByWeight(THTreeItem * pNewItem, unsigned
 }
 
 // Builds Huffman tree. Called with the first 8 bits loaded from input stream
-void THuffmannTree::BuildTree(unsigned int CompressionType)
+bool THuffmannTree::BuildTree(unsigned int CompressionType)
 {
     THTreeItem * pNewItem;
     THTreeItem * pChildLo;
@@ -530,7 +530,8 @@ void THuffmannTree::BuildTree(unsigned int CompressionType)
     MaxWeight = 0;
  
     // Ensure that the compression type is in range
-    assert((CompressionType & 0x0F) <= 0x08);
+    if((CompressionType & 0x0F) > 0x08)
+        return false;
     WeightTable  = WeightTables[CompressionType & 0x0F];
  
     // Build the linear list of entries that is sorted by byte weight
@@ -581,6 +582,7 @@ void THuffmannTree::BuildTree(unsigned int CompressionType)
 
     // Initialize the MinValidValue to 1, which invalidates all quick-link items
     MinValidValue = 1;
+    return true;
 }
 
 void THuffmannTree::IncWeightsAndRebalance(THTreeItem * pItem)
@@ -765,7 +767,8 @@ unsigned int THuffmannTree::Compress(TOutputStream * os, void * pvInBuffer, int
     unsigned char * pbOutBuff = os->pbOutBuffer;
     unsigned char InputByte;
  
-    BuildTree(CompressionType);
+    if(!BuildTree(CompressionType))
+        return 0;
     bIsCmp0 = (CompressionType == 0);
 
     // Store the compression type into output buffer
@@ -832,7 +835,8 @@ unsigned int THuffmannTree::Decompress(void * pvOutBuffer, unsigned int cbOutLen
     bIsCmp0 = (CompressionType == 0) ? 1 : 0;
  
     // Build the Huffman tree
-    BuildTree(CompressionType);    
+    if(!BuildTree(CompressionType))
+        return 0;
  
     // Process the entire input buffer until end of the stream
     while((DecompressedValue = DecodeOneByte(is)) != 0x100)
diff --git a/src/huffman/huff.h b/src/huffman/huff.h
index 2bd5abc..89993fd 100644
--- a/src/huffman/huff.h
+++ b/src/huffman/huff.h
@@ -115,7 +115,7 @@ class THuffmannTree
     THTreeItem * CreateNewItem(unsigned int DecompressedValue, unsigned int Weight, TInsertPoint InsertPoint);
 
     unsigned int FixupItemPosByWeight(THTreeItem * pItem, unsigned int MaxWeight);
-    void  BuildTree(unsigned int CompressionType);
+    bool  BuildTree(unsigned int CompressionType);
 
     void  IncWeightsAndRebalance(THTreeItem * pItem);
     void  InsertNewBranchAndRebalance(unsigned int Value1, unsigned int Value2);