Integrate fixes that make dutil a little more robust to failure

6 files changed, 265 insertions, 14 deletions

diff --git a/src/dutil/cabutil.cpp b/src/dutil/cabutil.cpp
index 4a6f7b7b..5d77e483 100644
--- a/src/dutil/cabutil.cpp
+++ b/src/dutil/cabutil.cpp
@@ -102,6 +102,29 @@ LExit:
 }
 
 
+static HANDLE OpenFileWithRetry(
+    __in LPCWSTR wzPath,
+    __in DWORD dwDesiredAccess,
+    __in DWORD dwCreationDisposition
+)
+{
+    HANDLE hFile = INVALID_HANDLE_VALUE;
+
+    for (DWORD i = 0; i < 30; ++i)
+    {
+        hFile = ::CreateFileW(wzPath, dwDesiredAccess, FILE_SHARE_READ, NULL, dwCreationDisposition, FILE_ATTRIBUTE_NORMAL, NULL);
+        if (INVALID_HANDLE_VALUE != hFile)
+        {
+            break;
+        }
+
+        ::Sleep(100);
+    }
+
+    return hFile;
+}
+
+
 /********************************************************************
  CabInitialize - initializes internal static variables
 
@@ -340,6 +363,7 @@ static __callback void DIAMONDAPI CabExtractFree(__in LPVOID pvData)
 static __callback INT_PTR FAR DIAMONDAPI CabExtractOpen(__in_z PSTR pszFile, __in int oflag, __in int pmode)
 {
     HRESULT hr = S_OK;
+    HANDLE hFile = INVALID_HANDLE_VALUE;
     INT_PTR pFile = -1;
     LPWSTR sczCabFile = NULL;
 
@@ -353,19 +377,24 @@ static __callback INT_PTR FAR DIAMONDAPI CabExtractOpen(__in_z PSTR pszFile, __i
     hr = StrAllocStringAnsi(&sczCabFile, pszFile, 0, CP_UTF8);
     CabExitOnFailure(hr, "Failed to convert UTF8 cab file name to wide character string");
 
-    pFile = reinterpret_cast<INT_PTR>(::CreateFileW(sczCabFile, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL));
-    if (INVALID_HANDLE_VALUE == reinterpret_cast<HANDLE>(pFile))
+    hFile = OpenFileWithRetry(sczCabFile, GENERIC_READ, OPEN_EXISTING);
+    if (INVALID_HANDLE_VALUE == hFile)
     {
         CabExitWithLastError(hr, "failed to open file: %ls", sczCabFile);
     }
 
+    pFile = reinterpret_cast<INT_PTR>(hFile);
+
     if (vdw64EmbeddedOffset)
     {
         hr = CabExtractSeek(pFile, 0, 0);
         CabExitOnFailure(hr, "Failed to seek to embedded offset %I64d", vdw64EmbeddedOffset);
     }
 
+    hFile = INVALID_HANDLE_VALUE;
+
 LExit:
+    ReleaseFileHandle(hFile);
     ReleaseStr(sczCabFile);
 
     return FAILED(hr) ? -1 : pFile;
@@ -460,6 +489,7 @@ static __callback INT_PTR DIAMONDAPI CabExtractCallback(__in FDINOTIFICATIONTYPE
     Assert(pFDINotify->pv);
 
     HRESULT hr = S_OK;
+    HANDLE hFile = INVALID_HANDLE_VALUE;
     INT_PTR ipResult = 0;   // result to return on success
 
     CAB_CALLBACK_STRUCT* pccs = static_cast<CAB_CALLBACK_STRUCT*>(pFDINotify->pv);
@@ -503,7 +533,6 @@ static __callback INT_PTR DIAMONDAPI CabExtractCallback(__in FDINOTIFICATIONTYPE
                 CabExitWithLastError(hr, "failed to get time for resource: %ls", wz);
             }
             ::LocalFileTimeToFileTime(&ftLocal, &ft);
-            
 
             WCHAR wzPath[MAX_PATH];
             hr = ::StringCchCopyW(wzPath, countof(wzPath), pccs->pwzExtractDir);
@@ -511,21 +540,24 @@ static __callback INT_PTR DIAMONDAPI CabExtractCallback(__in FDINOTIFICATIONTYPE
             hr = ::StringCchCatW(wzPath, countof(wzPath), wz);
             CabExitOnFailure(hr, "failed to concat onto path: %ls file: %ls", wzPath, wz);
 
-            ipResult = reinterpret_cast<INT_PTR>(::CreateFileW(wzPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL));
-            if (INVALID_HANDLE_VALUE == reinterpret_cast<HANDLE>(ipResult))
+            hFile = OpenFileWithRetry(wzPath, GENERIC_WRITE, CREATE_ALWAYS);
+            if (INVALID_HANDLE_VALUE == hFile)
             {
                 CabExitWithLastError(hr, "failed to create file: %ls", wzPath);
             }
 
-            ::SetFileTime(reinterpret_cast<HANDLE>(ipResult), &ft, &ft, &ft);   // try to set the file time (who cares if it fails)
+            ::SetFileTime(hFile, &ft, &ft, &ft);   // try to set the file time (who cares if it fails)
 
-            if (::SetFilePointer(reinterpret_cast<HANDLE>(ipResult), pFDINotify->cb, NULL, FILE_BEGIN))   // try to set the end of the file (don't worry if this fails)
+            if (::SetFilePointer(hFile, pFDINotify->cb, NULL, FILE_BEGIN))   // try to set the end of the file (don't worry if this fails)
             {
-                if (::SetEndOfFile(reinterpret_cast<HANDLE>(ipResult)))
+                if (::SetEndOfFile(hFile))
                 {
-                    ::SetFilePointer(reinterpret_cast<HANDLE>(ipResult), 0, NULL, FILE_BEGIN);  // reset the file pointer
+                    ::SetFilePointer(hFile, 0, NULL, FILE_BEGIN);  // reset the file pointer
                 }
             }
+
+            ipResult = reinterpret_cast<INT_PTR>(hFile);
+            hFile = INVALID_HANDLE_VALUE;
         }
         else  // resource wasn't requested, skip it
         {
@@ -579,5 +611,7 @@ static __callback INT_PTR DIAMONDAPI CabExtractCallback(__in FDINOTIFICATIONTYPE
     };
 
 LExit:
+    ReleaseFileHandle(hFile);
+
     return (S_OK == hr) ? ipResult : -1;
 }
diff --git a/src/dutil/fileutil.cpp b/src/dutil/fileutil.cpp
index c76017de..06a44b45 100644
--- a/src/dutil/fileutil.cpp
+++ b/src/dutil/fileutil.cpp
@@ -1106,6 +1106,134 @@ LExit:
 
 
 /*******************************************************************
+ FileCopyUsingHandlesWithProgress
+
+*******************************************************************/
+extern "C" HRESULT DAPI FileCopyUsingHandlesWithProgress(
+    __in HANDLE hSource,
+    __in HANDLE hTarget,
+    __in DWORD64 cbCopy,
+    __in_opt LPPROGRESS_ROUTINE lpProgressRoutine,
+    __in_opt LPVOID lpData,
+    __in_opt LPBOOL pbCancel,
+    __out_opt DWORD64* pcbCopied
+)
+{
+    HRESULT hr = S_OK;
+    BOOL fStop = FALSE;
+    BOOL fCanceled = FALSE;
+    DWORD64 cbTotalCopied = 0;
+    BYTE rgbData[64 * 1024];
+    DWORD cbRead = 0;
+
+    LARGE_INTEGER liSourceSize = { };
+    LARGE_INTEGER liTotalCopied = { };
+    LARGE_INTEGER liZero = { };
+    DWORD dwResult = 0;
+
+    hr = FileSizeByHandle(hSource, &liSourceSize.QuadPart);
+    FileExitOnFailure(hr, "Failed to get size of source.");
+
+    if (0 < cbCopy && cbCopy < (DWORD64)liSourceSize.QuadPart)
+    {
+        liSourceSize.QuadPart = cbCopy;
+    }
+
+    dwResult = lpProgressRoutine(liSourceSize, liTotalCopied, liZero, liZero, 0, CALLBACK_STREAM_SWITCH, hSource, hTarget, lpData);
+    switch (dwResult)
+    {
+    case PROGRESS_CONTINUE:
+        break;
+
+    case PROGRESS_CANCEL:
+        fCanceled = TRUE;
+        fStop = TRUE;
+        break;
+
+    case PROGRESS_STOP:
+        fStop = TRUE;
+        break;
+
+    case PROGRESS_QUIET:
+        lpProgressRoutine = NULL;
+        break;
+    }
+
+    // Set size of the target file.
+    ::SetFilePointerEx(hTarget, liSourceSize, NULL, FILE_BEGIN);
+
+    if (!::SetEndOfFile(hTarget))
+    {
+        FileExitWithLastError(hr, "Failed to set end of target file.");
+    }
+
+    if (!::SetFilePointerEx(hTarget, liZero, NULL, FILE_BEGIN))
+    {
+        FileExitWithLastError(hr, "Failed to reset target file pointer.");
+    }
+
+    // Copy with progress.
+    while (!fStop && (0 == cbCopy || cbTotalCopied < cbCopy))
+    {
+        cbRead = static_cast<DWORD>((0 == cbCopy) ? countof(rgbData) : min(countof(rgbData), cbCopy - cbTotalCopied));
+        if (!::ReadFile(hSource, rgbData, cbRead, &cbRead, NULL))
+        {
+            FileExitWithLastError(hr, "Failed to read from source.");
+        }
+
+        if (cbRead)
+        {
+            hr = FileWriteHandle(hTarget, rgbData, cbRead);
+            FileExitOnFailure(hr, "Failed to write to target.");
+
+            cbTotalCopied += cbRead;
+
+            if (lpProgressRoutine)
+            {
+                liTotalCopied.QuadPart = cbTotalCopied;
+                dwResult = lpProgressRoutine(liSourceSize, liTotalCopied, liZero, liZero, 0, CALLBACK_CHUNK_FINISHED, hSource, hTarget, lpData);
+                switch (dwResult)
+                {
+                case PROGRESS_CONTINUE:
+                    break;
+
+                case PROGRESS_CANCEL:
+                    fCanceled = TRUE;
+                    fStop = TRUE;
+                    break;
+
+                case PROGRESS_STOP:
+                    fStop = TRUE;
+                    break;
+
+                case PROGRESS_QUIET:
+                    lpProgressRoutine = NULL;
+                    break;
+                }
+            }
+        }
+        else
+        {
+            fStop = TRUE;
+        }
+    }
+
+LExit:
+    if (pbCancel)
+    {
+        *pbCancel = fCanceled;
+    }
+
+    if (pcbCopied)
+    {
+        *pcbCopied = cbTotalCopied;
+    }
+
+    return hr;
+}
+
+
+/*******************************************************************
  FileEnsureCopy
 
 *******************************************************************/
diff --git a/src/dutil/inc/fileutil.h b/src/dutil/inc/fileutil.h
index 7caa62b8..48830043 100644
--- a/src/dutil/inc/fileutil.h
+++ b/src/dutil/inc/fileutil.h
@@ -148,6 +148,15 @@ HRESULT DAPI FileCopyUsingHandles(
     __in DWORD64 cbCopy,
     __out_opt DWORD64* pcbCopied
     );
+HRESULT DAPI FileCopyUsingHandlesWithProgress(
+    __in HANDLE hSource,
+    __in HANDLE hTarget,
+    __in DWORD64 cbCopy,
+    __in_opt LPPROGRESS_ROUTINE lpProgressRoutine,
+    __in_opt LPVOID lpData,
+    __in_opt LPBOOL pbCancel,
+    __out_opt DWORD64* pcbCopied
+    );
 HRESULT DAPI FileEnsureCopy(
     __in_z LPCWSTR wzSource,
     __in_z LPCWSTR wzTarget,
diff --git a/src/dutil/inc/regutil.h b/src/dutil/inc/regutil.h
index 2f09d244..75284940 100644
--- a/src/dutil/inc/regutil.h
+++ b/src/dutil/inc/regutil.h
@@ -226,6 +226,19 @@ HRESULT DAPI RegQueryKey(
     __out_opt DWORD* pcSubKeys,
     __out_opt DWORD* pcValues
     );
+HRESULT DAPI RegKeyReadNumber(
+    __in HKEY hk,
+    __in_z LPCWSTR wzSubKey,
+    __in_z_opt LPCWSTR wzName,
+    __in BOOL f64Bit,
+    __out DWORD* pdwValue
+    );
+BOOL DAPI RegValueExists(
+    __in HKEY hk,
+    __in_z LPCWSTR wzSubKey,
+    __in_z_opt LPCWSTR wzName,
+    __in BOOL f64Bit
+    );
 
 #ifdef __cplusplus
 }
diff --git a/src/dutil/regutil.cpp b/src/dutil/regutil.cpp
index e1ef19e8..afd2d089 100644
--- a/src/dutil/regutil.cpp
+++ b/src/dutil/regutil.cpp
@@ -283,7 +283,7 @@ LExit:
 
 
 /********************************************************************
- RegKeyEnum - enumerates a registry key.
+ RegKeyEnum - enumerates child registry keys.
 
 *********************************************************************/
 extern "C" HRESULT DAPI RegKeyEnum(
@@ -340,7 +340,7 @@ LExit:
 
 
 /********************************************************************
- RegValueEnum - enumerates a registry value.
+ RegValueEnum - enumerates registry values.
 
 *********************************************************************/
 HRESULT DAPI RegValueEnum(
@@ -939,6 +939,61 @@ LExit:
     return hr;
 }
 
+/********************************************************************
+RegKeyReadNumber - reads a DWORD registry key value as a number from
+a specified subkey.
+
+*********************************************************************/
+extern "C" HRESULT DAPI RegKeyReadNumber(
+    __in HKEY hk,
+    __in_z LPCWSTR wzSubKey,
+    __in_z_opt LPCWSTR wzName,
+    __in BOOL f64Bit,
+    __out DWORD* pdwValue
+    )
+{
+    HRESULT hr = S_OK;
+    HKEY hkKey = NULL;
+
+    hr = RegOpen(hk, wzSubKey, KEY_READ | f64Bit ? KEY_WOW64_64KEY : 0, &hkKey);
+    RegExitOnFailure(hr, "Failed to open key: %ls", wzSubKey);
+
+    hr = RegReadNumber(hkKey, wzName, pdwValue);
+    RegExitOnFailure(hr, "Failed to read value: %ls/@%ls", wzSubKey, wzName);
+
+LExit:
+    ReleaseRegKey(hkKey);
+
+    return hr;
+}
+
+/********************************************************************
+RegValueExists - determines whether a named value exists in a 
+specified subkey.
+
+*********************************************************************/
+extern "C" BOOL DAPI RegValueExists(
+    __in HKEY hk,
+    __in_z LPCWSTR wzSubKey,
+    __in_z_opt LPCWSTR wzName,
+    __in BOOL f64Bit
+    )
+{
+    HRESULT hr = S_OK;
+    HKEY hkKey = NULL;
+    DWORD dwType = 0;
+
+    hr = RegOpen(hk, wzSubKey, KEY_READ | f64Bit ? KEY_WOW64_64KEY : 0, &hkKey);
+    RegExitOnFailure(hr, "Failed to open key: %ls", wzSubKey);
+
+    hr = RegGetType(hkKey, wzName, &dwType);
+    RegExitOnFailure(hr, "Failed to read value type: %ls/@%ls", wzSubKey, wzName);
+
+LExit:
+    ReleaseRegKey(hkKey);
+
+    return SUCCEEDED(hr);
+}
 
 static HRESULT WriteStringToRegistry(
     __in HKEY hk,
diff --git a/src/dutil/wiutil.cpp b/src/dutil/wiutil.cpp
index ffbfe85a..f1984266 100644
--- a/src/dutil/wiutil.cpp
+++ b/src/dutil/wiutil.cpp
@@ -4,6 +4,7 @@
 
 
 // Exit macros
+#define WiuExitTrace(x, s, ...) ExitTraceSource(DUTIL_SOURCE_WIUTIL, x, s, __VA_ARGS__)
 #define WiuExitOnLastError(x, s, ...) ExitOnLastErrorSource(DUTIL_SOURCE_WIUTIL, x, s, __VA_ARGS__)
 #define WiuExitOnLastErrorDebugTrace(x, s, ...) ExitOnLastErrorDebugTraceSource(DUTIL_SOURCE_WIUTIL, x, s, __VA_ARGS__)
 #define WiuExitWithLastError(x, s, ...) ExitWithLastErrorSource(DUTIL_SOURCE_WIUTIL, x, s, __VA_ARGS__)
@@ -692,12 +693,23 @@ extern "C" HRESULT DAPI WiuEnumRelatedProductCodes(
 
         if (fReturnHighestVersionOnly)
         {
-            // get the version
+            // try to get the version but if the product registration is broken
+            // (for whatever reason), skip this product
             hr = WiuGetProductInfo(wzCurrentProductCode, L"VersionString", &sczInstalledVersion);
-            WiuExitOnFailure(hr, "Failed to get version for product code: %ls", wzCurrentProductCode);
+            if (FAILED(hr))
+            {
+                WiuExitTrace(hr, "Could not get product version for product code: %ls, skipping...", wzCurrentProductCode);
+                continue;
+            }
 
+            // try to parse the product version but if it is corrupt (for whatever
+            // reason), skip it
             hr = FileVersionFromStringEx(sczInstalledVersion, 0, &qwCurrentVersion);
-            WiuExitOnFailure(hr, "Failed to convert version: %ls to DWORD64 for product code: %ls", sczInstalledVersion, wzCurrentProductCode);
+            if (FAILED(hr))
+            {
+                WiuExitTrace(hr, "Could not convert version: %ls to DWORD64 for product code: %ls, skipping...", sczInstalledVersion, wzCurrentProductCode);
+                continue;
+            }
 
             // if this is the first product found then it is the highest version (for now)
             if (0 == *pcRelatedProducts)