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