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
21using namespace NWindows;
22
23namespace NArchive {
24namespace NLzma {
25
26static 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
37static const Byte kProps[] =
38{
39  kpidSize,
40  kpidPackSize,
41  kpidMethod
42};
43
44static const Byte kArcProps[] =
45{
46  kpidNumStreams
47};
48
49struct 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
60bool 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
76class CDecoder
77{
78  CMyComPtr<ICompressCoder> _lzmaDecoder;
79  CMyComPtr<ISequentialOutStream> _bcjStream;
80public:
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
97static const UInt32 k_BCJ = 0x03030103;
98
99HRESULT 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
127CDecoder::~CDecoder()
128{
129  ReleaseInStream();
130}
131
132HRESULT 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
186class 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
216public:
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
232IMP_IInArchive_Props
233IMP_IInArchive_ArcProps
234
235STDMETHODIMP 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
258STDMETHODIMP CHandler::GetNumberOfItems(UInt32 *numItems)
259{
260  *numItems = 1;
261  return S_OK;
262}
263
264static 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
281STDMETHODIMP 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
305API_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
337API_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
348STDMETHODIMP 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
373STDMETHODIMP CHandler::OpenSeq(ISequentialInStream *stream)
374{
375  Close();
376  _isArc = true;
377  _seqStream = stream;
378  return S_OK;
379}
380
381STDMETHODIMP 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
402class CCompressProgressInfoImp:
403  public ICompressProgressInfo,
404  public CMyUnknownImp
405{
406  CMyComPtr<IArchiveOpenCallback> Callback;
407public:
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
415STDMETHODIMP 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
426STDMETHODIMP 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
568IMPL_ISetCompressCodecsInfo
569
570namespace NLzmaAr {
571
572IMP_CreateArcIn_2(CHandler(false))
573
574static 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
584REGISTER_ARC(Lzma)
585
586}
587
588namespace NLzma86Ar {
589
590IMP_CreateArcIn_2(CHandler(true))
591
592static CArcInfo g_ArcInfo =
593  { "lzma86", "lzma86", 0, 0xB,
594  0, { 0 },
595  0,
596  NArcInfoFlags::kKeepName,
597  CreateArc, NULL,
598  IsArc_Lzma86 };
599
600REGISTER_ARC(Lzma86)
601
602}
603
604}}
605