[PATCH] ntdll: Release IO thread pool object from ioqueue_thread_proc().

[Paul Gofman]

Signed-off-by: Paul Gofman <pgofman(a)codeweavers.com>
---
    When TpReleaseIoCompletion() is called it frees the TP_IO, but there might still be
    outstanding completions. When one arrives, ioqueue_thread_proc() may accesses
    TP_IO after free, and even if it doesn't crashes on "assert( !object->shutdown );"
    in tp_object_submit().
    Spotted in Forza Horizon 4.

 dlls/ntdll/threadpool.c | 38 +++++++++++++++++++-------------------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/dlls/ntdll/threadpool.c b/dlls/ntdll/threadpool.c
index 1b51a191979..6f03e96780c 100644
--- a/dlls/ntdll/threadpool.c
+++ b/dlls/ntdll/threadpool.c
@@ -1519,10 +1519,21 @@ static void CALLBACK ioqueue_thread_proc( void *param )
             ERR("NtRemoveIoCompletion failed, status %#x.\n", status);
         RtlEnterCriticalSection( &ioqueue.cs );
 
-        if (key)
-        {
-            io = (struct threadpool_object *)key;
+        io = (struct threadpool_object *)key;
 
+        if (io && io->shutdown)
+        {
+            if (iosb.u.Status != STATUS_THREADPOOL_RELEASED_DURING_OPERATION)
+            {
+                /* Skip remaining completions until the final one. */
+                continue;
+            }
+            --ioqueue.objcount;
+            TRACE( "Releasing io %p.\n", io );
+            tp_object_release( io );
+        }
+        else if (io)
+        {
             RtlEnterCriticalSection( &io->pool->cs );
 
             if (!array_reserve((void **)&io->u.io.completions, &io->u.io.completion_max,
@@ -1609,18 +1620,6 @@ static NTSTATUS tp_ioqueue_lock( struct threadpool_object *io, HANDLE file )
     return status;
 }
 
-static void tp_ioqueue_unlock( struct threadpool_object *io )
-{
-    assert( io->type == TP_OBJECT_TYPE_IO );
-
-    RtlEnterCriticalSection( &ioqueue.cs );
-
-    if (!--ioqueue.objcount)
-        NtSetIoCompletion( ioqueue.port, 0, 0, STATUS_SUCCESS, 0 );
-
-    RtlLeaveCriticalSection( &ioqueue.cs );
-}
-
 /***********************************************************************
  *           tp_threadpool_alloc    (internal)
  *
@@ -2051,8 +2050,6 @@ static void tp_object_prepare_shutdown( struct threadpool_object *object )
         tp_timerqueue_unlock( object );
     else if (object->type == TP_OBJECT_TYPE_WAIT)
         tp_waitqueue_unlock( object );
-    else if (object->type == TP_OBJECT_TYPE_IO)
-        tp_ioqueue_unlock( object );
 }
 
 /***********************************************************************
@@ -2796,9 +2793,12 @@ void WINAPI TpReleaseIoCompletion( TP_IO *io )
 
     TRACE( "%p\n", io );
 
-    tp_object_prepare_shutdown( this );
+    RtlEnterCriticalSection( &ioqueue.cs );
+
+    assert( ioqueue.objcount );
     this->shutdown = TRUE;
-    tp_object_release( this );
+    NtSetIoCompletion( ioqueue.port, (ULONG_PTR)this, 0, STATUS_THREADPOOL_RELEASED_DURING_OPERATION, 1 );
+    RtlLeaveCriticalSection( &ioqueue.cs );
 }
 
 /***********************************************************************
-- 
2.30.2