[2/2] ntoskrnl.exe: Fix management of input/output buffers in dispatch handlers.

[Sebastian Lackner]

Signed-off-by: Sebastian Lackner <sebastian at fds-team.de>
---

Initially I was planning to reuse the in_buff where possible, however I came to
the conclusion that this is not worth the effort. We could only do that for less-
frequently used requests, but usually not for dispatch_ioctl.

Now, all dispatch_*() functions have to release in_buff (either immediately or
deferred), if they return STATUS_SUCCESS. Read/Write uses IRP_DEALLOCATE_BUFFER to
let IoCompleteRequest do that. Ioctls are a bit more complicated - if a separate
output buffer (UserBuffer) is used it is deallocated in the completion routine.
After sending a request to the dispatcher, wine_ntoskrnl_main_loop now allocates
a fresh in_buff with default size (based on the assumption that most of the
packets will be rather small, and that there is no need to keep a huge buffer size).

This patch also fixes a small bug in wine_ntoskrnl_main_loop. in_size was assigned
even when the wineserver returned STATUS_PENDING. When the following request has
in_size != 0 this required an unnecessary additional roundtrip.

 dlls/ntoskrnl.exe/ntoskrnl.c |   47 +++++++++++++++++++++++++++++++------------
 1 file changed, 34 insertions(+), 13 deletions(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index f3b6913..974f946 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -202,7 +202,11 @@ static NTSTATUS WINAPI dispatch_irp_completion( DEVICE_OBJECT *device, IRP *irp,
         irp->Tail.Overlay.OriginalFileObject = NULL;
     }
 
-    HeapFree( GetProcessHeap(), 0, irp->UserBuffer );
+    if (irp->UserBuffer != irp->AssociatedIrp.SystemBuffer)
+    {
+        HeapFree( GetProcessHeap(), 0, irp->UserBuffer );
+        irp->UserBuffer = NULL;
+    }
     return STATUS_SUCCESS;
 }
 
@@ -260,6 +264,7 @@ static NTSTATUS dispatch_create( const irp_params_t *params, void *in_buff, ULON
     irp->Flags |= IRP_CREATE_OPERATION;
     dispatch_irp( device, irp, irp_handle );
 
+    HeapFree( GetProcessHeap(), 0, in_buff );
     return STATUS_SUCCESS;
 }
 
@@ -298,6 +303,7 @@ static NTSTATUS dispatch_close( const irp_params_t *params, void *in_buff, ULONG
     irp->Flags |= IRP_CLOSE_OPERATION;
     dispatch_irp( device, irp, irp_handle );
 
+    HeapFree( GetProcessHeap(), 0, in_buff );
     return STATUS_SUCCESS;
 }
 
@@ -336,8 +342,10 @@ static NTSTATUS dispatch_read( const irp_params_t *params, void *in_buff, ULONG
     irpsp->Parameters.Read.Key = params->read.key;
 
     irp->Flags |= IRP_READ_OPERATION;
+    irp->Flags |= IRP_DEALLOCATE_BUFFER;  /* deallocate out_buff */
     dispatch_irp( device, irp, irp_handle );
 
+    HeapFree( GetProcessHeap(), 0, in_buff );
     return STATUS_SUCCESS;
 }
 
@@ -370,6 +378,7 @@ static NTSTATUS dispatch_write( const irp_params_t *params, void *in_buff, ULONG
     irpsp->Parameters.Write.Key = params->write.key;
 
     irp->Flags |= IRP_WRITE_OPERATION;
+    irp->Flags |= IRP_DEALLOCATE_BUFFER;  /* deallocate in_buff */
     dispatch_irp( device, irp, irp_handle );
 
     return STATUS_SUCCESS;
@@ -398,6 +407,7 @@ static NTSTATUS dispatch_flush( const irp_params_t *params, void *in_buff, ULONG
 
     dispatch_irp( device, irp, irp_handle );
 
+    HeapFree( GetProcessHeap(), 0, in_buff );
     return STATUS_SUCCESS;
 }
 
@@ -425,6 +435,7 @@ static NTSTATUS dispatch_ioctl( const irp_params_t *params, void *in_buff, ULONG
         if ((params->ioctl.code & 3) == METHOD_BUFFERED)
         {
             memcpy( out_buff, in_buff, in_size );
+            HeapFree( GetProcessHeap(), 0, in_buff );
             in_buff = out_buff;
         }
     }
@@ -439,7 +450,9 @@ static NTSTATUS dispatch_ioctl( const irp_params_t *params, void *in_buff, ULONG
 
     irp->Tail.Overlay.OriginalFileObject = file;
     irp->RequestorMode = UserMode;
+    irp->AssociatedIrp.SystemBuffer = in_buff;
 
+    irp->Flags |= IRP_DEALLOCATE_BUFFER;  /* deallocate in_buff */
     dispatch_irp( device, irp, irp_handle );
 
     return STATUS_SUCCESS;
@@ -490,23 +503,23 @@ NTSTATUS CDECL wine_ntoskrnl_main_loop( HANDLE stop_event )
     HANDLE irp = 0;
     NTSTATUS status = STATUS_SUCCESS;
     irp_params_t irp_params;
-    void *in_buff;
     ULONG in_size = 4096, out_size = 0;
+    void *in_buff = NULL;
     HANDLE handles[2];
 
     request_thread = GetCurrentThreadId();
 
-    if (!(in_buff = HeapAlloc( GetProcessHeap(), 0, in_size )))
-    {
-        ERR( "failed to allocate buffer\n" );
-        return STATUS_NO_MEMORY;
-    }
-
     handles[0] = stop_event;
     handles[1] = manager;
 
     for (;;)
     {
+        if (!in_buff && !(in_buff = HeapAlloc( GetProcessHeap(), 0, in_size )))
+        {
+            ERR( "failed to allocate buffer\n" );
+            return STATUS_NO_MEMORY;
+        }
+
         SERVER_START_REQ( get_next_device_request )
         {
             req->manager = wine_server_obj_handle( manager );
@@ -525,13 +538,13 @@ NTSTATUS CDECL wine_ntoskrnl_main_loop( HANDLE stop_event )
             else
             {
                 irp = 0; /* no previous irp */
-                out_size = 0;
-                in_size = reply->in_size;
+                if (status == STATUS_BUFFER_OVERFLOW)
+                    in_size = reply->in_size;
             }
         }
         SERVER_END_REQ;
 
-        switch(status)
+        switch (status)
         {
         case STATUS_SUCCESS:
             if (irp_params.major > IRP_MJ_MAXIMUM_FUNCTION || !dispatch_funcs[irp_params.major])
@@ -541,11 +554,16 @@ NTSTATUS CDECL wine_ntoskrnl_main_loop( HANDLE stop_event )
                 break;
             }
             status = dispatch_funcs[irp_params.major]( &irp_params, in_buff, in_size, out_size, irp );
-            if (status == STATUS_SUCCESS) irp = 0;  /* status reported by IoCompleteRequest */
+            if (status == STATUS_SUCCESS)
+            {
+                irp = 0;  /* status reported by IoCompleteRequest */
+                in_size = 4096;
+                in_buff = NULL;
+            }
             break;
         case STATUS_BUFFER_OVERFLOW:
             HeapFree( GetProcessHeap(), 0, in_buff );
-            in_buff = HeapAlloc( GetProcessHeap(), 0, in_size );
+            in_buff = NULL;
             /* restart with larger buffer */
             break;
         case STATUS_PENDING:
@@ -1451,6 +1469,9 @@ VOID WINAPI IoCompleteRequest( IRP *irp, UCHAR priority_boost )
         }
     }
 
+    if (irp->Flags & IRP_DEALLOCATE_BUFFER)
+        HeapFree( GetProcessHeap(), 0, irp->AssociatedIrp.SystemBuffer );
+
     IoFreeIrp( irp );
 }
 
-- 
2.9.0