From: Aric Stewart Subject: [PATCH 02/03 v5] ntoskrnl.exe: Track drivers created with IoCreateDriver Message-Id: Date: Mon, 25 Jul 2016 07:24:50 -0500 v2: Suggestions from Sebastian Lackner including locking v3: be sure to free DriverName and remove duplicate driver_name v4: Use a rb_tree instead of a list, since this will be used for searching far more than iterating v5: Corrections found by Thomas Faber Signed-off-by: Aric Stewart --- dlls/ntoskrnl.exe/ntoskrnl.c | 108 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 90 insertions(+), 18 deletions(-) diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index 4ed06b8..2c79383 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -39,9 +39,10 @@ #include "ddk/wdm.h" #include "wine/unicode.h" #include "wine/server.h" -#include "wine/list.h" #include "wine/debug.h" +#include "wine/rbtree.h" + WINE_DEFAULT_DEBUG_CHANNEL(ntoskrnl); WINE_DECLARE_DEBUG_CHANNEL(relay); @@ -72,6 +73,37 @@ static DWORD request_thread; static DWORD client_tid; static DWORD client_pid; +struct wine_driver { + struct wine_rb_entry entry; + + DRIVER_OBJECT driver_obj; + DRIVER_EXTENSION driver_extension; +}; + +struct wine_rb_tree *wine_drivers = NULL; + +static CRITICAL_SECTION drivers_cs; +static CRITICAL_SECTION_DEBUG critsect_debug = +{ + 0, 0, &drivers_cs, + { &critsect_debug.ProcessLocksList, &critsect_debug.ProcessLocksList }, + 0, 0, { (DWORD_PTR)(__FILE__ ": drivers_cs") } +}; +static CRITICAL_SECTION drivers_cs = { &critsect_debug, -1, 0, 0, 0, 0 }; + +static void *wine_drivers_rb_alloc(size_t size); +static void *wine_drivers_rb_realloc(void *ptr, size_t size); +static void wine_drivers_rb_free(void *ptr); +static int wine_drivers_rb_compare(const void *key, const struct wine_rb_entry *entry); + +static const struct wine_rb_functions wine_drivers_rb_functions = +{ + wine_drivers_rb_alloc, + wine_drivers_rb_realloc, + wine_drivers_rb_free, + wine_drivers_rb_compare, +}; + #ifdef __i386__ #define DEFINE_FASTCALL1_ENTRYPOINT( name ) \ __ASM_STDCALL_FUNC( name, 4, \ @@ -850,41 +882,76 @@ static void build_driver_keypath(UNICODE_STRING *name, UNICODE_STRING *keypath) RtlInitUnicodeString( keypath, str ); } +static void *wine_drivers_rb_alloc(size_t size) +{ + return HeapAlloc(GetProcessHeap(), 0, size); +} + +static void *wine_drivers_rb_realloc(void *ptr, size_t size) +{ + return HeapReAlloc(GetProcessHeap(), 0, ptr, size); +} + +static void wine_drivers_rb_free(void *ptr) +{ + HeapFree(GetProcessHeap(), 0, ptr); +} + +static int wine_drivers_rb_compare(const void *key, const struct wine_rb_entry *entry) +{ + const UNICODE_STRING *k = key; + const struct wine_driver *driver = WINE_RB_ENTRY_VALUE(entry, + const struct wine_driver, entry); + + return RtlCompareUnicodeString(k, &driver->driver_obj.DriverName, FALSE); +} + /*********************************************************************** * IoCreateDriver (NTOSKRNL.EXE.@) */ NTSTATUS WINAPI IoCreateDriver( UNICODE_STRING *name, PDRIVER_INITIALIZE init ) { - DRIVER_OBJECT *driver; - DRIVER_EXTENSION *extension; + struct wine_driver *driver; NTSTATUS status; TRACE("(%s, %p)\n", debugstr_us(name), init); if (!(driver = RtlAllocateHeap( GetProcessHeap(), HEAP_ZERO_MEMORY, - sizeof(*driver) + sizeof(*extension) ))) + sizeof(*driver) ))) return STATUS_NO_MEMORY; - if ((status = RtlDuplicateUnicodeString( 1, name, &driver->DriverName ))) + if ((status = RtlDuplicateUnicodeString( 1, name, &driver->driver_obj.DriverName ))) { RtlFreeHeap( GetProcessHeap(), 0, driver ); return status; } - extension = (DRIVER_EXTENSION *)(driver + 1); - driver->Size = sizeof(*driver); - driver->DriverInit = init; - driver->DriverExtension = extension; - extension->DriverObject = driver; - build_driver_keypath( &driver->DriverName, &extension->ServiceKeyName ); + driver->driver_obj.Size = sizeof(driver->driver_obj); + driver->driver_obj.DriverInit = init; + driver->driver_obj.DriverExtension = &driver->driver_extension; + driver->driver_extension.DriverObject = &driver->driver_obj; - status = driver->DriverInit( driver, &extension->ServiceKeyName ); + build_driver_keypath( &driver->driver_obj.DriverName, &driver->driver_extension.ServiceKeyName ); + + status = driver->driver_obj.DriverInit( &driver->driver_obj, &driver->driver_extension.ServiceKeyName ); if (status) { - RtlFreeUnicodeString( &driver->DriverName ); + RtlFreeUnicodeString( &driver->driver_obj.DriverName); RtlFreeHeap( GetProcessHeap(), 0, driver ); } + else + { + EnterCriticalSection( &drivers_cs ); + if (!wine_drivers) + { + wine_drivers = HeapAlloc( GetProcessHeap(), 0, sizeof(*wine_drivers) ); + wine_rb_init(wine_drivers, &wine_drivers_rb_functions); + } + wine_rb_put( wine_drivers, &driver->driver_obj.DriverName, &driver->entry ); + LeaveCriticalSection( &drivers_cs ); + } + return status; } @@ -892,13 +959,18 @@ NTSTATUS WINAPI IoCreateDriver( UNICODE_STRING *name, PDRIVER_INITIALIZE init ) /*********************************************************************** * IoDeleteDriver (NTOSKRNL.EXE.@) */ -void WINAPI IoDeleteDriver( DRIVER_OBJECT *driver ) +void WINAPI IoDeleteDriver( DRIVER_OBJECT *driver_object ) { - TRACE("(%p)\n", driver); + TRACE("(%p)\n", driver_object); + + EnterCriticalSection( &drivers_cs ); + if (wine_drivers) + wine_rb_remove( wine_drivers, &driver_object->DriverName ); + LeaveCriticalSection( &drivers_cs ); - RtlFreeUnicodeString( &driver->DriverName ); - RtlFreeUnicodeString( &driver->DriverExtension->ServiceKeyName ); - RtlFreeHeap( GetProcessHeap(), 0, driver ); + RtlFreeUnicodeString( &driver_object->DriverName ); + RtlFreeUnicodeString( &driver_object->DriverExtension->ServiceKeyName ); + RtlFreeHeap( GetProcessHeap(), 0, CONTAINING_RECORD( driver_object, struct wine_driver, driver_obj ) ); }