wc3re/TreFile.cc

#include <cstdio>
#include <cstdint>
#include <cstring>

#include "common.hh"
#include "util.hh"
#include "decompress.hh"
#include "TreFile.hh"

struct TreHeader { // little endian
  char magic[4]; // == "XTRE"
  //uint32_t unknown;
  uint32_t table1Ofs;
  uint32_t table2Ofs;
  uint32_t table3Ofs;
  uint32_t dataOfs;
};

static const size_t headerSize = 24;
static const size_t table1EntrySize = 8;
static const size_t table3EntrySize = 8;

TreFile::TreFile(uint8_t const* base, size_t length)
  : base_(base), length_(length)
{
  if (length_ < 6*4)
    throw FormatException{"Size < header size"};

  TreHeader header;
  memcpy(header.magic, base, 4);

  if (memcmp(header.magic, "XTRE", 4) != 0)
    throw FormatException{"Wrong magic, not a TRE?"};

  header.table1Ofs = readU32LE(base_+8);
  header.table2Ofs = readU32LE(base_+12);
  header.table3Ofs = readU32LE(base_+16);
  header.dataOfs = readU32LE(base_+20);
  
  if ((header.table1Ofs > length_) ||
      (header.table2Ofs > length_) ||
      (header.table3Ofs > length_) ||
      (header.dataOfs > length_))
    throw FormatException{"Table offset exceeds length"};

  if ((header.table1Ofs > header.table2Ofs) ||
      (header.table2Ofs > header.table3Ofs) ||
      (header.table3Ofs > header.dataOfs))
    throw FormatException{"Table offsets not in ascending order"};

  const size_t numTable1Entries = (header.table2Ofs-header.table1Ofs)/table1EntrySize;
  const size_t numTable3Entries = (header.dataOfs-header.table3Ofs)/table3EntrySize;

#ifndef NDEBUG
  printf("Table 1: start %u, length %u (%lu entries)\n",
	 header.table1Ofs, header.table2Ofs-header.table1Ofs,
	 numTable1Entries);
  printf("Table 2: start %u, length %u\n",
	 header.table2Ofs, header.table3Ofs-header.table2Ofs);
  printf("Table 3: start %u, length %u (%lu entries)\n",
	 header.table3Ofs, header.dataOfs-header.table3Ofs,
	 numTable3Entries);
  printf("Data: start %u, length %lu\n",
	 header.dataOfs, length_ - header.dataOfs);
#endif

  // Read table 3
  for (size_t i = 0;i < numTable3Entries;++i) {
    uint8_t const* const entryBase = base_+header.table3Ofs+i*table3EntrySize;
    
    const uint32_t dataPtr = readU32LE(entryBase);
    const uint32_t length = readU32LE(entryBase+4)&0x0fffffffu;
    const uint8_t flags = *(entryBase+7)&0xf0;
    
    if ((dataPtr > length_) ||
	(dataPtr < header.dataOfs))
      throw FormatException{"Data pointer out of range"};

    if ((dataPtr + length) > length_) {
      fprintf(stderr, "Data length exceeds file length: %.8x %.8x\n", dataPtr, length);
    }

    if (!table3_.empty()) {
      uint32_t lastPtr, lastLen;
      uint8_t lastFlags;
      std::tie(lastPtr, std::ignore, lastLen, lastFlags) = table3_.back();
      if (lastPtr+lastLen > dataPtr) {
	if (lastFlags&0x80)
	  std::get<2>(table3_.back()) = dataPtr - lastPtr;
	else
	  fprintf(stderr, "Overlap? %u %u (%hhu) -> %u\n",lastPtr, lastLen, lastFlags, dataPtr);
      }
    }
    
    table3_.emplace_back(dataPtr, length, length, flags);
  }

  // Fixup last file clength
  {
    uint32_t lastPtr, lastLen;
    uint8_t lastFlags;
    std::tie(lastPtr, std::ignore, lastLen, lastFlags) = table3_.back();
    if (lastPtr+lastLen > length_) {
      if (lastFlags&0x80)
	std::get<2>(table3_.back()) = length_ - lastPtr;
      else
	fprintf(stderr, "Overrun? %u %u (%hhu) -> %lu\n",lastPtr, lastLen, lastFlags, length_);
    }
  }
  
  
  // Read Table 2
  std::map<size_t, uint32_t> table2Pos;
    
  size_t pos = header.table2Ofs;
  while((pos+5) < header.table3Ofs) {
    uint8_t nameLen = *(base+pos);

    if (pos+nameLen+5 > header.table3Ofs)
      throw FormatException{"Table 2 entry exceeds table " + std::to_string(nameLen)};

    for (unsigned i = 0;i < nameLen;++i)
      if (!isprint(base[pos+1+i]))
	throw FormatException{"Filename not printable"};
    
    std::string nameStr(reinterpret_cast<char const*>(base)+pos+1, nameLen);

    const uint32_t table3Ptr = readU32LE(base+pos+nameLen+1);
    
    if ((table3Ptr < header.table3Ofs) ||
	(table3Ptr >= header.dataOfs)) {
      throw FormatException{"Table3 pointer out of range"};
    }

    const size_t table3Index = (table3Ptr - header.table3Ofs)/table3EntrySize;

    auto ins = table2_.emplace(nameStr, table3Index);
    if (!ins.second) {
      if (ins.first->second == table3Index)
	printf("Duplicate entry for name %s\n", nameStr.c_str());
      else
	printf("Collision for name %s: prev %lu, new %lu\n", nameStr.c_str(),
	       ins.first->second, table3Index);
    }
    // Store an position-indexed copy of Table2 to resolve Table1->Table2 ptr's
    table2Pos.emplace(pos, table3Ptr);

    pos += 1+nameLen+4;
  }

  // Read Table 1
  for (size_t i = 0;i < numTable1Entries;++i) {
    uint8_t const* const entryBase = base+header.table1Ofs+i*table1EntrySize;

    const uint32_t crc = readU32LE(entryBase);
    uint32_t table3Ptr = readU32LE(entryBase+4);
    
    if ((crc == 0x00000000) &&
	(table3Ptr == 0xffffffff))
      continue; // Unused entry
      
    if ((table3Ptr < header.table2Ofs) ||
	(table3Ptr >= header.dataOfs)) {
      throw FormatException{"Table3 pointer out of range"};
    }

    if (table3Ptr < header.table3Ofs) {
      auto it = table2Pos.find(table3Ptr);
      if (it == table2Pos.end()) {
	fprintf(stderr, "Table1->Table2 pointer: Table 2 entry not found: crc %.8x ptr %.8x\n",
	        crc, table3Ptr);
	continue;
      }
      printf("Table1->Table2 pointer resolved: Table 2 %.8x -> Table 3 %.8x\n",
	     table3Ptr, it->second);
      table3Ptr = it->second;
    }

    const size_t table3Index = (table3Ptr - header.table3Ofs)/table3EntrySize;
    auto ins = table1_.emplace(crc, table3Index);
    if (!ins.second) {
      if (ins.first->second == table3Index)
	printf("Duplicate entry for CRC %.8x\n", crc);
      else
	printf("Collision for CRC %.8x: prev %lu, new %lu\n", crc, ins.first->second, table3Index);
    }
  }
}


TreFile::~TreFile()
{
}

std::vector<std::string> TreFile::getNames() const
{
  std::vector<std::string> ret;
  for (auto ent : table2_) {
    ret.push_back(ent.first);
  }

  return ret;
}

std::vector<uint32_t> TreFile::getCRCs() const
{
  std::vector<uint32_t> ret;
  for (auto ent : table1_) {
    ret.push_back(ent.first);
  }

  return ret;
}

TreFile::Object TreFile::openName(std::string const& name) const
{
  return openIdx_(findName_(name));
}

TreFile::Object TreFile::openCRC(uint32_t crc) const
{
  return openIdx_(findCRC_(crc));
}
  
TreFile::Stat TreFile::statName(std::string const& name) const
{
  return statIdx_(findName_(name));
}

TreFile::Stat TreFile::statCRC(uint32_t crc) const
{
  return statIdx_(findCRC_(crc));
}

void TreFile::dumpName(std::string path, std::string const& name) const
{
  auto idx = findName_(name);

  if (path.back() != '/')
    path.push_back('/');
  path.append(name);
  
  dumpIdx_(path, idx);
}

void TreFile::dumpCRC(std::string path, uint32_t crc) const
{
  auto idx = findCRC_(crc);
  char crcStr[9];
  snprintf(crcStr, 9, "%.8X", crc);
  
  if (path.back() != '/')
    path.push_back('/');
  path.append(crcStr);
  
  dumpIdx_(path, idx);
}

void TreFile::dumpAll(std::string path) const
{
  if (path.back() != '/')
    path.push_back('/');
  
  for (auto ent : table1_) {
    char crcStr[9];
    snprintf(crcStr, 9, "%.8X", ent.first);
    std::string name = path;
    name.append(crcStr);

    dumpIdx_(name, ent.second);
  }

  for (auto ent : table2_) {
    std::string name = path;
    name.append(ent.first);

    dumpIdx_(name, ent.second);
  } 
}
  
void TreFile::printStructure()
{
  printf("Table 1: %lu entries\n", table1_.size());
  printf("Table 2: %lu entries\n", table2_.size());
  printf("Table 3: %lu entries\n", table3_.size());

  printf("Files by CRC:\n");
  for(auto const& ent : table1_) {
    printf("\t%.8x -> (ofs %.8x, len %.8x, clen %.8x, flags %.2hhx)\n", ent.first,
	   std::get<0>(table3_[ent.second]),
	   std::get<1>(table3_[ent.second]),
	   std::get<2>(table3_[ent.second]),
	   std::get<3>(table3_[ent.second]));
  }

  printf("Files by Name:\n");
  for(auto const& ent : table2_) {
    printf("\t%s -> (ofs %.8x, len %.8x, clen %.8x, flags %.2hhx)\n", ent.first.c_str(),
	   std::get<0>(table3_[ent.second]),
	   std::get<1>(table3_[ent.second]),
	   std::get<2>(table3_[ent.second]),
	   std::get<3>(table3_[ent.second]));
  }
}

size_t TreFile::findName_(std::string const& name) const
{
  auto it = table2_.find(name);
  if (it == table2_.end())
    throw Exception{name + " not found"};

  return it->second;
}

size_t TreFile::findCRC_(uint32_t crc) const
{
  auto it = table1_.find(crc);
  if (it == table1_.end()) {
    char crcStr[9];
    snprintf(crcStr, 9, "%.8X", crc);
    throw Exception{std::string("CRC ") + crcStr + " not found"};
  }

  return it->second;
}

void TreFile::dumpIdx_(std::string const& name, size_t table3Idx) const
{
  auto obj = openIdx_(table3Idx);
  
  FILEUPtr outFile{fopen(name.c_str(), "wb")};
  if (!outFile)
    throw POSIXException{errno, "Could not open " + name};

  if (fwrite(obj.data(), obj.size(), 1, outFile.get()) != 1)
      throw POSIXException{errno, "Could not write"};
}

TreFile::Object TreFile::openIdx_(size_t table3Idx) const
{
  uint32_t dataPtr, length, clength;
  uint8_t flags;
  std::tie(dataPtr, length, clength, flags) = table3_[table3Idx];

  if ((dataPtr + clength) > length_)
    throw FormatException{"length exceeds file size"};
  
  if (flags&0x80) {
    if (flags&0x40) {
      auto dec = decompressLZ(base_+dataPtr, clength);
#ifndef NDEBUG
      if (dec.size() != length)
	printf("WARNING: Decompressed size != expected (%lu, %u)\n", dec.size(), length);
#endif
      return Object(std::move(dec));
    } else
      throw Exception{"Compression type 0 NYI"};
  } else {
    return Object(base_+dataPtr, length);
  }
}

TreFile::Stat TreFile::statIdx_(size_t table3Idx) const
{
  Stat ret;
  std::tie(std::ignore, ret.size, ret.csize, ret.flags) = table3_[table3Idx];

  return ret;
}