Merge pull request #76 from fancycode/binsearch_export_names

Implement binary search for looking up exports by name.

.gitignore

@@ -3,3 +3,6 @@
 *.exe
 tests/*.dll
 tests/*.res
+
+tests/SampleExports.cpp
+tests/SampleExports.h

MemoryModule.c

@@ -22,6 +22,12 @@
  * Portions created by Joachim Bauch are Copyright (C) 2004-2015
  * Joachim Bauch. All Rights Reserved.
  *
+ *
+ * THeller: Added binary search in MemoryGetProcAddress function
+ * (#define USE_BINARY_SEARCH to enable it).  This gives a very large
+ * speedup for libraries that exports lots of functions.
+ *
+ * These portions are Copyright (C) 2013 Thomas Heller.
  */
 
 #include <windows.h>
@@ -56,6 +62,11 @@
 
 #include "MemoryModule.h"
 
+struct ExportNameEntry {
+    LPCSTR name;
+    WORD idx;
+};
+
 typedef BOOL (WINAPI *DllEntryProc)(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved);
 typedef int (WINAPI *ExeEntryProc)(void);
 
@@ -72,6 +83,7 @@ typedef struct {
     CustomLoadLibraryFunc loadLibrary;
     CustomGetProcAddressFunc getProcAddress;
     CustomFreeLibraryFunc freeLibrary;
+    struct ExportNameEntry *nameExportsTable;
     void *userdata;
     ExeEntryProc exeEntry;
     DWORD pageSize;
@@ -688,12 +700,27 @@ error:
     return NULL;
 }
 
-FARPROC MemoryGetProcAddress(HMEMORYMODULE module, LPCSTR name)
+static int _compare(const void *a, const void *b)
 {
-    unsigned char *codeBase = ((PMEMORYMODULE)module)->codeBase;
+    const struct ExportNameEntry *p1 = (const struct ExportNameEntry*) a;
+    const struct ExportNameEntry *p2 = (const struct ExportNameEntry*) b;
+    return _stricmp(p1->name, p2->name);
+}
+
+static int _find(const void *a, const void *b)
+{
+    LPCSTR *name = (LPCSTR *) a;
+    const struct ExportNameEntry *p = (const struct ExportNameEntry*) b;
+    return _stricmp(*name, p->name);
+}
+
+FARPROC MemoryGetProcAddress(HMEMORYMODULE mod, LPCSTR name)
+{
+    PMEMORYMODULE module = (PMEMORYMODULE)mod;
+    unsigned char *codeBase = module->codeBase;
     DWORD idx = 0;
     PIMAGE_EXPORT_DIRECTORY exports;
-    PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY((PMEMORYMODULE)module, IMAGE_DIRECTORY_ENTRY_EXPORT);
+    PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_EXPORT);
     if (directory->Size == 0) {
         // no export table found
         SetLastError(ERROR_PROC_NOT_FOUND);
@@ -715,25 +742,44 @@ FARPROC MemoryGetProcAddress(HMEMORYMODULE module, LPCSTR name)
         }
 
         idx = LOWORD(name) - exports->Base;
+    } else if (!exports->NumberOfNames) {
+        SetLastError(ERROR_PROC_NOT_FOUND);
+        return NULL;
     } else {
-        // search function name in list of exported names
-        DWORD i;
-        DWORD *nameRef = (DWORD *) (codeBase + exports->AddressOfNames);
-        WORD *ordinal = (WORD *) (codeBase + exports->AddressOfNameOrdinals);
-        BOOL found = FALSE;
-        for (i=0; i<exports->NumberOfNames; i++, nameRef++, ordinal++) {
-            if (_stricmp(name, (const char *) (codeBase + (*nameRef))) == 0) {
-                idx = *ordinal;
-                found = TRUE;
-                break;
+        const struct ExportNameEntry *found;
+
+        // Lazily build name table and sort it by names
+        if (!module->nameExportsTable) {
+            DWORD i;
+            DWORD *nameRef = (DWORD *) (codeBase + exports->AddressOfNames);
+            WORD *ordinal = (WORD *) (codeBase + exports->AddressOfNameOrdinals);
+            struct ExportNameEntry *entry = (struct ExportNameEntry*) malloc(exports->NumberOfNames * sizeof(struct ExportNameEntry));
+            module->nameExportsTable = entry;
+            if (!entry) {
+                SetLastError(ERROR_OUTOFMEMORY);
+                return NULL;
             }
+            for (i=0; i<exports->NumberOfNames; i++, nameRef++, ordinal++, entry++) {
+                entry->name = (const char *) (codeBase + (*nameRef));
+                entry->idx = *ordinal;
+            }
+            qsort(module->nameExportsTable,
+                    exports->NumberOfNames,
+                    sizeof(struct ExportNameEntry), _compare);
         }
 
+        // search function name in list of exported names with binary search
+        found = (const struct ExportNameEntry*) bsearch(&name,
+                module->nameExportsTable,
+                exports->NumberOfNames,
+                sizeof(struct ExportNameEntry), _find);
         if (!found) {
             // exported symbol not found
             SetLastError(ERROR_PROC_NOT_FOUND);
             return NULL;
         }
+
+        idx = found->idx;
     }
 
     if (idx > exports->NumberOfFunctions) {
@@ -759,6 +805,7 @@ void MemoryFreeLibrary(HMEMORYMODULE mod)
         (*DllEntry)((HINSTANCE)module->codeBase, DLL_PROCESS_DETACH, 0);
     }
 
+    free(module->nameExportsTable);
     if (module->modules != NULL) {
         // free previously opened libraries
         int i;

tests/LoadDll.cpp

@@ -11,6 +11,7 @@
 
 #include "../MemoryModule.h"
 
+typedef int (*addProc)(int);
 typedef int (*addNumberProc)(int, int);
 
 // Thanks to Tim Cooper (from http://stackoverflow.com/a/8584708)
@@ -127,6 +128,11 @@ BOOL LoadFromMemory(char *filename)
     }
 
     addNumber = (addNumberProc)MemoryGetProcAddress(handle, "addNumbers");
+    if (!addNumber) {
+        _tprintf(_T("MemoryGetProcAddress(\"addNumber\") returned NULL\n"));
+        result = FALSE;
+        goto exit;
+    }
     _tprintf(_T("From memory: %d\n"), addNumber(1, 2));
 
     // the DLL only exports one function, try to load by ordinal value
@@ -218,6 +224,65 @@ exit:
     return result;
 }
 
+BOOL LoadExportsFromMemory(char *filename)
+{
+    FILE *fp;
+    unsigned char *data=NULL;
+    long size;
+    size_t read;
+    HMEMORYMODULE handle = NULL;
+    int i;
+    BOOL result = TRUE;
+
+    fp = fopen(filename, "rb");
+    if (fp == NULL)
+    {
+        printf("Can't open DLL file \"%s\".", filename);
+        result = FALSE;
+        goto exit;
+    }
+
+    fseek(fp, 0, SEEK_END);
+    size = ftell(fp);
+    assert(size > 0);
+    data = (unsigned char *)malloc(size);
+    assert(data != NULL);
+    fseek(fp, 0, SEEK_SET);
+    read = fread(data, 1, size, fp);
+    assert(read == static_cast<size_t>(size));
+    fclose(fp);
+
+    handle = MemoryLoadLibrary(data, size);
+    if (handle == NULL)
+    {
+        _tprintf(_T("Can't load library from memory.\n"));
+        result = FALSE;
+        goto exit;
+    }
+
+    for (i = 1; i <= 100; i++) {
+        char name[100];
+        sprintf(name, "add%d", i);
+        addProc addNumber = (addProc)MemoryGetProcAddress(handle, name);
+        if (!addNumber) {
+            _tprintf(_T("MemoryGetProcAddress(\"%s\") returned NULL\n"), name);
+            result = FALSE;
+            goto exit;
+        }
+        int result = addNumber(1);
+        if (result != 1 + i) {
+            _tprintf(_T("(\"%s\") returned %d, expected %d\n"), name, result, 1 + i);
+            result = FALSE;
+            goto exit;
+        }
+        _tprintf(_T("%s: %d\n"), name, result);
+    }
+exit:
+    MemoryFreeLibrary(handle);
+    free(data);
+    return result;
+}
+
 int main(int argc, char* argv[])
 {
     if (argc < 2) {
@@ -225,8 +290,14 @@ int main(int argc, char* argv[])
         return 1;
     }
 
-    if (!LoadFromMemory(argv[1])) {
-        return 2;
+    if (!strstr((const char *) argv[1], "exports")) {
+        if (!LoadFromMemory(argv[1])) {
+            return 2;
+        }
+    } else {
+        if (!LoadExportsFromMemory(argv[1])) {
+            return 2;
+        }
     }
 
     return 0;

tests/Makefile

@@ -47,6 +47,7 @@ TEST_DLLS = \
 	test-align-800.dll \
 	test-align-900.dll \
 	test-relocate.dll \
+	test-exports.dll
 
 LOADDLL_OBJ = LoadDll.o ../MemoryModule.o
 TESTSUITE_OBJ = TestSuite.o ../MemoryModule.o
@@ -72,6 +73,12 @@ test-align-%.dll: $(DLL_OBJ)
 test-relocate.dll: $(DLL_OBJ)
 	$(CXX) $(LDFLAGS_DLL) $(LDFLAGS) -Wl,--image-base -Wl,0x20000000 -o $@ $(DLL_OBJ)
 
+test-exports.dll: SampleExports.o
+	$(CXX) $(LDFLAGS_DLL) $(LDFLAGS) -o $@ SampleExports.o
+
+SampleExports.cpp: generate-exports.sh
+	./generate-exports.sh
+
 %.o: %.cpp
 	$(CXX) $(CFLAGS) $(CFLAGS_DLL) -c $<
 

tests/generate-exports.sh

@@ -0,0 +1,52 @@
+#!/bin/sh
+
+##
+## Generate header file.
+##
+
+cat > SampleExports.h << EOF
+extern "C" {
+
+#ifdef SAMPLEDLL_EXPORTS
+#define SAMPLEDLL_API __declspec(dllexport)
+#else
+#define SAMPLEDLL_API __declspec(dllimport)
+#endif
+
+EOF
+
+for i in `seq 1 100`;
+do
+cat >> SampleExports.h << EOF
+SAMPLEDLL_API int add$i(int a);
+EOF
+done
+
+cat >> SampleExports.h << EOF
+}
+EOF
+
+
+##
+## Generate source file.
+##
+
+cat > SampleExports.cpp << EOF
+#include "SampleExports.h"
+
+extern "C" {
+EOF
+
+for i in `seq 1 100 | sort -R`;
+do
+cat >> SampleExports.cpp << EOF
+SAMPLEDLL_API int add$i(int a)
+{
+    return a + $i;
+}
+EOF
+done
+
+cat >> SampleExports.cpp << EOF
+}
+EOF