// Copyright 2021 Google Inc. All Rights Reserved.

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//     http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

#include "absl/strings/substitute.h"

#include "bloaty.h"

#include "util.h"

using absl::string_view;

using std::string;

namespace bloaty {

namespace pe {

constexpr uint16_t dos_magic = 0x5A4D;  // MZ

// From LIEF/include/LIEF/PE/Structures.hpp.in

//! Sizes in bytes of various things in the COFF format.

constexpr size_t kHeader16Size = 20;

constexpr size_t kHeader32Size = 56;

constexpr size_t kNameSize = 8;

constexpr size_t kSymbol16Size = 18;

constexpr size_t kSymbol32Size = 20;

constexpr size_t kSectionSize = 40;

constexpr size_t kRelocationSize = 10;

constexpr size_t kBaseRelocationBlockSize = 8;

constexpr size_t kImportDirectoryTableEntrySize = 20;

constexpr size_t kResourceDirectoryTableSize = 16;

constexpr size_t kResourceDirectoryEntriesSize = 8;

constexpr size_t kResourceDataEntrySize = 16;

#include "third_party/lief_pe/pe_enums.h"

#include "third_party/lief_pe/pe_structures.h"

static_assert(kSectionSize == sizeof(pe_section),

              "Compiler options broke LIEF struct layout");

static_assert(kRelocationSize == sizeof(pe_relocation),

              "Compiler options broke LIEF struct layout");

static_assert(kBaseRelocationBlockSize == sizeof(pe_base_relocation_block),

              "Compiler options broke LIEF struct layout");

static_assert(kImportDirectoryTableEntrySize == sizeof(pe_import),

              "Compiler options broke LIEF struct layout");

static_assert(kResourceDirectoryTableSize ==

                  sizeof(pe_resource_directory_table),

              "Compiler options broke LIEF struct layout");

static_assert(kResourceDirectoryEntriesSize ==

                  sizeof(pe_resource_directory_entries),

              "Compiler options broke LIEF struct layout");

static_assert(kResourceDataEntrySize == sizeof(pe_resource_data_entry),

              "Compiler options broke LIEF struct layout");

static_assert(kSymbol16Size == sizeof(pe_symbol),

              "Compiler options broke LIEF struct layout");

static_assert(sizeof(PE_TYPE) == sizeof(uint16_t),

              "Compiler options broke PE_TYPE size");

class PeFile {

 public:

  PeFile(string_view data) : data_(data) { ok_ = Initialize(); }

  bool IsOpen() const { return ok_; }

  string_view entire_file() const { return data_; }

  string_view pe_headers() const { return pe_headers_; }

  string_view section_headers() const { return section_headers_; }

  uint32_t section_count() const { return section_count_; }

  string_view section_header(size_t n) const {

    return StrictSubstr(section_headers_, n * sizeof(pe_section),

                        sizeof(pe_section));

  }

 private:

  bool Initialize();

  string_view GetRegion(uint64_t start, uint64_t n) const {

    return StrictSubstr(data_, start, n);

  }

  bool ok_;

  bool is_64bit_;

  const string_view data_;

  pe_dos_header dos_header_;

  pe_header pe_header_;

  string_view pe_headers_;

  string_view section_headers_;

  uint32_t section_count_;

};

bool PeFile::Initialize() {

  if (data_.size() < sizeof(dos_header_)) {

    return false;

  }

  memcpy(&dos_header_, data_.data(), sizeof(dos_header_));

  if (dos_header_.Magic != dos_magic) {

    // Not a PE file.

    return false;

  }

  PE_TYPE Magic;

  auto pe_end =

      CheckedAdd(dos_header_.AddressOfNewExeHeader, sizeof(pe_header_));

  if (CheckedAdd(pe_end, sizeof(Magic)) > data_.size()) {

    // Cannot fit the headers / magic from optional header

    return false;

  }

  memcpy(&pe_header_, data_.data() + dos_header_.AddressOfNewExeHeader,

         sizeof(pe_header_));

  if (!std::equal(pe_header_.signature, pe_header_.signature + sizeof(PE_Magic),

                  std::begin(PE_Magic))) {

    // Not a PE file.

    return false;

  }

  memcpy(&Magic, data_.data() + pe_end, sizeof(Magic));

  if (Magic != PE_TYPE::PE32 && Magic != PE_TYPE::PE32_PLUS) {

    // Unknown PE magic

    return false;

  }

  is_64bit_ = Magic == PE_TYPE::PE32_PLUS;

  section_count_ = pe_header_.NumberOfSections;

  // TODO(mj): Figure out if we should trust SizeOfOptionalHeader here

  const uint32_t sections_offset = dos_header_.AddressOfNewExeHeader +

                                   sizeof(pe_header_) +

                                   pe_header_.SizeOfOptionalHeader;

  auto sections_size = CheckedMul(section_count_, sizeof(pe_section));

  if ((sections_offset + sections_size) > data_.size()) {

    // Cannot fit the headers

    return false;

  }

  pe_headers_ = GetRegion(0, sections_offset);

  section_headers_ = GetRegion(sections_offset, sections_size);

  return true;

}

class Section {

 public:

  string name;

  uint32_t raw_offset() const { return header_.PointerToRawData; }

  uint32_t raw_size() const { return header_.SizeOfRawData; }

  uint32_t virtual_addr() const { return header_.VirtualAddress; }

  uint32_t virtual_size() const { return header_.VirtualSize; }

  Section(string_view header_data) {

    assert(header_data.size() == sizeof(header_));

    memcpy(&header_, header_data.data(), sizeof(header_));

    // TODO(mj): Handle long section names:

    // For longer names, this member contains a forward slash (/) followed by an

    // ASCII representation of a decimal number that is an offset into the

    // string table.

    name = string(header_.Name, strnlen(header_.Name, kNameSize));

  }

 private:

  pe_section header_;

};

template <class Func>

void ForEachSection(const PeFile& pe, Func&& section_func) {

  for (auto n = 0; n < pe.section_count(); ++n) {

    Section section(pe.section_header(n));

    section_func(section);

  }

}

void ParseSections(const PeFile& pe, RangeSink* sink) {

  assert(pe.IsOpen());

  ForEachSection(pe, [sink, &pe](const Section& section) {

    uint64_t vmaddr = section.virtual_addr();

    uint64_t vmsize = section.virtual_size();

    absl::string_view section_data = StrictSubstr(

        pe.entire_file(), section.raw_offset(), section.raw_size());

    sink->AddRange("pe_sections", section.name, vmaddr, vmsize, section_data);

  });

}

void AddCatchAll(const PeFile& pe, RangeSink* sink) {

  assert(pe.IsOpen());

  auto begin = pe.pe_headers().data() - sink->input_file().data().data();

  sink->AddRange("pe_catchall", "[PE Headers]", begin, pe.pe_headers().size(),

                 pe.pe_headers());

  begin = pe.section_headers().data() - sink->input_file().data().data();

  sink->AddRange("pe_catchall", "[PE Section Headers]", begin,

                 pe.section_headers().size(), pe.section_headers());

  // The last-line fallback to make sure we cover the entire file.

  sink->AddFileRange("pe_catchall", "[Unmapped]", sink->input_file().data());

}

class PEObjectFile : public ObjectFile {

 public:

  PEObjectFile(std::unique_ptr<InputFile> file_data,

               std::unique_ptr<pe::PeFile> pe)

      : ObjectFile(std::move(file_data)), pe_file(std::move(pe)) {}

  std::string GetBuildId() const override {

    // TODO(mj): Read from pe_pdb_??

    return std::string();

  }

  void ProcessFile(const std::vector<RangeSink*>& sinks) const override {

    for (auto sink : sinks) {

      switch (sink->data_source()) {

        case DataSource::kSegments:

          // TODO(mj): sections: list out imports and other stuff!

        case DataSource::kSections:

          ParseSections(*pe_file, sink);

          break;

        case DataSource::kSymbols:

        case DataSource::kRawSymbols:

        case DataSource::kShortSymbols:

        case DataSource::kFullSymbols:

          // TODO(mj): Generate symbols from debug info, exports, imports, tls

          // data, relocations, resources ...

        case DataSource::kArchiveMembers:

        case DataSource::kCompileUnits:

        case DataSource::kInlines:

        default:

          THROW("PE doesn't support this data source");

      }

      AddCatchAll(*pe_file, sink);

    }

  }

  bool GetDisassemblyInfo(string_view /*symbol*/, DataSource /*symbol_source*/,

                          DisassemblyInfo* /*info*/) const override {

    WARN("PE files do not support disassembly yet");

    return false;

  }

 protected:

  std::unique_ptr<pe::PeFile> pe_file;

};

bool ReadMagic(const string_view& data) {

  // If the size is smaller than a dos header, it cannot be a PE file, right?

  if (data.size() < sizeof(pe_dos_header)) {

    return false;

  }

  uint16_t magic;

  memcpy(&magic, data.data(), sizeof(magic));

  return magic == dos_magic;

}

}  // namespace pe

std::unique_ptr<ObjectFile> TryOpenPEFile(std::unique_ptr<InputFile>& file) {

  // Do not bother creating an object if the first magic is not even there

  if (pe::ReadMagic(file->data())) {

    std::unique_ptr<pe::PeFile> pe(new pe::PeFile(file->data()));

    if (pe->IsOpen()) {

      return std::unique_ptr<ObjectFile>(

          new pe::PEObjectFile(std::move(file), std::move(pe)));

    }

  }

  return nullptr;

}

}  // namespace bloaty