Create InternalHeap.

This is a heap wrapper that takes any HeapInterface and a MemoryNotificationInterface, and makes the heap notify all allocations as being for internal use.

BUG=
R=sebmarchand@chromium.org

Review URL: https://codereview.appspot.com/136890043

git-svn-id: http://sawbuck.googlecode.com/svn/trunk@2278 15e8cca8-e42c-11de-a347-f34a4f72eb7d
diff --git a/syzygy/agent/asan/asan.gyp b/syzygy/agent/asan/asan.gyp
index 3dcbdf0..97d4dbf 100644
--- a/syzygy/agent/asan/asan.gyp
+++ b/syzygy/agent/asan/asan.gyp
@@ -73,6 +73,8 @@
         'windows_heap_adapter.h',
         'heaps/ctmalloc_heap.cc',
         'heaps/ctmalloc_heap.h',
+        'heaps/internal_heap.cc',
+        'heaps/internal_heap.h',
         'heaps/large_block_heap.cc',
         'heaps/large_block_heap.h',
         'heaps/simple_block_heap.cc',
@@ -225,6 +227,7 @@
         'unittest_util.h',
         'windows_heap_adapter_unittest.cc',
         'heaps/ctmalloc_heap_unittest.cc',
+        'heaps/internal_heap_unittest.cc',
         'heaps/large_block_heap_unittest.cc',
         'heaps/simple_block_heap_unittest.cc',
         'heap_managers/block_heap_manager_unittest.cc',
diff --git a/syzygy/agent/asan/heaps/internal_heap.cc b/syzygy/agent/asan/heaps/internal_heap.cc
new file mode 100644
index 0000000..1518d15
--- /dev/null
+++ b/syzygy/agent/asan/heaps/internal_heap.cc
@@ -0,0 +1,102 @@
+// Copyright 2014 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "syzygy/agent/asan/heaps/internal_heap.h"
+
+#include "syzygy/common/align.h"
+
+namespace agent {
+namespace asan {
+namespace heaps {
+
+namespace {
+
+struct InternalHeapEntry {
+  uint32 size;
+  // Actually of a size such that the whole InternalHeapAlloc is of size
+  // |size|.
+  uint8 body[1];
+};
+
+const size_t kBodyOffset = offsetof(InternalHeapEntry, body);
+
+}  // namespace
+
+InternalHeap::InternalHeap(MemoryNotifierInterface* memory_notifier,
+                           HeapInterface* heap)
+    : memory_notifier_(memory_notifier), heap_(heap) {
+  DCHECK_NE(static_cast<MemoryNotifierInterface*>(NULL), memory_notifier);
+  DCHECK_NE(static_cast<HeapInterface*>(NULL), heap);
+  notifying_heap_ =
+      heap_->GetHeapFeatures() & HeapInterface::kHeapReportsReservations;
+}
+
+uint32 InternalHeap::GetHeapFeatures() const {
+  return heap_->GetHeapFeatures();
+}
+
+void* InternalHeap::Allocate(size_t bytes) {
+  size_t size = common::AlignUp(bytes + kBodyOffset, kShadowRatio);
+  void* alloc = heap_->Allocate(size);
+  if (alloc == NULL)
+    return NULL;
+
+  InternalHeapEntry* entry = reinterpret_cast<InternalHeapEntry*>(alloc);
+  entry->size = size;
+  memory_notifier_->NotifyInternalUse(entry, size);
+
+  return entry->body;
+}
+
+bool InternalHeap::Free(void* alloc) {
+  if (alloc != NULL) {
+    uint8* bytes = reinterpret_cast<uint8*>(alloc);
+    InternalHeapEntry* entry = reinterpret_cast<InternalHeapEntry*>(
+        bytes - kBodyOffset);
+    if (notifying_heap_) {
+      // A notifying heap redzones the memory from which allocations are made.
+      // We return the redzone to its initial state.
+      memory_notifier_->NotifyFutureHeapUse(entry, entry->size);
+    } else {
+      // A non-notifying heap serves memory from greenzoned pages, so indicate
+      // the memory has returned to the OS.
+      memory_notifier_->NotifyReturnedToOS(entry, entry->size);
+    }
+
+    // Adjust the allocation pointer to that of the wrapped heap.
+    alloc = entry;
+  }
+
+  return heap_->Free(alloc);
+}
+
+bool InternalHeap::IsAllocated(void* alloc) {
+  if (alloc != NULL) {
+    uint32* header = reinterpret_cast<uint32*>(alloc) - 1;
+    alloc = header;
+  }
+  return heap_->IsAllocated(alloc);
+}
+
+void InternalHeap::Lock() {
+  heap_->Lock();
+}
+
+void InternalHeap::Unlock() {
+  heap_->Unlock();
+}
+
+}  // namespace heaps
+}  // namespace asan
+}  // namespace agent
diff --git a/syzygy/agent/asan/heaps/internal_heap.h b/syzygy/agent/asan/heaps/internal_heap.h
new file mode 100644
index 0000000..dee75ea
--- /dev/null
+++ b/syzygy/agent/asan/heaps/internal_heap.h
@@ -0,0 +1,78 @@
+// Copyright 2014 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Defines InternalHeap, a simple wrapper of any other HeapInterface that
+// adds internal-use notifications via a MemoryNotifierInterface.
+
+#ifndef SYZYGY_AGENT_ASAN_HEAPS_INTERNAL_HEAP_H_
+#define SYZYGY_AGENT_ASAN_HEAPS_INTERNAL_HEAP_H_
+
+#include "base/logging.h"
+#include "syzygy/agent/asan/heap.h"
+#include "syzygy/agent/asan/memory_notifier.h"
+
+namespace agent {
+namespace asan {
+namespace heaps {
+
+// An implementation of HeapInterface that wraps another HeapInterface and
+// a MemoryNotificationInterface. It subsequently will notify all allocations
+// as being for internal use. This incurs a small amount of memory overhead
+// per allocation to store the original size of the allocation. This heap
+// does *not* return allocations that are kShadowRatio aligned. Rather, it
+// returns allocations that sizeof(uint32) % kShadowRatio aligned, due to the
+// extra incurred header. This is not an issue as the allocations are only
+// for internal use and no shadow memory notations will be applied to them.
+class InternalHeap : public HeapInterface {
+ public:
+  // Constructor.
+  // @param memory_notifier The notifier that will be used to inform the
+  //     runtime of all allocations.
+  // @param heap The underlying heap that is being wrapped.
+  InternalHeap(MemoryNotifierInterface* memory_notifier,
+               HeapInterface* heap);
+
+  // Destructor.
+  virtual ~InternalHeap() { }
+
+  // @name HeapInterface functions.
+  // @{
+  virtual uint32 GetHeapFeatures() const;
+  virtual void* Allocate(size_t bytes);
+  virtual bool Free(void* alloc);
+  virtual bool IsAllocated(void* alloc);
+  virtual void Lock();
+  virtual void Unlock();
+  // @}
+
+ protected:
+  // The interface that will be notified of all memory use. Has its own
+  // locking.
+  MemoryNotifierInterface* memory_notifier_;
+
+  // The underlying heap interface. Provides locking for us.
+  HeapInterface* heap_;
+
+  // This is true if the wrapped heap is a notifying heap.
+  bool notifying_heap_;
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(InternalHeap);
+};
+
+}  // namespace heaps
+}  // namespace asan
+}  // namespace agent
+
+#endif  // SYZYGY_AGENT_ASAN_HEAPS_INTERNAL_HEAP_H_
diff --git a/syzygy/agent/asan/heaps/internal_heap_unittest.cc b/syzygy/agent/asan/heaps/internal_heap_unittest.cc
new file mode 100644
index 0000000..135201b
--- /dev/null
+++ b/syzygy/agent/asan/heaps/internal_heap_unittest.cc
@@ -0,0 +1,101 @@
+// Copyright 2014 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "syzygy/agent/asan/heaps/internal_heap.h"
+
+#include "gtest/gtest.h"
+#include "syzygy/agent/asan/unittest_util.h"
+#include "syzygy/agent/asan/heaps/ctmalloc_heap.h"
+#include "syzygy/agent/asan/heaps/win_heap.h"
+
+namespace agent {
+namespace asan {
+namespace heaps {
+
+namespace {
+
+using testing::_;
+using testing::Return;
+
+}  // namespace
+
+TEST(InternalHeapTest, EndToEnd) {
+  memory_notifiers::NullMemoryNotifier mock_notifier;
+  WinHeap win_heap;
+  InternalHeap h(&mock_notifier, &win_heap);
+
+  // Allocate and free a zero-sized allocation. This should succeed
+  // by definition.
+  void* alloc = h.Allocate(0);
+  EXPECT_TRUE(h.Free(alloc));
+
+  // Make a bunch of different sized allocations.
+  std::set<void*> allocs;
+  for (size_t i = 1; i < 1024 * 1024; i <<= 1) {
+    void* alloc = h.Allocate(i);
+    allocs.insert(alloc);
+  }
+
+  // Now free them.
+  std::set<void*>::const_iterator it = allocs.begin();
+  for (; it != allocs.end(); ++it)
+    EXPECT_TRUE(h.Free(*it));
+}
+
+TEST(InternalHeapTest, NotificationsWorkWithNonNotifyingHeap) {
+  testing::MockMemoryNotifier mock_notifier;
+  WinHeap win_heap;
+  InternalHeap h(&mock_notifier, &win_heap);
+
+  EXPECT_CALL(mock_notifier, NotifyInternalUse(_, 16)).Times(1);
+  EXPECT_CALL(mock_notifier, NotifyReturnedToOS(_, 16)).Times(1);
+  void* alloc = h.Allocate(8);
+  h.Free(alloc);
+}
+
+TEST(InternalHeapTest, NotificationsWorkWithNotifyingHeap) {
+  memory_notifiers::NullMemoryNotifier null_notifier;
+  testing::MockMemoryNotifier mock_notifier;
+  CtMallocHeap ctmalloc_heap(&null_notifier);
+  InternalHeap h(&mock_notifier, &ctmalloc_heap);
+
+  EXPECT_CALL(mock_notifier, NotifyInternalUse(_, 16)).Times(1);
+  EXPECT_CALL(mock_notifier, NotifyFutureHeapUse(_, 16)).Times(1);
+  void* alloc = h.Allocate(8);
+  h.Free(alloc);
+}
+
+TEST(InternalHeapTest, HeaderIsAllocated) {
+  memory_notifiers::NullMemoryNotifier null_notifier;
+  testing::MockHeap mock_heap;
+  uint8 dummy_allocation[16] = {};
+
+  EXPECT_CALL(mock_heap, GetHeapFeatures()).Times(1).WillOnce(Return(0));
+  InternalHeap h(&null_notifier, &mock_heap);
+
+  void* header = dummy_allocation;
+  void* expected_alloc = dummy_allocation + sizeof(uint32);
+
+  EXPECT_CALL(mock_heap, Allocate(16)).Times(1).WillOnce(Return(header));
+  void* alloc = h.Allocate(8);
+  EXPECT_EQ(expected_alloc, alloc);
+  EXPECT_EQ(16, *reinterpret_cast<uint32*>(dummy_allocation));
+
+  EXPECT_CALL(mock_heap, Free(header)).Times(1).WillOnce(Return(true));
+  EXPECT_TRUE(h.Free(alloc));
+}
+
+}  // namespace heaps
+}  // namespace asan
+}  // namespace agent