// Copyright 2018 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome_elf/pe_image_safe/pe_image_safe.h"

#include "base/files/file_util.h"
#include "base/path_service.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace pe_image_safe {
namespace {

constexpr wchar_t kPeFile[] = L"chrome_elf.dll";
constexpr DWORD kPageSize = 4096;

struct CompareData {
  bool bits_64;
  DWORD size_of_image;
  DWORD time_date_stamp;
  PIMAGE_DOS_HEADER dos_header;
  PIMAGE_NT_HEADERS nt_headers;
  PIMAGE_FILE_HEADER file_header;
  PIMAGE_OPTIONAL_HEADER optional_header;
};

// Raw collection of some PE header data, without using pe_image_safe.
// This function assumes full headers from a legitimate PE image.
bool GetComparisonData(char* buffer, CompareData* data) {
  data->dos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(buffer);
  if (data->dos_header->e_magic != IMAGE_DOS_SIGNATURE)
    return false;

  data->nt_headers = reinterpret_cast<PIMAGE_NT_HEADERS>(
      reinterpret_cast<char*>(data->dos_header) + data->dos_header->e_lfanew);
  if (data->nt_headers->Signature != IMAGE_NT_SIGNATURE)
    return false;

  data->file_header = &data->nt_headers->FileHeader;
  data->time_date_stamp = data->nt_headers->FileHeader.TimeDateStamp;
  data->optional_header = &data->nt_headers->OptionalHeader;

  if (data->optional_header->Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
    PIMAGE_OPTIONAL_HEADER64 optional_header =
        reinterpret_cast<PIMAGE_OPTIONAL_HEADER64>(
            &data->nt_headers->OptionalHeader);
    data->bits_64 = true;
    data->size_of_image = optional_header->SizeOfImage;
  } else {
    PIMAGE_OPTIONAL_HEADER32 optional_header =
        reinterpret_cast<PIMAGE_OPTIONAL_HEADER32>(
            &data->nt_headers->OptionalHeader);
    data->bits_64 = false;
    data->size_of_image = optional_header->SizeOfImage;
  }

  return true;
}

//------------------------------------------------------------------------------

TEST(PEImageSafe, SanityTest) {
  // Open and read in a PE file.
  base::FilePath pe_path;
  EXPECT_TRUE(base::PathService::Get(base::DIR_EXE, &pe_path));
  pe_path = pe_path.Append(kPeFile);

  base::File file(pe_path, base::File::FLAG_OPEN | base::File::FLAG_READ);
  ASSERT_TRUE(file.IsValid());

  std::vector<char> buffer;
  buffer.resize(kPageSize);
  ASSERT_EQ(file.Read(0, &buffer[0], kPageSize), static_cast<int>(kPageSize));
  file.Close();

  // Grab some key data out of the pe headers first, NOT using pe_image_safe.
  CompareData data_for_comparison = {};
  ASSERT_TRUE(GetComparisonData(buffer.data(), &data_for_comparison));

  // Apply scaffolding.
  PEImageSafe pe_image(buffer.data(), static_cast<DWORD>(buffer.size()));
  PIMAGE_DOS_HEADER dos_header = pe_image.GetDosHeader();
  EXPECT_TRUE(dos_header);
  EXPECT_EQ(dos_header, data_for_comparison.dos_header);

  PIMAGE_FILE_HEADER file_header = pe_image.GetFileHeader();
  EXPECT_TRUE(file_header);
  EXPECT_EQ(file_header, data_for_comparison.file_header);
  EXPECT_EQ(file_header->TimeDateStamp, data_for_comparison.time_date_stamp);

  BYTE* optional_header = pe_image.GetOptionalHeader();
  EXPECT_TRUE(optional_header);
  EXPECT_EQ(optional_header,
            reinterpret_cast<BYTE*>(data_for_comparison.optional_header));

  pe_image_safe::ImageBitness bitness = pe_image.GetImageBitness();
  EXPECT_NE(bitness, pe_image_safe::ImageBitness::kUnknown);
  EXPECT_TRUE(((bitness == pe_image_safe::ImageBitness::k64) &&
               data_for_comparison.bits_64) ||
              ((bitness == pe_image_safe::ImageBitness::k32) &&
               !data_for_comparison.bits_64));

  if (bitness == pe_image_safe::ImageBitness::k64) {
    PIMAGE_OPTIONAL_HEADER64 optional_header64 =
        reinterpret_cast<PIMAGE_OPTIONAL_HEADER64>(optional_header);
    EXPECT_EQ(optional_header64->SizeOfImage,
              data_for_comparison.size_of_image);
  } else {
    PIMAGE_OPTIONAL_HEADER32 optional_header32 =
        reinterpret_cast<PIMAGE_OPTIONAL_HEADER32>(optional_header);
    EXPECT_EQ(optional_header32->SizeOfImage,
              data_for_comparison.size_of_image);
  }
}

}  // namespace
}  // namespace pe_image_safe