Home | History | Annotate | Download | only in Archive
      1 // LzmaHandler.cpp
      2 
      3 #include "StdAfx.h"
      4 
      5 #include "../../../C/CpuArch.h"
      6 
      7 #include "../../Common/ComTry.h"
      8 #include "../../Common/IntToString.h"
      9 
     10 #include "../../Windows/PropVariant.h"
     11 
     12 #include "../Common/CreateCoder.h"
     13 #include "../Common/ProgressUtils.h"
     14 #include "../Common/RegisterArc.h"
     15 #include "../Common/StreamUtils.h"
     16 
     17 #include "../Compress/LzmaDecoder.h"
     18 
     19 #include "Common/DummyOutStream.h"
     20 
     21 using namespace NWindows;
     22 
     23 namespace NArchive {
     24 namespace NLzma {
     25 
     26 static bool CheckDicSize(const Byte *p)
     27 {
     28   UInt32 dicSize = GetUi32(p);
     29   if (dicSize == 1)
     30     return true;
     31   for (unsigned i = 0; i <= 30; i++)
     32     if (dicSize == ((UInt32)2 << i) || dicSize == ((UInt32)3 << i))
     33       return true;
     34   return (dicSize == 0xFFFFFFFF);
     35 }
     36 
     37 static const Byte kProps[] =
     38 {
     39   kpidSize,
     40   kpidPackSize,
     41   kpidMethod
     42 };
     43 
     44 static const Byte kArcProps[] =
     45 {
     46   kpidNumStreams
     47 };
     48 
     49 struct CHeader
     50 {
     51   UInt64 Size;
     52   Byte FilterID;
     53   Byte LzmaProps[5];
     54 
     55   UInt32 GetDicSize() const { return GetUi32(LzmaProps + 1); }
     56   bool HasSize() const { return (Size != (UInt64)(Int64)-1); }
     57   bool Parse(const Byte *buf, bool isThereFilter);
     58 };
     59 
     60 bool CHeader::Parse(const Byte *buf, bool isThereFilter)
     61 {
     62   FilterID = 0;
     63   if (isThereFilter)
     64     FilterID = buf[0];
     65   const Byte *sig = buf + (isThereFilter ? 1 : 0);
     66   for (int i = 0; i < 5; i++)
     67     LzmaProps[i] = sig[i];
     68   Size = GetUi64(sig + 5);
     69   return
     70     LzmaProps[0] < 5 * 5 * 9 &&
     71     FilterID < 2 &&
     72     (!HasSize() || Size < ((UInt64)1 << 56))
     73     && CheckDicSize(LzmaProps + 1);
     74 }
     75 
     76 class CDecoder
     77 {
     78   CMyComPtr<ICompressCoder> _lzmaDecoder;
     79   CMyComPtr<ISequentialOutStream> _bcjStream;
     80 public:
     81   NCompress::NLzma::CDecoder *_lzmaDecoderSpec;
     82 
     83   ~CDecoder();
     84   HRESULT Create(DECL_EXTERNAL_CODECS_LOC_VARS
     85       bool filtered, ISequentialInStream *inStream);
     86 
     87   HRESULT Code(const CHeader &header, ISequentialOutStream *outStream, ICompressProgressInfo *progress);
     88 
     89   UInt64 GetInputProcessedSize() const { return _lzmaDecoderSpec->GetInputProcessedSize(); }
     90 
     91   void ReleaseInStream() { if (_lzmaDecoder) _lzmaDecoderSpec->ReleaseInStream(); }
     92 
     93   HRESULT ReadInput(Byte *data, UInt32 size, UInt32 *processedSize)
     94     { return _lzmaDecoderSpec->ReadFromInputStream(data, size, processedSize); }
     95 };
     96 
     97 static const UInt32 k_BCJ = 0x03030103;
     98 
     99 HRESULT CDecoder::Create(
    100     DECL_EXTERNAL_CODECS_LOC_VARS
    101     bool filteredMode, ISequentialInStream *inStream)
    102 {
    103   if (!_lzmaDecoder)
    104   {
    105     _lzmaDecoderSpec = new NCompress::NLzma::CDecoder;
    106     _lzmaDecoderSpec->FinishStream = true;
    107     _lzmaDecoder = _lzmaDecoderSpec;
    108   }
    109 
    110   if (filteredMode)
    111   {
    112     if (!_bcjStream)
    113     {
    114       CMyComPtr<ICompressCoder> coder;
    115       RINOK(CreateCoder(EXTERNAL_CODECS_LOC_VARS k_BCJ, coder, false));
    116       if (!coder)
    117         return E_NOTIMPL;
    118       coder.QueryInterface(IID_ISequentialOutStream, &_bcjStream);
    119       if (!_bcjStream)
    120         return E_NOTIMPL;
    121     }
    122   }
    123 
    124   return _lzmaDecoderSpec->SetInStream(inStream);
    125 }
    126 
    127 CDecoder::~CDecoder()
    128 {
    129   ReleaseInStream();
    130 }
    131 
    132 HRESULT CDecoder::Code(const CHeader &header, ISequentialOutStream *outStream,
    133     ICompressProgressInfo *progress)
    134 {
    135   if (header.FilterID > 1)
    136     return E_NOTIMPL;
    137 
    138   {
    139     CMyComPtr<ICompressSetDecoderProperties2> setDecoderProperties;
    140     _lzmaDecoder.QueryInterface(IID_ICompressSetDecoderProperties2, &setDecoderProperties);
    141     if (!setDecoderProperties)
    142       return E_NOTIMPL;
    143     RINOK(setDecoderProperties->SetDecoderProperties2(header.LzmaProps, 5));
    144   }
    145 
    146   CMyComPtr<ICompressSetOutStream> setOutStream;
    147 
    148   bool filteredMode = (header.FilterID == 1);
    149 
    150   if (filteredMode)
    151   {
    152     _bcjStream.QueryInterface(IID_ICompressSetOutStream, &setOutStream);
    153     if (!setOutStream)
    154       return E_NOTIMPL;
    155     RINOK(setOutStream->SetOutStream(outStream));
    156     outStream = _bcjStream;
    157   }
    158 
    159   const UInt64 *Size = header.HasSize() ? &header.Size : NULL;
    160   HRESULT res = _lzmaDecoderSpec->CodeResume(outStream, Size, progress);
    161 
    162   if (filteredMode)
    163   {
    164     CMyComPtr<IOutStreamFlush> flush;
    165     _bcjStream.QueryInterface(IID_IOutStreamFlush, &flush);
    166     if (flush)
    167     {
    168       HRESULT res2 = flush->Flush();
    169       if (res == S_OK)
    170         res = res2;
    171     }
    172     HRESULT res2 = setOutStream->ReleaseOutStream();
    173     if (res == S_OK)
    174       res = res2;
    175   }
    176   RINOK(res);
    177 
    178   if (header.HasSize())
    179     if (_lzmaDecoderSpec->GetOutputProcessedSize() != header.Size)
    180       return S_FALSE;
    181 
    182   return S_OK;
    183 }
    184 
    185 
    186 class CHandler:
    187   public IInArchive,
    188   public IArchiveOpenSeq,
    189   PUBLIC_ISetCompressCodecsInfo
    190   public CMyUnknownImp
    191 {
    192   CHeader _header;
    193   bool _lzma86;
    194   CMyComPtr<IInStream> _stream;
    195   CMyComPtr<ISequentialInStream> _seqStream;
    196 
    197   bool _isArc;
    198   bool _needSeekToStart;
    199   bool _dataAfterEnd;
    200   bool _needMoreInput;
    201 
    202   bool _packSize_Defined;
    203   bool _unpackSize_Defined;
    204   bool _numStreams_Defined;
    205 
    206   bool _unsupported;
    207   bool _dataError;
    208 
    209   UInt64 _packSize;
    210   UInt64 _unpackSize;
    211   UInt64 _numStreams;
    212 
    213   DECL_EXTERNAL_CODECS_VARS
    214   DECL_ISetCompressCodecsInfo
    215 
    216 public:
    217   MY_QUERYINTERFACE_BEGIN2(IInArchive)
    218   MY_QUERYINTERFACE_ENTRY(IArchiveOpenSeq)
    219   QUERY_ENTRY_ISetCompressCodecsInfo
    220   MY_QUERYINTERFACE_END
    221   MY_ADDREF_RELEASE
    222 
    223   INTERFACE_IInArchive(;)
    224   STDMETHOD(OpenSeq)(ISequentialInStream *stream);
    225 
    226   CHandler(bool lzma86) { _lzma86 = lzma86; }
    227 
    228   unsigned GetHeaderSize() const { return 5 + 8 + (_lzma86 ? 1 : 0); }
    229 
    230 };
    231 
    232 IMP_IInArchive_Props
    233 IMP_IInArchive_ArcProps
    234 
    235 STDMETHODIMP CHandler::GetArchiveProperty(PROPID propID, PROPVARIANT *value)
    236 {
    237   NCOM::CPropVariant prop;
    238   switch (propID)
    239   {
    240     case kpidPhySize: if (_packSize_Defined) prop = _packSize; break;
    241     case kpidNumStreams: if (_numStreams_Defined) prop = _numStreams; break;
    242     case kpidUnpackSize: if (_unpackSize_Defined) prop = _unpackSize; break;
    243     case kpidErrorFlags:
    244     {
    245       UInt32 v = 0;
    246       if (!_isArc) v |= kpv_ErrorFlags_IsNotArc;;
    247       if (_needMoreInput) v |= kpv_ErrorFlags_UnexpectedEnd;
    248       if (_dataAfterEnd) v |= kpv_ErrorFlags_DataAfterEnd;
    249       if (_unsupported) v |= kpv_ErrorFlags_UnsupportedMethod;
    250       if (_dataError) v |= kpv_ErrorFlags_DataError;
    251       prop = v;
    252     }
    253   }
    254   prop.Detach(value);
    255   return S_OK;
    256 }
    257 
    258 STDMETHODIMP CHandler::GetNumberOfItems(UInt32 *numItems)
    259 {
    260   *numItems = 1;
    261   return S_OK;
    262 }
    263 
    264 static void DictSizeToString(UInt32 value, char *s)
    265 {
    266   for (int i = 0; i <= 31; i++)
    267     if (((UInt32)1 << i) == value)
    268     {
    269       ::ConvertUInt32ToString(i, s);
    270       return;
    271     }
    272   char c = 'b';
    273        if ((value & ((1 << 20) - 1)) == 0) { value >>= 20; c = 'm'; }
    274   else if ((value & ((1 << 10) - 1)) == 0) { value >>= 10; c = 'k'; }
    275   ::ConvertUInt32ToString(value, s);
    276   s += MyStringLen(s);
    277   *s++ = c;
    278   *s = 0;
    279 }
    280 
    281 STDMETHODIMP CHandler::GetProperty(UInt32 /* index */, PROPID propID, PROPVARIANT *value)
    282 {
    283   NCOM::CPropVariant prop;
    284   switch (propID)
    285   {
    286     case kpidSize: if (_stream && _header.HasSize()) prop = _header.Size; break;
    287     case kpidPackSize: if (_packSize_Defined) prop = _packSize; break;
    288     case kpidMethod:
    289       if (_stream)
    290       {
    291         char sz[64];
    292         char *s = sz;
    293         if (_header.FilterID != 0)
    294           s = MyStpCpy(s, "BCJ ");
    295         s = MyStpCpy(s, "LZMA:");
    296         DictSizeToString(_header.GetDicSize(), s);
    297         prop = sz;
    298       }
    299       break;
    300   }
    301   prop.Detach(value);
    302   return S_OK;
    303 }
    304 
    305 API_FUNC_static_IsArc IsArc_Lzma(const Byte *p, size_t size)
    306 {
    307   const UInt32 kHeaderSize = 1 + 4 + 8;
    308   if (size < kHeaderSize)
    309     return k_IsArc_Res_NEED_MORE;
    310   if (p[0] >= 5 * 5 * 9)
    311     return k_IsArc_Res_NO;
    312   UInt64 unpackSize = GetUi64(p + 1 + 4);
    313   if (unpackSize != (UInt64)(Int64)-1)
    314   {
    315     if (size >= ((UInt64)1 << 56))
    316       return k_IsArc_Res_NO;
    317   }
    318   if (unpackSize != 0)
    319   {
    320     if (size < kHeaderSize + 2)
    321       return k_IsArc_Res_NEED_MORE;
    322     if (p[kHeaderSize] != 0)
    323       return k_IsArc_Res_NO;
    324     if (unpackSize != (UInt64)(Int64)-1)
    325     {
    326       if ((p[kHeaderSize + 1] & 0x80) != 0)
    327         return k_IsArc_Res_NO;
    328     }
    329   }
    330   if (!CheckDicSize(p + 1))
    331     // return k_IsArc_Res_YES_LOW_PROB;
    332     return k_IsArc_Res_NO;
    333   return k_IsArc_Res_YES;
    334 }
    335 }
    336 
    337 API_FUNC_static_IsArc IsArc_Lzma86(const Byte *p, size_t size)
    338 {
    339   if (size < 1)
    340     return k_IsArc_Res_NEED_MORE;
    341   Byte filterID = p[0];
    342   if (filterID != 0 && filterID != 1)
    343     return k_IsArc_Res_NO;
    344   return IsArc_Lzma(p + 1, size - 1);
    345 }
    346 }
    347 
    348 STDMETHODIMP CHandler::Open(IInStream *inStream, const UInt64 *, IArchiveOpenCallback *)
    349 {
    350   Close();
    351 
    352   const UInt32 kBufSize = 1 + 5 + 8 + 2;
    353   Byte buf[kBufSize];
    354 
    355   RINOK(ReadStream_FALSE(inStream, buf, kBufSize));
    356 
    357   if (!_header.Parse(buf, _lzma86))
    358     return S_FALSE;
    359   const Byte *start = buf + GetHeaderSize();
    360   if (start[0] != 0 /* || (start[1] & 0x80) != 0 */ ) // empty stream with EOS is not 0x80
    361     return S_FALSE;
    362 
    363   RINOK(inStream->Seek(0, STREAM_SEEK_END, &_packSize));
    364   if (_packSize >= 24 && _header.Size == 0 && _header.FilterID == 0 && _header.LzmaProps[0] == 0)
    365     return S_FALSE;
    366   _isArc = true;
    367   _stream = inStream;
    368   _seqStream = inStream;
    369   _needSeekToStart = true;
    370   return S_OK;
    371 }
    372 
    373 STDMETHODIMP CHandler::OpenSeq(ISequentialInStream *stream)
    374 {
    375   Close();
    376   _isArc = true;
    377   _seqStream = stream;
    378   return S_OK;
    379 }
    380 
    381 STDMETHODIMP CHandler::Close()
    382 {
    383   _isArc = false;
    384   _packSize_Defined = false;
    385   _unpackSize_Defined = false;
    386   _numStreams_Defined = false;
    387 
    388   _dataAfterEnd = false;
    389   _needMoreInput = false;
    390   _unsupported = false;
    391   _dataError = false;
    392 
    393   _packSize = 0;
    394 
    395   _needSeekToStart = false;
    396 
    397   _stream.Release();
    398   _seqStream.Release();
    399    return S_OK;
    400 }
    401 
    402 class CCompressProgressInfoImp:
    403   public ICompressProgressInfo,
    404   public CMyUnknownImp
    405 {
    406   CMyComPtr<IArchiveOpenCallback> Callback;
    407 public:
    408   UInt64 Offset;
    409 
    410   MY_UNKNOWN_IMP1(ICompressProgressInfo)
    411   STDMETHOD(SetRatioInfo)(const UInt64 *inSize, const UInt64 *outSize);
    412   void Init(IArchiveOpenCallback *callback) { Callback = callback; }
    413 };
    414 
    415 STDMETHODIMP CCompressProgressInfoImp::SetRatioInfo(const UInt64 *inSize, const UInt64 * /* outSize */)
    416 {
    417   if (Callback)
    418   {
    419     UInt64 files = 0;
    420     UInt64 value = Offset + *inSize;
    421     return Callback->SetCompleted(&files, &value);
    422   }
    423   return S_OK;
    424 }
    425 
    426 STDMETHODIMP CHandler::Extract(const UInt32 *indices, UInt32 numItems,
    427     Int32 testMode, IArchiveExtractCallback *extractCallback)
    428 {
    429   COM_TRY_BEGIN
    430   if (numItems == 0)
    431     return S_OK;
    432   if (numItems != (UInt32)(Int32)-1 && (numItems != 1 || indices[0] != 0))
    433     return E_INVALIDARG;
    434 
    435   if (_packSize_Defined)
    436     extractCallback->SetTotal(_packSize);
    437 
    438 
    439   CMyComPtr<ISequentialOutStream> realOutStream;
    440   Int32 askMode = testMode ?
    441       NExtract::NAskMode::kTest :
    442       NExtract::NAskMode::kExtract;
    443   RINOK(extractCallback->GetStream(0, &realOutStream, askMode));
    444   if (!testMode && !realOutStream)
    445     return S_OK;
    446 
    447   extractCallback->PrepareOperation(askMode);
    448 
    449   CDummyOutStream *outStreamSpec = new CDummyOutStream;
    450   CMyComPtr<ISequentialOutStream> outStream(outStreamSpec);
    451   outStreamSpec->SetStream(realOutStream);
    452   outStreamSpec->Init();
    453   realOutStream.Release();
    454 
    455   CLocalProgress *lps = new CLocalProgress;
    456   CMyComPtr<ICompressProgressInfo> progress = lps;
    457   lps->Init(extractCallback, true);
    458 
    459   if (_needSeekToStart)
    460   {
    461     if (!_stream)
    462       return E_FAIL;
    463     RINOK(_stream->Seek(0, STREAM_SEEK_SET, NULL));
    464   }
    465   else
    466     _needSeekToStart = true;
    467 
    468   CDecoder decoder;
    469   HRESULT result = decoder.Create(
    470       EXTERNAL_CODECS_VARS
    471       _lzma86, _seqStream);
    472   RINOK(result);
    473 
    474   bool firstItem = true;
    475 
    476   UInt64 packSize = 0;
    477   UInt64 unpackSize = 0;
    478   UInt64 numStreams = 0;
    479 
    480   bool dataAfterEnd = false;
    481 
    482   for (;;)
    483   {
    484     lps->InSize = packSize;
    485     lps->OutSize = unpackSize;
    486     RINOK(lps->SetCur());
    487 
    488     const UInt32 kBufSize = 1 + 5 + 8;
    489     Byte buf[kBufSize];
    490     const UInt32 headerSize = GetHeaderSize();
    491     UInt32 processed;
    492     RINOK(decoder.ReadInput(buf, headerSize, &processed));
    493     if (processed != headerSize)
    494     {
    495       if (processed != 0)
    496         dataAfterEnd = true;
    497       break;
    498     }
    499 
    500     CHeader st;
    501     if (!st.Parse(buf, _lzma86))
    502     {
    503       dataAfterEnd = true;
    504       break;
    505     }
    506     numStreams++;
    507     firstItem = false;
    508 
    509     result = decoder.Code(st, outStream, progress);
    510 
    511     packSize = decoder.GetInputProcessedSize();
    512     unpackSize = outStreamSpec->GetSize();
    513 
    514     if (result == E_NOTIMPL)
    515     {
    516       _unsupported = true;
    517       result = S_FALSE;
    518       break;
    519     }
    520     if (result == S_FALSE)
    521       break;
    522     RINOK(result);
    523   }
    524 
    525   if (firstItem)
    526   {
    527     _isArc = false;
    528     result = S_FALSE;
    529   }
    530   else if (result == S_OK || result == S_FALSE)
    531   {
    532     if (dataAfterEnd)
    533       _dataAfterEnd = true;
    534     else if (decoder._lzmaDecoderSpec->NeedMoreInput)
    535       _needMoreInput = true;
    536 
    537     _packSize = packSize;
    538     _unpackSize = unpackSize;
    539     _numStreams = numStreams;
    540 
    541     _packSize_Defined = true;
    542     _unpackSize_Defined = true;
    543     _numStreams_Defined = true;
    544   }
    545 
    546   Int32 opResult = NExtract::NOperationResult::kOK;
    547 
    548   if (!_isArc)
    549     opResult = NExtract::NOperationResult::kIsNotArc;
    550   else if (_needMoreInput)
    551     opResult = NExtract::NOperationResult::kUnexpectedEnd;
    552   else if (_unsupported)
    553     opResult = NExtract::NOperationResult::kUnsupportedMethod;
    554   else if (_dataAfterEnd)
    555     opResult = NExtract::NOperationResult::kDataAfterEnd;
    556   else if (result == S_FALSE)
    557     opResult = NExtract::NOperationResult::kDataError;
    558   else if (result == S_OK)
    559     opResult = NExtract::NOperationResult::kOK;
    560   else
    561     return result;
    562 
    563   outStream.Release();
    564   return extractCallback->SetOperationResult(opResult);
    565   COM_TRY_END
    566 }
    567 
    568 IMPL_ISetCompressCodecsInfo
    569 
    570 namespace NLzmaAr {
    571 
    572 IMP_CreateArcIn_2(CHandler(false))
    573 
    574 static CArcInfo g_ArcInfo =
    575   { "lzma", "lzma", 0, 0xA,
    576   0, { 0 },
    577   // 2, { 0x5D, 0x00 },
    578   0,
    579   NArcInfoFlags::kStartOpen |
    580   NArcInfoFlags::kKeepName,
    581   CreateArc, NULL,
    582   IsArc_Lzma };
    583 
    584 REGISTER_ARC(Lzma)
    585 
    586 }
    587 
    588 namespace NLzma86Ar {
    589 
    590 IMP_CreateArcIn_2(CHandler(true))
    591 
    592 static CArcInfo g_ArcInfo =
    593   { "lzma86", "lzma86", 0, 0xB,
    594   0, { 0 },
    595   0,
    596   NArcInfoFlags::kKeepName,
    597   CreateArc, NULL,
    598   IsArc_Lzma86 };
    599 
    600 REGISTER_ARC(Lzma86)
    601 
    602 }
    603 
    604 }}
    605