wc3re/TreFile.cc

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

#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#include "common.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;
} __attribute__((__packed__));

struct Table1Entry { // little endian
  uint32_t crc;
  uint32_t table3Ptr;
} __attribute__((__packed__));

struct Table3Entry { // little endian
  unsigned dataPtr:32;
  //  unsigned unknown1:8;
  unsigned length:24;
  unsigned unknown2:8;
} __attribute__((__packed__));


TreFile::TreFile(std::string const& filename)
  : file_(nullptr), start_(0), length_(0)
{
  file_ = fopen(filename.c_str(), "rb");
  if (!file_)
    throw POSIXException{errno, "Could not open "s + filename};

  {
    struct stat statBuf;
    
    if (fstat(fileno(file_), &statBuf) != 0) {
      fclose(file_);
      throw POSIXException(errno, "Could not stat");
    }

    length_ = statBuf.st_size;
  }
  try {
    construct_();
  } catch(...) {
    fclose(file_);
    throw;
  }
}

TreFile::TreFile(FILE* file, off_t start, off_t length)
  : file_(file), start_(start), length_(length)
{
  construct_();
}

TreFile::~TreFile()
{
  if (file_) {
    fclose(file_);
    file_ = nullptr;
  }
}

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::File TreFile::openName(std::string const& name) const
{
  return openIdx_(findName_(name));
}

TreFile::File 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)
{
  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)
{
  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)
{
  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, flags %.2hhx)\n", ent.first,
	   std::get<0>(table3_[ent.second]),
	   std::get<1>(table3_[ent.second]),
	   std::get<2>(table3_[ent.second]));
  }

  printf("Files by Name:\n");
  for(auto const& ent : table2_) {
    printf("\t%s -> (ofs %.8x, len %.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]));
  }
}
  
void TreFile::construct_()
{
  if (length_ < sizeof(TreHeader))
    throw FormatException{"Size < header size"};
    
  if (fseeko(file_, start_, SEEK_SET) != 0) {
    throw POSIXException(errno, "Could not seek");
  }
 
  TreHeader header;
  if (fread(&header, sizeof(TreHeader), 1, file_) != 1)
    throw POSIXException{errno, "Could not read header"};

  if (memcmp(header.magic, "XTRE", 4) != 0)
    throw FormatException{"Wrong magic, not a TRE?"};
    
  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)/sizeof(Table1Entry);
  const size_t numTable3Entries = (header.dataOfs-header.table3Ofs)/sizeof(Table3Entry);
    
  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);
    
  // Read Table 3
  if (fseeko(file_, start_+header.table3Ofs, SEEK_SET) != 0) {
    throw POSIXException{errno, "Could not seek"};
  }

  table3_.reserve(numTable3Entries);
  for (size_t i = 0;i < numTable3Entries;++i) {
    Table3Entry entry;
    if (fread(&entry, sizeof(Table3Entry), 1, file_) != 1)
      throw POSIXException{errno, "Could not read"};

    if ((entry.dataPtr > length_) ||
	(entry.dataPtr < header.dataOfs))
      throw FormatException{"Data pointer out of range"};

    if ((entry.dataPtr + entry.length) > length_) {
      fprintf(stderr, "Data length exceeds file length: %.8x %.8x\n", entry.dataPtr, entry.length);
      //throw FormatException{"Data length exceeds file length"};
    }
      
    table3_.push_back(std::make_tuple(static_cast<uint32_t>(entry.dataPtr),
				      static_cast<uint32_t>(entry.length),
				      static_cast<uint8_t>(entry.unknown2)));
  }

  // Read Table 2
  if (fseeko(file_, start_+header.table2Ofs, SEEK_SET) != 0) {
    throw POSIXException{errno, "Could not seek"};
  }

  std::map<size_t, uint32_t> table2Pos;
    
  size_t pos = header.table2Ofs;
  while((pos+5) < header.table3Ofs) {
    uint8_t nameLen;

    if (fread(&nameLen, 1, 1, file_) != 1)
      throw POSIXException{errno, "Could not read"};

    if (pos+nameLen+5 > header.table3Ofs)
      throw FormatException{"Table 2 entry exceeds table "s + std::to_string(nameLen)};
      
    std::vector<char> nameData(nameLen);
    if (fread(nameData.data(), nameLen, 1, file_) != 1)
      throw POSIXException{errno, "Could not read"};

    uint32_t table3Ptr;
    if (fread(&table3Ptr, 4, 1, file_) != 1)
      throw POSIXException{errno, "Could not read"};

    if ((table3Ptr < header.table3Ofs) ||
	(table3Ptr >= header.dataOfs)) {
      throw FormatException{"Table3 pointer out of range"};
    }
      
    std::string nameStr{nameData.data(), nameLen};
    size_t table3Index = (table3Ptr - header.table3Ofs)/sizeof(Table3Entry);
      
    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
  if (fseeko(file_, start_+header.table1Ofs, SEEK_SET) != 0) {
    throw POSIXException{errno, "Could not seek"};
  }

  for (size_t i = 0;i < numTable1Entries;++i) {
    Table1Entry entry;
    if (fread(&entry, sizeof(Table1Entry), 1, file_) != 1)
      throw POSIXException{errno, "Could not read"};

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

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

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

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{"CRC "s + crcStr + " not found"s};
  }

  return it->second;
}

void TreFile::dumpIdx_(std::string const& name, size_t table3Idx)
{
  uint32_t ofs, len;
  std::tie(ofs, len, std::ignore) = table3_[table3Idx];

  // if (ofs%2 != 0) {
  //   ++ofs;
  //   len;
  // }
  
  std::vector<char> buf(len);

  if (fseeko(file_, start_+ofs, SEEK_SET) != 0)
    throw POSIXException{errno, "Could not seek"};

  if (fread(buf.data(), len, 1, file_) != 1)
    throw POSIXException{errno, "Could not read data"};

  FILEUPtr outFile{fopen(name.c_str(), "wb")};
  if (!outFile)
    throw POSIXException{errno, "Could not open " + name};

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

TreFile::File TreFile::openIdx_(size_t table3Idx) const
{
  uint32_t ofs, len;
  uint8_t flags;
  std::tie(ofs, len, flags) = table3_[table3Idx];

  if(flags&0x80)
    throw Exception{"Compressed TRE objects NYI"};
  
  std::vector<char> buf(len);

  if (fseeko(file_, start_+ofs, SEEK_SET) != 0)
    throw POSIXException{errno, "Could not seek"};

  if (fread(buf.data(), len, 1, file_) != 1)
    throw POSIXException{errno, "Could not read data"};

  return File{buf};
}

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

  return ret;
}