[PATCH v3 3/3] odbccp32: SQLWrite/RemoveDSNFromIniW correctly handle config_mode

[Alistair Leslie-Hughes]

From: Alistair Leslie-Hughes

---
 dlls/odbccp32/odbccp32.c   | 36 +++++++++++++++++++++++++-----
 dlls/odbccp32/tests/misc.c | 45 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/dlls/odbccp32/odbccp32.c b/dlls/odbccp32/odbccp32.c
index b77c5698c07..1676cfa5746 100644
--- a/dlls/odbccp32/odbccp32.c
+++ b/dlls/odbccp32/odbccp32.c
@@ -1582,7 +1582,7 @@ BOOL WINAPI SQLRemoveDriverManager(LPDWORD pdwUsageCount)

BOOL WINAPI SQLRemoveDSNFromIniW(LPCWSTR lpszDSN)
 {
-    HKEY hkey;
+    HKEY hkey, hkeyroot = HKEY_CURRENT_USER;

TRACE("%s\n", debugstr_w(lpszDSN));

@@ -1594,13 +1594,36 @@ BOOL WINAPI SQLRemoveDSNFromIniW(LPCWSTR lpszDSN)

clear_errors();

-    if (RegOpenKeyW(HKEY_LOCAL_MACHINE, L"Software\ODBC\ODBC.INI\ODBC Data Sources", &hkey) == ERROR_SUCCESS)
+    if (config_mode == ODBC_SYSTEM_DSN)
+        hkeyroot = HKEY_LOCAL_MACHINE;
+    else if (config_mode == ODBC_BOTH_DSN)
+    {
+        WCHAR *regpath = malloc( (wcslen(L"Software\ODBC\ODBC.INI\") + wcslen(lpszDSN) + 1) * sizeof(WCHAR) );
+        if (!regpath)
+        {
+            push_error(ODBC_ERROR_OUT_OF_MEM, L"Out of memory");
+            return FALSE;
+        }
+        wcscpy(regpath, L"Software\ODBC\ODBC.INI\");
+        wcscat(regpath, lpszDSN);
+
+        /* ONLY removes one DSN, USER or SYSTEM */
+        if (RegOpenKeyW(HKEY_CURRENT_USER, regpath, &hkey) == ERROR_SUCCESS)
+            hkeyroot = HKEY_CURRENT_USER;
+        else
+            hkeyroot = HKEY_LOCAL_MACHINE;
+
+        RegCloseKey(hkey);
+        free(regpath);
+    }
+
+    if (RegOpenKeyW(hkeyroot, L"Software\ODBC\ODBC.INI\ODBC Data Sources", &hkey) == ERROR_SUCCESS)
     {
         RegDeleteValueW(hkey, lpszDSN);
         RegCloseKey(hkey);
     }

-    if (RegOpenKeyW(HKEY_LOCAL_MACHINE, L"Software\ODBC\ODBC.INI", &hkey) == ERROR_SUCCESS)
+    if (RegOpenKeyW(hkeyroot, L"Software\ODBC\ODBC.INI", &hkey) == ERROR_SUCCESS)
     {
         RegDeleteTreeW(hkey, lpszDSN);
         RegCloseKey(hkey);
@@ -1756,7 +1779,7 @@ BOOL WINAPI SQLValidDSN(LPCSTR lpszDSN)
 BOOL WINAPI SQLWriteDSNToIniW(LPCWSTR lpszDSN, LPCWSTR lpszDriver)
 {
     DWORD ret;
-    HKEY hkey, hkeydriver;
+    HKEY hkey, hkeydriver, hkeyroot = HKEY_CURRENT_USER;
     WCHAR filename[MAX_PATH];

TRACE("%s %s\n", debugstr_w(lpszDSN), debugstr_w(lpszDriver));
@@ -1784,7 +1807,10 @@ BOOL WINAPI SQLWriteDSNToIniW(LPCWSTR lpszDSN, LPCWSTR lpszDriver)
         RegCloseKey(hkey);
     }

-    if ((ret = RegCreateKeyW(HKEY_LOCAL_MACHINE, L"SOFTWARE\ODBC\ODBC.INI", &hkey)) == ERROR_SUCCESS)
+    if (config_mode == ODBC_SYSTEM_DSN)
+        hkeyroot = HKEY_LOCAL_MACHINE;
+
+    if ((ret = RegCreateKeyW(hkeyroot, L"SOFTWARE\ODBC\ODBC.INI", &hkey)) == ERROR_SUCCESS)
     {
         HKEY sources;

diff --git a/dlls/odbccp32/tests/misc.c b/dlls/odbccp32/tests/misc.c
index af149e6c973..b294e465cc7 100644
--- a/dlls/odbccp32/tests/misc.c
+++ b/dlls/odbccp32/tests/misc.c
@@ -876,6 +876,22 @@ static void test_SQLConfigDataSource(void)
     check_error(ODBC_ERROR_COMPONENT_NOT_FOUND);
 }

+static BOOL check_dsn_exists(HKEY key, const WCHAR *dsn)
+{
+    HKEY hkey;
+    WCHAR buffer[256];
+    LONG res;
+
+    wcscpy(buffer, L"Software\ODBC\ODBC.INI\");
+    wcscat(buffer, dsn);
+
+    res = RegOpenKeyExW(key, buffer, 0, KEY_READ, &hkey);
+    if (res)
+        RegCloseKey(hkey);
+
+    return res == ERROR_SUCCESS;
+}
+
 static void test_SQLWriteDSNToIni(void)
 {
     BOOL ret;
@@ -891,6 +907,9 @@ static void test_SQLWriteDSNToIni(void)
     ret = SQLRemoveDSNFromIniW(L"");
     ok(!ret, "got %d\n", ret);

+    ret = check_dsn_exists(HKEY_LOCAL_MACHINE, L"wine_dbs");
+    ok(!ret, "Found registry key\n");
+
     ret = SQLWriteDSNToIni("wine_dbs", "SQL Server");
     if (!ret)
     {
@@ -898,6 +917,18 @@ static void test_SQLWriteDSNToIni(void)
         return;
     }

+    ret = check_dsn_exists(HKEY_LOCAL_MACHINE, L"wine_dbs");
+    ok(ret, "Failed to find registry key\n");
+
+    SQLSetConfigMode(ODBC_USER_DSN);
+    ret = SQLWriteDSNToIni("wine_dbs", "SQL Server");
+    ok(ret, "got %d\n", ret);
+
+    ret = check_dsn_exists(HKEY_CURRENT_USER, L"wine_dbs");
+    ok(ret, "Failed to find registry key\n");
+
+    SQLSetConfigMode(ODBC_SYSTEM_DSN);
+
     if(ret)
     {
         HKEY hkey;
@@ -920,6 +951,18 @@ static void test_SQLWriteDSNToIni(void)
             RegCloseKey(hkey);
         }

+        SQLSetConfigMode(ODBC_BOTH_DSN);
+
+        /* ODBC_BOTH_DSN set and has both System/User DSN but only removes USER. */
+        ret = SQLRemoveDSNFromIni("wine_dbs");
+        ok(ret, "got %d\n", ret);
+
+        ret = check_dsn_exists(HKEY_CURRENT_USER, L"wine_dbs");
+        ok(!ret, "Found registry key\n");
+
+        ret = check_dsn_exists(HKEY_LOCAL_MACHINE, L"wine_dbs");
+        ok(ret, "Failed to find registry key\n");
+
         res = RegOpenKeyExA(HKEY_LOCAL_MACHINE, "Software\ODBC\ODBC.INI\wine_dbs", 0,
                             KEY_READ, &hkey);
         ok(res == ERROR_SUCCESS, "RegOpenKeyExW failed\n");
@@ -938,6 +981,8 @@ static void test_SQLWriteDSNToIni(void)
             RegCloseKey(hkey);
         }

+        SQLSetConfigMode(ODBC_SYSTEM_DSN);
+
         ret = SQLRemoveDSNFromIni("wine_dbs");
         ok(ret, "got %d\n", ret);
     }