SERVER-2942 Hook functions in the Windows Import Address Table

In the Windows version of mongod and test, hook the RtlCreateHeap and
NtAllocateVirtualMemory functions in ntdll.dll through the Import
Address Tables of kernel32.dll and/or kernelbase.dll.  Use RemapLock
to block memory allocations during the period when remapPrivateView()
has the private view unmapped so that other threads can't grab any of
the address space required for the private view.

12 files changed, 384 insertions, 4 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index b8dc7be2882..bc827293da0 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -137,7 +137,6 @@ coreServerFiles = [ "util/version.cpp",
                     "db/security_commands.cpp",
                     ]
 
-
 env.StaticLibrary('ntservice', ['util/ntservice.cpp'])
 
 scripting_common_files = [ "scripting/engine.cpp",
@@ -272,6 +271,9 @@ serverOnlyFiles = [ "db/curop.cpp",
                     "db/commands/document_source_cursor.cpp",
                     "db/driverHelpers.cpp" ]
 
+if os.sys.platform == 'win32':
+    serverOnlyFiles += [ "util/hook_win32.cpp", "util/hook_windows_memory.cpp" ]
+
 env.Library( "dbcmdline", "db/cmdline.cpp" )
 
 env.Library('stacktrace', 'util/stacktrace.cpp')
diff --git a/src/mongo/db/db.cpp b/src/mongo/db/db.cpp
index 274672b6b7d..e323e00b19f 100644
--- a/src/mongo/db/db.cpp
+++ b/src/mongo/db/db.cpp
@@ -52,7 +52,8 @@
 #include "ttl.h"
 
 #if defined(_WIN32)
-# include "../util/ntservice.h"
+# include "mongo/util/ntservice.h"
+# include "mongo/util/hook_windows_memory.h"
 # include <DbgHelp.h>
 #else
 # include <sys/file.h>
@@ -1001,6 +1002,15 @@ int main(int argc, char* argv[]) {
         if( repairpath.empty() )
             repairpath = dbpath;
 
+#if defined(_WIN32)
+        if ( cmdLine.dur ) {
+            // Hook Windows APIs that can allocate memory so that we can RemapLock them out while
+            //  remapPrivateView() has a data file unmapped (so only needed when journaling)
+            // This is the last point where we are still single-threaded, makes hooking simpler
+            hookWindowsMemory();
+        }
+#endif
+
         Module::configAll( params );
         dataFileSync.go();
 
diff --git a/src/mongo/db/mongod.vcxproj b/src/mongo/db/mongod.vcxproj
index e623bf0c51c..02ccd3c4097 100755
--- a/src/mongo/db/mongod.vcxproj
+++ b/src/mongo/db/mongod.vcxproj
@@ -1055,6 +1055,8 @@ cscript //Nologo ..\shell\createCPPfromJavaScriptFiles.js "$(ProjectDir).."
     <ClCompile Include="..\util\file_allocator.cpp" />

     <ClCompile Include="..\util\histogram.cpp" />

     <ClCompile Include="..\util\intrusive_counter.cpp" />

+    <ClCompile Include="..\util\hook_win32.cpp" />

+    <ClCompile Include="..\util\hook_windows_memory.cpp" />

     <ClCompile Include="..\util\log.cpp" />

     <ClCompile Include="..\util\logfile.cpp" />

     <ClCompile Include="..\util\net\listen.cpp" />

@@ -1199,6 +1201,8 @@ cscript //Nologo ..\shell\createCPPfromJavaScriptFiles.js "$(ProjectDir).."
     <ClInclude Include="..\..\third_party\js-1.7\prmjtime.h" />

     <ClInclude Include="..\..\third_party\js-1.7\resource.h" />

     <ClInclude Include="..\util\histogram.h" />

+    <ClInclude Include="..\util\hook_win32.h" />

+    <ClInclude Include="..\util\hook_windows_memory.h" />

     <ClInclude Include="..\util\progress_meter.h" />

     <ClInclude Include="..\util\concurrency\remap_lock.h" />

     <ClInclude Include="..\util\signal_handlers.h" />

diff --git a/src/mongo/db/mongod.vcxproj.filters b/src/mongo/db/mongod.vcxproj.filters
index 92c014bd5f6..823ad5c8c1f 100755
--- a/src/mongo/db/mongod.vcxproj.filters
+++ b/src/mongo/db/mongod.vcxproj.filters
@@ -896,6 +896,12 @@
     <ClCompile Include="..\util\touch_pages.cpp">

       <Filter>util\Source Files</Filter>

     </ClCompile>

+    <ClCompile Include="..\util\hook_win32.cpp">

+      <Filter>util\Source Files</Filter>

+    </ClCompile>

+    <ClCompile Include="..\util\hook_windows_memory.cpp">

+      <Filter>util\Source Files</Filter>

+    </ClCompile>

   </ItemGroup>

   <ItemGroup>

     <ClInclude Include="..\targetver.h" />

@@ -1620,6 +1626,12 @@
     <ClInclude Include="..\util\touch_pages.h">

       <Filter>util\Header Files</Filter>

     </ClInclude>

+    <ClInclude Include="..\util\hook_win32.h">

+      <Filter>util\Header Files</Filter>

+    </ClInclude>

+    <ClInclude Include="..\util\hook_windows_memory.h">

+      <Filter>util\Header Files</Filter>

+    </ClInclude>

   </ItemGroup>

   <ItemGroup>

     <None Include="mongo.ico">

diff --git a/src/mongo/dbtests/framework.cpp b/src/mongo/dbtests/framework.cpp
index 6675b990399..975212430b8 100644
--- a/src/mongo/dbtests/framework.cpp
+++ b/src/mongo/dbtests/framework.cpp
@@ -25,8 +25,8 @@
 #include <sys/file.h>
 #endif
 
-#include <boost/program_options.hpp>
 #include <boost/filesystem/operations.hpp>
+#include <boost/program_options.hpp>
 
 #include "mongo/db/client.h"
 #include "mongo/db/cmdline.h"
@@ -37,6 +37,10 @@
 #include "mongo/util/file_allocator.h"
 #include "mongo/util/version.h"
 
+#if defined(_WIN32)
+#include "mongo/util/hook_windows_memory.h"
+#endif
+
 namespace po = boost::program_options;
 
 namespace mongo {
@@ -215,6 +219,15 @@ namespace mongo {
                 log() << "****************" << endl;
             }
 
+#if defined(_WIN32)
+            if ( cmdLine.dur ) {
+                // Hook Windows APIs that allocate memory so that we can RemapLock them out while
+                //  remapPrivateView() has a data file unmapped (so only needed when journaling)
+                // This is the last point where we are still single-threaded, makes hooking simpler
+                hookWindowsMemory();
+            }
+#endif
+
             FileAllocator::get()->start();
 
             vector<string> suites;
diff --git a/src/mongo/dbtests/test.vcxproj b/src/mongo/dbtests/test.vcxproj
index 49f6334be88..84c9c568d7c 100644
--- a/src/mongo/dbtests/test.vcxproj
+++ b/src/mongo/dbtests/test.vcxproj
@@ -556,6 +556,8 @@ cscript //Nologo ..\shell\createCPPfromJavaScriptFiles.js "$(ProjectDir).."
     <ClInclude Include="..\util\goodies.h" />

     <ClInclude Include="..\util\hashtab.h" />

     <ClInclude Include="..\db\lasterror.h" />

+    <ClInclude Include="..\util\hook_win32.h" />

+    <ClInclude Include="..\util\hook_windows_memory.h" />

     <ClInclude Include="..\util\log.h" />

     <ClInclude Include="..\util\logfile.h" />

     <ClInclude Include="..\util\lruishmap.h" />

@@ -1279,6 +1281,8 @@ cscript //Nologo ..\shell\createCPPfromJavaScriptFiles.js "$(ProjectDir).."
     <ClCompile Include="..\util\concurrency\vars.cpp" />

     <ClCompile Include="..\util\debug_util.cpp" />

     <ClCompile Include="..\util\file_allocator.cpp" />

+    <ClCompile Include="..\util\hook_win32.cpp" />

+    <ClCompile Include="..\util\hook_windows_memory.cpp" />

     <ClCompile Include="..\util\log.cpp" />

     <ClCompile Include="..\util\logfile.cpp" />

     <ClCompile Include="..\util\mmap_win.cpp" />

diff --git a/src/mongo/dbtests/test.vcxproj.filters b/src/mongo/dbtests/test.vcxproj.filters
index ea067d5dad7..1c4bc1bdbab 100755
--- a/src/mongo/dbtests/test.vcxproj.filters
+++ b/src/mongo/dbtests/test.vcxproj.filters
@@ -582,6 +582,12 @@
     <ClInclude Include="..\util\touch_pages.h">

       <Filter>util\Header Files</Filter>

     </ClInclude>

+    <ClInclude Include="..\util\hook_win32.h">

+      <Filter>util\Header Files</Filter>

+    </ClInclude>

+    <ClInclude Include="..\util\hook_windows_memory.h">

+      <Filter>util\Header Files</Filter>

+    </ClInclude>

   </ItemGroup>

   <ItemGroup>

     <ClCompile Include="..\client\connpool.cpp">

@@ -1417,6 +1423,12 @@
     <ClCompile Include="..\util\touch_pages.cpp">

       <Filter>util\Source Files</Filter>

     </ClCompile>

+    <ClCompile Include="..\util\hook_win32.cpp">

+      <Filter>util\Source Files</Filter>

+    </ClCompile>

+    <ClCompile Include="..\util\hook_windows_memory.cpp">

+      <Filter>util\Source Files</Filter>

+    </ClCompile>

   </ItemGroup>

   <ItemGroup>

     <None Include="btreetests.inl">

diff --git a/src/mongo/util/hook_win32.cpp b/src/mongo/util/hook_win32.cpp
new file mode 100644
index 00000000000..09c81c65cd7
--- /dev/null
+++ b/src/mongo/util/hook_win32.cpp
@@ -0,0 +1,107 @@
+// @file hook_win32.cpp : Used to hook Windows functions imported through the Import Address Table

+

+/**

+*    Copyright (C) 2012 10gen Inc.

+*

+*    This program is free software: you can redistribute it and/or modify

+*    it under the terms of the GNU Affero General Public License, version 3,

+*    as published by the Free Software Foundation.

+*

+*    This program is distributed in the hope that it will be useful,

+*    but WITHOUT ANY WARRANTY; without even the implied warranty of

+*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

+*    GNU Affero General Public License for more details.

+*

+*    You should have received a copy of the GNU Affero General Public License

+*    along with this program.  If not, see <http://www.gnu.org/licenses/>.

+*/

+

+#if defined(_WIN32)

+

+#include "mongo/util/hook_win32.h"

+

+#include "mongo/platform/windows_basic.h"

+

+#include <DbgHelp.h>

+

+#include "mongo/util/assert_util.h"

+

+#pragma comment(lib, "dbghelp.lib")

+

+namespace mongo {

+

+    /**

+     * Hook a Windows API through the Import Address Table of a calling module

+     *

+     * @param hookModuleAddress     ptr to start of importing executable

+     * @param functionModuleName    name of module containing function to hook

+     * @param functionName          name of function to hook

+     * @param hookFunction          ptr to replacement (hook) function

+     * @return                      ptr to original (hooked) function

+     */

+    void* hookWin32(

+            char* hookModuleAddress,

+            char* functionModuleName,

+            char* functionName,

+            void* hookFunction ) {

+

+        if ( hookModuleAddress == 0 ) {

+            return NULL;

+        }

+

+        // look up the original function by module and function name

+        HMODULE moduleHandle = GetModuleHandleA( functionModuleName );

+        void* originalFunction = GetProcAddress( moduleHandle, functionName );

+        if ( originalFunction == 0 ) {

+            return NULL;

+        }

+

+        // get a pointer to this module's Import Table

+        ULONG entrySize;

+        void* imageEntry = ImageDirectoryEntryToDataEx(

+                hookModuleAddress,              // address of module with IAT to patch

+                TRUE,                           // mapped as image

+                IMAGE_DIRECTORY_ENTRY_IMPORT,   // desired directory

+                &entrySize,                     // returned directory size

+                NULL                            // returned section ptr (always zero for IAT)

+        );

+        IMAGE_IMPORT_DESCRIPTOR* importModule =

+                reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>( imageEntry );

+

+        // walk the imported module list until we find the desired module

+        for ( ; importModule->Name; ++importModule ) {

+            char* foundModuleName = hookModuleAddress + importModule->Name;

+            if ( _strcmpi( foundModuleName, functionModuleName ) == 0 ) {

+

+                // search the list of imported functions for the address of the one we want to hook

+                IMAGE_THUNK_DATA* thunk = reinterpret_cast<IMAGE_THUNK_DATA*>(

+                        hookModuleAddress + importModule->FirstThunk );

+                for ( ; thunk->u1.Function; ++thunk ) {

+                    if ( reinterpret_cast<void*>( thunk->u1.Function ) == originalFunction ) {

+

+                        // we found our function, remember its location and then hook it

+                        void** _tablePointer = reinterpret_cast<void**>( &thunk->u1.Function );

+                        MEMORY_BASIC_INFORMATION mbi;

+                        VirtualQuery( _tablePointer, &mbi, sizeof(mbi) );

+                        fassert( 16239, VirtualProtect( mbi.BaseAddress,

+                                                        mbi.RegionSize,

+                                                        PAGE_EXECUTE_READWRITE,

+                                                        &mbi.Protect ) );

+                        *_tablePointer = hookFunction;

+                        DWORD unused;

+                        fassert( 16240, VirtualProtect( mbi.BaseAddress,

+                                                        mbi.RegionSize,

+                                                        mbi.Protect,

+                                                        &unused ) );

+                        return originalFunction;

+                    }

+                }

+                break;

+            }

+        }

+        return NULL;

+    }

+

+} // namespace mongo

+

+#endif // #if defined(_WIN32)

diff --git a/src/mongo/util/hook_win32.h b/src/mongo/util/hook_win32.h
new file mode 100644
index 00000000000..8be163173db
--- /dev/null
+++ b/src/mongo/util/hook_win32.h
@@ -0,0 +1,42 @@
+// @file hook_win32.h : Used to hook Windows functions imported through the Import Address Table
+
+/**
+*    Copyright (C) 2012 10gen Inc.
+*
+*    This program is free software: you can redistribute it and/or modify
+*    it under the terms of the GNU Affero General Public License, version 3,
+*    as published by the Free Software Foundation.
+*
+*    This program is distributed in the hope that it will be useful,
+*    but WITHOUT ANY WARRANTY; without even the implied warranty of
+*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+*    GNU Affero General Public License for more details.
+*
+*    You should have received a copy of the GNU Affero General Public License
+*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#pragma once
+
+#if defined(_WIN32)
+
+namespace mongo {
+
+    /**
+     * Hook a Windows API through the Import Address Table of a calling module
+     *
+     * @param hookModuleAddress     ptr to start of importing executable
+     * @param functionModuleName    name of module containing function to hook
+     * @param functionName          name of function to hook
+     * @param hookFunction          ptr to replacement (hook) function
+     * @return                      ptr to original (hooked) function
+     */
+    void* hookWin32(
+        char* hookModuleAddress,
+        char* functionModuleName,
+        char* functionName,
+        void* hookFunction );
+
+} // namespace mongo
+
+#endif // #if defined(_WIN32)
diff --git a/src/mongo/util/hook_windows_memory.cpp b/src/mongo/util/hook_windows_memory.cpp
new file mode 100644
index 00000000000..423d87defb9
--- /dev/null
+++ b/src/mongo/util/hook_windows_memory.cpp
@@ -0,0 +1,149 @@
+// @file hook_windows_memory.cpp : Used to hook Windows functions that allocate virtual memory
+
+/**
+*    Copyright (C) 2012 10gen Inc.
+*
+*    This program is free software: you can redistribute it and/or modify
+*    it under the terms of the GNU Affero General Public License, version 3,
+*    as published by the Free Software Foundation.
+*
+*    This program is distributed in the hope that it will be useful,
+*    but WITHOUT ANY WARRANTY; without even the implied warranty of
+*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+*    GNU Affero General Public License for more details.
+*
+*    You should have received a copy of the GNU Affero General Public License
+*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#if defined(_WIN32)
+
+#include "mongo/util/hook_windows_memory.h"
+
+#include <vector>
+
+#include "mongo/platform/windows_basic.h"
+#include "mongo/util/concurrency/remap_lock.h"
+#include "mongo/util/hook_win32.h"
+
+namespace {
+
+    // hooked function typedefs and routines
+
+    typedef PVOID ( WINAPI *RtlCreateHeap_t ) ( ULONG, PVOID, SIZE_T, SIZE_T, PVOID, void* );
+
+    RtlCreateHeap_t originalRtlCreateHeap_kernel32;
+
+    static PVOID WINAPI myRtlCreateHeap_kernel32(
+        ULONG Flags,
+        PVOID HeapBase,
+        SIZE_T ReserveSize,
+        SIZE_T CommitSize,
+        PVOID Lock,
+        void* /* PRTL_HEAP_PARAMETERS */ Parameters
+    ) {
+        mongo::RemapLock remapLock;
+        return originalRtlCreateHeap_kernel32( Flags,
+                                               HeapBase,
+                                               ReserveSize,
+                                               CommitSize,
+                                               Lock,
+                                               Parameters );
+    }
+
+    RtlCreateHeap_t originalRtlCreateHeap_kernelbase;
+
+    static PVOID WINAPI myRtlCreateHeap_kernelbase(
+        ULONG Flags,
+        PVOID HeapBase,
+        SIZE_T ReserveSize,
+        SIZE_T CommitSize,
+        PVOID Lock,
+        void* /* PRTL_HEAP_PARAMETERS */ Parameters
+    ) {
+        mongo::RemapLock remapLock;
+        return originalRtlCreateHeap_kernelbase( Flags,
+                                                 HeapBase,
+                                                 ReserveSize,
+                                                 CommitSize,
+                                                 Lock,
+                                                 Parameters );
+    }
+
+    typedef LONG ( WINAPI *NtAllocateVirtualMemory_t )
+            ( HANDLE, PVOID*, ULONG_PTR, PSIZE_T, ULONG, ULONG );
+
+    NtAllocateVirtualMemory_t originalNtAllocateVirtualMemory_kernel32;
+
+    static LONG WINAPI myNtAllocateVirtualMemory_kernel32(
+        HANDLE ProcessHandle,
+        PVOID* BaseAddress,
+        ULONG_PTR ZeroBits,
+        PSIZE_T RegionSize,
+        ULONG AllocationType,
+        ULONG Protect
+    ) {
+        mongo::RemapLock remapLock;
+        return originalNtAllocateVirtualMemory_kernel32( ProcessHandle,
+                                                         BaseAddress,
+                                                         ZeroBits,
+                                                         RegionSize,
+                                                         AllocationType,
+                                                         Protect );
+    }
+
+    NtAllocateVirtualMemory_t originalNtAllocateVirtualMemory_kernelbase;
+
+    static LONG WINAPI myNtAllocateVirtualMemory_kernelbase(
+        HANDLE ProcessHandle,
+        PVOID* BaseAddress,
+        ULONG_PTR ZeroBits,
+        PSIZE_T RegionSize,
+        ULONG AllocationType,
+        ULONG Protect
+    ) {
+        mongo::RemapLock remapLock;
+        return originalNtAllocateVirtualMemory_kernelbase( ProcessHandle,
+                                                           BaseAddress,
+                                                           ZeroBits,
+                                                           RegionSize,
+                                                           AllocationType,
+                                                           Protect );
+    }
+
+} // namespace
+
+namespace mongo {
+
+    void hookWindowsMemory( void ) {
+
+        // kernel32.dll calls directly prior to Windows 7
+        char* hookModuleAddress = reinterpret_cast<char*>( GetModuleHandleA( "kernel32.dll" ) );
+        originalRtlCreateHeap_kernel32 = reinterpret_cast<RtlCreateHeap_t>(
+                        hookWin32( hookModuleAddress,
+                                   "ntdll.dll",
+                                   "RtlCreateHeap",
+                                   myRtlCreateHeap_kernel32 ) );
+        originalNtAllocateVirtualMemory_kernel32 = reinterpret_cast<NtAllocateVirtualMemory_t>(
+                        hookWin32( hookModuleAddress,
+                                   "ntdll.dll",
+                                   "NtAllocateVirtualMemory",
+                                    myNtAllocateVirtualMemory_kernel32 ) );
+
+        // in Windows 7 and Server 2008 R2, calls are through kernelbase.dll
+        hookModuleAddress = reinterpret_cast<char*>( GetModuleHandleA( "kernelbase.dll" ) );
+        originalRtlCreateHeap_kernelbase = reinterpret_cast<RtlCreateHeap_t>(
+                hookWin32( hookModuleAddress,
+                           "ntdll.dll",
+                           "RtlCreateHeap",
+                           myRtlCreateHeap_kernelbase ) );
+        originalNtAllocateVirtualMemory_kernelbase = reinterpret_cast<NtAllocateVirtualMemory_t>(
+                hookWin32( hookModuleAddress,
+                           "ntdll.dll",
+                           "NtAllocateVirtualMemory",
+                           myNtAllocateVirtualMemory_kernelbase ) );
+    }
+
+} //namespace mongo
+
+#endif // #if defined(_WIN32)
diff --git a/src/mongo/util/hook_windows_memory.h b/src/mongo/util/hook_windows_memory.h
new file mode 100644
index 00000000000..b478b6b086d
--- /dev/null
+++ b/src/mongo/util/hook_windows_memory.h
@@ -0,0 +1,25 @@
+// @file hook_windows_memory.h : Used to hook Windows functions that allocate virtual memory
+
+/**
+*    Copyright (C) 2012 10gen Inc.
+*
+*    This program is free software: you can redistribute it and/or modify
+*    it under the terms of the GNU Affero General Public License, version 3,
+*    as published by the Free Software Foundation.
+*
+*    This program is distributed in the hope that it will be useful,
+*    but WITHOUT ANY WARRANTY; without even the implied warranty of
+*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+*    GNU Affero General Public License for more details.
+*
+*    You should have received a copy of the GNU Affero General Public License
+*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#if defined(_WIN32)
+
+namespace mongo {
+    void hookWindowsMemory( void );
+}
+
+#endif // #if defined(_WIN32)
diff --git a/src/mongo/util/mmap_win.cpp b/src/mongo/util/mmap_win.cpp
index d9f8b2809a9..a66b017a370 100644
--- a/src/mongo/util/mmap_win.cpp
+++ b/src/mongo/util/mmap_win.cpp
@@ -263,7 +263,7 @@ namespace mongo {
 
         LockMongoFilesExclusive lockMongoFiles;
 
-        RemapLock lk;   // Interlock with PortMessageServer::acceptedMP() to stop thread creation
+        RemapLock remapLock;
 
         clearWritableBits(oldPrivateAddr);
         if( !UnmapViewOfFile(oldPrivateAddr) ) {