pe_open_image_from_addr(): use pe_get_ldr_entry_from_addr() as the backend.

2 files changed, 38 insertions, 37 deletions

diff --git a/src/internal/pe_os.h b/src/internal/pe_os.h
index 1544e44..11e334b 100644
--- a/src/internal/pe_os.h
+++ b/src/internal/pe_os.h
@@ -13,6 +13,7 @@
 #define OS_STATUS_COULD_NOT_INTERPRET		(int32_t)0xC00000B9
 #define OS_STATUS_NOT_SUPPORTED			(int32_t)0xC00000BB
 #define OS_STATUS_NAME_TOO_LONG			(int32_t)0xC0000106
+#define OS_STATUS_BUFFER_TOO_SMALL		(int32_t)0xC0000023
 #define OS_STATUS_INTERNAL_ERROR		(int32_t)0xC00000E5
 #define OS_STATUS_BAD_FILE_TYPE			(int32_t)0xC0000903
 #define OS_STATUS_OBJECT_NAME_NOT_FOUND 	(int32_t)0xC0000034
@@ -133,15 +134,6 @@ typedef int32_t __stdcall os_zw_query_object(
 	__out	uint32_t *		returned_length		__optional);
 
 
-typedef int32_t __stdcall os_zw_query_virtual_memory(
-	__in	void *			hprocess,
-	__in	void *			base_address,
-	__in	int			mem_info_class,
-	__out	void *			mem_info,
-	__in	size_t			mem_info_length,
-	__out	uint32_t *		returned_length	__optional);
-
-
 typedef int32_t __stdcall os_zw_read_virtual_memory(
 	__in	void *			hprocess,
 	__in	void *			base_address,
diff --git a/src/ldso/pe_open_image_from_addr.c b/src/ldso/pe_open_image_from_addr.c
index b8a1845..e241fac 100644
--- a/src/ldso/pe_open_image_from_addr.c
+++ b/src/ldso/pe_open_image_from_addr.c
@@ -9,11 +9,6 @@
 #include <pemagine/pe_structs.h>
 #include "pe_os.h"
 
-struct os_memory_section_name {
-	struct pe_unicode_str	section_name;
-	wchar16_t		section_name_buffer[];
-};
-
 pe_api int32_t pe_open_image_from_addr(
 	__out	void **			himage,
 	__in	void *			addr,
@@ -24,47 +19,61 @@ pe_api int32_t pe_open_image_from_addr(
 	__in	uint32_t		share_access,
 	__in	uint32_t		open_options)
 {
-	int32_t				status;
 	struct os_oa			oa;
 	struct os_iosb			iosb;
-	struct os_memory_section_name *	path;
-	uint32_t			len;
 	void *				hntdll;
-	os_zw_query_virtual_memory *	zw_query_virtual_memory;
 	os_zw_open_file *		zw_open_file;
-
+	struct pe_unicode_str		path;
+	struct pe_ldr_tbl_entry *	lentry;
+	wchar16_t *			name;
+	wchar16_t *			cap;
+	wchar16_t *			src;
+	wchar16_t *			dst;
+	wchar16_t *			mark;
 
 	/* init */
-	path = (struct os_memory_section_name *)buffer;
-	path->section_name.strlen = 0;
-	path->section_name.maxlen = (uint16_t)(bufsize - sizeof(struct pe_unicode_str));
-	path->section_name.buffer = path->section_name_buffer;
-
 	if (!(hntdll = pe_get_ntdll_module_handle()))
 		return OS_STATUS_INTERNAL_ERROR;
 
-	if (!(zw_query_virtual_memory = (os_zw_query_virtual_memory *)pe_get_procedure_address(
-			hntdll,"ZwQueryVirtualMemory")))
-		return OS_STATUS_INTERNAL_ERROR;
-
 	if (!(zw_open_file = (os_zw_open_file *)pe_get_procedure_address(
 			hntdll,"ZwOpenFile")))
 		return OS_STATUS_INTERNAL_ERROR;
 
 	/* native path of image containing addr */
-	if ((status = zw_query_virtual_memory(
-			OS_CURRENT_PROCESS_HANDLE,
-			addr,
-			OS_MEMORY_SECTION_NAME,
-			buffer,
-			bufsize,
-			&len)))
-		return status;
+	if (!(lentry = pe_get_ldr_entry_from_addr(addr)))
+		return OS_STATUS_INTERNAL_ERROR;
+
+	if (bufsize - 4*sizeof(wchar16_t) < lentry->full_dll_name.strlen)
+		return OS_STATUS_BUFFER_TOO_SMALL;
+
+	name = lentry->full_dll_name.buffer;
+	mark = (wchar16_t *)buffer;
+	dst  = mark;
+
+	if (name[1] == ':') {
+		*dst++ = '\\';
+		*dst++ = '?';
+		*dst++ = '?';
+		*dst++ = '\\';
+	}
+
+	src = name;
+	dst = mark;
+	cap  = &dst[lentry->full_dll_name.strlen / sizeof(wchar16_t)];
+
+	for (; dst<cap; )
+		*dst++ = *src++;
+
+	*dst = 0;
+
+	path.strlen = sizeof(wchar16_t) * (cap - mark);
+	path.maxlen = 0;
+	path.buffer = mark;
 
 	/* oa */
 	oa.len      = sizeof(struct os_oa);
 	oa.root_dir = 0;
-	oa.obj_name = &path->section_name;
+	oa.obj_name = &path;
 	oa.obj_attr = oattr;
 	oa.sec_desc = 0;
 	oa.sec_qos  = 0;