Add magic number and canary checks to malloc and free implementation

include/malloc_from_scratch/memory_internal.h

@@ -11,6 +11,7 @@ namespace internal
 
 struct MemoryBlock
 {
+    size_t magic_;
     size_t size_;
     bool allocated_;
     MemoryBlock* next_;
@@ -22,20 +23,24 @@ inline size_t total_memory_allocated = 0;
 inline pthread_mutex_t allocator_mutex = PTHREAD_MUTEX_INITIALIZER;
 
 constexpr size_t BLOCK_METADATA_SIZE = sizeof(MemoryBlock);
-constexpr size_t BLOCK_SIZE = 16;
 constexpr size_t CHUNK_SIZE = 65536;
+constexpr size_t MAGIC_NUMBER = 0xDEADBEEF;
+constexpr size_t CANARY_VALUE = 0xC0FFEE00;
+constexpr size_t CANARY_SIZE = sizeof(size_t);
 
 // malloc
 void* increaseHeap(size_t size);
 MemoryBlock* findLargeEnoughFreeMemoryBlock(MemoryBlock** block_list_head, size_t size);
 void* splitFreeMemoryBlockIfPossible(MemoryBlock* new_block, size_t size);
 void* getMemoryBlockSplitAddress(MemoryBlock* new_block, size_t size);
 void insertMemoryBlockAtEnd(MemoryBlock** block_list_head, MemoryBlock* new_block);
+void setCanary(MemoryBlock* block);
 
 // free
 void decreaseHeap(MemoryBlock* block_heap_end);
 void mergeFreeMemoryBlocks();
 void getLastMemoryBlock(MemoryBlock** block_list_tail, MemoryBlock** block_previous_from_last);
+bool isBlockCorrupted(MemoryBlock* block);
 
 // helpers
 void* getPayloadAddress(MemoryBlock* block);
@@ -44,6 +49,8 @@ MemoryBlock* getMemoryBlockFromAddress(void* address);
 size_t getSizeOfAllocatedMemoryBlock(MemoryBlock* block);
 bool isPointerInHeap(void* ptr);
 void* getErrorCodeInVoidPtr(size_t error_code);
+bool isValidBlock(MemoryBlock* block);
+bool checkCanary(MemoryBlock* block);
 
 // Test helper functions to inspect allocator state
 inline size_t getTotalUsedMemory() { return total_memory_allocated; }

src/free.cpp

@@ -20,15 +20,20 @@ void free(void* ptr)
     }
 
     internal::MemoryBlock* block_to_free = internal::getMetadata(ptr);
-    if (block_to_free == nullptr || block_to_free->allocated_ == false)
+    if (!internal::isValidBlock(block_to_free))
     {
-        if (block_to_free)
+        if (internal::isBlockCorrupted(block_to_free))
         {
             exit(-1);
         }
         return;
     }
 
+    if (!internal::checkCanary(block_to_free))
+    {
+        exit(-1);
+    }
+
     internal::total_memory_allocated -= block_to_free->size_;
     block_to_free->allocated_ = false;
 
@@ -91,7 +96,7 @@ void mergeFreeMemoryBlocks()
     {
         if (current->allocated_ == false && current->next_->allocated_ == false)
         {
-            current->size_ += BLOCK_METADATA_SIZE + current->next_->size_;
+            current->size_ += BLOCK_METADATA_SIZE + CANARY_SIZE + current->next_->size_;
             current->next_ = current->next_->next_;
         }
         else
@@ -116,5 +121,7 @@ void getLastMemoryBlock(MemoryBlock** block_list_tail, MemoryBlock** block_previ
     *block_previous_from_last = previous;
 }
 
+bool isBlockCorrupted(MemoryBlock* block) { return (block && block->magic_ != MAGIC_NUMBER); }
+
 } // namespace internal
 } // namespace mem

src/malloc.cpp

@@ -37,7 +37,7 @@ void* malloc(size_t size)
     }
     else
     {
-        total_size = internal::BLOCK_METADATA_SIZE + size;
+        total_size = internal::BLOCK_METADATA_SIZE + size + internal::CANARY_SIZE;
         is_frame_used = false;
     }
 
@@ -50,6 +50,7 @@ void* malloc(size_t size)
     internal::MemoryBlock* new_block = internal::getMemoryBlockFromAddress(new_memory_allocation);
     if (is_frame_used)
     {
+        new_block->magic_ = internal::MAGIC_NUMBER;
         new_block->size_ = internal::CHUNK_SIZE;
         new_block->allocated_ = false;
         new_block->next_ = nullptr;
@@ -59,10 +60,12 @@ void* malloc(size_t size)
     }
     else
     {
+        new_block->magic_ = internal::MAGIC_NUMBER;
         new_block->size_ = size;
         new_block->allocated_ = true;
         new_block->next_ = nullptr;
         internal::insertMemoryBlockAtEnd(&internal::block_list_head, new_block);
+        internal::setCanary(new_block);
         internal::total_memory_allocated += size;
     }
 
@@ -86,14 +89,16 @@ void* increaseHeap(size_t size)
 MemoryBlock* findLargeEnoughFreeMemoryBlock(MemoryBlock** block_list_head, size_t size)
 {
     MemoryBlock* current = *block_list_head;
+    size_t size_needed = size + CANARY_SIZE;
     while (current != nullptr)
     {
-        if (current->allocated_ == false && current->size_ >= size)
+        if (current->allocated_ == false && current->size_ >= size_needed)
         {
             return current;
         }
         current = current->next_;
     }
+
     return nullptr;
 }
 
@@ -104,12 +109,14 @@ void* splitFreeMemoryBlockIfPossible(MemoryBlock* new_block, size_t size)
         return nullptr;
     }
 
-    size_t remaining_size = new_block->size_ - size;
+    size_t size_with_canary = size + CANARY_SIZE;
+    size_t remaining_size = new_block->size_ - size_with_canary;
     char one_byte_payload_size_requirement = 1;
     if (remaining_size >= BLOCK_METADATA_SIZE + one_byte_payload_size_requirement)
     {
         void* split_address = getMemoryBlockSplitAddress(new_block, size);
         MemoryBlock* new_temp_block = getMemoryBlockFromAddress(split_address);
+        new_temp_block->magic_ = MAGIC_NUMBER;
         new_temp_block->size_ = remaining_size - BLOCK_METADATA_SIZE;
         new_temp_block->allocated_ = false;
         new_temp_block->next_ = new_block->next_;
@@ -120,17 +127,19 @@ void* splitFreeMemoryBlockIfPossible(MemoryBlock* new_block, size_t size)
     }
     else
     {
-        new_block->size_ = new_block->size_;
+        new_block->size_ = new_block->size_ - CANARY_SIZE;
         new_block->allocated_ = true;
     }
 
+    setCanary(new_block);
     total_memory_allocated += new_block->size_;
+
     return getPayloadAddress(new_block);
 }
 
 void* getMemoryBlockSplitAddress(MemoryBlock* new_block, size_t size)
 {
-    return (reinterpret_cast<char*>(new_block) + BLOCK_METADATA_SIZE + size);
+    return (reinterpret_cast<char*>(new_block) + BLOCK_METADATA_SIZE + size + CANARY_SIZE);
 }
 
 void insertMemoryBlockAtEnd(MemoryBlock** block_list_head, MemoryBlock* new_block)
@@ -150,6 +159,18 @@ void insertMemoryBlockAtEnd(MemoryBlock** block_list_head, MemoryBlock* new_bloc
     }
 }
 
+void setCanary(MemoryBlock* block)
+{
+    if (!block || block->allocated_ == false)
+    {
+        return;
+    }
+
+    size_t* canary = reinterpret_cast<size_t*>(reinterpret_cast<char*>(block) +
+                                               BLOCK_METADATA_SIZE + block->size_);
+    *canary = CANARY_VALUE;
+}
+
 // helpers
 void* getPayloadAddress(MemoryBlock* block)
 {
@@ -180,15 +201,28 @@ size_t getSizeOfAllocatedMemoryBlock(MemoryBlock* block)
 bool isPointerInHeap(void* ptr)
 {
     void* current_program_break = sbrk(0);
-    if (ptr < heap_start || ptr >= current_program_break)
+    return (ptr >= heap_start && ptr < current_program_break);
+}
+
+void* getErrorCodeInVoidPtr(size_t error_code) { return reinterpret_cast<void*>(error_code); }
+
+bool isValidBlock(MemoryBlock* block)
+{
+    return block != nullptr && block->magic_ == MAGIC_NUMBER && block->allocated_ == true;
+}
+
+bool checkCanary(MemoryBlock* block)
+{
+    if (!block)
     {
         return false;
     }
 
-    return true;
-}
+    size_t* canary = reinterpret_cast<size_t*>(reinterpret_cast<char*>(block) +
+                                               BLOCK_METADATA_SIZE + block->size_);
 
-void* getErrorCodeInVoidPtr(size_t error_code) { return reinterpret_cast<void*>(error_code); }
+    return (*canary == CANARY_VALUE);
+}
 
 } // namespace internal
 } // namespace mem