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  for (int i = 1; i <= 30; i++)
30    if (dicSize == ((UInt32)2 << i) || dicSize == ((UInt32)3 << i))
31      return true;
32  return (dicSize == 0xFFFFFFFF);
33}
34
35STATPROPSTG kProps[] =
36{
37  { NULL, kpidSize, VT_UI8},
38  { NULL, kpidPackSize, VT_UI8},
39  { NULL, kpidMethod, VT_BSTR}
40};
41
42struct CHeader
43{
44  UInt64 Size;
45  Byte FilterID;
46  Byte LzmaProps[5];
47
48  UInt32 GetDicSize() const { return GetUi32(LzmaProps + 1); }
49  bool HasSize() const { return (Size != (UInt64)(Int64)-1); }
50  bool Parse(const Byte *buf, bool isThereFilter);
51};
52
53bool CHeader::Parse(const Byte *buf, bool isThereFilter)
54{
55  FilterID = 0;
56  if (isThereFilter)
57    FilterID = buf[0];
58  const Byte *sig = buf + (isThereFilter ? 1 : 0);
59  for (int i = 0; i < 5; i++)
60    LzmaProps[i] = sig[i];
61  Size = GetUi64(sig + 5);
62  return
63    LzmaProps[0] < 5 * 5 * 9 &&
64    FilterID < 2 &&
65    (!HasSize() || Size < ((UInt64)1 << 56)) &&
66    CheckDicSize(LzmaProps + 1);
67}
68
69class CDecoder
70{
71  NCompress::NLzma::CDecoder *_lzmaDecoderSpec;
72  CMyComPtr<ICompressCoder> _lzmaDecoder;
73  CMyComPtr<ISequentialOutStream> _bcjStream;
74public:
75  ~CDecoder();
76  HRESULT Create(DECL_EXTERNAL_CODECS_LOC_VARS
77      bool filtered, ISequentialInStream *inStream);
78
79  HRESULT Code(const CHeader &header, ISequentialOutStream *outStream, ICompressProgressInfo *progress);
80
81  UInt64 GetInputProcessedSize() const { return _lzmaDecoderSpec->GetInputProcessedSize(); }
82
83  void ReleaseInStream() { if (_lzmaDecoder) _lzmaDecoderSpec->ReleaseInStream(); }
84
85  HRESULT ReadInput(Byte *data, UInt32 size, UInt32 *processedSize)
86    { return _lzmaDecoderSpec->ReadFromInputStream(data, size, processedSize); }
87};
88
89static const UInt64 k_BCJ = 0x03030103;
90
91HRESULT CDecoder::Create(
92    DECL_EXTERNAL_CODECS_LOC_VARS
93    bool filteredMode, ISequentialInStream *inStream)
94{
95  if (!_lzmaDecoder)
96  {
97    _lzmaDecoderSpec = new NCompress::NLzma::CDecoder;
98    _lzmaDecoder = _lzmaDecoderSpec;
99  }
100
101  if (filteredMode)
102  {
103    if (!_bcjStream)
104    {
105      CMyComPtr<ICompressCoder> coder;
106      RINOK(CreateCoder(EXTERNAL_CODECS_LOC_VARS k_BCJ, coder, false));
107      if (!coder)
108        return E_NOTIMPL;
109      coder.QueryInterface(IID_ISequentialOutStream, &_bcjStream);
110      if (!_bcjStream)
111        return E_NOTIMPL;
112    }
113  }
114
115  return _lzmaDecoderSpec->SetInStream(inStream);
116}
117
118CDecoder::~CDecoder()
119{
120  ReleaseInStream();
121}
122
123HRESULT CDecoder::Code(const CHeader &header, ISequentialOutStream *outStream,
124    ICompressProgressInfo *progress)
125{
126  if (header.FilterID > 1)
127    return E_NOTIMPL;
128
129  {
130    CMyComPtr<ICompressSetDecoderProperties2> setDecoderProperties;
131    _lzmaDecoder.QueryInterface(IID_ICompressSetDecoderProperties2, &setDecoderProperties);
132    if (!setDecoderProperties)
133      return E_NOTIMPL;
134    RINOK(setDecoderProperties->SetDecoderProperties2(header.LzmaProps, 5));
135  }
136
137  CMyComPtr<ICompressSetOutStream> setOutStream;
138
139  bool filteredMode = (header.FilterID == 1);
140
141  if (filteredMode)
142  {
143    _bcjStream.QueryInterface(IID_ICompressSetOutStream, &setOutStream);
144    if (!setOutStream)
145      return E_NOTIMPL;
146    RINOK(setOutStream->SetOutStream(outStream));
147    outStream = _bcjStream;
148  }
149
150  const UInt64 *Size = header.HasSize() ? &header.Size : NULL;
151  HRESULT res = _lzmaDecoderSpec->CodeResume(outStream, Size, progress);
152
153  if (filteredMode)
154  {
155    CMyComPtr<IOutStreamFlush> flush;
156    _bcjStream.QueryInterface(IID_IOutStreamFlush, &flush);
157    if (flush)
158    {
159      HRESULT res2 = flush->Flush();
160      if (res == S_OK)
161        res = res2;
162    }
163    HRESULT res2 = setOutStream->ReleaseOutStream();
164    if (res == S_OK)
165      res = res2;
166  }
167  RINOK(res);
168
169  return S_OK;
170}
171
172
173class CHandler:
174  public IInArchive,
175  public IArchiveOpenSeq,
176  PUBLIC_ISetCompressCodecsInfo
177  public CMyUnknownImp
178{
179  CHeader _header;
180  bool _lzma86;
181  UInt64 _startPosition;
182  UInt64 _packSize;
183  bool _packSizeDefined;
184  CMyComPtr<IInStream> _stream;
185  CMyComPtr<ISequentialInStream> _seqStream;
186
187  DECL_EXTERNAL_CODECS_VARS
188  DECL_ISetCompressCodecsInfo
189
190public:
191  MY_QUERYINTERFACE_BEGIN2(IInArchive)
192  MY_QUERYINTERFACE_ENTRY(IArchiveOpenSeq)
193  QUERY_ENTRY_ISetCompressCodecsInfo
194  MY_QUERYINTERFACE_END
195  MY_ADDREF_RELEASE
196
197  INTERFACE_IInArchive(;)
198  STDMETHOD(OpenSeq)(ISequentialInStream *stream);
199
200  CHandler(bool lzma86) { _lzma86 = lzma86; }
201
202  unsigned GetHeaderSize() const { return 5 + 8 + (_lzma86 ? 1 : 0); }
203
204};
205
206IMP_IInArchive_Props
207IMP_IInArchive_ArcProps_NO_Table
208
209STDMETHODIMP CHandler::GetArchiveProperty(PROPID propID, PROPVARIANT *value)
210{
211  NCOM::CPropVariant prop;
212  switch(propID)
213  {
214    case kpidPhySize: if (_packSizeDefined) prop = _packSize; break;
215  }
216  prop.Detach(value);
217  return S_OK;
218}
219
220STDMETHODIMP CHandler::GetNumberOfItems(UInt32 *numItems)
221{
222  *numItems = 1;
223  return S_OK;
224}
225
226static void DictSizeToString(UInt32 value, char *s)
227{
228  for (int i = 0; i <= 31; i++)
229    if ((UInt32(1) << i) == value)
230    {
231      ::ConvertUInt32ToString(i, s);
232      return;
233    }
234  char c = 'b';
235  if ((value & ((1 << 20) - 1)) == 0)
236  {
237    value >>= 20;
238    c = 'm';
239  }
240  else if ((value & ((1 << 10) - 1)) == 0)
241  {
242    value >>= 10;
243    c = 'k';
244  }
245  ::ConvertUInt32ToString(value, s);
246  int p = MyStringLen(s);
247  s[p++] = c;
248  s[p++] = '\0';
249}
250
251static void MyStrCat(char *d, const char *s)
252{
253  MyStringCopy(d + MyStringLen(d), s);
254}
255
256STDMETHODIMP CHandler::GetProperty(UInt32 /* index */, PROPID propID,  PROPVARIANT *value)
257{
258  NCOM::CPropVariant prop;
259  switch(propID)
260  {
261    case kpidSize: if (_stream && _header.HasSize()) prop = _header.Size; break;
262    case kpidPackSize: if (_packSizeDefined) prop = _packSize; break;
263    case kpidMethod:
264      if (_stream)
265      {
266        char s[64];
267        s[0] = '\0';
268        if (_header.FilterID != 0)
269          MyStrCat(s, "BCJ ");
270        MyStrCat(s, "LZMA:");
271        DictSizeToString(_header.GetDicSize(), s + MyStringLen(s));
272        prop = s;
273      }
274      break;
275  }
276  prop.Detach(value);
277  return S_OK;
278}
279
280STDMETHODIMP CHandler::Open(IInStream *inStream, const UInt64 *, IArchiveOpenCallback *)
281{
282  RINOK(inStream->Seek(0, STREAM_SEEK_CUR, &_startPosition));
283
284  const UInt32 kBufSize = 1 + 5 + 8 + 1;
285  Byte buf[kBufSize];
286
287  RINOK(ReadStream_FALSE(inStream, buf, kBufSize));
288
289  if (!_header.Parse(buf, _lzma86))
290    return S_FALSE;
291  const Byte *start = buf + GetHeaderSize();
292  if (start[0] != 0)
293    return S_FALSE;
294
295  UInt64 endPos;
296  RINOK(inStream->Seek(0, STREAM_SEEK_END, &endPos));
297  _packSize = endPos - _startPosition;
298  _packSizeDefined = true;
299
300  _stream = inStream;
301  _seqStream = inStream;
302  return S_OK;
303}
304
305STDMETHODIMP CHandler::OpenSeq(ISequentialInStream *stream)
306{
307  Close();
308  _seqStream = stream;
309  return S_OK;
310}
311
312STDMETHODIMP CHandler::Close()
313{
314  _packSizeDefined = false;
315  _stream.Release();
316  _seqStream.Release();
317   return S_OK;
318}
319
320
321STDMETHODIMP CHandler::Extract(const UInt32 *indices, UInt32 numItems,
322    Int32 testMode, IArchiveExtractCallback *extractCallback)
323{
324  COM_TRY_BEGIN
325  if (numItems == 0)
326    return S_OK;
327  if (numItems != (UInt32)-1 && (numItems != 1 || indices[0] != 0))
328    return E_INVALIDARG;
329
330  if (_stream)
331    extractCallback->SetTotal(_packSize);
332
333
334  CMyComPtr<ISequentialOutStream> realOutStream;
335  Int32 askMode = testMode ?
336      NExtract::NAskMode::kTest :
337      NExtract::NAskMode::kExtract;
338  RINOK(extractCallback->GetStream(0, &realOutStream, askMode));
339  if (!testMode && !realOutStream)
340    return S_OK;
341
342  extractCallback->PrepareOperation(askMode);
343
344  CDummyOutStream *outStreamSpec = new CDummyOutStream;
345  CMyComPtr<ISequentialOutStream> outStream(outStreamSpec);
346  outStreamSpec->SetStream(realOutStream);
347  outStreamSpec->Init();
348  realOutStream.Release();
349
350  CLocalProgress *lps = new CLocalProgress;
351  CMyComPtr<ICompressProgressInfo> progress = lps;
352  lps->Init(extractCallback, true);
353
354  if (_stream)
355  {
356    RINOK(_stream->Seek(_startPosition, STREAM_SEEK_SET, NULL));
357  }
358
359  CDecoder decoder;
360  HRESULT result = decoder.Create(
361      EXTERNAL_CODECS_VARS
362      _lzma86, _seqStream);
363  RINOK(result);
364
365  Int32 opRes = NExtract::NOperationResult::kOK;
366  bool firstItem = true;
367
368  for (;;)
369  {
370    lps->OutSize = outStreamSpec->GetSize();
371    lps->InSize = _packSize = decoder.GetInputProcessedSize();
372    _packSizeDefined = true;
373    RINOK(lps->SetCur());
374
375    CHeader st;
376
377    const UInt32 kBufSize = 1 + 5 + 8;
378    Byte buf[kBufSize];
379    const UInt32 headerSize = GetHeaderSize();
380    UInt32 processed;
381    RINOK(decoder.ReadInput(buf, headerSize, &processed));
382    if (processed != headerSize)
383      break;
384
385    if (!st.Parse(buf, _lzma86))
386      break;
387    firstItem = false;
388
389    result = decoder.Code(st, outStream, progress);
390    if (result == E_NOTIMPL)
391    {
392      opRes = NExtract::NOperationResult::kUnSupportedMethod;
393      break;
394    }
395    if (result == S_FALSE)
396    {
397      opRes = NExtract::NOperationResult::kDataError;
398      break;
399    }
400    RINOK(result);
401  }
402  if (firstItem)
403    return E_FAIL;
404  outStream.Release();
405  return extractCallback->SetOperationResult(opRes);
406  COM_TRY_END
407}
408
409IMPL_ISetCompressCodecsInfo
410
411static IInArchive *CreateArc() { return new CHandler(false); }
412static IInArchive *CreateArc86() { return new CHandler(true); }
413
414namespace NLzmaAr {
415
416static CArcInfo g_ArcInfo =
417  { L"lzma", L"lzma", 0, 0xA, { 0 }, 0, true, CreateArc, NULL };
418REGISTER_ARC(Lzma)
419
420}
421
422namespace NLzma86Ar {
423
424static CArcInfo g_ArcInfo =
425  { L"lzma86", L"lzma86", 0, 0xB, { 0 }, 0, true, CreateArc86, NULL };
426REGISTER_ARC(Lzma86)
427
428}
429
430}}
431