Reimplement AllocateNearTo for 64bit.

sandbox/win/src/sandbox_nt_util_unittest.cc

Index: sandbox/win/src/sandbox_nt_util_unittest.cc
diff --git a/sandbox/win/src/sandbox_nt_util_unittest.cc b/sandbox/win/src/sandbox_nt_util_unittest.cc
index 0fbea6680247323a351574b8f9dfdb31f00c7685..b916e3d5ea247c181fb55d9cfa4405dfb3264759 100644
--- a/sandbox/win/src/sandbox_nt_util_unittest.cc
+++ b/sandbox/win/src/sandbox_nt_util_unittest.cc
@@ -2,7 +2,9 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include <memory>
 #include <windows.h>
+#include <vector>
 
 #include "base/win/scoped_handle.h"
 #include "base/win/scoped_process_information.h"
@@ -43,5 +45,151 @@ TEST(SandboxNtUtil, IsSameProcessDifferentProcess) {
   EXPECT_TRUE(TerminateProcess(process_info.process_handle(), 0));
 }
 
+#if defined(_WIN64)
+struct VirtualMemDeleter {
+  void operator()(char* p) { ::VirtualFree(p, 0, MEM_RELEASE); }
+};
+
+typedef std::unique_ptr<char, VirtualMemDeleter> unique_ptr_vmem;
+
+void AllocateBlock(SIZE_T size,
+                   SIZE_T free_size,
+                   char** base_address,
+                   std::vector<unique_ptr_vmem>* mem_range) {
+  unique_ptr_vmem ptr(static_cast<char*>(::VirtualAlloc(
+      *base_address, size - free_size, MEM_RESERVE, PAGE_READWRITE)));
+  ASSERT_NE(nullptr, ptr.get());
+  mem_range->push_back(std::move(ptr));
+  *base_address += size;
+}
+
+#define KIB(x) ((x)*1024ULL)
+#define MIB(x) (KIB(x) * 1024ULL)
+#define GIB(x) (MIB(x) * 1024ULL)
+// Construct a basic memory layout to do the test. We reserve first to get a
+// base address then reallocate with the following pattern.
+// |512MiB-64KiB Free|512MiB-128Kib Free|512MiB-256Kib Free|512MiB+512KiB Free|
+// The purpose of this is leave a couple of free memory regions within a 2GiB
+// block of reserved memory that we can test the searching allocator.
+void AllocateTestRange(std::vector<unique_ptr_vmem>* mem_range) {
+  // Ensure we preallocate enough space in the vector to prevent unexpected
+  // allocations.
+  mem_range->reserve(5);
+  SIZE_T total_size =
+      MIB(512) + MIB(512) + MIB(512) + MIB(512) + KIB(512) + KIB(64);
+  unique_ptr_vmem ptr(static_cast<char*>(
+      ::VirtualAlloc(nullptr, total_size, MEM_RESERVE, PAGE_READWRITE)));
+  ASSERT_NE(nullptr, ptr.get());
+  char* base_address = ptr.get();
+  char* orig_base = base_address;
+  ptr.reset();
+  AllocateBlock(MIB(512), KIB(64), &base_address, mem_range);
+  AllocateBlock(MIB(512), KIB(128), &base_address, mem_range);
+  AllocateBlock(MIB(512), KIB(256), &base_address, mem_range);
+  AllocateBlock(MIB(512) + KIB(512), KIB(512), &base_address, mem_range);
+  // Allocate a memory block at end to act as an upper bound.
+  AllocateBlock(KIB(64), 0, &base_address, mem_range);
+  ASSERT_EQ(total_size, static_cast<SIZE_T>(base_address - orig_base));
+}
+
+// Test we can allocate appropriate blocks.
+void TestAlignedRange(char* base_address) {
+  unique_ptr_vmem ptr_256k(new (sandbox::NT_PAGE, base_address) char[KIB(256)]);
+  EXPECT_EQ(base_address + GIB(1) + MIB(512) - KIB(256), ptr_256k.get());
+  unique_ptr_vmem ptr_64k(new (sandbox::NT_PAGE, base_address) char[KIB(64)]);
+  EXPECT_EQ(base_address + MIB(512) - KIB(64), ptr_64k.get());
+  unique_ptr_vmem ptr_128k(new (sandbox::NT_PAGE, base_address) char[KIB(128)]);
+  EXPECT_EQ(base_address + GIB(1) - KIB(128), ptr_128k.get());
+  // We will have run out of space here so should also fail.
+  unique_ptr_vmem ptr_64k_noalloc(
+      new (sandbox::NT_PAGE, base_address) char[KIB(64)]);
+  EXPECT_EQ(nullptr, ptr_64k_noalloc.get());
+}
+
+// Test the 512k block which exists at the end of the maximum allocation
+// boundary.
+void Test512kBlock(char* base_address) {
+  // This should fail as it'll just be out of range.
+  unique_ptr_vmem ptr_512k_noalloc(
+      new (sandbox::NT_PAGE, base_address) char[KIB(512)]);
+  EXPECT_EQ(nullptr, ptr_512k_noalloc.get());
+  // Check that moving base address we can allocate the 512k block.
+  unique_ptr_vmem ptr_512k(
+      new (sandbox::NT_PAGE, base_address + GIB(1)) char[KIB(512)]);
+  EXPECT_EQ(base_address + GIB(2), ptr_512k.get());
+  // Free pointer first.
+  ptr_512k.reset();
+  ptr_512k.reset(new (sandbox::NT_PAGE, base_address + GIB(2)) char[KIB(512)]);
+  EXPECT_EQ(base_address + GIB(2), ptr_512k.get());
+}
+
+// Test we can allocate appropriate blocks even when starting at an unaligned
+// address.
+void TestUnalignedRange(char* base_address) {
+  char* unaligned_base = base_address + 123456;
+  unique_ptr_vmem ptr_256k(
+      new (sandbox::NT_PAGE, unaligned_base) char[KIB(256)]);
+  EXPECT_EQ(base_address + GIB(1) + MIB(512) - KIB(256), ptr_256k.get());
+  unique_ptr_vmem ptr_64k(new (sandbox::NT_PAGE, unaligned_base) char[KIB(64)]);
+  EXPECT_EQ(base_address + MIB(512) - KIB(64), ptr_64k.get());
+  unique_ptr_vmem ptr_128k(
+      new (sandbox::NT_PAGE, unaligned_base) char[KIB(128)]);
+  EXPECT_EQ(base_address + GIB(1) - KIB(128), ptr_128k.get());
+}
+
+// Test maximum number of available allocations within the predefined pattern.
+void TestMaxAllocations(char* base_address) {
+  // There's only 7 64k blocks in the first 2g which we can fill.
+  unique_ptr_vmem ptr_1(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_1.get());
+  unique_ptr_vmem ptr_2(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_2.get());
+  unique_ptr_vmem ptr_3(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_3.get());
+  unique_ptr_vmem ptr_4(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_4.get());
+  unique_ptr_vmem ptr_5(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_5.get());
+  unique_ptr_vmem ptr_6(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_6.get());
+  unique_ptr_vmem ptr_7(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_NE(nullptr, ptr_7.get());
+  unique_ptr_vmem ptr_8(new (sandbox::NT_PAGE, base_address) char[1]);
+  EXPECT_EQ(nullptr, ptr_8.get());
+}
+
+// Test extreme allocations we know should fail.
+void TestExtremes() {
+  unique_ptr_vmem ptr_null(new (sandbox::NT_PAGE, nullptr) char[1]);
+  EXPECT_EQ(nullptr, ptr_null.get());
+  unique_ptr_vmem ptr_too_large(
+      new (sandbox::NT_PAGE, reinterpret_cast<void*>(0x1000000)) char[GIB(4)]);
+  EXPECT_EQ(nullptr, ptr_too_large.get());
+  unique_ptr_vmem ptr_overflow(
+      new (sandbox::NT_PAGE, reinterpret_cast<void*>(SIZE_MAX)) char[1]);
+  EXPECT_EQ(nullptr, ptr_overflow.get());
+  unique_ptr_vmem ptr_invalid(new (
+      sandbox::NT_PAGE, reinterpret_cast<void*>(SIZE_MAX - 0x1000000)) char[1]);
+  EXPECT_EQ(nullptr, ptr_invalid.get());
+}
+
+// Test nearest allocator, only do this for 64 bit. We test through the exposed
+// new operator as we can't call the AllocateNearTo function directly.
+TEST(SandboxNtUtil, NearestAllocator) {
+  InitGlobalNt();
+  std::vector<unique_ptr_vmem> mem_range;
+  AllocateTestRange(&mem_range);
+  ASSERT_LT(0U, mem_range.size());
+  char* base_address = static_cast<char*>(mem_range[0].get());
+
+  TestAlignedRange(base_address);
+  Test512kBlock(base_address);
+  TestUnalignedRange(base_address);
+  TestMaxAllocations(base_address);
+  TestExtremes();
+}
+
+#endif  // defined(_WIN64)
+
 }  // namespace
 }  // namespace sandbox