1 files changed, 90 insertions, 45 deletions

diff --git a/CPP/7zip/Compress/DeflateDecoder.cpp b/CPP/7zip/Compress/DeflateDecoder.cpp
index 2848cd81..5285e143 100755..100644
--- a/CPP/7zip/Compress/DeflateDecoder.cpp
+++ b/CPP/7zip/Compress/DeflateDecoder.cpp
@@ -8,24 +8,27 @@ namespace NCompress {
 namespace NDeflate {
 namespace NDecoder {
 
-static const int kLenIdFinished = -1;
-static const int kLenIdNeedInit = -2;
-
 CCoder::CCoder(bool deflate64Mode, bool deflateNSIS):
     _deflate64Mode(deflate64Mode),
     _deflateNSIS(deflateNSIS),
     _keepHistory(false),
+    _needFinishInput(false),
     _needInitInStream(true),
     ZlibMode(false) {}
 
-UInt32 CCoder::ReadBits(int numBits)
+UInt32 CCoder::ReadBits(unsigned numBits)
 {
   return m_InBitStream.ReadBits(numBits);
 }
 
-bool CCoder::DeCodeLevelTable(Byte *values, int numSymbols)
+Byte CCoder::ReadAlignedByte()
+{
+  return m_InBitStream.ReadAlignedByte();
+}
+
+bool CCoder::DeCodeLevelTable(Byte *values, unsigned numSymbols)
 {
-  int i = 0;
+  unsigned i = 0;
   do
   {
     UInt32 number = m_LevelDecoder.DecodeSymbol(&m_InBitStream);
@@ -37,25 +40,25 @@ bool CCoder::DeCodeLevelTable(Byte *values, int numSymbols)
       {
         if (i == 0)
           return false;
-        int num = ReadBits(2) + 3;
+        unsigned num = ReadBits(2) + 3;
         for (; num > 0 && i < numSymbols; num--, i++)
           values[i] = values[i - 1];
       }
       else
       {
-        int num;
+        unsigned num;
         if (number == kTableLevel0Number)
           num = ReadBits(3) + 3;
         else
           num = ReadBits(7) + 11;
-        for (;num > 0 && i < numSymbols; num--)
+        for (; num > 0 && i < numSymbols; num--)
           values[i++] = 0;
       }
     }
     else
       return false;
   }
-  while(i < numSymbols);
+  while (i < numSymbols);
   return true;
 }
 
@@ -64,18 +67,22 @@ bool CCoder::DeCodeLevelTable(Byte *values, int numSymbols)
 bool CCoder::ReadTables(void)
 {
   m_FinalBlock = (ReadBits(kFinalBlockFieldSize) == NFinalBlockField::kFinalBlock);
+  if (m_InBitStream.ExtraBitsWereRead())
+    return false;
   UInt32 blockType = ReadBits(kBlockTypeFieldSize);
   if (blockType > NBlockType::kDynamicHuffman)
     return false;
+  if (m_InBitStream.ExtraBitsWereRead())
+    return false;
 
   if (blockType == NBlockType::kStored)
   {
     m_StoredMode = true;
     m_InBitStream.AlignToByte();
-    m_StoredBlockSize = ReadBits(kStoredBlockLengthFieldSize);
+    m_StoredBlockSize = ReadAligned_UInt16(); // ReadBits(kStoredBlockLengthFieldSize)
     if (_deflateNSIS)
       return true;
-    return (m_StoredBlockSize == (UInt16)~ReadBits(kStoredBlockLengthFieldSize));
+    return (m_StoredBlockSize == (UInt16)~ReadAligned_UInt16());
   }
 
   m_StoredMode = false;
@@ -88,29 +95,35 @@ bool CCoder::ReadTables(void)
   }
   else
   {
-    int numLitLenLevels = ReadBits(kNumLenCodesFieldSize) + kNumLitLenCodesMin;
+    unsigned numLitLenLevels = ReadBits(kNumLenCodesFieldSize) + kNumLitLenCodesMin;
     _numDistLevels = ReadBits(kNumDistCodesFieldSize) + kNumDistCodesMin;
-    int numLevelCodes = ReadBits(kNumLevelCodesFieldSize) + kNumLevelCodesMin;
+    unsigned numLevelCodes = ReadBits(kNumLevelCodesFieldSize) + kNumLevelCodesMin;
 
     if (!_deflate64Mode)
       if (_numDistLevels > kDistTableSize32)
         return false;
     
     Byte levelLevels[kLevelTableSize];
-    for (int i = 0; i < kLevelTableSize; i++)
+    for (unsigned i = 0; i < kLevelTableSize; i++)
     {
       int position = kCodeLengthAlphabetOrder[i];
-      if(i < numLevelCodes)
+      if (i < numLevelCodes)
         levelLevels[position] = (Byte)ReadBits(kLevelFieldSize);
       else
         levelLevels[position] = 0;
     }
     
+    if (m_InBitStream.ExtraBitsWereRead())
+      return false;
+
     RIF(m_LevelDecoder.SetCodeLengths(levelLevels));
     
     Byte tmpLevels[kFixedMainTableSize + kFixedDistTableSize];
     if (!DeCodeLevelTable(tmpLevels, numLitLenLevels + _numDistLevels))
       return false;
+    
+    if (m_InBitStream.ExtraBitsWereRead())
+      return false;
 
     levels.SubClear();
     memcpy(levels.litLenLevels, tmpLevels, numLitLenLevels);
@@ -120,7 +133,7 @@ bool CCoder::ReadTables(void)
   return m_DistDecoder.SetCodeLengths(levels.distLevels);
 }
 
-HRESULT CCoder::CodeSpec(UInt32 curSize)
+HRESULT CCoder::CodeSpec(UInt32 curSize, bool finishInputStream)
 {
   if (_remainLen == kLenIdFinished)
     return S_OK;
@@ -136,10 +149,7 @@ HRESULT CCoder::CodeSpec(UInt32 curSize)
     _needReadTable = true;
   }
 
-  if (curSize == 0)
-    return S_OK;
-
-  while(_remainLen > 0 && curSize > 0)
+  while (_remainLen > 0 && curSize > 0)
   {
     _remainLen--;
     Byte b = m_OutWindowStream.GetByte(_rep0);
@@ -147,8 +157,10 @@ HRESULT CCoder::CodeSpec(UInt32 curSize)
     curSize--;
   }
 
-  while(curSize > 0)
+  while (curSize > 0 || finishInputStream)
   {
+    if (m_InBitStream.ExtraBitsWereRead())
+      return S_FALSE;
     if (_needReadTable)
     {
       if (m_FinalBlock)
@@ -158,19 +170,28 @@ HRESULT CCoder::CodeSpec(UInt32 curSize)
       }
       if (!ReadTables())
         return S_FALSE;
+      if (m_InBitStream.ExtraBitsWereRead())
+        return S_FALSE;
       _needReadTable = false;
     }
 
-    if(m_StoredMode)
+    if (m_StoredMode)
     {
+      if (finishInputStream && curSize == 0 && m_StoredBlockSize != 0)
+        return S_FALSE;
+      /* NSIS version contains some bits in bitl bits buffer.
+         So we must read some first bytes via ReadAlignedByte */
+      for (; m_StoredBlockSize > 0 && curSize > 0 && m_InBitStream.ThereAreDataInBitsBuffer(); m_StoredBlockSize--, curSize--)
+        m_OutWindowStream.PutByte(ReadAlignedByte());
       for (; m_StoredBlockSize > 0 && curSize > 0; m_StoredBlockSize--, curSize--)
-        m_OutWindowStream.PutByte(m_InBitStream.ReadByte());
+        m_OutWindowStream.PutByte(m_InBitStream.ReadDirectByte());
       _needReadTable = (m_StoredBlockSize == 0);
       continue;
     }
-    while(curSize > 0)
+    
+    while (curSize > 0)
     {
-      if (m_InBitStream.NumExtraBytes > 4)
+      if (m_InBitStream.ExtraBitsWereRead_Fast())
         return S_FALSE;
 
       UInt32 number = m_MainDecoder.DecodeSymbol(&m_InBitStream);
@@ -190,7 +211,7 @@ HRESULT CCoder::CodeSpec(UInt32 curSize)
         number -= kSymbolMatch;
         UInt32 len;
         {
-          int numBits;
+          unsigned numBits;
           if (_deflate64Mode)
           {
             len = kLenStart64[number];
@@ -224,28 +245,40 @@ HRESULT CCoder::CodeSpec(UInt32 curSize)
       else
         return S_FALSE;
     }
+    
+    if (finishInputStream && curSize == 0)
+    {
+      if (m_MainDecoder.DecodeSymbol(&m_InBitStream) != kSymbolEndOfBlock)
+        return S_FALSE;
+      _needReadTable = true;
+    }
   }
+
+  if (m_InBitStream.ExtraBitsWereRead())
+    return S_FALSE;
+
   return S_OK;
 }
 
 #ifdef _NO_EXCEPTIONS
 
 #define DEFLATE_TRY_BEGIN
-#define DEFLATE_TRY_END
+#define DEFLATE_TRY_END(res)
 
 #else
 
 #define DEFLATE_TRY_BEGIN try {
-#define DEFLATE_TRY_END } \
-  catch(const CInBufferException &e)  { return e.ErrorCode; } \
-  catch(const CLzOutWindowException &e)  { return e.ErrorCode; } \
-  catch(...) { return S_FALSE; }
+#define DEFLATE_TRY_END(res) } \
+  catch(const CInBufferException &e)  { res = e.ErrorCode; } \
+  catch(const CLzOutWindowException &e)  { res = e.ErrorCode; } \
+  catch(...) { res = S_FALSE; }
 
 #endif
 
 HRESULT CCoder::CodeReal(ISequentialOutStream *outStream,
       const UInt64 *outSize, ICompressProgressInfo *progress)
 {
+  HRESULT res;
   DEFLATE_TRY_BEGIN
   m_OutWindowStream.SetStream(outStream);
   CCoderReleaser flusher(this);
@@ -255,18 +288,23 @@ HRESULT CCoder::CodeReal(ISequentialOutStream *outStream,
   for (;;)
   {
     UInt32 curSize = 1 << 18;
-    if (outSize != 0)
+    bool finishInputStream = false;
+    if (outSize)
     {
       const UInt64 rem = *outSize - (m_OutWindowStream.GetProcessedSize() - start);
-      if (curSize > rem)
+      if (curSize >= rem)
+      {
         curSize = (UInt32)rem;
+        if (ZlibMode || _needFinishInput)
+          finishInputStream = true;
+      }
     }
     if (curSize == 0)
       break;
-    RINOK(CodeSpec(curSize));
+    RINOK(CodeSpec(curSize, finishInputStream));
     if (_remainLen == kLenIdFinished)
       break;
-    if (progress != NULL)
+    if (progress)
     {
       const UInt64 inSize = m_InBitStream.GetProcessedSize() - inStart;
       const UInt64 nowPos64 = m_OutWindowStream.GetProcessedSize() - start;
@@ -277,14 +315,14 @@ HRESULT CCoder::CodeReal(ISequentialOutStream *outStream,
   {
     m_InBitStream.AlignToByte();
     for (int i = 0; i < 4; i++)
-      ZlibFooter[i] = m_InBitStream.ReadByte();
+      ZlibFooter[i] = ReadAlignedByte();
   }
   flusher.NeedFlush = false;
-  HRESULT res = Flush();
+  res = Flush();
   if (res == S_OK && InputEofError())
     return S_FALSE;
+  DEFLATE_TRY_END(res)
   return res;
-  DEFLATE_TRY_END
 }
 
 HRESULT CCoder::Code(ISequentialInStream *inStream, ISequentialOutStream *outStream,
@@ -307,13 +345,14 @@ STDMETHODIMP CCoder::GetInStreamProcessedSize(UInt64 *value)
 
 STDMETHODIMP CCoder::SetInStream(ISequentialInStream *inStream)
 {
+  m_InStreamRef = inStream;
   m_InBitStream.SetStream(inStream);
   return S_OK;
 }
 
 STDMETHODIMP CCoder::ReleaseInStream()
 {
-  m_InBitStream.ReleaseStream();
+  m_InStreamRef.Release();
   return S_OK;
 }
 
@@ -329,16 +368,22 @@ STDMETHODIMP CCoder::SetOutStreamSize(const UInt64 * /* outSize */)
 
 STDMETHODIMP CCoder::Read(void *data, UInt32 size, UInt32 *processedSize)
 {
+  HRESULT res;
   DEFLATE_TRY_BEGIN
   if (processedSize)
     *processedSize = 0;
   const UInt64 startPos = m_OutWindowStream.GetProcessedSize();
   m_OutWindowStream.SetMemStream((Byte *)data);
-  RINOK(CodeSpec(size));
-  if (processedSize)
-    *processedSize = (UInt32)(m_OutWindowStream.GetProcessedSize() - startPos);
-  return Flush();
-  DEFLATE_TRY_END
+  res = CodeSpec(size, false);
+  if (res == S_OK)
+  {
+    res = Flush();
+    if (processedSize)
+      *processedSize = (UInt32)(m_OutWindowStream.GetProcessedSize() - startPos);
+  }
+  DEFLATE_TRY_END(res)
+  m_OutWindowStream.SetMemStream(NULL);
+  return res;
 }
 
 #endif