// Copyright 2012 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 "syzygy/agent/asan/rtl_impl.h"

#include <windows.h>  // NOLINT

#include "base/rand_util.h"
#include "gtest/gtest.h"
#include "syzygy/agent/asan/runtime.h"
#include "syzygy/agent/asan/shadow.h"
#include "syzygy/agent/asan/unittest_util.h"
#include "syzygy/core/unittest_util.h"

namespace agent {
namespace asan {

namespace {

class AsanRtlImplTest : public testing::TestWithAsanLogger {
 public:
  AsanRtlImplTest() : heap_(NULL) {
  }

  void SetUp() override {
    testing::TestWithAsanLogger::SetUp();
    asan_runtime_.SetUp(runtime_command_line_);
    agent::asan::SetUpRtl(&asan_runtime_);
    heap_ = asan_HeapCreate(0, 0, 0);
    ASSERT_TRUE(heap_ != NULL);
  }

  void TearDown() override {
    if (heap_ != NULL) {
      asan_HeapDestroy(heap_);
      heap_ = NULL;
    }
    agent::asan::TearDownRtl();
    asan_runtime_.TearDown();
    testing::TestWithAsanLogger::TearDown();
  }

 protected:
  agent::asan::AsanRuntime asan_runtime_;

  // Arbitrary constant for all size limit.
  static const size_t kMaxAllocSize = 134584;

  // Scratch heap handle valid from SetUp to TearDown.
  HANDLE heap_;

  // The command line used to initialize the runtime. Empty by default but
  // can be override by the derived classes to enable some runtime flags
  // during setup.
  std::wstring runtime_command_line_;
};

}  // namespace

TEST_F(AsanRtlImplTest, CreateDestroy) {
  HANDLE heap = asan_HeapCreate(0, 0, 0);
  ASSERT_TRUE(heap != NULL);
  ASSERT_TRUE(asan_HeapDestroy(heap));
}

TEST_F(AsanRtlImplTest, Alloc) {
  for (uint32_t size = 10; size < kMaxAllocSize; size = size * 5 + 123) {
    void* mem = asan_HeapAlloc(heap_, 0, size);
    ASSERT_TRUE(mem != NULL);
    ::memset(mem, '\0', size);

    uint32_t new_size = size;
    while (new_size == size)
      new_size = base::RandInt(size / 2, size * 2);

    void* new_mem = asan_HeapReAlloc(heap_, 0, mem, new_size);
    ASSERT_TRUE(new_mem != NULL);
    ASSERT_NE(mem, new_mem);

    ASSERT_TRUE(asan_HeapFree(heap_, 0, new_mem));
  }
}

TEST_F(AsanRtlImplTest, Size) {
  for (size_t size = 10; size < kMaxAllocSize; size = size * 5 + 123) {
    void* mem = asan_HeapAlloc(heap_, 0, size);
    ASSERT_TRUE(mem != NULL);
    ASSERT_EQ(size, asan_HeapSize(heap_, 0, mem));
    ASSERT_TRUE(asan_HeapFree(heap_, 0, mem));
  }
}

TEST_F(AsanRtlImplTest, Validate) {
  for (size_t size = 10; size < kMaxAllocSize; size = size * 5 + 123) {
    void* mem = asan_HeapAlloc(heap_, 0, size);
    ASSERT_TRUE(mem != NULL);
    ASSERT_TRUE(asan_HeapValidate(heap_, 0, mem));
    ASSERT_TRUE(asan_HeapFree(heap_, 0, mem));
  }
}

TEST_F(AsanRtlImplTest, Compact) {
  // Compact isn't supported by the current heap implementation and should
  // always return 0.
  ASSERT_EQ(0U, asan_HeapCompact(heap_, 0));
}

TEST_F(AsanRtlImplTest, LockUnlock) {
  // We can't really test these, aside from not crashing.
  ASSERT_TRUE(asan_HeapLock(heap_));
  ASSERT_TRUE(asan_HeapUnlock(heap_));
}

TEST_F(AsanRtlImplTest, Walk) {
  // Walk isn't supported by the current heap implementation.
  PROCESS_HEAP_ENTRY entry = {};
  ASSERT_FALSE(asan_HeapWalk(heap_, &entry));
}

TEST_F(AsanRtlImplTest, SetQueryInformation) {
  ULONG compat_flag = ULONG_MAX;
  SIZE_T ret = 0;
  // QueryInformation isn't supported by the current heap implementation and
  // should always return false.
  ASSERT_FALSE(
      asan_HeapQueryInformation(heap_, HeapCompatibilityInformation,
                                &compat_flag, sizeof(compat_flag), &ret));

  // Put the heap in LFH, which should always succeed, except when a debugger
  // is attached. When a debugger is attached, the heap is wedged in certain
  // debug settings.
  if (base::debug::BeingDebugged()) {
    LOG(WARNING) << "Can't test HeapProxy::SetInformation under debugger.";
    return;
  }

  // SetInformation isn't supported by the current heap implementation and
  // should always return true.
  ASSERT_TRUE(
      asan_HeapSetInformation(heap_, HeapCompatibilityInformation,
                              &compat_flag, sizeof(compat_flag)));
}

TEST_F(AsanRtlImplTest, SetInformationWithNullHeapPtr) {
  // The documentation of HeapSetInformation specify that the heap handle is
  // optional.
  ASSERT_TRUE(
      asan_HeapSetInformation(NULL, HeapEnableTerminationOnCorruption,
                              NULL, 0));
}

namespace {

// Derived class that defers the crash reporter initialization.
class AsanRtlImplTestCrashReporterInitialization : public AsanRtlImplTest {
 public:
  AsanRtlImplTestCrashReporterInitialization() {
    runtime_command_line_ = L"--defer_crash_reporter_initialization";
  }
};

}  // namespace

TEST_F(AsanRtlImplTestCrashReporterInitialization, InitializeCrashReporter) {
  EXPECT_FALSE(asan_runtime_.crash_reporter_initialized());
  asan_InitializeCrashReporter();
  EXPECT_TRUE(asan_runtime_.crash_reporter_initialized());
}

}  // namespace asan
}  // namespace agent