From 033911d0ffa830ad5f6de4ebb657405d411ac7e4 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Sat, 7 Nov 2009 03:33:00 +0000
Subject: [PATCH] [RTL] Simplify RtlLookupFunctionTable

svn path=/branches/ros-amd64-bringup/; revision=43993
---
 reactos/lib/rtl/amd64/unwind.c | 48 ++++++++++++++++++----------------
 1 file changed, 26 insertions(+), 22 deletions(-)

diff --git a/reactos/lib/rtl/amd64/unwind.c b/reactos/lib/rtl/amd64/unwind.c
index 5d1597f9b72..276981df6cb 100644
--- a/reactos/lib/rtl/amd64/unwind.c
+++ b/reactos/lib/rtl/amd64/unwind.c
@@ -60,6 +60,17 @@ typedef struct _UNWIND_INFO
 
 /* FUNCTIONS *****************************************************************/
 
+/*! RtlLookupFunctionTable
+ * \brief Locates the table of RUNTIME_FUNCTION entries for a code address.
+ * \param ControlPc
+ *            Address of the code, for which the table should be searched.
+ * \param ImageBase 
+ *            Pointer to a DWORD64 that receives the base address of the
+ *            corresponding executable image.
+ * \param Length
+ *            Pointer to an ULONG that receives the number of table entries
+ *            present in the table.
+ */
 PRUNTIME_FUNCTION
 NTAPI
 RtlLookupFunctionTable(
@@ -67,34 +78,27 @@ RtlLookupFunctionTable(
     OUT PDWORD64 ImageBase,
     OUT PULONG Length)
 {
-    PIMAGE_DOS_HEADER DosHeader;
-    PIMAGE_NT_HEADERS NtHeader;
-    PIMAGE_DATA_DIRECTORY Directory;
+    PVOID Table;
+    ULONG Size;
 
-    /* Find ModuleBase */
-    if (!RtlPcToFileHeader((PVOID)ControlPc, (PVOID*)&DosHeader))
-    {
-        return NULL;
-    }
-
-    /* Locate NT header and check number of directories */
-    NtHeader = (PVOID)((ULONG64)DosHeader + DosHeader->e_lfanew);
-    if (NtHeader->OptionalHeader.NumberOfRvaAndSizes 
-         < IMAGE_DIRECTORY_ENTRY_EXCEPTION)
+    /* Find corresponding file header from code address */
+    if (!RtlPcToFileHeader((PVOID)ControlPc, (PVOID*)ImageBase))
     {
+        /* Nothing found */
         return NULL;
     }
 
     /* Locate the exception directory */
-    Directory = &NtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
-    *Length = Directory->Size / sizeof(RUNTIME_FUNCTION);
-    *ImageBase = (ULONG64)DosHeader;
-    if (!Directory->VirtualAddress)
-    {
-        return NULL;
-    }
+    Table = RtlImageDirectoryEntryToData((PVOID)*ImageBase,
+                                         TRUE,
+                                         IMAGE_DIRECTORY_ENTRY_EXCEPTION,
+                                         &Size);
 
-    return (PVOID)((ULONG64)DosHeader + Directory->VirtualAddress);
+    /* Return the number of entries */
+    *Length = Size / sizeof(RUNTIME_FUNCTION);
+
+    /* Return the address of the table */
+    return Table;
 }
 
 /*! RtlLookupFunctionEntry
@@ -119,7 +123,7 @@ RtlLookupFunctionEntry(
     /* Fail, if no table is found */
     if (!FunctionTable)
     {
-        return (PVOID)1;
+        return NULL;
     }
 
     /* Use relative virtual address */