From 1fa9d7752c58dd9696e9ee69b786fa11be1317eb Mon Sep 17 00:00:00 2001
From: Teon Banek <teon.banek@memgraph.io>
Date: Mon, 26 Aug 2019 13:35:28 +0200
Subject: [PATCH] Fix a bug in Pool::Deallocate

Reviewers: mferencevic, ipaljak

Reviewed By: mferencevic

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D2325
---
 src/utils/memory.cpp        |  5 +++++
 tests/unit/utils_memory.cpp | 23 +++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/utils/memory.cpp b/src/utils/memory.cpp
index f9d894a72..b14e58e9e 100644
--- a/src/utils/memory.cpp
+++ b/src/utils/memory.cpp
@@ -220,10 +220,15 @@ void Pool::Deallocate(void *p) {
            ptr < reinterpret_cast<uintptr_t>(chunk.data + data_size);
   };
   auto deallocate_block_from_chunk = [this, p](Chunk *chunk) {
+    // NOTE: This check is not enough to cover all double-free issues.
+    CHECK(chunk->blocks_available < blocks_per_chunk_)
+        << "Deallocating more blocks than a chunk can contain, possibly a "
+           "double-free situation or we have a bug in the allocator.";
     // Link the block into the free-list
     auto *block = reinterpret_cast<unsigned char *>(p);
     *block = chunk->first_available_block_ix;
     chunk->first_available_block_ix = (block - chunk->data) / block_size_;
+    chunk->blocks_available++;
   };
   if (is_in_chunk(*last_dealloc_chunk_)) {
     deallocate_block_from_chunk(last_dealloc_chunk_);
diff --git a/tests/unit/utils_memory.cpp b/tests/unit/utils_memory.cpp
index 417408370..64fdd5e1a 100644
--- a/tests/unit/utils_memory.cpp
+++ b/tests/unit/utils_memory.cpp
@@ -303,6 +303,29 @@ TEST(PoolResource, AllocationWithOverflow) {
   }
 }
 
+TEST(PoolResource, BlockDeallocation) {
+  TestMemory test_mem;
+  const size_t max_blocks_per_chunk = 2U;
+  const size_t max_block_size = 64U;
+  utils::PoolResource mem(max_blocks_per_chunk, max_block_size, &test_mem);
+  auto *ptr = CheckAllocation(&mem, max_block_size);
+  test_mem.new_count_ = 0U;
+  // Do another allocation before deallocating `ptr`, so that we are sure that
+  // the chunk of 2 blocks is still alive and therefore `ptr` may be reused when
+  // it's deallocated. If we deallocate now, the implementation may choose to
+  // free the whole chunk, and we do not want that for the purposes of this
+  // test.
+  CheckAllocation(&mem, max_block_size);
+  EXPECT_EQ(test_mem.new_count_, 0U);
+  EXPECT_EQ(test_mem.delete_count_, 0U);
+  mem.Deallocate(ptr, max_block_size);
+  EXPECT_EQ(test_mem.delete_count_, 0U);
+  // CheckAllocation(&mem, max_block_size) will fail as PoolResource should
+  // reuse free blocks.
+  EXPECT_EQ(ptr, mem.Allocate(max_block_size));
+  EXPECT_EQ(test_mem.new_count_, 0U);
+}
+
 // NOLINTNEXTLINE(hicpp-special-member-functions)
 class ContainerWithAllocatorLast final {
  public: