1 files changed, 246 insertions, 0 deletions

diff --git a/src/api/burn/balutil/balretry.cpp b/src/api/burn/balutil/balretry.cpp
new file mode 100644
index 00000000..9d8abd6d
--- /dev/null
+++ b/src/api/burn/balutil/balretry.cpp
@@ -0,0 +1,246 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+#include "precomp.h"
+
+typedef enum BALRETRY_TYPE
+{
+    BALRETRY_TYPE_CACHE_CONTAINER,
+    BALRETRY_TYPE_CACHE_PAYLOAD,
+    BALRETRY_TYPE_EXECUTE,
+} BALRETRY_TYPE;
+
+struct BALRETRY_INFO
+{
+    LPWSTR sczId;
+    DWORD cRetries;
+    DWORD dwLastError;
+};
+
+static DWORD vdwMaxRetries = 0;
+static DWORD vdwTimeout = 0;
+static BALRETRY_INFO vrgRetryInfo[3];
+
+// prototypes
+static BOOL IsActiveRetryEntry(
+    __in BALRETRY_TYPE type,
+    __in_z LPCWSTR sczId
+    );
+
+static HRESULT StartActiveRetryEntry(
+    __in BALRETRY_TYPE type,
+    __in_z LPCWSTR sczId
+    );
+
+
+DAPI_(void) BalRetryInitialize(
+    __in DWORD dwMaxRetries,
+    __in DWORD dwTimeout
+    )
+{
+    BalRetryUninitialize(); // clean everything out.
+
+    vdwMaxRetries = dwMaxRetries;
+    vdwTimeout = dwTimeout;
+}
+
+
+DAPI_(void) BalRetryUninitialize()
+{
+    for (DWORD i = 0; i < countof(vrgRetryInfo); ++i)
+    {
+        ReleaseStr(vrgRetryInfo[i].sczId);
+        memset(vrgRetryInfo + i, 0, sizeof(BALRETRY_INFO));
+    }
+
+    vdwMaxRetries = 0;
+    vdwTimeout = 0;
+}
+
+
+DAPI_(void) BalRetryStartContainerOrPayload(
+    __in_z_opt LPCWSTR wzContainerOrPackageId,
+    __in_z_opt LPCWSTR wzPayloadId
+    )
+{
+    if (!wzContainerOrPackageId && !wzPayloadId)
+    {
+        ReleaseNullStr(vrgRetryInfo[BALRETRY_TYPE_CACHE_CONTAINER].sczId);
+        ReleaseNullStr(vrgRetryInfo[BALRETRY_TYPE_CACHE_PAYLOAD].sczId);
+    }
+    else if (wzPayloadId)
+    {
+        StartActiveRetryEntry(BALRETRY_TYPE_CACHE_PAYLOAD, wzPayloadId);
+    }
+    else
+    {
+        StartActiveRetryEntry(BALRETRY_TYPE_CACHE_CONTAINER, wzContainerOrPackageId);
+    }
+}
+
+
+DAPI_(void) BalRetryStartPackage(
+    __in_z LPCWSTR wzPackageId
+    )
+{
+    StartActiveRetryEntry(BALRETRY_TYPE_EXECUTE, wzPackageId);
+}
+
+
+DAPI_(void) BalRetryErrorOccurred(
+    __in_z LPCWSTR wzPackageId,
+    __in DWORD dwError
+    )
+{
+    if (IsActiveRetryEntry(BALRETRY_TYPE_EXECUTE, wzPackageId))
+    {
+        vrgRetryInfo[BALRETRY_TYPE_EXECUTE].dwLastError = dwError;
+    }
+}
+
+
+DAPI_(HRESULT) BalRetryEndContainerOrPayload(
+    __in_z_opt LPCWSTR wzContainerOrPackageId,
+    __in_z_opt LPCWSTR wzPayloadId,
+    __in HRESULT hrError,
+    __inout BOOL* pfRetry
+    )
+{
+    HRESULT hr = S_OK;
+    BALRETRY_TYPE type = BALRETRY_TYPE_CACHE_PAYLOAD;
+    LPCWSTR wzId = NULL;
+
+    if (!wzContainerOrPackageId && !wzPayloadId)
+    {
+        ReleaseNullStr(vrgRetryInfo[BALRETRY_TYPE_CACHE_CONTAINER].sczId);
+        ReleaseNullStr(vrgRetryInfo[BALRETRY_TYPE_CACHE_PAYLOAD].sczId);
+        ExitFunction();
+    }
+    else if (wzPayloadId)
+    {
+        type = BALRETRY_TYPE_CACHE_PAYLOAD;
+        wzId = wzPayloadId;
+    }
+    else
+    {
+        type = BALRETRY_TYPE_CACHE_CONTAINER;
+        wzId = wzContainerOrPackageId;
+    }
+
+    if (FAILED(hrError) && vrgRetryInfo[type].cRetries < vdwMaxRetries && IsActiveRetryEntry(type, wzId))
+    {
+        // Retry on all errors except the following.
+        if (HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT) != hrError &&
+            BG_E_NETWORK_DISCONNECTED != hrError &&
+            HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND) != hrError &&
+            HRESULT_FROM_WIN32(ERROR_INTERNET_NAME_NOT_RESOLVED) != hrError)
+        {
+            *pfRetry = TRUE;
+        }
+    }
+
+LExit:
+    return hr;
+}
+
+
+DAPI_(HRESULT) BalRetryEndPackage(
+    __in_z LPCWSTR wzPackageId,
+    __in HRESULT hrError,
+    __inout BOOL* pfRetry
+    )
+{
+    HRESULT hr = S_OK;
+    BALRETRY_TYPE type = BALRETRY_TYPE_EXECUTE;
+
+    if (!wzPackageId || !*wzPackageId)
+    {
+        ReleaseNullStr(vrgRetryInfo[type].sczId);
+    }
+    else if (FAILED(hrError) && vrgRetryInfo[type].cRetries < vdwMaxRetries && IsActiveRetryEntry(type, wzPackageId))
+    {
+        // If the service is out of whack, just try again.
+        if (HRESULT_FROM_WIN32(ERROR_INSTALL_SERVICE_FAILURE) == hrError)
+        {
+            *pfRetry = TRUE;
+        }
+        else if (HRESULT_FROM_WIN32(ERROR_INSTALL_FAILURE) == hrError)
+        {
+            DWORD dwError = vrgRetryInfo[type].dwLastError;
+
+            // If we failed with one of these specific error codes, then retry since
+            // we've seen these have a high success of succeeding on retry.
+            if (1303 == dwError ||
+                1304 == dwError ||
+                1306 == dwError ||
+                1307 == dwError ||
+                1309 == dwError ||
+                1310 == dwError ||
+                1311 == dwError ||
+                1312 == dwError ||
+                1316 == dwError ||
+                1317 == dwError ||
+                1321 == dwError ||
+                1335 == dwError ||
+                1402 == dwError ||
+                1406 == dwError ||
+                1606 == dwError ||
+                1706 == dwError ||
+                1719 == dwError ||
+                1723 == dwError ||
+                1923 == dwError ||
+                1931 == dwError)
+            {
+                *pfRetry = TRUE;
+            }
+        }
+        else if (HRESULT_FROM_WIN32(ERROR_INSTALL_ALREADY_RUNNING) == hrError)
+        {
+            *pfRetry = TRUE;
+        }
+    }
+
+    return hr;
+}
+
+
+// Internal functions.
+
+static BOOL IsActiveRetryEntry(
+    __in BALRETRY_TYPE type,
+    __in_z LPCWSTR sczId
+    )
+{
+    BOOL fActive = FALSE;
+
+    fActive = vrgRetryInfo[type].sczId && CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, sczId, -1, vrgRetryInfo[type].sczId, -1);
+
+    return fActive;
+}
+
+static HRESULT StartActiveRetryEntry(
+    __in BALRETRY_TYPE type,
+    __in_z LPCWSTR sczId
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (!sczId || !*sczId)
+    {
+        ReleaseNullStr(vrgRetryInfo[type].sczId);
+    }
+    else if (IsActiveRetryEntry(type, sczId))
+    {
+        ++vrgRetryInfo[type].cRetries;
+        ::Sleep(vdwTimeout);
+    }
+    else
+    {
+        hr = StrAllocString(&vrgRetryInfo[type].sczId, sczId, 0);
+
+        vrgRetryInfo[type].cRetries = 0;
+    }
+
+    vrgRetryInfo[type].dwLastError = ERROR_SUCCESS;
+
+    return hr;
+}