diff --git a/include/malloc_from_scratch/memory_internal.h b/include/malloc_from_scratch/memory_internal.h index 53466a0..e72a982 100644 --- a/include/malloc_from_scratch/memory_internal.h +++ b/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,8 +23,10 @@ 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); @@ -31,11 +34,13 @@ MemoryBlock* findLargeEnoughFreeMemoryBlock(MemoryBlock** block_list_head, 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; } diff --git a/src/free.cpp b/src/free.cpp index 0f4e2a7..396d21f 100644 --- a/src/free.cpp +++ b/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 diff --git a/src/malloc.cpp b/src/malloc.cpp index 40e66dc..e5d7827 100644 --- a/src/malloc.cpp +++ b/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(new_block) + BLOCK_METADATA_SIZE + size); + return (reinterpret_cast(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(reinterpret_cast(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(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(reinterpret_cast(block) + + BLOCK_METADATA_SIZE + block->size_); -void* getErrorCodeInVoidPtr(size_t error_code) { return reinterpret_cast(error_code); } + return (*canary == CANARY_VALUE); +} } // namespace internal } // namespace mem