1// Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4//
5// Static class for hooking Win32 API routines.  For now,
6// we only add one watcher at a time.
7//
8// TODO(mbelshe):  Support multiple watchers.
9
10#ifndef MEMORY_WATCHER_MEMORY_HOOK_
11#define MEMORY_WATCHER_MEMORY_HOOK_
12
13#include "base/logging.h"
14
15// When allocating memory for internal use with the MemoryHook,
16// we must always use the MemoryHook's heap; otherwise, the memory
17// gets tracked, and it becomes an infinite loop (allocation() calls
18// MemoryHook() which calls allocation(), etc).
19//
20// PrivateHookAllocator is an STL-friendly Allocator so that STL lists,
21// maps, etc can be used on the global MemoryHook's heap.
22template <class T>
23class PrivateHookAllocator {
24 public:
25  // These type definitions are needed for stl allocators.
26  typedef size_t    size_type;
27  typedef ptrdiff_t difference_type;
28  typedef T*        pointer;
29  typedef const T*  const_pointer;
30  typedef T&        reference;
31  typedef const T&  const_reference;
32  typedef T         value_type;
33
34  PrivateHookAllocator() {}
35
36  // Allocate memory for STL.
37  pointer allocate(size_type n, const void * = 0) {
38    return reinterpret_cast<T*>(MemoryHook::Alloc(n * sizeof(T)));
39  }
40
41  // Deallocate memory for STL.
42  void deallocate(void* p, size_type) {
43    if (p)
44      MemoryHook::Free(p);
45  }
46
47  // Construct the object
48  void construct(pointer p, const T& val) {
49    new (reinterpret_cast<T*>(p))T(val);
50  }
51
52  // Destruct an object
53  void destroy(pointer p) { p->~T(); }
54
55  size_type max_size() const { return size_t(-1); }
56
57  template <class U>
58  struct rebind { typedef PrivateHookAllocator<U> other; };
59
60  template <class U>
61  PrivateHookAllocator(const PrivateHookAllocator<U>&) {}
62};
63
64template<class T, class U> inline
65bool operator==(const PrivateHookAllocator<T>&,
66                const PrivateHookAllocator<U>&) {
67  return (true);
68}
69
70template<class T, class U> inline
71bool operator!=(const PrivateHookAllocator<T>& left,
72                const PrivateHookAllocator<U>& right) {
73  return (!(left == right));
74}
75
76
77// Classes which monitor memory from these hooks implement
78// the MemoryObserver interface.
79class MemoryObserver {
80 public:
81  virtual ~MemoryObserver() {}
82
83  // Track a pointer.  Will capture the current StackTrace.
84  virtual void OnTrack(HANDLE heap, int32 id, int32 size) = 0;
85
86  // Untrack a pointer, removing it from our list.
87  virtual void OnUntrack(HANDLE heap, int32 id, int32 size) = 0;
88};
89
90class MemoryHook : MemoryObserver {
91 public:
92  // Initialize the MemoryHook.  Must be called before
93  // registering watchers.  This can be called repeatedly,
94  // but is not thread safe.
95  static bool Initialize();
96
97  // Returns true is memory allocations and deallocations
98  // are being traced.
99  static bool hooked() { return hooked_ != NULL; }
100
101  // Register a class to receive memory allocation & deallocation
102  // callbacks.  If we haven't hooked memory yet, this call will
103  // force memory hooking to start.
104  static bool RegisterWatcher(MemoryObserver* watcher);
105
106  // Register a class to stop receiving callbacks.  If there are
107  // no more watchers, this call will unhook memory.
108  static bool UnregisterWatcher(MemoryObserver* watcher);
109
110  // MemoryHook provides a private heap for allocating
111  // unwatched memory.
112  static void* Alloc(size_t size) {
113    DCHECK(global_hook_ && global_hook_->heap_);
114    return HeapAlloc(global_hook_->heap_, 0, size);
115  }
116  static void Free(void* ptr) {
117    DCHECK(global_hook_ && global_hook_->heap_);
118    HeapFree(global_hook_->heap_, 0, ptr);
119  }
120
121  // Access the global hook.  For internal use only from static "C"
122  // hooks.
123  static MemoryHook* hook() { return global_hook_; }
124
125  // MemoryObserver interface.
126  virtual void OnTrack(HANDLE hHeap, int32 id, int32 size);
127  virtual void OnUntrack(HANDLE hHeap, int32 id, int32 size);
128
129 private:
130  MemoryHook();
131  ~MemoryHook();
132
133  // Enable memory tracing.  When memory is 'hooked',
134  // MemoryWatchers which have registered will be called
135  // as memory is allocated and deallocated.
136  static bool Hook();
137
138  // Disables memory tracing.
139  static bool Unhook();
140
141  // Create our private heap
142  bool CreateHeap();
143
144  // Close our private heap.
145  bool CloseHeap();
146
147  MemoryObserver* watcher_;
148  HANDLE heap_;   // An internal accounting heap.
149  static bool hooked_;
150  static MemoryHook* global_hook_;
151};
152
153#endif  // MEMORY_WATCHER_MEMORY_HOOK_
154