From 7cafb0e9737f81d6930b8c52804d55c0ff773907 Mon Sep 17 00:00:00 2001
From: Bryan Bernhart <bryan.bernhart@intel.com>
Date: Mon, 25 Jul 2022 16:23:43 -0700
Subject: [PATCH] Fix overflow of slabs during sub-allocation.

Also, enables asserts for standalone fuzzing.
---
 .github/workflows/win_clang_rel_x64.yaml         |  2 +-
 build_overrides/gpgmm_features.gni               |  4 ++++
 src/fuzzers/D3D12ResourceAllocatorFuzzer.cpp     |  2 ++
 src/gpgmm/common/BUILD.gn                        |  2 +-
 src/gpgmm/common/SlabMemoryAllocator.cpp         |  2 ++
 src/tests/unittests/SlabMemoryAllocatorTests.cpp | 11 +++++++++++
 6 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/win_clang_rel_x64.yaml b/.github/workflows/win_clang_rel_x64.yaml
index 81e0d05ec..10ad53b1e 100644
--- a/.github/workflows/win_clang_rel_x64.yaml
+++ b/.github/workflows/win_clang_rel_x64.yaml
@@ -161,7 +161,7 @@ jobs:
         set "PATH=%CD%\..\depot_tools;%PATH%"
         set "DEPOT_TOOLS_WIN_TOOLCHAIN=0"
         cd test
-        gn gen out\Fuzzer --args="is_debug=false use_libfuzzer=true is_asan=true gpgmm_enable_assert_on_warning=true gpgmm_enable_device_leak_checks=true"
+        gn gen out\Fuzzer --args="is_debug=false gpgmm_use_fuzzing_engine=true use_libfuzzer=true is_asan=true gpgmm_enable_assert_on_warning=true gpgmm_enable_device_leak_checks=true"
 
     - name: Build fuzzer for main branch (with patch)
       shell: cmd
diff --git a/build_overrides/gpgmm_features.gni b/build_overrides/gpgmm_features.gni
index c65355944..56a915b61 100644
--- a/build_overrides/gpgmm_features.gni
+++ b/build_overrides/gpgmm_features.gni
@@ -30,6 +30,10 @@ declare_args() {
   # Sets -dGPGMM_ENABLE_ASSERTS
   gpgmm_always_assert = false
 
+  # Enable GPGMM for fuzzing in standalone builds.
+  # Sets -dGPGMM_USE_FUZZING_ENGINE
+  gpgmm_use_fuzzing_engine = false
+
   # Enables event tracing even in release builds.
   # Sets -dGPGMM_FORCE_TRACING
   gpgmm_force_tracing = false
diff --git a/src/fuzzers/D3D12ResourceAllocatorFuzzer.cpp b/src/fuzzers/D3D12ResourceAllocatorFuzzer.cpp
index 526d1150c..1c6abc740 100644
--- a/src/fuzzers/D3D12ResourceAllocatorFuzzer.cpp
+++ b/src/fuzzers/D3D12ResourceAllocatorFuzzer.cpp
@@ -68,6 +68,8 @@ extern "C" int LLVMFuzzerInitialize(int* argc, char*** argv) {
     }
     allocatorDesc.ResourceHeapTier = options.ResourceHeapTier;
 
+    allocatorDesc.MinLogLevel = D3D12_MESSAGE_SEVERITY_MESSAGE;
+
     if (FAILED(
             gpgmm::d3d12::ResourceAllocator::CreateAllocator(allocatorDesc, &gResourceAllocator))) {
         return 0;
diff --git a/src/gpgmm/common/BUILD.gn b/src/gpgmm/common/BUILD.gn
index 3a2fae222..ff4463b71 100644
--- a/src/gpgmm/common/BUILD.gn
+++ b/src/gpgmm/common/BUILD.gn
@@ -29,7 +29,7 @@ if (build_with_chromium) {
 if (build_with_chromium) {
   import("//build/config/sanitizers/sanitizers.gni")
 } else {
-  use_fuzzing_engine = false
+  use_fuzzing_engine = gpgmm_use_fuzzing_engine
 }
 
 ###############################################################################
diff --git a/src/gpgmm/common/SlabMemoryAllocator.cpp b/src/gpgmm/common/SlabMemoryAllocator.cpp
index df84ceaf5..e894cba62 100644
--- a/src/gpgmm/common/SlabMemoryAllocator.cpp
+++ b/src/gpgmm/common/SlabMemoryAllocator.cpp
@@ -99,6 +99,7 @@ namespace gpgmm {
         ASSERT(mSlabGrowthFactor >= 1);
         ASSERT(IsAligned(mMaxSlabSize, mSlabAlignment));
         ASSERT(IsAligned(mMinSlabSize, mSlabAlignment));
+        ASSERT(blockSize <= mMaxSlabSize);
     }
 
     SlabMemoryAllocator::~SlabMemoryAllocator() {
@@ -461,6 +462,7 @@ namespace gpgmm {
         GPGMM_INVALID_IF(!ValidateRequest(request));
 
         const uint64_t blockSize = AlignTo(request.SizeInBytes, request.Alignment);
+        GPGMM_INVALID_IF(blockSize > mMaxSlabSize);
 
         // Create a slab allocator for the new entry.
         auto entry =
diff --git a/src/tests/unittests/SlabMemoryAllocatorTests.cpp b/src/tests/unittests/SlabMemoryAllocatorTests.cpp
index a7599d706..ea3641375 100644
--- a/src/tests/unittests/SlabMemoryAllocatorTests.cpp
+++ b/src/tests/unittests/SlabMemoryAllocatorTests.cpp
@@ -684,6 +684,17 @@ TEST_F(SlabMemoryAllocatorTests, SlabGrowthLimit) {
 
 class SlabCacheAllocatorTests : public SlabMemoryAllocatorTests {};
 
+// Attempting to allocate a block greater then the slab should always fail.
+TEST_F(SlabCacheAllocatorTests, SlabOversized) {
+    constexpr uint64_t kMaxSlabSize = 256;
+    constexpr uint64_t kMinSlabSize = 16;
+    SlabCacheAllocator allocator(kMaxSlabSize, kMinSlabSize, kDefaultSlabAlignment,
+                                 kDefaultSlabFragmentationLimit, kNoSlabPrefetchAllowed,
+                                 kNoSlabGrowthFactor, std::make_unique<DummyMemoryAllocator>());
+
+    EXPECT_EQ(allocator.TryAllocateMemory(CreateBasicRequest(kMaxSlabSize + 1, 1)), nullptr);
+}
+
 TEST_F(SlabCacheAllocatorTests, SingleSlabMultipleSize) {
     constexpr uint64_t kMaxSlabSize = 256;
     constexpr uint64_t kSlabSize = 0;  // deduce slab size from allocation size.