1 files changed, 601 insertions, 0 deletions

diff --git a/src/engine/variant.cpp b/src/engine/variant.cpp
new file mode 100644
index 00000000..2a9f08ed
--- /dev/null
+++ b/src/engine/variant.cpp
@@ -0,0 +1,601 @@
+// 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"
+
+#define VARIANT_ENCRYPTION_SCOPE CRYPTPROTECTMEMORY_SAME_PROCESS
+
+// internal function declarations
+
+static HRESULT BVariantEncryptNumeric(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    );
+
+static HRESULT BVariantEncryptString(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    );
+
+static HRESULT BVariantEncryptVersion(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    );
+
+static HRESULT BVariantRetrieveDecryptedNumeric(
+    __in BURN_VARIANT* pVariant,
+    __out LONGLONG* pllValue
+    );
+
+static HRESULT BVariantRetrieveDecryptedString(
+    __in BURN_VARIANT* pVariant,
+    __out LPWSTR* psczValue
+    );
+
+static HRESULT BVariantRetrieveDecryptedVersion(
+    __in BURN_VARIANT* pVariant,
+    __out DWORD64* pqwValue
+    );
+
+// function definitions
+
+extern "C" void BVariantUninitialize(
+    __in BURN_VARIANT* pVariant
+    )
+{
+    if (BURN_VARIANT_TYPE_STRING == pVariant->Type)
+    {
+        StrSecureZeroFreeString(pVariant->sczValue);
+    }
+    SecureZeroMemory(pVariant, sizeof(BURN_VARIANT));
+}
+
+// The contents of pllValue may be sensitive, should keep encrypted and SecureZeroMemory.
+extern "C" HRESULT BVariantGetNumeric(
+    __in BURN_VARIANT* pVariant,
+    __out LONGLONG* pllValue
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczValue = NULL;
+
+    switch (pVariant->Type)
+    {
+    case BURN_VARIANT_TYPE_NUMERIC:
+        BVariantRetrieveDecryptedNumeric(pVariant, pllValue);
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantRetrieveDecryptedString(pVariant, &sczValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = StrStringToInt64(sczValue, 0, pllValue);
+            if (FAILED(hr))
+            {
+                hr = DISP_E_TYPEMISMATCH;
+            }
+        }
+        StrSecureZeroFreeString(sczValue);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        BVariantRetrieveDecryptedVersion(pVariant, (DWORD64*)pllValue);
+        break;
+    default:
+        hr = E_INVALIDARG;
+        break;
+    }
+
+    return hr;
+}
+
+// The contents of psczValue may be sensitive, should keep encrypted and SecureZeroFree.
+extern "C" HRESULT BVariantGetString(
+    __in BURN_VARIANT* pVariant,
+    __out_z LPWSTR* psczValue
+    )
+{
+    HRESULT hr = S_OK;
+    LONGLONG llValue = 0;
+    DWORD64 qwValue = 0;
+
+    switch (pVariant->Type)
+    {
+    case BURN_VARIANT_TYPE_NUMERIC:
+        hr = BVariantRetrieveDecryptedNumeric(pVariant, &llValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = StrAllocFormattedSecure(psczValue, L"%I64d", llValue);
+            ExitOnFailure(hr, "Failed to convert int64 to string.");
+        }
+        SecureZeroMemory(&llValue, sizeof(llValue));
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantRetrieveDecryptedString(pVariant, psczValue);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        hr = BVariantRetrieveDecryptedVersion(pVariant, &qwValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = StrAllocFormattedSecure(psczValue, L"%hu.%hu.%hu.%hu",
+                (WORD)(qwValue >> 48),
+                (WORD)(qwValue >> 32),
+                (WORD)(qwValue >> 16),
+                (WORD)qwValue);
+            ExitOnFailure(hr, "Failed to convert version to string.");
+        }
+        SecureZeroMemory(&qwValue, sizeof(qwValue));
+        break;
+    default:
+        hr = E_INVALIDARG;
+        break;
+    }
+
+LExit:
+    return hr;
+}
+
+// The contents of pqwValue may be sensitive, should keep encrypted and SecureZeroMemory.
+extern "C" HRESULT BVariantGetVersion(
+    __in BURN_VARIANT* pVariant,
+    __out DWORD64* pqwValue
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczValue = NULL;
+
+    switch (pVariant->Type)
+    {
+    case BURN_VARIANT_TYPE_NUMERIC:
+        BVariantRetrieveDecryptedNumeric(pVariant, (LONGLONG*)pqwValue);
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantRetrieveDecryptedString(pVariant, &sczValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = FileVersionFromStringEx(sczValue, 0, pqwValue);
+            if (FAILED(hr))
+            {
+                hr = DISP_E_TYPEMISMATCH;
+            }
+        }
+        StrSecureZeroFreeString(sczValue);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        BVariantRetrieveDecryptedVersion(pVariant, pqwValue);
+        break;
+    default:
+        hr = E_INVALIDARG;
+        break;
+    }
+
+    return hr;
+}
+
+extern "C" HRESULT BVariantSetNumeric(
+    __in BURN_VARIANT* pVariant,
+    __in LONGLONG llValue
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fEncryptValue = pVariant->fEncryptValue;
+
+    if (BURN_VARIANT_TYPE_STRING == pVariant->Type)
+    {
+        StrSecureZeroFreeString(pVariant->sczValue);
+    }
+    memset(pVariant, 0, sizeof(BURN_VARIANT));
+    pVariant->llValue = llValue;
+    pVariant->Type = BURN_VARIANT_TYPE_NUMERIC;
+    BVariantSetEncryption(pVariant, fEncryptValue);
+
+    return hr;
+}
+
+extern "C" HRESULT BVariantSetString(
+    __in BURN_VARIANT* pVariant,
+    __in_z_opt LPCWSTR wzValue,
+    __in DWORD_PTR cchValue
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fEncryptValue = pVariant->fEncryptValue;
+
+    if (!wzValue) // if we're nulling out the string, make the variable NONE.
+    {
+        BVariantUninitialize(pVariant);
+    }
+    else // assign the value.
+    {
+        if (BURN_VARIANT_TYPE_STRING != pVariant->Type)
+        {
+            memset(pVariant, 0, sizeof(BURN_VARIANT));
+        }
+        else
+        {
+            // We're about to copy an unencrypted value.
+            pVariant->fEncryptValue = FALSE;
+        }
+
+        hr = StrAllocStringSecure(&pVariant->sczValue, wzValue, cchValue);
+        ExitOnFailure(hr, "Failed to copy string.");
+
+        pVariant->Type = BURN_VARIANT_TYPE_STRING;
+    }
+
+LExit:
+    BVariantSetEncryption(pVariant, fEncryptValue);
+    return hr;
+}
+
+extern "C" HRESULT BVariantSetVersion(
+    __in BURN_VARIANT* pVariant,
+    __in DWORD64 qwValue
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fEncryptValue = pVariant->fEncryptValue;
+
+    if (BURN_VARIANT_TYPE_STRING == pVariant->Type)
+    {
+        StrSecureZeroFreeString(pVariant->sczValue);
+    }
+    memset(pVariant, 0, sizeof(BURN_VARIANT));
+    pVariant->qwValue = qwValue;
+    pVariant->Type = BURN_VARIANT_TYPE_VERSION;
+    BVariantSetEncryption(pVariant, fEncryptValue);
+
+    return hr;
+}
+
+extern "C" HRESULT BVariantSetValue(
+    __in BURN_VARIANT* pVariant,
+    __in BURN_VARIANT* pValue
+    )
+{
+    HRESULT hr = S_OK;
+    LONGLONG llValue = 0;
+    LPWSTR sczValue = NULL;
+    DWORD64 qwValue = 0;
+    BOOL fEncrypt = pVariant->fEncryptValue;
+
+    switch (pValue->Type)
+    {
+    case BURN_VARIANT_TYPE_NONE:
+        BVariantUninitialize(pVariant);
+        break;
+    case BURN_VARIANT_TYPE_NUMERIC:
+        hr = BVariantGetNumeric(pValue, &llValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = BVariantSetNumeric(pVariant, llValue);
+        }
+        SecureZeroMemory(&llValue, sizeof(llValue));
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantGetString(pValue, &sczValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = BVariantSetString(pVariant, sczValue, 0);
+        }
+        StrSecureZeroFreeString(sczValue);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        hr = BVariantGetVersion(pValue, &qwValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = BVariantSetVersion(pVariant, qwValue);
+        }
+        SecureZeroMemory(&qwValue, sizeof(qwValue));
+        break;
+    default:
+        hr = E_INVALIDARG;
+    }
+    ExitOnFailure(hr, "Failed to copy variant.");
+
+    hr = BVariantSetEncryption(pVariant, fEncrypt);
+
+LExit:
+    return hr;
+}
+
+extern "C" HRESULT BVariantCopy(
+    __in BURN_VARIANT* pSource,
+    __out BURN_VARIANT* pTarget
+    )
+{
+    HRESULT hr = S_OK;
+    LONGLONG llValue = 0;
+    LPWSTR sczValue = NULL;
+    DWORD64 qwValue = 0;
+
+    BVariantUninitialize(pTarget);
+
+    switch (pSource->Type)
+    {
+    case BURN_VARIANT_TYPE_NONE:
+        break;
+    case BURN_VARIANT_TYPE_NUMERIC:
+        hr = BVariantGetNumeric(pSource, &llValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = BVariantSetNumeric(pTarget, llValue);
+        }
+        SecureZeroMemory(&llValue, sizeof(llValue));
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantGetString(pSource, &sczValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = BVariantSetString(pTarget, sczValue, 0);
+        }
+        StrSecureZeroFreeString(sczValue);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        hr = BVariantGetVersion(pSource, &qwValue);
+        if (SUCCEEDED(hr))
+        {
+            hr = BVariantSetVersion(pTarget, qwValue);
+        }
+        SecureZeroMemory(&qwValue, sizeof(qwValue));
+        break;
+    default:
+        hr = E_INVALIDARG;
+    }
+    ExitOnFailure(hr, "Failed to copy variant.");
+
+    hr = BVariantSetEncryption(pTarget, pSource->fEncryptValue);
+
+LExit:
+    return hr;
+}
+
+extern "C" HRESULT BVariantChangeType(
+    __in BURN_VARIANT* pVariant,
+    __in BURN_VARIANT_TYPE type
+    )
+{
+    HRESULT hr = S_OK;
+    BURN_VARIANT variant = { };
+    BOOL fEncryptValue = pVariant->fEncryptValue;
+
+    if (pVariant->Type == type)
+    {
+        ExitFunction(); // variant already is of the requested type
+    }
+
+    switch (type)
+    {
+    case BURN_VARIANT_TYPE_NONE:
+        hr = S_OK;
+        break;
+    case BURN_VARIANT_TYPE_NUMERIC:
+        hr = BVariantGetNumeric(pVariant, &variant.llValue);
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantGetString(pVariant, &variant.sczValue);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        hr = BVariantGetVersion(pVariant, &variant.qwValue);
+        break;
+    default:
+        ExitFunction1(hr = E_INVALIDARG);
+    }
+    ExitOnFailure(hr, "Failed to copy variant value.");
+    variant.Type = type;
+
+    BVariantUninitialize(pVariant);
+    memcpy_s(pVariant, sizeof(BURN_VARIANT), &variant, sizeof(BURN_VARIANT));
+    SecureZeroMemory(&variant, sizeof(BURN_VARIANT));
+    BVariantSetEncryption(pVariant, fEncryptValue);
+
+LExit:
+    return hr;
+}
+
+extern "C" HRESULT BVariantSetEncryption(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (pVariant->fEncryptValue == fEncrypt)
+    {
+        // The requested encryption state is already applied.
+        ExitFunction();
+    }
+    
+    switch (pVariant->Type)
+    {
+    case BURN_VARIANT_TYPE_NONE:
+        hr = S_OK;
+        break;
+    case BURN_VARIANT_TYPE_NUMERIC:
+        hr = BVariantEncryptNumeric(pVariant, fEncrypt);
+        break;
+    case BURN_VARIANT_TYPE_STRING:
+        hr = BVariantEncryptString(pVariant, fEncrypt);
+        break;
+    case BURN_VARIANT_TYPE_VERSION:
+        hr = BVariantEncryptVersion(pVariant, fEncrypt);
+        break;
+    default:
+        hr = E_INVALIDARG;
+    }
+    ExitOnFailure(hr, "Failed to set the variant's encryption state");
+    pVariant->fEncryptValue = fEncrypt;
+
+LExit:
+    return hr;
+}
+
+static HRESULT BVariantEncryptNumeric(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (fEncrypt)
+    {
+        hr = CrypEncryptMemory(&pVariant->llValue, sizeof(pVariant->encryptionPadding), VARIANT_ENCRYPTION_SCOPE);
+    }
+    else
+    {
+        hr = CrypDecryptMemory(&pVariant->llValue, sizeof(pVariant->encryptionPadding), VARIANT_ENCRYPTION_SCOPE);
+    }
+
+//LExit:
+    return hr;
+}
+
+static HRESULT BVariantEncryptString(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T cbData = 0;
+
+    if (NULL == pVariant->sczValue)
+    {
+        ExitFunction();
+    }
+
+    cbData = MemSize(pVariant->sczValue);
+    if (-1 == cbData)
+    {
+        hr = E_INVALIDARG;
+        ExitOnFailure(hr, "Failed to get the size of the string");
+    }
+    
+    DWORD remainder = fEncrypt ? cbData % CRYP_ENCRYPT_MEMORY_SIZE : 0;
+    DWORD extraNeeded = 0 < remainder ? CRYP_ENCRYPT_MEMORY_SIZE - remainder : 0;
+    if ((MAXDWORD - extraNeeded) < cbData)
+    {
+        hr = E_INVALIDDATA;
+        ExitOnFailure(hr, "The string is too big: size %u", cbData);
+    }
+    else if (0 < extraNeeded)
+    {
+        cbData += extraNeeded;
+        LPVOID pvNew = NULL;
+        hr = MemReAllocSecure(static_cast<LPVOID>(pVariant->sczValue), cbData, TRUE, &pvNew);
+        ExitOnFailure(hr, "Failed to resize the string so it could be encrypted");
+        pVariant->sczValue = static_cast<LPWSTR>(pvNew);
+    }
+    
+    if (fEncrypt)
+    {
+        hr = CrypEncryptMemory(pVariant->sczValue, static_cast<DWORD>(cbData), VARIANT_ENCRYPTION_SCOPE);
+    }
+    else
+    {
+        hr = CrypDecryptMemory(pVariant->sczValue, static_cast<DWORD>(cbData), VARIANT_ENCRYPTION_SCOPE);
+    }
+
+LExit:
+    return hr;
+}
+
+static HRESULT BVariantEncryptVersion(
+    __in BURN_VARIANT* pVariant,
+    __in BOOL fEncrypt
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (fEncrypt)
+    {
+        hr = CrypEncryptMemory(&pVariant->qwValue, sizeof(pVariant->encryptionPadding), VARIANT_ENCRYPTION_SCOPE);
+    }
+    else
+    {
+        hr = CrypDecryptMemory(&pVariant->qwValue, sizeof(pVariant->encryptionPadding), VARIANT_ENCRYPTION_SCOPE);
+    }
+
+//LExit:
+    return hr;
+}
+
+// The contents of pllValue may be sensitive, should keep encrypted and SecureZeroMemory.
+static HRESULT BVariantRetrieveDecryptedNumeric(
+    __in BURN_VARIANT* pVariant,
+    __out LONGLONG* pllValue
+    )
+{
+    HRESULT hr = S_OK;
+
+    Assert(NULL != pllValue);
+    if (pVariant->fEncryptValue)
+    {
+        hr = BVariantEncryptNumeric(pVariant, FALSE);
+        ExitOnFailure(hr, "Failed to decrypt numeric");
+    }
+
+    *pllValue = pVariant->llValue;
+
+    if (pVariant->fEncryptValue)
+    {
+        hr = BVariantEncryptNumeric(pVariant, TRUE);
+    }
+
+LExit:
+    return hr;
+}
+
+// The contents of psczValue may be sensitive, should keep encrypted and SecureZeroFree.
+static HRESULT BVariantRetrieveDecryptedString(
+    __in BURN_VARIANT* pVariant,
+    __out LPWSTR* psczValue
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (NULL == pVariant->sczValue)
+    {
+        *psczValue = NULL;
+        ExitFunction();
+    }
+
+    if (pVariant->fEncryptValue)
+    {
+        hr = BVariantEncryptString(pVariant, FALSE);
+        ExitOnFailure(hr, "Failed to decrypt string");
+    }
+
+    hr = StrAllocStringSecure(psczValue, pVariant->sczValue, 0);
+    ExitOnFailure(hr, "Failed to copy value.");
+
+    if (pVariant->fEncryptValue)
+    {
+        hr = BVariantEncryptString(pVariant, TRUE);
+    }
+
+LExit:
+    return hr;
+}
+
+// The contents of pqwValue may be sensitive, should keep encrypted and SecureZeroMemory.
+static HRESULT BVariantRetrieveDecryptedVersion(
+    __in BURN_VARIANT* pVariant,
+    __out DWORD64* pqwValue
+    )
+{
+    HRESULT hr = S_OK;
+
+    Assert(NULL != pqwValue);
+    if (pVariant->fEncryptValue)
+    {
+        hr = BVariantEncryptVersion(pVariant, FALSE);
+        ExitOnFailure(hr, "Failed to decrypt version");
+    }
+
+    *pqwValue = pVariant->qwValue;
+
+    if (pVariant->fEncryptValue)
+    {
+        hr = BVariantEncryptVersion(pVariant, TRUE);
+    }
+
+LExit:
+    return hr;
+}