Create a gist now

Instantly share code, notes, and snippets.

PartitionAlloc randomization of the freelist and double free protection
diff --git a/PartitionAlloc.cpp b/PartitionAlloc.cpp
index 1215713..3b249e2 100644
--- a/PartitionAlloc.cpp
+++ b/PartitionAlloc.cpp
@@ -32,6 +32,7 @@
#include "PartitionAlloc.h"
#include <string.h>
+#include <vector>
#ifndef NDEBUG
#include <stdio.h>
@@ -142,6 +143,10 @@ void partitionAllocInit(PartitionRoot* root, size_t numBuckets, size_t maxAlloca
{
parititonAllocBaseInit(root);
+ // Reseed the RNG using time() with some bits
+ // from the current random pool
+ srand(time(NULL) + (rand() % 100000));
+
root->numBuckets = numBuckets;
root->maxAllocation = maxAllocation;
size_t i;
@@ -564,16 +569,31 @@ static ALWAYS_INLINE char* partitionPageAllocAndFillFreelist(PartitionPage* page
page->numUnprovisionedSlots = numSlots;
page->numAllocatedSlots++;
+ std::vector<PartitionFreelistEntry *> vfreelist;
+
if (LIKELY(numNewFreelistEntries)) {
char* freelistPointer = firstFreelistPointer;
PartitionFreelistEntry* entry = reinterpret_cast<PartitionFreelistEntry*>(freelistPointer);
page->freelistHead = entry;
+ vfreelist.push_back(entry);
+
while (--numNewFreelistEntries) {
freelistPointer += size;
- PartitionFreelistEntry* nextEntry = reinterpret_cast<PartitionFreelistEntry*>(freelistPointer);
+ vfreelist.push_back(reinterpret_cast<PartitionFreelistEntry*>(freelistPointer));
+ }
+
+ std::random_shuffle(vfreelist.begin(), vfreelist.end());
+
+ entry = vfreelist[vfreelist.size()-1];
+ vfreelist.pop_back();
+ page->freelistHead = entry;
+
+ for(int i=0;i<vfreelist.size();i++) {
+ PartitionFreelistEntry* nextEntry = vfreelist[i];
entry->next = partitionFreelistMask(nextEntry);
entry = nextEntry;
}
+
entry->next = partitionFreelistMask(0);
} else {
page->freelistHead = 0;
@@ -842,9 +862,30 @@ void* partitionAllocSlowPath(PartitionRootBase* root, int flags, size_t size, Pa
// usable freelist head.
if (LIKELY(newPage->freelistHead != nullptr)) {
PartitionFreelistEntry* entry = newPage->freelistHead;
- PartitionFreelistEntry* newHead = partitionFreelistMask(entry->next);
- newPage->freelistHead = newHead;
+ PartitionFreelistEntry* z = entry;
+
+ while(entry) {
+ if((rand() % 10) == 1) {
+ break;
+ }
+
+ z = entry;
+ entry = partitionFreelistMask(entry->next);
+ }
+
+ if(entry == nullptr) {
+ entry = newPage->freelistHead;
+ }
+
+ if(entry == newPage->freelistHead) {
+ PartitionFreelistEntry* newHead = partitionFreelistMask(entry->next);
+ newPage->freelistHead = newHead;
+ } else {
+ z->next = entry->next;
+ }
+
newPage->numAllocatedSlots++;
+
return entry;
}
// Otherwise, we need to build the freelist.
diff --git a/PartitionAlloc.h b/PartitionAlloc.h
index 2458d5a..7b13e40 100644
--- a/PartitionAlloc.h
+++ b/PartitionAlloc.h
@@ -1,3 +1,6 @@
+#include <stdio.h>
+#include <stdlib.h>
+
/*
* Copyright (C) 2013 Google Inc. All rights reserved.
*
@@ -579,13 +582,37 @@ ALWAYS_INLINE void* partitionBucketAlloc(PartitionRootBase* root, int flags, siz
ASSERT(page->numAllocatedSlots >= 0);
void* ret = page->freelistHead;
if (LIKELY(ret != 0)) {
+ PartitionFreelistEntry* t = page->freelistHead;
+ PartitionFreelistEntry* z = t;
+
+ while(t) {
+ if((rand() % 10) == 1) {
+ break;
+ }
+
+ z = t;
+ t = partitionFreelistMask(t->next);
+ }
+
+ if(t) {
+ ret = t;
+ } else {
+ ret = page->freelistHead;
+ }
+
// If these asserts fire, you probably corrupted memory.
ASSERT(partitionPointerIsValid(ret));
// All large allocations must go through the slow path to correctly
// update the size metadata.
ASSERT(partitionPageGetRawSize(page) == 0);
- PartitionFreelistEntry* newHead = partitionFreelistMask(static_cast<PartitionFreelistEntry*>(ret)->next);
- page->freelistHead = newHead;
+
+ if(ret == page->freelistHead) {
+ PartitionFreelistEntry* newHead = partitionFreelistMask(static_cast<PartitionFreelistEntry*>(ret)->next);
+ page->freelistHead = newHead;
+ } else {
+ z->next = t->next;
+ }
+
page->numAllocatedSlots++;
} else {
ret = partitionAllocSlowPath(root, flags, size, bucket);
@@ -647,6 +674,14 @@ ALWAYS_INLINE void partitionFreeWithPage(void* ptr, PartitionPage* page)
ASSERT(!freelistHead || partitionPointerIsValid(freelistHead));
RELEASE_ASSERT_WITH_SECURITY_IMPLICATION(ptr != freelistHead); // Catches an immediate double free.
ASSERT_WITH_SECURITY_IMPLICATION(!freelistHead || ptr != partitionFreelistMask(freelistHead->next)); // Look for double free one level deeper in debug.
+
+ PartitionFreelistEntry* f = freelistHead;
+
+ while(f != NULL) {
+ RELEASE_ASSERT_WITH_SECURITY_IMPLICATION(f != ptr);
+ f = partitionFreelistMask(f->next);
+ }
+
PartitionFreelistEntry* entry = static_cast<PartitionFreelistEntry*>(ptr);
entry->next = partitionFreelistMask(freelistHead);
page->freelistHead = entry;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment