[PATCH v10 8/9] kernel32: Rework ExpandEnvironmentStringsA to return ansi size and fix corner cases.

[Fabian Maurer]

From: Fabian Maurer

Wine-Bug: https://bugs.winehq.org/show_bug.cgi?id=54295
---
 dlls/advapi32/tests/registry.c |  2 +-
 dlls/kernel32/tests/environ.c  | 92 ++++++++++++++++------------------
 dlls/kernelbase/process.c      | 30 +++++++----
 dlls/shcore/tests/shcore.c     |  3 --
 dlls/shlwapi/tests/shreg.c     |  9 ++--
 dlls/userenv/tests/userenv.c   |  1 +
 6 files changed, 69 insertions(+), 68 deletions(-)

diff --git a/dlls/advapi32/tests/registry.c b/dlls/advapi32/tests/registry.c
index 0f77b2fa066..2e4195e723f 100644
--- a/dlls/advapi32/tests/registry.c
+++ b/dlls/advapi32/tests/registry.c
@@ -934,7 +934,7 @@ static void test_get_value(void)
     size = 0;
     ret = pRegGetValueA(hkey_main, NULL, "TP2_EXP_SZ", RRF_RT_REG_SZ, NULL, NULL, &size);
     ok(ret == ERROR_SUCCESS, "ret=%ld\n", ret);
-    todo_wine ok(size == strlen(expanded2)+2,
+    ok(size == strlen(expanded2)+2,
        "strlen(expanded2)=%d, strlen(sTestpath2)=%d, size=%ld\n", lstrlenA(expanded2), lstrlenA(sTestpath2), size);

/* Query REG_EXPAND_SZ using RRF_RT_REG_SZ (ok, expands) */
diff --git a/dlls/kernel32/tests/environ.c b/dlls/kernel32/tests/environ.c
index ec765758369..bd088d1670f 100644
--- a/dlls/kernel32/tests/environ.c
+++ b/dlls/kernel32/tests/environ.c
@@ -492,54 +492,54 @@ static void test_ExpandEnvironmentStringsA(void)
     BOOL todo_len;
     } tests[] =
     {
-        /*  0 */ { "Long long value",       "",                   0, 17, FALSE, TRUE, TRUE },
-        /*  1 */ { "Long long value",       "",                   1, 17, TRUE, TRUE, TRUE },
-        /*  2 */ { "Long long value",       "",                   2, 17, TRUE, TRUE, TRUE },
-        /*  3 */ { "Long long value",       "",                   3, 17, TRUE, TRUE, TRUE },
-        /*  4 */ { "Long long value",       "",                  15, 17, TRUE, TRUE, TRUE },
-        /*  5 */ { "Long long value",       "",                  16, 17, FALSE, TRUE, TRUE },
+        /*  0 */ { "Long long value",       "",                   0, 17 },
+        /*  1 */ { "Long long value",       "",                   1, 17 },
+        /*  2 */ { "Long long value",       "",                   2, 17 },
+        /*  3 */ { "Long long value",       "",                   3, 17 },
+        /*  4 */ { "Long long value",       "",                  15, 17 },
+        /*  5 */ { "Long long value",       "",                  16, 17 },
         /*  6 */ { "Long long value",       "Long long value",   17, 16 },
-        /*  7 */ { "%TVAR% long long",      "",                   0, 16, FALSE, TRUE, TRUE },
-        /*  8 */ { "%TVAR% long long",      "",                   1, 16, TRUE, FALSE, TRUE },
-        /*  9 */ { "%TVAR% long long",      "",                   2, 16, TRUE, FALSE, TRUE },
-        /* 10 */ { "%TVAR% long long",      "",                   4, 16, TRUE, FALSE, TRUE },
-        /* 11 */ { "%TVAR% long long",      "",                   5, 16, TRUE, TRUE, TRUE },
-        /* 12 */ { "%TVAR% long long",      "",                   6, 16, TRUE, TRUE, TRUE },
-        /* 13 */ { "%TVAR% long long",      "",                   7, 16, TRUE, TRUE, TRUE },
-        /* 14 */ { "%TVAR% long long",      "",                  15, 16, FALSE, TRUE, TRUE },
+        /*  7 */ { "%TVAR% long long",      "",                   0, 16 },
+        /*  8 */ { "%TVAR% long long",      "",                   1, 16 },
+        /*  9 */ { "%TVAR% long long",      "",                   2, 16 },
+        /* 10 */ { "%TVAR% long long",      "",                   4, 16 },
+        /* 11 */ { "%TVAR% long long",      "",                   5, 16 },
+        /* 12 */ { "%TVAR% long long",      "",                   6, 16 },
+        /* 13 */ { "%TVAR% long long",      "",                   7, 16 },
+        /* 14 */ { "%TVAR% long long",      "",                  15, 16 },
         /* 15 */ { "%TVAR% long long",      "WINE long long",    16, 15 },
-        /* 16 */ { "%TVAR%%TVAR% long",     "",                   4, 15, TRUE, FALSE, TRUE },
-        /* 17 */ { "%TVAR%%TVAR% long",     "",                   5, 15, TRUE, TRUE, TRUE },
-        /* 18 */ { "%TVAR%%TVAR% long",     "",                   6, 15, TRUE, TRUE, TRUE },
-        /* 19 */ { "%TVAR%%TVAR% long",     "",                   8, 15, TRUE, TRUE, TRUE },
-        /* 20 */ { "%TVAR%%TVAR% long",     "",                   9, 15, TRUE, TRUE, TRUE },
-        /* 21 */ { "%TVAR%%TVAR% long",     "",                  10, 15, TRUE, TRUE, TRUE },
-        /* 22 */ { "%TVAR%%TVAR% long",     "",                  14, 15, FALSE, TRUE, TRUE },
+        /* 16 */ { "%TVAR%%TVAR% long",     "",                   4, 15 },
+        /* 17 */ { "%TVAR%%TVAR% long",     "",                   5, 15 },
+        /* 18 */ { "%TVAR%%TVAR% long",     "",                   6, 15 },
+        /* 19 */ { "%TVAR%%TVAR% long",     "",                   8, 15 },
+        /* 20 */ { "%TVAR%%TVAR% long",     "",                   9, 15 },
+        /* 21 */ { "%TVAR%%TVAR% long",     "",                  10, 15 },
+        /* 22 */ { "%TVAR%%TVAR% long",     "",                  14, 15 },
         /* 23 */ { "%TVAR%%TVAR% long",     "WINEWINE long",     15, 14 },
-        /* 24 */ { "%TVAR% %TVAR% long",    "",                   5, 16, TRUE, TRUE, TRUE },
-        /* 25 */ { "%TVAR% %TVAR% long",    "",                   6, 16, TRUE, TRUE, TRUE },
-        /* 26 */ { "%TVAR% %TVAR% long",    "",                   8, 16, TRUE, TRUE, TRUE },
-        /* 27 */ { "%TVAR% %TVAR% long",    "",                   9, 16, TRUE, TRUE, TRUE },
-        /* 28 */ { "%TVAR% %TVAR% long",    "",                  10, 16, TRUE, TRUE, TRUE },
-        /* 29 */ { "%TVAR% %TVAR% long",    "",                  11, 16, TRUE, TRUE, TRUE },
-        /* 30 */ { "%TVAR% %TVAR% long",    "",                  12, 16, TRUE, TRUE, TRUE },
-        /* 31 */ { "%TVAR% %TVAR% long",    "",                  14, 16, TRUE, TRUE, TRUE },
-        /* 32 */ { "%TVAR% %TVAR% long",    "",                  15, 16, FALSE, TRUE, TRUE },
+        /* 24 */ { "%TVAR% %TVAR% long",    "",                   5, 16 },
+        /* 25 */ { "%TVAR% %TVAR% long",    "",                   6, 16 },
+        /* 26 */ { "%TVAR% %TVAR% long",    "",                   8, 16 },
+        /* 27 */ { "%TVAR% %TVAR% long",    "",                   9, 16 },
+        /* 28 */ { "%TVAR% %TVAR% long",    "",                  10, 16 },
+        /* 29 */ { "%TVAR% %TVAR% long",    "",                  11, 16 },
+        /* 30 */ { "%TVAR% %TVAR% long",    "",                  12, 16 },
+        /* 31 */ { "%TVAR% %TVAR% long",    "",                  14, 16 },
+        /* 32 */ { "%TVAR% %TVAR% long",    "",                  15, 16 },
         /* 33 */ { "%TVAR% %TVAR% long",    "WINE WINE long",    16, 15 },
-        /* 34 */ { "%TVAR2% long long",     "",                   1, 19, TRUE, TRUE, TRUE },
-        /* 35 */ { "%TVAR2% long long",     "",                   2, 19, TRUE, TRUE, TRUE },
-        /* 36 */ { "%TVAR2% long long",     "",                   4, 19, TRUE, TRUE, TRUE },
-        /* 37 */ { "%TVAR2% long long",     "",                   7, 19, TRUE, TRUE, TRUE },
-        /* 38 */ { "%TVAR2% long long",     "",                   8, 19, TRUE, TRUE, TRUE },
-        /* 39 */ { "%TVAR2% long long",     "",                   9, 19, TRUE, TRUE, TRUE },
-        /* 40 */ { "%TVAR2% long long",     "",                  10, 19, TRUE, TRUE, TRUE },
-        /* 41 */ { "%TVAR2% long long",     "",                  18, 19, FALSE, TRUE, TRUE },
+        /* 34 */ { "%TVAR2% long long",     "",                   1, 19 },
+        /* 35 */ { "%TVAR2% long long",     "",                   2, 19 },
+        /* 36 */ { "%TVAR2% long long",     "",                   4, 19 },
+        /* 37 */ { "%TVAR2% long long",     "",                   7, 19 },
+        /* 38 */ { "%TVAR2% long long",     "",                   8, 19 },
+        /* 39 */ { "%TVAR2% long long",     "",                   9, 19 },
+        /* 40 */ { "%TVAR2% long long",     "",                  10, 19 },
+        /* 41 */ { "%TVAR2% long long",     "",                  18, 19 },
         /* 42 */ { "%TVAR2% long long",     "%TVAR2% long long", 19, 18 },
-        /* 43 */ { "%TVAR long long",       "",                   1, 17, TRUE, TRUE, TRUE },
-        /* 44 */ { "%TVAR long long",       "",                   2, 17, TRUE, TRUE, TRUE },
-        /* 45 */ { "%TVAR long long",       "",                   3, 17, TRUE, TRUE, TRUE },
-        /* 46 */ { "%TVAR long long",       "",                  15, 17, TRUE, TRUE, TRUE },
-        /* 47 */ { "%TVAR long long",       "",                  16, 17, FALSE, TRUE, TRUE },
+        /* 43 */ { "%TVAR long long",       "",                   1, 17 },
+        /* 44 */ { "%TVAR long long",       "",                   2, 17 },
+        /* 45 */ { "%TVAR long long",       "",                   3, 17 },
+        /* 46 */ { "%TVAR long long",       "",                  15, 17 },
+        /* 47 */ { "%TVAR long long",       "",                  16, 17 },
         /* 48 */ { "%TVAR long long",       "%TVAR long long",   17, 16 },
     };

@@ -653,23 +653,17 @@ static void test_ExpandEnvironmentStringsA(void)
         SetLastError(0xdeadbeef);
         ret_size = ExpandEnvironmentStringsA(japanese_test, NULL, 0);
         ok(GetLastError() == 0xdeadbeef, "got last error %ld\n", GetLastError());
-        todo_wine
         ok(ret_size >= japanese_len, "Needed at least %d, got %lu\n", japanese_len, ret_size);

SetLastError(0xdeadbeef);
         ret_size = ExpandEnvironmentStringsA(japanese_test, buf, ret_size);
-        todo_wine
         ok(GetLastError() == 0xdeadbeef, "got last error %ld\n", GetLastError());
-        todo_wine
         ok(ret_size >= japanese_len, "Needed at least %d, got %lu\n", japanese_len, ret_size);
-        todo_wine
         ok(!strcmp(buf, japanese_test), "Got %s\n", debugstr_a(buf));

SetLastError(0xdeadbeef);
         ret_size = ExpandEnvironmentStringsA(japanese_test, buf, japanese_len / 2); /* Buffer too small */
-        todo_wine
         ok(GetLastError() == 0xdeadbeef, "got last error %ld\n", GetLastError());
-        todo_wine
         ok(ret_size >= japanese_len, "Needed at least %d, got %lu\n", japanese_len, ret_size);
     }
     else
diff --git a/dlls/kernelbase/process.c b/dlls/kernelbase/process.c
index c6f2f30b2d8..09e6547245d 100644
--- a/dlls/kernelbase/process.c
+++ b/dlls/kernelbase/process.c
@@ -1447,20 +1447,32 @@ DWORD WINAPI DECLSPEC_HOTPATCH ExpandEnvironmentStringsA( LPCSTR src, LPSTR dst,
 {
     UNICODE_STRING us_src;
     PWSTR dstW = NULL;
-    DWORD ret;
+    DWORD count_neededW;
+    DWORD count_neededA = 0;

RtlCreateUnicodeStringFromAsciiz( &us_src, src );
-    if (count)
-    {
-        if (!(dstW = HeapAlloc(GetProcessHeap(), 0, count * sizeof(WCHAR)))) return 0;
-        ret = ExpandEnvironmentStringsW( us_src.Buffer, dstW, count);
-        if (ret) WideCharToMultiByte( CP_ACP, 0, dstW, ret, dst, count, NULL, NULL );
-    }
-    else ret = ExpandEnvironmentStringsW( us_src.Buffer, NULL, 0 );

+    /* We always need to call ExpandEnvironmentStringsW, since we need the result to calculate the needed buffer size */
+    count_neededW = ExpandEnvironmentStringsW( us_src.Buffer, NULL, 0 );
+    if (!(dstW = HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY, count_neededW * sizeof(WCHAR) ))) goto cleanup;
+    count_neededW = ExpandEnvironmentStringsW( us_src.Buffer, dstW, count_neededW );
+
+    /* Calculate needed buffer */
+    count_neededA = WideCharToMultiByte( CP_ACP, 0, dstW, count_neededW, NULL, 0, NULL, NULL );
+
+    /* If provided buffer is enough, do actual conversion */
+    if (count > count_neededA)
+        count_neededA = WideCharToMultiByte( CP_ACP, 0, dstW, count_neededW, dst, count, NULL, NULL );
+    else if(dst)
+        *dst = 0;
+
+cleanup:
     RtlFreeUnicodeString( &us_src );
     HeapFree( GetProcessHeap(), 0, dstW );
-    return ret;
+
+    if (count_neededA >= count) /* When the buffer is too small, native over-reports by one byte */
+        return count_neededA + 1;
+    return count_neededA;
 }

diff --git a/dlls/shcore/tests/shcore.c b/dlls/shcore/tests/shcore.c
index 731fa340e8b..3e621c8b950 100644
--- a/dlls/shcore/tests/shcore.c
+++ b/dlls/shcore/tests/shcore.c
@@ -584,11 +584,8 @@ static void test_SHQueryValueEx(void)
     ret = pSHQueryValueExA( hKey, "Test3", NULL, &type, buf, &size);
     ok(ret == ERROR_MORE_DATA, "Unexpected retval %ld.\n", ret);

-    todo_wine
-    {
     ok( (0 == strcmp("", buf)) || (0 == strcmp(test_envvar2, buf)),
     "Expected empty or first part of the string "%s", got "%s"\n", test_envvar2, buf);
-    }

ok(size >= buffer_len2, "Buffer size %lu should be >= %lu.\n", size, buffer_len2);
     ok(type == REG_SZ, "Unexpected type %ld.\n", type);
diff --git a/dlls/shlwapi/tests/shreg.c b/dlls/shlwapi/tests/shreg.c
index 75d4dfae7eb..72ca470856e 100644
--- a/dlls/shlwapi/tests/shreg.c
+++ b/dlls/shlwapi/tests/shreg.c
@@ -313,12 +313,9 @@ static void test_SHQueryValueEx(void)
             broken(ERROR_SUCCESS == dwRet), /* < IE5.5 */
             "Expected ERROR_MORE_DATA, got (%lu)\n", dwRet);

-        todo_wine
-        {
-            ok( (0 == strcmp("", buf)) || (0 == strcmp(sEnvvar2, buf)) ||
-                broken(0 == strcmp(sTestpath2, buf)), /* IE 5.5 */
-                "Expected empty or first part of the string "%s", got "%s"\n", sEnvvar2, buf);
-        }
+        ok( (0 == strcmp("", buf)) || (0 == strcmp(sEnvvar2, buf)) ||
+            broken(0 == strcmp(sTestpath2, buf)), /* IE 5.5 */
+            "Expected empty or first part of the string "%s", got "%s"\n", sEnvvar2, buf);

ok( dwSize >= nUsedBuffer2 ||
             broken(dwSize == (strlen(sEnvvar2) + 1)) || /* IE4.01 SP1 (W98) and IE5 (W98SE) */
diff --git a/dlls/userenv/tests/userenv.c b/dlls/userenv/tests/userenv.c
index 3ddb893c18c..7f69a329c36 100644
--- a/dlls/userenv/tests/userenv.c
+++ b/dlls/userenv/tests/userenv.c
@@ -294,6 +294,7 @@ static void test_get_profiles_dir(void)
     /* Rather than a BOOL, the return value is also the number of characters
      * stored in the buffer.
      */
+    todo_wine
     ok(profiles_len - 1 == r, "expected %ld, got %d\n", profiles_len - 1, r);
     ok(!strcmp(buf, profiles_dir), "expected %s, got %s\n", profiles_dir, buf);
     ok(strlen(buf) + 1 == cch, "String length is %Iu, but cch is %lu\n", strlen(buf), cch);