shlwapi: Fixed error handling in StrRetToBufW.

Signed-off-by: Jacek Caban <jacek@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

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)
+    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';
+
+    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;
+    }
+
+    case STRRET_CSTR:
+        if (!MultiByteToWideChar( CP_ACP, 0, src->u.cStr, -1, dest, len ))
+            dest[len-1] = 0;
+        break;
+
+    case STRRET_OFFSET:
+        if (pidl)
 	{
-	  WARN("Invalid lpStrRet would crash under Win32!\n");
+            if (!MultiByteToWideChar( CP_ACP, 0, ((LPCSTR)&pidl->mkid)+src->u.uOffset, -1,
-	  if (dest)
+                                      dest, len ))
-	    *dest = '\0';
+                dest[len-1] = 0;
-	  return E_FAIL;
+        }
-	}
+        break;
 
-	if (!dest || !len)
+    default:
-	  return E_FAIL;
+        FIXME("unknown type!\n");
+        return E_NOTIMPL;
+    }
 
-	*dest = '\0';
+    return S_OK;
-
-	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_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;
-	}
-	return S_OK;
 }
 
 /*************************************************************************

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");