From: Jacek Caban <jacek@codeweavers.com>
Subject: [PATCH 2/3] shlwapi: Fixed error handling in StrRetToBufW.
Message-Id: <7afceb5d-c026-6a59-7d5f-57041caeaa33@codeweavers.com>
Date: Thu, 23 Feb 2017 15:42:41 +0100

Signed-off-by: Jacek Caban <jacek@codeweavers.com>
---
 dlls/shlwapi/string.c       | 79
++++++++++++++++++++++++++-------------------
 dlls/shlwapi/tests/string.c | 21 +++++++++++-
 2 files changed, 65 insertions(+), 35 deletions(-)

diff --git a/dlls/shlwapi/string.c b/dlls/shlwapi/string.c
index 8c4cb71..0b34d43 100644
--- a/dlls/shlwapi/string.c
+++ b/dlls/shlwapi/string.c
@@ -1521,47 +1521,58 @@ HRESULT WINAPI StrRetToBufA (LPSTRRET src, const ITEMIDLIST *pidl, LPSTR dest, U
  */
 HRESULT WINAPI StrRetToBufW (LPSTRRET src, const ITEMIDLIST *pidl, LPWSTR dest, UINT len)
 {
-        TRACE("dest=%p len=0x%x strret=%p pidl=%p\n", dest, len, src, pidl);
+    TRACE("dest=%p len=0x%x strret=%p pidl=%p\n", dest, len, src, pidl);
 
-	if (!src)
-	{
-	  WARN("Invalid lpStrRet would crash under Win32!\n");
-	  if (dest)
-	    *dest = '\0';
-	  return E_FAIL;
-	}
+    if (!dest || !len)
+        return E_FAIL;
 
-	if (!dest || !len)
-	  return E_FAIL;
+    if (!src)
+    {
+        WARN("Invalid lpStrRet would crash under Win32!\n");
+        if (dest)
+            *dest = '\0';
+        return E_FAIL;
+    }
 
-	*dest = '\0';
+    *dest = '\0';
+
+    switch (src->uType) {
+    case STRRET_WSTR: {
+        size_t dst_len;
+        if (!src->u.pOleStr)
+            return E_FAIL;
+        dst_len = strlenW(src->u.pOleStr);
+        memcpy(dest, src->u.pOleStr, min(dst_len, len-1) * sizeof(WCHAR));
+        dest[min(dst_len, len-1)] = 0;
+        CoTaskMemFree(src->u.pOleStr);
+        if (len <= dst_len)
+        {
+            dest[0] = 0;
+            return E_NOT_SUFFICIENT_BUFFER;
+        }
+        break;
+    }
 
-	switch (src->uType)
-	{
-	  case STRRET_WSTR:
-            lstrcpynW(dest, src->u.pOleStr, len);
-	    CoTaskMemFree(src->u.pOleStr);
-	    break;
+    case STRRET_CSTR:
+        if (!MultiByteToWideChar( CP_ACP, 0, src->u.cStr, -1, dest, len ))
+            dest[len-1] = 0;
+        break;
 
-	  case STRRET_CSTR:
-              if (!MultiByteToWideChar( CP_ACP, 0, src->u.cStr, -1, dest, len ))
-                  dest[len-1] = 0;
-	    break;
+    case STRRET_OFFSET:
+        if (pidl)
+	{
+            if (!MultiByteToWideChar( CP_ACP, 0, ((LPCSTR)&pidl->mkid)+src->u.uOffset, -1,
+                                      dest, len ))
+                dest[len-1] = 0;
+        }
+        break;
 
-	  case STRRET_OFFSET:
-	    if (pidl)
-	    {
-              if (!MultiByteToWideChar( CP_ACP, 0, ((LPCSTR)&pidl->mkid)+src->u.uOffset, -1,
-                                        dest, len ))
-                  dest[len-1] = 0;
-	    }
-	    break;
+    default:
+        FIXME("unknown type!\n");
+        return E_NOTIMPL;
+    }
 
-	  default:
-	    FIXME("unknown type!\n");
-	    return E_NOTIMPL;
-	}
-	return S_OK;
+    return S_OK;
 }
 
 /*************************************************************************
diff --git a/dlls/shlwapi/tests/string.c b/dlls/shlwapi/tests/string.c
index ac14db0..f23ae3f 100644
--- a/dlls/shlwapi/tests/string.c
+++ b/dlls/shlwapi/tests/string.c
@@ -988,6 +988,7 @@ static void test_StrXXX_overflows(void)
     WCHAR wstr1[2*MAX_PATH+1], wbuf[2*MAX_PATH];
     const WCHAR fmt[] = {'%','s',0};
     STRRET strret;
+    HRESULT hres;
     int ret;
     int i;
 
@@ -1059,9 +1060,27 @@ if (0)
         memset(wbuf, 0xbf, sizeof(wbuf));
         strret.uType = STRRET_WSTR;
         U(strret).pOleStr = StrDupW(wstr1);
-        expect_eq2(pStrRetToBufW(&strret, NULL, wbuf, 10), S_OK, E_NOT_SUFFICIENT_BUFFER /* Vista */, HRESULT, "%x");
+        hres = pStrRetToBufW(&strret, NULL, wbuf, 10);
+        ok(hres == E_NOT_SUFFICIENT_BUFFER || broken(hres == S_OK) /* winxp */,
+           "StrRetToBufW returned %08x\n", hres);
+        if (hres == E_NOT_SUFFICIENT_BUFFER)
+            expect_eq(wbuf[0], 0, WCHAR, "%x");
         expect_eq(wbuf[9], 0, WCHAR, "%x");
         expect_eq(wbuf[10], (WCHAR)0xbfbf, WCHAR, "%x");
+
+        memset(wbuf, 0xbf, sizeof(wbuf));
+        strret.uType = STRRET_CSTR;
+        StrCpyNA(U(strret).cStr, str1, MAX_PATH);
+        hres = pStrRetToBufW(&strret, NULL, wbuf, 10);
+        ok(hres == S_OK, "StrRetToBufW returned %08x\n", hres);
+        ok(!memcmp(wbuf, wstr1, 9*sizeof(WCHAR)) && !wbuf[9], "StrRetToBuf returned %s\n", wine_dbgstr_w(wbuf));
+
+        memset(wbuf, 0xbf, sizeof(wbuf));
+        strret.uType = STRRET_WSTR;
+        U(strret).pOleStr = NULL;
+        hres = pStrRetToBufW(&strret, NULL, wbuf, 10);
+        ok(hres == E_FAIL, "StrRetToBufW returned %08x\n", hres);
+        ok(!wbuf[0], "StrRetToBuf returned %s\n", wine_dbgstr_w(wbuf));
     }
     else
         win_skip("StrRetToBufW() is not available\n");