CPP/7zip/Archive/LzmaHandler.cpp

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