[PATCH v10 0/3] MR4243: rtworkq: free cancelled work items

[Yuxuan Shui (＠yshui)]

--
  v10: rtworkq: Release cancelled work items.
       mfplat/tests: Validate MFCancelWorkItem releases the callback.

https://gitlab.winehq.org/wine/wine/-/merge_requests/4243

[PATCH v10 1/3] rtworkq: Avoid use-after-free.

[Yuxuan Shui]

From: Yuxuan Shui

queue_release_pending_item releases the work_item reference but later accesses `item->queue`, which
is a potential use-after-free.
---
 dlls/rtworkq/queue.c | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/dlls/rtworkq/queue.c b/dlls/rtworkq/queue.c
index eebb096ad31..00b77bf6953 100644
--- a/dlls/rtworkq/queue.c
+++ b/dlls/rtworkq/queue.c
@@ -734,9 +734,10 @@ static HRESULT invoke_async_callback(IRtwqAsyncResult *result)
  * removed from pending items when it got canceled. */
 static BOOL queue_release_pending_item(struct work_item *item)
 {
+    struct queue *queue = item->queue;
     BOOL ret = FALSE;

-    EnterCriticalSection(&item->queue->cs);
+    EnterCriticalSection(&queue->cs);
     if (item->key)
     {
         list_remove(&item->entry);
@@ -744,7 +745,7 @@ static BOOL queue_release_pending_item(struct work_item *item)
         item->key = 0;
         IUnknown_Release(&item->IUnknown_iface);
     }
-    LeaveCriticalSection(&item->queue->cs);
+    LeaveCriticalSection(&queue->cs);
     return ret;
 }

[PATCH v10 2/3] mfplat/tests: Validate MFCancelWorkItem releases the callback.

[Yuxuan Shui]

From: Yuxuan Shui

---
 dlls/mfplat/tests/mfplat.c | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/dlls/mfplat/tests/mfplat.c b/dlls/mfplat/tests/mfplat.c
index 6897036e9dd..363c240d13d 100644
--- a/dlls/mfplat/tests/mfplat.c
+++ b/dlls/mfplat/tests/mfplat.c
@@ -3598,6 +3598,7 @@ static void test_scheduled_items(void)
     IMFAsyncResult *result;
     MFWORKITEM_KEY key, key2;
     HRESULT hr;
+    ULONG refcount;

callback = create_test_callback(NULL);

@@ -3610,6 +3611,9 @@ static void test_scheduled_items(void)
     hr = MFCancelWorkItem(key);
     ok(hr == S_OK, "Failed to cancel item, hr %#lx.\n", hr);

+    refcount = IMFAsyncCallback_Release(&callback->IMFAsyncCallback_iface);
+    todo_wine ok(refcount == 0, "Unexpected refcount %lu.\n", refcount);
+
     hr = MFCancelWorkItem(key);
     ok(hr == MF_E_NOT_FOUND || broken(hr == S_OK) /* < win10 */, "Unexpected hr %#lx.\n", hr);

@@ -3619,6 +3623,8 @@ static void test_scheduled_items(void)
         return;
     }

+    callback = create_test_callback(NULL);
+
     hr = MFCreateAsyncResult(NULL, &callback->IMFAsyncCallback_iface, NULL, &result);
     ok(hr == S_OK, "Failed to create result, hr %#lx.\n", hr);

[PATCH v10 3/3] rtworkq: Release cancelled work items.

[Yuxuan Shui]

From: Yuxuan Shui

And avoiding race condition while doing so.

Usually the threadpool holds a reference to the `work_item`, which is released when the `work_item`'s
callback is invoked. However `queue_cancel_item` closes the threadpool object without releasing it,
leading to a leak. The fix is not as simple as adding a `IUnknown_Release` to `queue_cancel_item`,
because the `work_item`'s callback can still be called after CloseThreadpoolTimer/Wait has returned.
In fact its callback might currently be running. In which case the callback will access freed memory
if `queue_cancel_item` frees the `work_item`.

We have to stop any further callbacks to be queued, wait for any currently running callbacks to
finish, then finally we can release the `work_item` if it hasn't already been freed during the wait.
---
 dlls/mfplat/tests/mfplat.c |  2 +-
 dlls/rtworkq/queue.c       | 53 +++++++++++++++++++++++++++++++-------
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/dlls/mfplat/tests/mfplat.c b/dlls/mfplat/tests/mfplat.c
index 363c240d13d..1e751b425fa 100644
--- a/dlls/mfplat/tests/mfplat.c
+++ b/dlls/mfplat/tests/mfplat.c
@@ -3612,7 +3612,7 @@ static void test_scheduled_items(void)
     ok(hr == S_OK, "Failed to cancel item, hr %#lx.\n", hr);

refcount = IMFAsyncCallback_Release(&callback->IMFAsyncCallback_iface);
-    todo_wine ok(refcount == 0, "Unexpected refcount %lu.\n", refcount);
+    ok(refcount == 0, "Unexpected refcount %lu.\n", refcount);

hr = MFCancelWorkItem(key);
     ok(hr == MF_E_NOT_FOUND || broken(hr == S_OK) /* < win10 */, "Unexpected hr %#lx.\n", hr);
diff --git a/dlls/rtworkq/queue.c b/dlls/rtworkq/queue.c
index 00b77bf6953..1fb367da03b 100644
--- a/dlls/rtworkq/queue.c
+++ b/dlls/rtworkq/queue.c
@@ -768,8 +768,9 @@ static void CALLBACK waiting_item_cancelable_callback(TP_CALLBACK_INSTANCE *inst

TRACE("result object %p.\n", item->result);

-    if (queue_release_pending_item(item))
-        invoke_async_callback(item->result);
+    queue_release_pending_item(item);
+
+    invoke_async_callback(item->result);

IUnknown_Release(&item->IUnknown_iface);
 }
@@ -791,8 +792,9 @@ static void CALLBACK scheduled_item_cancelable_callback(TP_CALLBACK_INSTANCE *in

TRACE("result object %p.\n", item->result);

-    if (queue_release_pending_item(item))
-        invoke_async_callback(item->result);
+    queue_release_pending_item(item);
+
+    invoke_async_callback(item->result);

IUnknown_Release(&item->IUnknown_iface);
 }
@@ -884,6 +886,8 @@ static HRESULT queue_submit_timer(struct queue *queue, IRtwqAsyncResult *result,
 static HRESULT queue_cancel_item(struct queue *queue, RTWQWORKITEM_KEY key)
 {
     HRESULT hr = RTWQ_E_NOT_FOUND;
+    TP_WAIT *wait_object;
+    TP_TIMER *timer_object;
     struct work_item *item;

EnterCriticalSection(&queue->cs);
@@ -891,22 +895,53 @@ static HRESULT queue_cancel_item(struct queue *queue, RTWQWORKITEM_KEY key)
     {
         if (item->key == key)
         {
+            /* We can't immediately release the item here, because the callback could already be
+             * running somewhere else. And if we release it here, the callback will access freed memory.
+             * So instead we have to make sure the callback is really stopped, or has really finished
+             * running before we do that. And we can't do that in this critical section, which would be a
+             * deadlock. So we first keep an extra reference to it, then leave the critical section to
+             * wait for the thread-pool objects, finally we re-enter critical section to release it. */
             key >>= 32;
+            IUnknown_AddRef(&item->IUnknown_iface);
             if ((key & WAIT_ITEM_KEY_MASK) == WAIT_ITEM_KEY_MASK)
             {
-                IRtwqAsyncResult_SetStatus(item->result, RTWQ_E_OPERATION_CANCELLED);
-                invoke_async_callback(item->result);
-                CloseThreadpoolWait(item->u.wait_object);
+                wait_object = item->u.wait_object;
                 item->u.wait_object = NULL;
+                LeaveCriticalSection(&queue->cs);
+
+                SetThreadpoolWait(wait_object, NULL, NULL);
+                WaitForThreadpoolWaitCallbacks(wait_object, TRUE);
+                CloseThreadpoolWait(wait_object);
+
+                EnterCriticalSection(&queue->cs);
             }
             else if ((key & SCHEDULED_ITEM_KEY_MASK) == SCHEDULED_ITEM_KEY_MASK)
             {
-                CloseThreadpoolTimer(item->u.timer_object);
+                timer_object = item->u.timer_object;
                 item->u.timer_object = NULL;
+                LeaveCriticalSection(&queue->cs);
+
+                SetThreadpoolTimer(timer_object, NULL, 0, 0);
+                WaitForThreadpoolTimerCallbacks(timer_object, TRUE);
+                CloseThreadpoolTimer(timer_object);
+
+                EnterCriticalSection(&queue->cs);
             }
             else
                 WARN("Unknown item key mask %#I64x.\n", key);
-            queue_release_pending_item(item);
+
+            if (queue_release_pending_item(item))
+            {
+                /* This means the callback wasn't run during our wait, so we can invoke the
+                 * callback with a canceled status, and release the work item. */
+                if ((key & WAIT_ITEM_KEY_MASK) == WAIT_ITEM_KEY_MASK)
+                {
+                    IRtwqAsyncResult_SetStatus(item->result, RTWQ_E_OPERATION_CANCELLED);
+                    invoke_async_callback(item->result);
+                }
+                IUnknown_Release(&item->IUnknown_iface);
+            }
+            IUnknown_Release(&item->IUnknown_iface);
             hr = S_OK;
             break;
         }