[try2] ntoskrnl.exe: Implement Io{Allocate, Get}DriverObjectExtension.

[Alexander Morozov]

Please, use this patch instead of previous "ntoskrnl.exe: Implement 
Io{Allocate,Get}DriverObjectExtension."

From b5441d22fe51640849436b563ca3ab6c9a80c1a5 Mon Sep 17 00:00:00 2001
From: Alexander Morozov 
Date: Wed, 17 Dec 2008 12:44:46 +0300
Subject: [PATCH] ntoskrnl.exe: Implement Io{Allocate,Get}DriverObjectExtension.

---
 dlls/ntoskrnl.exe/ntoskrnl.c |   53 +++++++++++++++++++++++++++++++++++++++--
 1 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index a669d3b..73054e3 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -63,6 +63,16 @@ struct IrpInstance
     IRP *irp;
 };

+static struct list DriverObjExtensions = LIST_INIT(DriverObjExtensions);
+
+struct DriverObjExtension
+{
+    struct list entry;
+    void *ptr;
+    DRIVER_OBJECT *driver;
+    void *id_addr;
+};
+
 #ifdef __i386__
 #define DEFINE_FASTCALL1_ENTRYPOINT( name ) \
     __ASM_GLOBAL_FUNC( name, \
@@ -267,9 +277,28 @@ NTSTATUS WINAPI IoAllocateDriverObjectExtension( PDRIVER_OBJECT DriverObject,
                                                  ULONG DriverObjectExtensionSize,
                                                  PVOID *DriverObjectExtension )
 {
-    FIXME( "%p, %p, %u, %p\n", DriverObject, ClientIdentificationAddress,
+    struct DriverObjExtension *ext;
+
+    TRACE( "%p, %p, %u, %p\n", DriverObject, ClientIdentificationAddress,
             DriverObjectExtensionSize, DriverObjectExtension );
-    return STATUS_NOT_IMPLEMENTED;
+
+    *DriverObjectExtension = NULL;
+    if (IoGetDriverObjectExtension( DriverObject, ClientIdentificationAddress ))
+        return STATUS_OBJECT_NAME_COLLISION;
+    ext = ExAllocatePool( NonPagedPool, sizeof(*ext) );
+    if (ext == NULL)
+        return STATUS_INSUFFICIENT_RESOURCES;
+    ext->ptr = ExAllocatePool( NonPagedPool, DriverObjectExtensionSize );
+    if (ext->ptr == NULL)
+    {
+        ExFreePool( ext );
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+    ext->driver = DriverObject;
+    ext->id_addr = ClientIdentificationAddress;
+    list_add_tail( &DriverObjExtensions, &ext->entry );
+    *DriverObjectExtension = ext->ptr;
+    return STATUS_SUCCESS;
 }

@@ -279,7 +308,16 @@ NTSTATUS WINAPI IoAllocateDriverObjectExtension( PDRIVER_OBJECT DriverObject,
 PVOID WINAPI IoGetDriverObjectExtension( PDRIVER_OBJECT DriverObject,
                                          PVOID ClientIdentificationAddress )
 {
-    FIXME( "%p, %p\n", DriverObject, ClientIdentificationAddress );
+    struct DriverObjExtension *ext;
+
+    TRACE( "%p, %p\n", DriverObject, ClientIdentificationAddress );
+
+    LIST_FOR_EACH_ENTRY( ext, &DriverObjExtensions, struct DriverObjExtension, entry )
+    {
+        if (DriverObject == ext->driver &&
+            ClientIdentificationAddress == ext->id_addr)
+            return ext->ptr;
+    }
     return NULL;
 }

@@ -1138,6 +1176,7 @@ PVOID WINAPI MmGetSystemRoutineAddress(PUNICODE_STRING SystemRoutineName)
 BOOL WINAPI DllMain( HINSTANCE inst, DWORD reason, LPVOID reserved )
 {
     LARGE_INTEGER count;
+    struct DriverObjExtension *ext, *ext2;

switch(reason)
     {
@@ -1146,6 +1185,14 @@ BOOL WINAPI DllMain( HINSTANCE inst, DWORD reason, LPVOID reserved )
         RtlAddVectoredExceptionHandler( TRUE, vectored_handler );
         KeQueryTickCount( &count );  /* initialize the global KeTickCount */
         break;
+    case DLL_PROCESS_DETACH:
+        LIST_FOR_EACH_ENTRY_SAFE( ext, ext2, &DriverObjExtensions,
+                struct DriverObjExtension, entry )
+        {
+            list_remove( &ext->entry );
+            ExFreePool( ext->ptr );
+            ExFreePool( ext );
+        }
     }
     return TRUE;
 }