[PATCH v2 3/5] ntoskrnl.exe: Make the driver object accessible by name while in its entry point.

[Zebediah Figura]

Signed-off-by: Zebediah Figura 
---
 dlls/ntoskrnl.exe/ntoskrnl.c     | 11 ++++++-----
 dlls/ntoskrnl.exe/tests/driver.c | 19 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index 17731e7aca..210f42b6b9 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -1537,10 +1537,15 @@ NTSTATUS WINAPI IoCreateDriver( UNICODE_STRING *name, PDRIVER_INITIALIZE init )
     for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
         driver->driver_obj.MajorFunction[i] = unhandled_irp;

+    EnterCriticalSection( &drivers_cs );
+    if (wine_rb_put( &wine_drivers, &driver->driver_obj.DriverName, &driver->entry ))
+        ERR( "failed to insert driver %s in tree\n", debugstr_us(name) );
+    LeaveCriticalSection( &drivers_cs );
+
     status = driver->driver_obj.DriverInit( &driver->driver_obj, &driver->driver_extension.ServiceKeyName );
     if (status)
     {
-        ObDereferenceObject( driver );
+        IoDeleteDriver( &driver->driver_obj );
         return status;
     }

@@ -1550,10 +1555,6 @@ NTSTATUS WINAPI IoCreateDriver( UNICODE_STRING *name, PDRIVER_INITIALIZE init )
         driver->driver_obj.MajorFunction[i] = unhandled_irp;
     }

-    EnterCriticalSection( &drivers_cs );
-    if (wine_rb_put( &wine_drivers, &driver->driver_obj.DriverName, &driver->entry ))
-        ERR( "failed to insert driver %s in tree\n", debugstr_us(name) );
-    LeaveCriticalSection( &drivers_cs );
     return STATUS_SUCCESS;
 }

diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index 9f55111636..c275e67ae5 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -56,7 +56,7 @@ static int running_under_wine;
 static int winetest_debug;
 static int winetest_report_success;

-static POBJECT_TYPE *pExEventObjectType, *pIoFileObjectType, *pPsThreadType;
+static POBJECT_TYPE *pExEventObjectType, *pIoFileObjectType, *pPsThreadType, *pIoDriverObjectType;
 static PEPROCESS *pPsInitialSystemProcess;
 static void *create_caller_thread;

@@ -1887,8 +1887,12 @@ static VOID WINAPI driver_Unload(DRIVER_OBJECT *driver)

NTSTATUS WINAPI DriverEntry(DRIVER_OBJECT *driver, PUNICODE_STRING registry)
 {
+    static const WCHAR IoDriverObjectTypeW[] = {'I','o','D','r','i','v','e','r','O','b','j','e','c','t','T','y','p','e',0};
+    static const WCHAR driver_nameW[] = {'\','D','r','i','v','e','r',
+            '\','W','i','n','e','T','e','s','t','D','r','i','v','e','r',0};
     UNICODE_STRING nameW, linkW;
     NTSTATUS status;
+    void *obj;

DbgPrint("loading driver\n");

@@ -1903,6 +1907,19 @@ NTSTATUS WINAPI DriverEntry(DRIVER_OBJECT *driver, PUNICODE_STRING registry)
     driver->MajorFunction[IRP_MJ_FLUSH_BUFFERS]     = driver_FlushBuffers;
     driver->MajorFunction[IRP_MJ_CLOSE]             = driver_Close;

+    RtlInitUnicodeString(&nameW, IoDriverObjectTypeW);
+    pIoDriverObjectType = MmGetSystemRoutineAddress(&nameW);
+
+    RtlInitUnicodeString(&nameW, driver_nameW);
+    if ((status = ObReferenceObjectByName(&nameW, 0, NULL, 0, *pIoDriverObjectType, KernelMode, NULL, &obj)))
+        return status;
+    if (obj != driver)
+    {
+        ObDereferenceObject(obj);
+        return STATUS_UNSUCCESSFUL;
+    }
+    ObDereferenceObject(obj);
+
     RtlInitUnicodeString(&nameW, device_name);
     RtlInitUnicodeString(&linkW, driver_link);