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..a55f75dac3b7c77e9fa9f7d4a5396f3ecbf63c8d
--- /dev/null
+++ b/syzygy/integration_tests/allocator_shim.cc
@@ -0,0 +1,137 @@
+// Copyright 2016 The Chromium Authors. All rights reserved.

[Sébastien Marchand, 2017/07/25 16:21:53]

2017

[njanevsk, 2017/07/25 17:43:02]

Done.

+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+// This header defines symbols to override the same functions in the Visual C++

[Sébastien Marchand, 2017/07/25 16:21:53]

This isn't a header :)

[njanevsk, 2017/07/25 17:43:01]

Done.

+// CRT implementation.
+
+#ifdef BASE_ALLOCATOR_ALLOCATOR_SHIM_OVERRIDE_UCRT_SYMBOLS_WIN_H_
+#error This header is meant to be included only once by allocator_shim.cc

[Sébastien Marchand, 2017/07/25 16:21:53]

We don't need this anymore.

[njanevsk, 2017/07/25 17:43:02]

Done.

+#endif
+#define BASE_ALLOCATOR_ALLOCATOR_SHIM_OVERRIDE_UCRT_SYMBOLS_WIN_H_
+
+#include <windows.h>
+#include <malloc.h>
+#include <mutex>
+#include "base/logging.h"

[Sébastien Marchand, 2017/07/25 16:21:53]

Do you really need logging?

[njanevsk, 2017/07/25 19:49:17]

Done.

+
+extern "C" {
+
+namespace {
+
+typedef HANDLE(WINAPI* HeapCreatePtr)(DWORD, SIZE_T, SIZE_T);

[Sébastien Marchand, 2017/07/25 16:21:53]

Add a comment to describe what these typedefs are.

[njanevsk, 2017/07/25 17:43:02]

Done.

+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);
+
+int win_new_mode = 0;

[Sébastien Marchand, 2017/07/25 16:21:53]

You don't need this.

[njanevsk, 2017/07/25 17:43:01]

Done.

+std::mutex m;

[Sébastien Marchand, 2017/07/25 16:21:53]

Rename this to something more meaningful, add a comment to describe what it's
for. Move the declaration next to the asan_shim one.

[njanevsk, 2017/07/25 17:43:02]

Done.

+

[Sébastien Marchand, 2017/07/25 16:21:52]

Remove one of these BLs

[njanevsk, 2017/07/25 17:43:01]

Done.

+
+struct asan_shim_struct {

[Sébastien Marchand, 2017/07/25 16:21:53]

Use a CamelCase name.

[njanevsk, 2017/07/25 17:43:02]

Done.

[chrisha, 2017/07/25 18:47:14]

Also, _struct / Struct isn't adding any additional information. This function is
a collection of points to runtime library heap functions, so maybe a better name
would be RuntimePointers?

[njanevsk, 2017/07/25 19:49:17]

I changed it to AsanRuntimePointers. I want to have Asan in the name to make it
more specific and descriptive. RuntimePointers is general.

+  HANDLE asan_heap = nullptr;
+  HMODULE asan_module;
+  HeapCreatePtr heap_create;
+  HeapAllocPtr heap_alloc;
+  HeapFreePtr heap_free;
+  HeapReAllocPtr heap_realloc;

[chrisha, 2017/07/25 18:47:14]

All of these should have = nullptr as well.

[njanevsk, 2017/07/25 19:49:17]

Done.

+};
+
+asan_shim_struct asan_shim;
+
+void load_asan_module() {
+  std::unique_lock<std::mutex> lock(m);

[Sébastien Marchand, 2017/07/25 16:21:52]

I think that we could avoid all this by moving this code to the constructor of
the asan_shim struct, Chris, wdyt?

[njanevsk, 2017/07/25 17:43:02]

That sounds like a really good idea. That way the variable will be initialized
when created and we can get rid of the load_asan_module calls in each function.

[chrisha, 2017/07/25 18:47:14]

That means that we rely on the static initializers of this module having run
before any heap allocations can be made. But, other static initializers running
before ours might need to make heap allocations. I'd stick with the current
mechanism (a gateway function for the initialization).

I'd also move the lock acquisition inside of the check for asan_module ==
nullptr.

In fact, you could get rid of the lock entirely, as we know we are running on
x86 systems where 32-bit reads and writes are atomic. The worst case scenario is
that multiple threads do the resolving at the same time, each coming up with the
same answer.

Finally, I'd call this functions something more meaningful, like
ResolveRuntimePointers, maybe?

[njanevsk, 2017/07/25 19:49:17]

Thanks. I removed the lock. Seb are you satisified with this solution?
I renamed the function to:
resolve_asan_runtime_pointers()
Again I am using asan to make it more specific and in agreement with the data
structure's name.

[Sébastien Marchand, 2017/07/25 20:07:06]

Yep, this approach sgtm.

+  if (asan_shim.asan_module == nullptr) {
+    asan_shim.asan_module = GetModuleHandle(L"syzyasan_rtl.dll");
+    asan_shim.heap_create = reinterpret_cast<HeapCreatePtr>(
+        ::GetProcAddress(asan_shim.asan_module, "asan_HeapCreate"));
+    asan_shim.heap_alloc = reinterpret_cast<HeapAllocPtr>(
+        ::GetProcAddress(asan_shim.asan_module, "asan_HeapAlloc"));
+    asan_shim.heap_free = reinterpret_cast<HeapFreePtr>(
+        ::GetProcAddress(asan_shim.asan_module, "asan_HeapFree"));
+    asan_shim.heap_realloc = reinterpret_cast<HeapReAllocPtr>(
+        ::GetProcAddress(asan_shim.asan_module, "asan_HeapReAlloc"));
+  }
+  lock.unlock();
+}
+
+inline HANDLE get_heap_handle() {
+  load_asan_module();

[Sébastien Marchand, 2017/07/25 16:21:53]

put this behind a "if (asan_shim.asan_module == nullptr)", otherwise you'll grab
the lock for nothing.

[njanevsk, 2017/07/25 19:49:17]

Since we are getting rid of the lock this is N\A.

+  return asan_shim.asan_heap;
+}
+
+}  // namespace
+
+// This function behaves similarly to MSVC's _set_new_mode.
+// If flag is 0 (default), calls to malloc will behave normally.
+// If flag is 1, calls to malloc will behave like calls to new,
+// and the std_new_handler will be invoked on failure.
+// Returns the previous mode.
+//
+// Replaces _set_new_mode in ucrt\heap\new_mode.cpp
+int _set_new_mode(int flag) {
+  return 0;
+}
+
+// Replaces _query_new_mode in ucrt\heap\new_mode.cpp
+int _query_new_mode() {
+  return 0;
+}
+
+// These symbols override the CRT's implementation of the same functions.
+__declspec(restrict) void* malloc(size_t size) {
+  

[Sébastien Marchand, 2017/07/25 16:21:53]

Remove this BL.

[njanevsk, 2017/07/25 17:43:02]

Done.

+  return asan_shim.heap_alloc(get_heap_handle(), 0, size);
+}
+
+void free(void* ptr) {
+  load_asan_module();

[Sébastien Marchand, 2017/07/25 16:21:52]

This is confusing, the load_asan_module is in some functions but not in the
others, you could add a "GetAllocShim" function that returns a pointer to the
initialized struct instead.

[njanevsk, 2017/07/25 17:43:02]

If we use a constructor for the structure then the asan_shim structure will
always be initialized and there will not be a need to call it in any of the
methods.

+  asan_shim.heap_free(get_heap_handle(), 0, ptr);
+}
+
+__declspec(restrict) void* realloc(void* ptr, size_t size) {
+  load_asan_module();
+  return asan_shim.heap_realloc(get_heap_handle(), 0, ptr, size);
+}
+
+__declspec(restrict) void* calloc(size_t n, size_t size) {
+  load_asan_module();
+  void* ptr = malloc(size * n);
+  ::memset(ptr, 0, size * n);

[Sébastien Marchand, 2017/07/25 16:21:53]

Check if ptr == nullptr before this, the malloc call might fail in which case
this function should return nullptr
(http://www.cplusplus.com/reference/cstdlib/calloc/)

[njanevsk, 2017/07/25 17:43:02]

Done.

+  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() {
+  __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);
+}
+
+// The default dispatch translation unit has to define also the following

[Sébastien Marchand, 2017/07/25 16:21:53]

You don't need this comment afaik.

[njanevsk, 2017/07/25 17:43:02]

Done.

+// symbols (unless they are ultimately routed to the system symbols):
+//   void malloc_stats(void);
+//   int mallopt(int, int);
+//   struct mallinfo mallinfo(void);
+//   size_t malloc_size(void*);
+//   size_t malloc_usable_size(const void*);
+
+}  // extern "C"