From 42f51185709e86f804c2babf259d0b1816b33d30 Mon Sep 17 00:00:00 2001
From: Andrew Eikum <aeikum@codeweavers.com>
Date: Fri, 13 Jun 2014 13:00:46 -0500
Subject: [PATCH] ole32: Implement CoTreatAsClass for classes without
 AutoTreatAs.

---
 dlls/ole32/compobj.c       | 26 ++++++++++++++++++------
 dlls/ole32/tests/compobj.c | 41 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/dlls/ole32/compobj.c b/dlls/ole32/compobj.c
index d117c0206b9..98f64399f6f 100644
--- a/dlls/ole32/compobj.c
+++ b/dlls/ole32/compobj.c
@@ -3731,7 +3731,8 @@ HRESULT WINAPI CoTreatAsClass(REFCLSID clsidOld, REFCLSID clsidNew)
     res = COM_OpenKeyForCLSID(clsidOld, NULL, KEY_READ | KEY_WRITE, &hkey);
     if (FAILED(res))
         goto done;
-    if (!memcmp( clsidOld, clsidNew, sizeof(*clsidOld) ))
+
+    if (IsEqualGUID( clsidOld, clsidNew ))
     {
        if (!RegQueryValueW(hkey, wszAutoTreatAs, auto_treat_as, &auto_treat_as_size) &&
            CLSIDFromString(auto_treat_as, &id) == S_OK)
@@ -3744,15 +3745,28 @@ HRESULT WINAPI CoTreatAsClass(REFCLSID clsidOld, REFCLSID clsidNew)
        }
        else
        {
-           RegDeleteKeyW(hkey, wszTreatAs);
+           if(RegDeleteKeyW(hkey, wszTreatAs))
+               res = REGDB_E_WRITEREGDB;
            goto done;
        }
     }
-    else if (!StringFromGUID2(clsidNew, szClsidNew, ARRAYSIZE(szClsidNew)) &&
-             !RegSetValueW(hkey, wszTreatAs, REG_SZ, szClsidNew, sizeof(szClsidNew)))
+    else
     {
-        res = REGDB_E_WRITEREGDB;
-	goto done;
+        if(IsEqualGUID(clsidNew, &CLSID_NULL)){
+           RegDeleteKeyW(hkey, wszTreatAs);
+        }else{
+            if(!StringFromGUID2(clsidNew, szClsidNew, ARRAYSIZE(szClsidNew))){
+                WARN("StringFromGUID2 failed\n");
+                res = E_FAIL;
+                goto done;
+            }
+
+            if(RegSetValueW(hkey, wszTreatAs, REG_SZ, szClsidNew, sizeof(szClsidNew)) != ERROR_SUCCESS){
+                WARN("RegSetValue failed\n");
+                res = REGDB_E_WRITEREGDB;
+                goto done;
+            }
+        }
     }
 
 done:
diff --git a/dlls/ole32/tests/compobj.c b/dlls/ole32/tests/compobj.c
index 16ae5012790..219d2f9db55 100644
--- a/dlls/ole32/tests/compobj.c
+++ b/dlls/ole32/tests/compobj.c
@@ -41,6 +41,7 @@ static HRESULT (WINAPI * pCoInitializeEx)(LPVOID lpReserved, DWORD dwCoInit);
 static HRESULT (WINAPI * pCoGetObjectContext)(REFIID riid, LPVOID *ppv);
 static HRESULT (WINAPI * pCoSwitchCallContext)(IUnknown *pObject, IUnknown **ppOldObject);
 static HRESULT (WINAPI * pCoGetTreatAsClass)(REFCLSID clsidOld, LPCLSID pClsidNew);
+static HRESULT (WINAPI * pCoTreatAsClass)(REFCLSID clsidOld, REFCLSID pClsidNew);
 static HRESULT (WINAPI * pCoGetContextToken)(ULONG_PTR *token);
 static LONG (WINAPI * pRegOverridePredefKey)(HKEY key, HKEY override);
 
@@ -1896,11 +1897,14 @@ static void test_CoGetContextToken(void)
     CoUninitialize();
 }
 
-static void test_CoGetTreatAsClass(void)
+static void test_TreatAsClass(void)
 {
     HRESULT hr;
     CLSID out;
     static GUID deadbeef = {0xdeadbeef,0xdead,0xbeef,{0xde,0xad,0xbe,0xef,0xde,0xad,0xbe,0xef}};
+    static const char deadbeefA[] = "{DEADBEEF-DEAD-BEEF-DEAD-BEEFDEADBEEF}";
+    HKEY clsidkey, deadbeefkey;
+    LONG lr;
 
     if (!pCoGetTreatAsClass)
     {
@@ -1910,6 +1914,38 @@ static void test_CoGetTreatAsClass(void)
     hr = pCoGetTreatAsClass(&deadbeef,&out);
     ok (hr == S_FALSE, "expected S_FALSE got %x\n",hr);
     ok (IsEqualGUID(&out,&deadbeef), "expected to get same clsid back\n");
+
+    lr = RegOpenKeyExA(HKEY_CLASSES_ROOT, "CLSID", 0, KEY_READ, &clsidkey);
+    ok(lr == ERROR_SUCCESS, "Couldn't open CLSID key\n");
+
+    lr = RegCreateKeyExA(clsidkey, deadbeefA, 0, NULL, 0, KEY_WRITE, NULL, &deadbeefkey, NULL);
+    ok(lr == ERROR_SUCCESS, "Couldn't create class key\n");
+
+    hr = pCoTreatAsClass(&deadbeef, &deadbeef);
+    ok(hr == REGDB_E_WRITEREGDB, "CoTreatAsClass gave wrong error: %08x\n", hr);
+
+    hr = pCoTreatAsClass(&deadbeef, &CLSID_FileProtocol);
+    if(hr == REGDB_E_WRITEREGDB){
+        win_skip("Insufficient privileges to use CoTreatAsClass\n");
+        goto exit;
+    }
+    ok(hr == S_OK, "CoTreatAsClass failed: %08x\n", hr);
+
+    hr = pCoGetTreatAsClass(&deadbeef, &out);
+    ok(hr == S_OK, "CoGetTreatAsClass failed: %08x\n",hr);
+    ok(IsEqualGUID(&out, &CLSID_FileProtocol), "expected to get substituted clsid\n");
+
+    hr = pCoTreatAsClass(&deadbeef, &CLSID_NULL);
+    ok(hr == S_OK, "CoTreatAsClass failed: %08x\n", hr);
+
+    hr = pCoGetTreatAsClass(&deadbeef, &out);
+    ok(hr == S_FALSE, "expected S_FALSE got %08x\n", hr);
+    ok(IsEqualGUID(&out, &deadbeef), "expected to get same clsid back\n");
+
+exit:
+    RegCloseKey(deadbeefkey);
+    RegDeleteKeyA(clsidkey, deadbeefA);
+    RegCloseKey(clsidkey);
 }
 
 static void test_CoInitializeEx(void)
@@ -1994,6 +2030,7 @@ static void init_funcs(void)
     pCoGetObjectContext = (void*)GetProcAddress(hOle32, "CoGetObjectContext");
     pCoSwitchCallContext = (void*)GetProcAddress(hOle32, "CoSwitchCallContext");
     pCoGetTreatAsClass = (void*)GetProcAddress(hOle32,"CoGetTreatAsClass");
+    pCoTreatAsClass = (void*)GetProcAddress(hOle32,"CoTreatAsClass");
     pCoGetContextToken = (void*)GetProcAddress(hOle32, "CoGetContextToken");
     pRegOverridePredefKey = (void*)GetProcAddress(hAdvapi32, "RegOverridePredefKey");
     pCoInitializeEx = (void*)GetProcAddress(hOle32, "CoInitializeEx");
@@ -2038,7 +2075,7 @@ START_TEST(compobj)
     test_CoGetObjectContext();
     test_CoGetCallContext();
     test_CoGetContextToken();
-    test_CoGetTreatAsClass();
+    test_TreatAsClass();
     test_CoInitializeEx();
     test_OleRegGetMiscStatus();
     test_CoCreateGuid();