From 609c311a81b87ed351a4cf02478552c737fab716 Mon Sep 17 00:00:00 2001
From: Nikolay Sivov <nsivov@codeweavers.com>
Date: Wed, 16 Sep 2020 16:02:40 +0300
Subject: [PATCH] combase: Manage per-thread call cancellation counter.

Signed-off-by: Nikolay Sivov <nsivov@codeweavers.com>
Signed-off-by: Huw Davies <huw@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
---
 dlls/combase/combase.c         | 27 +++++++++++++---
 dlls/combase/combase_private.h |  1 +
 dlls/ole32/compobj_private.h   |  1 +
 dlls/ole32/tests/compobj.c     | 59 ++++++++++++++++++++++++++++++++++
 include/winerror.h             |  1 +
 5 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/dlls/combase/combase.c b/dlls/combase/combase.c
index c66f0bbf8e2..7a85b95a19b 100644
--- a/dlls/combase/combase.c
+++ b/dlls/combase/combase.c
@@ -3104,9 +3104,20 @@ HRESULT WINAPI CoRegisterChannelHook(REFGUID guidExtension, IChannelHook *channe
  */
 HRESULT WINAPI CoDisableCallCancellation(void *reserved)
 {
-    FIXME("%p stub\n", reserved);
+    struct tlsdata *tlsdata;
+    HRESULT hr;
 
-    return E_NOTIMPL;
+    TRACE("%p\n", reserved);
+
+    if (FAILED(hr = com_get_tlsdata(&tlsdata)))
+        return hr;
+
+    if (!tlsdata->cancelcount)
+        return CO_E_CANCEL_DISABLED;
+
+    tlsdata->cancelcount--;
+
+    return S_OK;
 }
 
 /***********************************************************************
@@ -3114,9 +3125,17 @@ HRESULT WINAPI CoDisableCallCancellation(void *reserved)
  */
 HRESULT WINAPI CoEnableCallCancellation(void *reserved)
 {
-    FIXME("%p stub\n", reserved);
+    struct tlsdata *tlsdata;
+    HRESULT hr;
 
-    return E_NOTIMPL;
+    TRACE("%p\n", reserved);
+
+    if (FAILED(hr = com_get_tlsdata(&tlsdata)))
+        return hr;
+
+    tlsdata->cancelcount++;
+
+    return S_OK;
 }
 
 /***********************************************************************
diff --git a/dlls/combase/combase_private.h b/dlls/combase/combase_private.h
index 2c938e6ff49..55ce02d06f0 100644
--- a/dlls/combase/combase_private.h
+++ b/dlls/combase/combase_private.h
@@ -84,6 +84,7 @@ struct tlsdata
     IUnknown         *state;         /* see CoSetState */
     struct list       spies;         /* Spies installed with CoRegisterInitializeSpy */
     DWORD             spies_lock;
+    DWORD             cancelcount;
 };
 
 extern HRESULT WINAPI InternalTlsAllocData(struct tlsdata **data);
diff --git a/dlls/ole32/compobj_private.h b/dlls/ole32/compobj_private.h
index f10b41907a1..4ec22d51f45 100644
--- a/dlls/ole32/compobj_private.h
+++ b/dlls/ole32/compobj_private.h
@@ -62,6 +62,7 @@ struct oletls
     IUnknown        *state;       /* see CoSetState */
     struct list      spies;         /* Spies installed with CoRegisterInitializeSpy */
     DWORD            spies_lock;
+    DWORD            cancelcount;
 };
 
 /* Global Interface Table Functions */
diff --git a/dlls/ole32/tests/compobj.c b/dlls/ole32/tests/compobj.c
index e627521fa59..5ed3f9638b1 100644
--- a/dlls/ole32/tests/compobj.c
+++ b/dlls/ole32/tests/compobj.c
@@ -4186,6 +4186,64 @@ todo_wine
     CoUninitialize();
 }
 
+static void test_call_cancellation(void)
+{
+    HRESULT hr;
+
+    /* Cancellation is disabled initially. */
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == CO_E_CANCEL_DISABLED, "Unexpected hr %#x.\n", hr);
+
+    hr = CoEnableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == CO_E_CANCEL_DISABLED, "Unexpected hr %#x.\n", hr);
+
+    hr = CoEnableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    /* Counter is not affected by initialization. */
+    hr = CoInitialize(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == CO_E_CANCEL_DISABLED, "Unexpected hr %#x.\n", hr);
+
+    hr = CoEnableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    CoUninitialize();
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == CO_E_CANCEL_DISABLED, "Unexpected hr %#x.\n", hr);
+
+    /* It's cumulative. */
+    hr = CoEnableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoEnableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = CoDisableCallCancellation(NULL);
+    ok(hr == CO_E_CANCEL_DISABLED, "Unexpected hr %#x.\n", hr);
+}
+
 START_TEST(compobj)
 {
     init_funcs();
@@ -4236,6 +4294,7 @@ START_TEST(compobj)
     test_CoGetCurrentProcess();
     test_mta_usage();
     test_CoCreateInstanceFromApp();
+    test_call_cancellation();
 
     DeleteFileA( testlib );
 }
diff --git a/include/winerror.h b/include/winerror.h
index 4ebb48c1249..86f7aace68e 100644
--- a/include/winerror.h
+++ b/include/winerror.h
@@ -2679,6 +2679,7 @@ static inline HRESULT HRESULT_FROM_WIN32(unsigned int x)
 #define RPC_E_NO_CONTEXT                                   _HRESULT_TYPEDEF_(0x8001011E)
 #define RPC_E_TIMEOUT                                      _HRESULT_TYPEDEF_(0x8001011F)
 #define RPC_E_NO_SYNC                                      _HRESULT_TYPEDEF_(0x80010120)
+#define CO_E_CANCEL_DISABLED                               _HRESULT_TYPEDEF_(0x80010140)
 #define RPC_E_UNEXPECTED                                   _HRESULT_TYPEDEF_(0x8001FFFF)
 
 #define DISP_E_UNKNOWNINTERFACE                            _HRESULT_TYPEDEF_(0x80020001)