Allocator shims design prototype.

syzygy/integration_tests/allocator_shim.cc

Index: syzygy/integration_tests/allocator_shim.cc
diff --git a/syzygy/integration_tests/allocator_shim.cc b/syzygy/integration_tests/allocator_shim.cc
new file mode 100644
index 0000000000000000000000000000000000000000..c6e0bc1be8c8a070162378350645c20f1f2aad0a
--- /dev/null
+++ b/syzygy/integration_tests/allocator_shim.cc
@@ -0,0 +1,112 @@
+// Copyright 2017 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+// This file defines symbols to override the same functions in the Visual C++
+// CRT implementation.
+
+#include <windows.h>
+#include <malloc.h>
+
+extern "C" {
+
+namespace {
+
+// Definitions of the following heap management functions provided by
+// the SyzyAsan library:
+// heap_create, heap_destroy, heap_alloc, heap_realloc, and heap_free
+typedef HANDLE(WINAPI* HeapCreatePtr)(DWORD, SIZE_T, SIZE_T);
+typedef BOOL(WINAPI* HeapDestroyPtr)(HANDLE);
+typedef LPVOID(WINAPI* HeapAllocPtr)(HANDLE, DWORD, SIZE_T);
+typedef LPVOID(WINAPI* HeapReAllocPtr)(HANDLE, DWORD, LPVOID, SIZE_T);
+typedef BOOL(WINAPI* HeapFreePtr)(HANDLE, DWORD, LPVOID);
+
+struct AsanRuntimePointers {
+  HANDLE asan_heap = nullptr;
+  HMODULE asan_module = nullptr;
+  HeapCreatePtr heap_create = nullptr;
+  HeapAllocPtr heap_alloc = nullptr;
+  HeapFreePtr heap_free = nullptr;
+  HeapReAllocPtr heap_realloc = nullptr;
+};
+
+// If the syzyasan_rtl.dll is not loaded it loads it and the heap functions

[Sébastien Marchand, 2017/07/31 15:01:40]

This comment isn't true, it only retrieves the module handle but doesn't load it

[njanevsk, 2017/07/31 18:57:28]

Done.

+// that it provides.
+void ResolveAsanRuntimePointers(AsanRuntimePointers* asan_runtime_ptrs) {
+  if (asan_runtime_ptrs->asan_module == nullptr) {
+    asan_runtime_ptrs->asan_module = GetModuleHandle(L"syzyasan_rtl.dll");

[Sigurður Ásgeirsson, 2017/07/28 17:03:14]

if this returns nullptr, then the calls below may malfunction
(GetProcAddress(nullptr, ...) has defined behavior.

Note also that GetModuleHandle does not load the module, it merely retrieves it
if it's already loaded into the process.

[Sébastien Marchand, 2017/07/31 15:01:40]

Ha good point, it's probably safe as-is for the integration tests but adding
some safety checks will help to prevent some flakiness

[njanevsk, 2017/07/31 18:57:28]

I added a check to make sure that the module is retrieved. However, this does
not exit the code. How do you handle this in here? Do you throw an exception or
modify all the methods to check that GetAsanHeapShim does not return nullptr?

+    asan_runtime_ptrs->heap_create = reinterpret_cast<HeapCreatePtr>(
+        ::GetProcAddress(asan_runtime_ptrs->asan_module, "asan_HeapCreate"));
+    asan_runtime_ptrs->heap_alloc = reinterpret_cast<HeapAllocPtr>(
+        ::GetProcAddress(asan_runtime_ptrs->asan_module, "asan_HeapAlloc"));
+    asan_runtime_ptrs->heap_free = reinterpret_cast<HeapFreePtr>(
+        ::GetProcAddress(asan_runtime_ptrs->asan_module, "asan_HeapFree"));
+    asan_runtime_ptrs->heap_realloc = reinterpret_cast<HeapReAllocPtr>(
+        ::GetProcAddress(asan_runtime_ptrs->asan_module, "asan_HeapReAlloc"));
+    asan_runtime_ptrs->asan_heap = asan_runtime_ptrs->heap_create(0, 0, 0);
+  }
+}
+
+// It returns a pointer to a static data structure containing the syzyasan
+// module and
+// the heap methods.

[Sébastien Marchand, 2017/07/31 15:01:40]

Move this to the previous line.

[njanevsk, 2017/07/31 18:57:28]

Done.

+const AsanRuntimePointers* GetAsanHeapShim() {
+  static AsanRuntimePointers asan_runtime_ptrs = {};

[Sigurður Ásgeirsson, 2017/07/28 17:03:14]

In Chrome land, we've now started to rely on thread-safe static initialization.
Maybe that's safer and faster here?

[Sébastien Marchand, 2017/07/31 15:01:40]

Ha, that's cool. In this case you could add a constructor to AsanRuntimePtr and
do:
static AsanRuntimePointers* asan_runtime_ptrs = new AsanRuntimePointers();

[njanevsk, 2017/07/31 18:57:28]

Done.

+  if (asan_runtime_ptrs.asan_module == nullptr) {
+    ResolveAsanRuntimePointers(&asan_runtime_ptrs);
+  }
+  return &asan_runtime_ptrs;
+}
+
+inline HANDLE get_heap_handle() {
+  return GetAsanHeapShim()->asan_heap;
+}
+
+}  // namespace
+
+// These symbols override the CRT's implementation of the same functions.
+__declspec(restrict) void* malloc(size_t size) {
+  return GetAsanHeapShim()->heap_alloc(get_heap_handle(), 0, size);

[Sigurður Ásgeirsson, 2017/07/28 17:03:14]

does compiler do a decent job of eliminating the duplicate call to
GetAsanHeapShim? Traversing through a function that contains a statically
initialized variable might be expensive (depending on how the compiler handles
it), so could be worth explicitly retrieving the shim and using it twice...

[Sébastien Marchand, 2017/07/31 15:01:40]

Not that this shim is only intended to be used in the context of the integration
tests so performance isn't a huge criteria but it'd be interesting to look at
the generated code to see how Clang handle this. 

[njanevsk, 2017/07/31 18:57:28]

Leaving it as it is for now.

+}
+
+void free(void* ptr) {
+  GetAsanHeapShim()->heap_free(get_heap_handle(), 0, ptr);
+}
+
+__declspec(restrict) void* realloc(void* ptr, size_t size) {
+  return GetAsanHeapShim()->heap_realloc(get_heap_handle(), 0, ptr, size);
+}
+
+__declspec(restrict) void* calloc(size_t n, size_t size) {
+  void* ptr = malloc(size * n);
+  if (ptr != nullptr)
+    ::memset(ptr, 0, size * n);
+  return ptr;
+}
+
+// The symbols
+//   * __acrt_heap
+//   * __acrt_initialize_heap
+//   * __acrt_uninitialize_heap
+//   * _get_heap_handle
+// must be overridden all or none, as they are otherwise supplied
+// by heap_handle.obj in the ucrt.lib file.
+HANDLE __acrt_heap = nullptr;
+
+bool __acrt_initialize_heap() {

[Sigurður Ásgeirsson, 2017/07/28 17:03:14]

if you intend this for use in Chrome, then it won't work as-is. There are cases
where Chrome calls _get_heap_handle, IIRC, and you're going to return the wrong
one, methinks?

[Sébastien Marchand, 2017/07/31 15:01:40]

This shim is not meant to be used in Chrome.

[njanevsk, 2017/07/31 18:57:28]

Not fixing cause it is not meant to be use with Chrome.

[Sigurður Ásgeirsson, 2017/07/31 19:05:52]

I don't understand why you'd want the integration test to implement incorrect
behavior, though. Why not use the same heap here and in the shim functions?

[njanevsk, 2017/08/01 15:46:43]

I tried using the same heap but the Clang test cases are crushing in that case.
Seb might be able to provide a better answer.

[Sébastien Marchand, 2017/08/03 21:49:45]

The problem here is that we can't use a heap that hasn't been created by
::HeapCreate to initialize __acrt_heap. The DLL main code ends up calling
_calloc_base directly, which calls ::HeapAlloc(__acrt_heap,...) and we're not
patching HeapAlloc here.

Note that this code comes from Chrome:
https://cs.chromium.org/chromium/src/base/allocator/allocator_shim_override_u...

This means that we won't have coverage for all the initialization made directly
by the CRT via HeapAlloc, this is probably fine because we don't instrument this
code with Clang anyway.

[Sigurður Ásgeirsson, 2017/08/04 13:55:13]

No, this isn't OK, and this WILL fail in Chrome, so your integration tests here
aren't going to be any good for what you want to use this for eventually. What's
going to happen when the pointer out of _calloc_base is freed?

+  // TODO(sebmarchand): Check if it is possible to use asan_heap here.
+  __acrt_heap = ::HeapCreate(0, 0, 0);
+  return true;
+}
+
+bool __acrt_uninitialize_heap() {
+  ::HeapDestroy(__acrt_heap);
+  __acrt_heap = nullptr;
+  return true;
+}
+
+intptr_t _get_heap_handle(void) {
+  return reinterpret_cast<intptr_t>(__acrt_heap);
+}
+
+}  // extern "C"