dmghandler.cpp

7z一个高压缩比的压缩程序源代码,重要的是里面的算法值得学习

// DmgHandler.cpp

#include "StdAfx.h"

#include "../../../C/CpuArch.h"

#include "Common/Buffer.h"
#include "Common/ComTry.h"
#include "Common/IntToString.h"
#include "Common/MyXml.h"
#include "Common/UTFConvert.h"

#include "Windows/PropVariant.h"

#include "../Common/LimitedStreams.h"
#include "../Common/ProgressUtils.h"
#include "../Common/RegisterArc.h"
#include "../Common/StreamUtils.h"

#include "../Compress/BZip2Decoder.h"
#include "../Compress/CopyCoder.h"
#include "../Compress/ZlibDecoder.h"

// #define DMG_SHOW_RAW

// #include <stdio.h>
#define PRF(x) // x

#define Get32(p) GetBe32(p)
#define Get64(p) GetBe64(p)

static int Base64ToByte(char c)
{
  if (c >= 'A' && c <= 'Z') return c - 'A';
  if (c >= 'a' && c <= 'z') return c - 'a' + 26;
  if (c >= '0' && c <= '9') return c - '0' + 52;
  if (c == '+') return 62;
  if (c == '/') return 63;
  if (c == '=') return 0;
  return -1;
}

static int Base64ToBin(Byte *dest, const char *src, int srcLen)
{
  int srcPos = 0;
  int destPos = 0;
  while (srcPos < srcLen)
  {
    Byte buf[4];
    int filled = 0;
    while (srcPos < srcLen)
    {
      int n = Base64ToByte(src[srcPos++]);
      if (n >= 0)
      {
        buf[filled++] = (Byte)n;
        if (filled == 4)
          break;
      }
    }
    if (filled >= 2) { if (dest) dest[destPos] = (buf[0] << 2) | (buf[1] >> 4); destPos++; }
    if (filled >= 3) { if (dest) dest[destPos] = (buf[1] << 4) | (buf[2] >> 2); destPos++; }
    if (filled >= 4) { if (dest) dest[destPos] = (buf[2] << 6) | (buf[3]     ); destPos++; }
  }
  return destPos;
}

static UString GetSizeString(UInt64 value)
{
  wchar_t s[32];
  wchar_t c;
  if (value < (UInt64)20000) c = 0;
  else if (value < ((UInt64)20000 << 10)) { value >>= 10; c = L'K'; }
  else if (value < ((UInt64)20000 << 20)) { value >>= 20; c = L'M'; }
  else                                    { value >>= 30; c = L'G'; }
  ConvertUInt64ToString(value, s);
  int p = MyStringLen(s);
  s[p++] = c;
  s[p++] = L'\0';
  return s;
}

namespace NArchive {
namespace NDmg {

struct CBlock
{
  UInt32 Type;
  UInt64 UnpPos;
  UInt64 UnpSize;
  UInt64 PackPos;
  UInt64 PackSize;
};

struct CFile
{
  CByteBuffer Raw;
  // UInt64 StartPos;
  CRecordVector<CBlock> Blocks;
  UInt64 GetUnpackSize() const
  {
    UInt64 size = 0;
    for (int i = 0; i < Blocks.Size(); i++)
      size += Blocks[i].UnpSize;
    return size;
  };
  UInt64 GetPackSize() const
  {
    UInt64 size = 0;
    for (int i = 0; i < Blocks.Size(); i++)
      size += Blocks[i].PackSize;
    return size;
  };

  AString Name;
};

class CHandler:
  public IInArchive,
  public CMyUnknownImp
{
  CMyComPtr<IInStream> _inStream;

  AString _xml;
  CObjectVector<CFile> _files;
  CRecordVector<int> _fileIndices;

  HRESULT Open2(IInStream *stream);
  HRESULT Extract(IInStream *stream);
public:
  MY_UNKNOWN_IMP1(IInArchive)
  INTERFACE_IInArchive(;)
};

const UInt32 kXmlSizeMax = ((UInt32)1 << 31) - (1 << 14);

enum
{
  METHOD_ZERO_0 = 0,
  METHOD_COPY   = 1,
  METHOD_ZERO_2 = 2,
  METHOD_ZLIB   = 0x80000005,
  METHOD_BZIP2  = 0x80000006,
  METHOD_DUMMY  = 0x7FFFFFFE,
  METHOD_END    = 0xFFFFFFFF
};

struct CMethodStat
{
  UInt32 NumBlocks;
  UInt64 PackSize;
  UInt64 UnpSize;
  CMethodStat(): NumBlocks(0), PackSize(0), UnpSize(0) {}
};

struct CMethods
{
  CRecordVector<CMethodStat> Stats;
  CRecordVector<UInt32> Types;
  void Update(const CFile &file);
  UString GetString() const;
};

void CMethods::Update(const CFile &file)
{
  for (int i = 0; i < file.Blocks.Size(); i++)
  {
    const CBlock &b = file.Blocks[i];
    int index = Types.FindInSorted(b.Type);
    if (index < 0)
    {
      index = Types.AddToUniqueSorted(b.Type);
      Stats.Insert(index, CMethodStat());
    }
    CMethodStat &m = Stats[index];
    m.PackSize += b.PackSize;
    m.UnpSize += b.UnpSize;
    m.NumBlocks++;
  }
}

UString CMethods::GetString() const
{
  UString res;
  for (int i = 0; i < Types.Size(); i++)
  {
    if (i != 0)
      res += L' ';
    wchar_t buf[32];
    const wchar_t *s;
    const CMethodStat &m = Stats[i];
    bool showPack = true;
    UInt32 type = Types[i];
    switch(type)
    {
      case METHOD_ZERO_0: s = L"zero0"; showPack = (m.PackSize != 0); break;
      case METHOD_ZERO_2: s = L"zero2"; showPack = (m.PackSize != 0); break;
      case METHOD_COPY:   s = L"copy"; showPack = (m.UnpSize != m.PackSize); break;
      case METHOD_ZLIB:   s = L"zlib"; break;
      case METHOD_BZIP2:  s = L"bzip2"; break;
      default: ConvertUInt64ToString(type, buf); s = buf;
    }
    res += s;
    if (m.NumBlocks != 1)
    {
      res += L'[';
      ConvertUInt64ToString(m.NumBlocks, buf);
      res += buf;
      res += L']';
    }
    res += L'-';
    res += GetSizeString(m.UnpSize);
    if (showPack)
    {
      res += L'-';
      res += GetSizeString(m.PackSize);
    }
  }
  return res;
}

STATPROPSTG kProps[] =
{
  { NULL, kpidPath, VT_BSTR},
  { NULL, kpidSize, VT_UI8},
  { NULL, kpidPackSize, VT_UI8},
  { NULL, kpidComment, VT_BSTR},
  { NULL, kpidMethod, VT_BSTR}
};

IMP_IInArchive_Props

STATPROPSTG kArcProps[] =
{
  { NULL, kpidMethod, VT_BSTR},
  { NULL, kpidNumBlocks, VT_UI4}
};

STDMETHODIMP CHandler::GetArchiveProperty(PROPID propID, PROPVARIANT *value)
{
  COM_TRY_BEGIN
  NWindows::NCOM::CPropVariant prop;
  switch(propID)
  {
    case kpidMethod:
    {
      CMethods m;
      for (int i = 0; i < _files.Size(); i++)
        m.Update(_files[i]);
      prop = m.GetString();
      break;
    }
    case kpidNumBlocks:
    {
      UInt64 numBlocks = 0;
      for (int i = 0; i < _files.Size(); i++)
        numBlocks += _files[i].Blocks.Size();
      prop = numBlocks;
      break;
    }
  }
  prop.Detach(value);
  return S_OK;
  COM_TRY_END
}

IMP_IInArchive_ArcProps

static int FindKeyPair(const CXmlItem &item, const AString &key, const AString &nextTag)
{
  for (int i = 0; i + 1 < item.SubItems.Size(); i++)
  {
    const CXmlItem &si = item.SubItems[i];
    if (si.IsTagged("key") && si.GetSubString() == key && item.SubItems[i + 1].IsTagged(nextTag))
      return i + 1;
  }
  return -1;
}

static AString GetStringFromKeyPair(const CXmlItem &item, const AString &key, const AString &nextTag)
{
  int index = FindKeyPair(item, key, nextTag);
  if (index >= 0)
    return item.SubItems[index].GetSubString();
  return AString();
}

HRESULT CHandler::Open2(IInStream *stream)
{
  const int HEADER_SIZE = 0x1E0;

  UInt64 headerPos;
  RINOK(stream->Seek(-HEADER_SIZE, STREAM_SEEK_END, &headerPos));
  Byte buf[HEADER_SIZE];
  RINOK(ReadStream_FALSE(stream, buf, HEADER_SIZE));
  UInt64 address1 = Get64(buf + 0);
  UInt64 address2 = Get64(buf + 0xB8);
  UInt64 size64 = Get64(buf + 0xC0);
  if (address1 != address2 || size64 >= kXmlSizeMax || size64 == 0 ||
      address1 >= headerPos || address1 + size64 > headerPos)
    return S_FALSE;
  RINOK(stream->Seek(address1, STREAM_SEEK_SET, NULL));
  size_t size = (size_t)size64;

  char *ss = _xml.GetBuffer((int)size + 1);
  RINOK(ReadStream_FALSE(stream, ss, size));
  ss[size] = 0;
  _xml.ReleaseBuffer();

  CXml xml;
  if (!xml.Parse(_xml))
    return S_FALSE;
  if (xml.Root.Name != "plist")
    return S_FALSE;
  
  int dictIndex = xml.Root.FindSubTag("dict");
  if (dictIndex < 0)
    return S_FALSE;
  
  const CXmlItem &dictItem = xml.Root.SubItems[dictIndex];
  int rfDictIndex = FindKeyPair(dictItem, "resource-fork", "dict");
  if (rfDictIndex < 0)
    return S_FALSE;
  
  const CXmlItem &rfDictItem = dictItem.SubItems[rfDictIndex];
  int arrIndex = FindKeyPair(rfDictItem, "blkx", "array");
  if (arrIndex < 0)
    return S_FALSE;

  const CXmlItem &arrItem = rfDictItem.SubItems[arrIndex];

  /* some DMG file has BUG:
  PackPos for each new file is 0.
  So we use additional "StartPos" to fix that BUG */

  /*
  UInt64 startPos = 0;
  bool startPosIsDefined = false;
  */


  for (int i = 0; i < arrItem.SubItems.Size(); i++)
  {
    const CXmlItem &item = arrItem.SubItems[i];
    if (!item.IsTagged("dict"))
      continue;

    CFile file;
    // file.StartPos = startPos;

    int destLen;
    {
      AString dataString;
      AString name = GetStringFromKeyPair(item, "Name", "string");
      if (name.IsEmpty())
        name = GetStringFromKeyPair(item, "CFName", "string");
      file.Name = name;
      dataString = GetStringFromKeyPair(item, "Data", "data");
     
      destLen = Base64ToBin(NULL, dataString, dataString.Length());
      file.Raw.SetCapacity(destLen);
      Base64ToBin(file.Raw, dataString, dataString.Length());
    }

    if (destLen > 0xCC && Get32(file.Raw) == 0x6D697368)
    {
      PRF(printf("\n\n index = %d", _files.Size()));
      const int kRecordSize = 40;
      for (int offset = 0xCC; offset + kRecordSize <= destLen; offset += kRecordSize)
      {
        const Byte *p = (const Byte *)file.Raw + offset;
        CBlock b;
        b.Type = Get32(p);
        if (b.Type == METHOD_END)
          break;
        if (b.Type == METHOD_DUMMY)
          continue;

        b.UnpPos   = Get64(p + 0x08) << 9;
        b.UnpSize  = Get64(p + 0x10) << 9;
        b.PackPos  = Get64(p + 0x18);
        b.PackSize = Get64(p + 0x20);

        /*
        if (startPosIsdefined)
        {
        }
        else
        {
          startPosIsdefined = true;