1    :  // Copyright 2015 Google Inc. All Rights Reserved.
   2    :  //
   3    :  // Licensed under the Apache License, Version 2.0 (the "License");
   4    :  // you may not use this file except in compliance with the License.
   5    :  // You may obtain a copy of the License at
   6    :  //
   7    :  //     http://www.apache.org/licenses/LICENSE-2.0
   8    :  //
   9    :  // Unless required by applicable law or agreed to in writing, software
  10    :  // distributed under the License is distributed on an "AS IS" BASIS,
  11    :  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12    :  // See the License for the specific language governing permissions and
  13    :  // limitations under the License.
  14    :  
  15    :  #include "syzygy/bard/backdrops/heap_backdrop.h"
  16    :  
  17    :  #include "base/logging.h"
  18    :  
  19    :  namespace bard {
  20    :  namespace backdrops {
  21    :  
  22  E :  HeapBackdrop::HeapBackdrop() {
  23  E :  }
  24    :  
  25  E :  LPVOID HeapBackdrop::HeapAlloc(HANDLE heap, DWORD flags, SIZE_T bytes) {
  26  E :    DCHECK(!heap_alloc_.is_null());
  27  E :    return heap_alloc_.Run(heap, flags, bytes);
  28  E :  }
  29    :  
  30    :  HANDLE HeapBackdrop::HeapCreate(DWORD options,
  31    :                                  SIZE_T initial_size,
  32  E :                                  SIZE_T maximum_size) {
  33  E :    DCHECK(!heap_create_.is_null());
  34  E :    return heap_create_.Run(options, initial_size, maximum_size);
  35  E :  }
  36    :  
  37  E :  BOOL HeapBackdrop::HeapDestroy(HANDLE heap) {
  38  E :    DCHECK(!heap_destroy_.is_null());
  39  E :    return heap_destroy_.Run(heap);
  40  E :  }
  41    :  
  42  E :  BOOL HeapBackdrop::HeapFree(HANDLE heap, DWORD flags, LPVOID mem) {
  43  E :    DCHECK(!heap_free_.is_null());
  44  E :    return heap_free_.Run(heap, flags, mem);
  45  E :  }
  46    :  
  47    :  LPVOID HeapBackdrop::HeapReAlloc(HANDLE heap,
  48    :                                   DWORD flags,
  49    :                                   LPVOID mem,
  50  E :                                   SIZE_T bytes) {
  51  E :    DCHECK(!heap_realloc_.is_null());
  52  E :    return heap_realloc_.Run(heap, flags, mem, bytes);
  53  E :  }
  54    :  
  55    :  BOOL HeapBackdrop::HeapSetInformation(HANDLE heap,
  56    :                                        HEAP_INFORMATION_CLASS info_class,
  57    :                                        PVOID info,
  58  E :                                        SIZE_T info_length) {
  59  E :    DCHECK(!heap_set_information_.is_null());
  60  E :    return heap_set_information_.Run(heap, info_class, info, info_length);
  61  E :  }
  62    :  
  63  E :  SIZE_T HeapBackdrop::HeapSize(HANDLE heap, DWORD flags, LPCVOID mem) {
  64  E :    DCHECK(!heap_size_.is_null());
  65  E :    return heap_size_.Run(heap, flags, mem);
  66  E :  }
  67    :  
  68  E :  void HeapBackdrop::UpdateStats(EventType type, uint64_t time) {
  69  E :    base::AutoLock auto_lock(lock_);
  70    :  
  71  E :    auto stats = total_stats_.insert(std::make_pair(type, struct Stats())).first;
  72  E :    stats->second.calls++;
  73  E :    stats->second.time += time;
  74  E :  }
  75    :  
  76  E :  bool HeapBackdrop::TearDown() {
  77    :    // Destroy heaps created via AddExistingHeap.
  78  E :    for (auto live_heap : existing_heaps_) {
  79  E :      HANDLE trace_heap = nullptr;
  80  E :      if (!heap_map_.GetTraceFromLive(live_heap, &trace_heap))
  81  i :        return false;
  82  E :      ::HeapDestroy(live_heap);
  83    :      // This can only fail under racy use of this class.
  84  E :      CHECK(heap_map_.RemoveMapping(trace_heap, live_heap));
  85  E :    }
  86    :  
  87    :    // Remove the heap created by SetProcessHeap. This can only fail under racy
  88    :    // use of this class.
  89  E :    HANDLE live_ph = ::GetProcessHeap();
  90  E :    HANDLE trace_ph = nullptr;
  91  E :    if (heap_map_.GetTraceFromLive(live_ph, &trace_ph))
  92  E :      CHECK(heap_map_.RemoveMapping(trace_ph, live_ph));
  93    :  
  94    :    // Handle any other heaps that were created via the HeapDestroy callback.
  95  E :    for (auto heap_pair : heap_map_.live_trace()) {
  96  i :      DCHECK(!heap_destroy_.is_null());
  97  i :      if (!heap_destroy_.Run(heap_pair.first))
  98  i :        return false;
  99  i :    }
 100  E :    heap_map_.Clear();
 101    :  
 102    :    // Since the heaps are clear, also clear the maps.
 103  E :    alloc_map_.Clear();
 104  E :    existing_heaps_.clear();
 105    :  
 106  E :    return true;
 107  E :  }
 108    :  
 109  E :  bool HeapBackdrop::SetProcessHeap(void* proc_heap) {
 110  E :    HANDLE ph = ::GetProcessHeap();
 111  E :    return heap_map_.AddMapping(proc_heap, ph);
 112  E :  }
 113    :  
 114  E :  bool HeapBackdrop::AddExistingHeap(void* heap) {
 115  E :    HANDLE h = ::HeapCreate(0, 0, 0);
 116  E :    if (!h)
 117  i :      return false;
 118  E :    existing_heaps_.push_back(h);
 119  E :    return heap_map_.AddMapping(heap, h);
 120  E :  }
 121    :  
 122    :  }  // namespace backdrops
 123    :  }  // namespace bard