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 
CheckDicSize(const Byte * p)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 
GetDicSizeNArchive::NLzma::CHeader55   UInt32 GetDicSize() const { return GetUi32(LzmaProps + 1); }
HasSizeNArchive::NLzma::CHeader56   bool HasSize() const { return (Size != (UInt64)(Int64)-1); }
57   bool Parse(const Byte *buf, bool isThereFilter);
58 };
59 
Parse(const Byte * buf,bool isThereFilter)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 
GetInputProcessedSize() const89   UInt64 GetInputProcessedSize() const { return _lzmaDecoderSpec->GetInputProcessedSize(); }
90 
ReleaseInStream()91   void ReleaseInStream() { if (_lzmaDecoder) _lzmaDecoderSpec->ReleaseInStream(); }
92 
ReadInput(Byte * data,UInt32 size,UInt32 * processedSize)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 
Create(DECL_EXTERNAL_CODECS_LOC_VARS bool filteredMode,ISequentialInStream * inStream)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 
~CDecoder()127 CDecoder::~CDecoder()
128 {
129   ReleaseInStream();
130 }
131 
Code(const CHeader & header,ISequentialOutStream * outStream,ICompressProgressInfo * progress)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 
CHandler(bool lzma86)226   CHandler(bool lzma86) { _lzma86 = lzma86; }
227 
GetHeaderSize() const228   unsigned GetHeaderSize() const { return 5 + 8 + (_lzma86 ? 1 : 0); }
229 
230 };
231 
232 IMP_IInArchive_Props
233 IMP_IInArchive_ArcProps
234 
GetArchiveProperty(PROPID propID,PROPVARIANT * value)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 
GetNumberOfItems(UInt32 * numItems)258 STDMETHODIMP CHandler::GetNumberOfItems(UInt32 *numItems)
259 {
260   *numItems = 1;
261   return S_OK;
262 }
263 
DictSizeToString(UInt32 value,char * s)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 
GetProperty(UInt32,PROPID propID,PROPVARIANT * value)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 
IsArc_Lzma(const Byte * p,size_t size)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 
IsArc_Lzma86(const Byte * p,size_t size)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 
Open(IInStream * inStream,const UInt64 *,IArchiveOpenCallback *)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 
OpenSeq(ISequentialInStream * stream)373 STDMETHODIMP CHandler::OpenSeq(ISequentialInStream *stream)
374 {
375   Close();
376   _isArc = true;
377   _seqStream = stream;
378   return S_OK;
379 }
380 
Close()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);
Init(IArchiveOpenCallback * callback)412   void Init(IArchiveOpenCallback *callback) { Callback = callback; }
413 };
414 
SetRatioInfo(const UInt64 * inSize,const UInt64 *)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 
Extract(const UInt32 * indices,UInt32 numItems,Int32 testMode,IArchiveExtractCallback * extractCallback)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