Add fixed buffer butil APIs.

3 files changed, 203 insertions, 21 deletions

diff --git a/src/burn/test/BurnUnitTest/RegistrationTest.cpp b/src/burn/test/BurnUnitTest/RegistrationTest.cpp
index a8319bb2..b1c911f7 100644
--- a/src/burn/test/BurnUnitTest/RegistrationTest.cpp
+++ b/src/burn/test/BurnUnitTest/RegistrationTest.cpp
@@ -593,9 +593,10 @@ namespace Bootstrapper
                 Assert::Empty((System::Collections::IEnumerable ^)Registry::GetValue(gcnew String(TEST_VARIABLE_KEY), gcnew String(L"WixBundleForcedRestartPackage"), nullptr));
 
                 hr = StrAlloc(&sczRelatedBundleId, MAX_GUID_CHARS + 1);
+                NativeAssert::Succeeded(hr, "Failed to allocate buffer for related bundle id.");
 
                 // Verify we can find ourself via the UpgradeCode
-                hr = BundleEnumRelatedBundle(TEST_BUNDLE_UPGRADE_CODE, BUNDLE_INSTALL_CONTEXT_USER, &dwRelatedBundleIndex, sczRelatedBundleId);
+                hr = BundleEnumRelatedBundleFixed(TEST_BUNDLE_UPGRADE_CODE, BUNDLE_INSTALL_CONTEXT_USER, &dwRelatedBundleIndex, sczRelatedBundleId);
                 TestThrowOnFailure(hr, L"Failed to enumerate related bundle.");
 
                 NativeAssert::StringEqual(TEST_BUNDLE_ID, sczRelatedBundleId);
diff --git a/src/libs/dutil/WixToolset.DUtil/butil.cpp b/src/libs/dutil/WixToolset.DUtil/butil.cpp
index 5e980699..4c96dfc1 100644
--- a/src/libs/dutil/WixToolset.DUtil/butil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/butil.cpp
@@ -57,6 +57,11 @@ static HRESULT OpenBundleKey(
     __in_opt LPCWSTR wzSubKey,
     __inout HKEY* phKey
     );
+static HRESULT CopyStringToBuffer(
+    __in_z LPWSTR wzValue,
+    __in_z_opt LPWSTR wzBuffer,
+    __inout SIZE_T* pcchBuffer
+    );
 
 DAPI_(HRESULT) BundleGetBundleInfo(
     __in_z LPCWSTR wzBundleId,
@@ -111,11 +116,39 @@ LExit:
 }
 
 
+DAPI_(HRESULT) BundleGetBundleInfoFixed(
+    __in_z LPCWSTR wzBundleId,
+    __in_z LPCWSTR wzAttribute,
+    __out_ecount_opt(*pcchValue) LPWSTR wzValue,
+    __inout SIZE_T* pcchValue
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczValue = NULL;
+
+    if (!pcchValue)
+    {
+        ButilExitWithRootFailure(hr, E_INVALIDARG, "An invalid parameter was passed to the function.");
+    }
+
+    hr = BundleGetBundleInfo(wzBundleId, wzAttribute, &sczValue);
+    if (SUCCEEDED(hr))
+    {
+        hr = CopyStringToBuffer(sczValue, wzValue, pcchValue);
+    }
+
+LExit:
+    ReleaseStr(sczValue);
+
+    return hr;
+}
+
+
 DAPI_(HRESULT) BundleEnumRelatedBundle(
-  __in_z LPCWSTR wzUpgradeCode,
-  __in BUNDLE_INSTALL_CONTEXT context,
-  __inout PDWORD pdwStartIndex,
-  __out_ecount(MAX_GUID_CHARS+1) LPWSTR lpBundleIdBuf
+    __in_z LPCWSTR wzUpgradeCode,
+    __in BUNDLE_INSTALL_CONTEXT context,
+    __inout PDWORD pdwStartIndex,
+    __deref_out_z LPWSTR* psczBundleId
     )
 {
     HRESULT hr = S_OK;
@@ -123,7 +156,6 @@ DAPI_(HRESULT) BundleEnumRelatedBundle(
     HKEY hkUninstall = NULL;
     HKEY hkBundle = NULL;
     LPWSTR sczUninstallSubKey = NULL;
-    size_t cchUninstallSubKey = 0;
     LPWSTR sczUninstallSubKeyPath = NULL;
     LPWSTR sczValue = NULL;
     DWORD dwType = 0;
@@ -132,7 +164,7 @@ DAPI_(HRESULT) BundleEnumRelatedBundle(
     DWORD cBundleUpgradeCodes = 0;
     BOOL fUpgradeCodeFound = FALSE;
 
-    if (!wzUpgradeCode || !lpBundleIdBuf || !pdwStartIndex)
+    if (!wzUpgradeCode || !pdwStartIndex)
     {
         ButilExitOnFailure(hr = E_INVALIDARG, "An invalid parameter was passed to the function.");
     }
@@ -205,13 +237,10 @@ DAPI_(HRESULT) BundleEnumRelatedBundle(
 
         if (fUpgradeCodeFound)
         {
-            if (lpBundleIdBuf)
+            if (psczBundleId)
             {
-                hr = ::StringCchLengthW(sczUninstallSubKey, STRSAFE_MAX_CCH, &cchUninstallSubKey);
-                ButilExitOnRootFailure(hr, "Failed to calculate length of string.");
-
-                hr = ::StringCchCopyNExW(lpBundleIdBuf, MAX_GUID_CHARS + 1, sczUninstallSubKey, cchUninstallSubKey, NULL, NULL, STRSAFE_FILL_BEHIND_NULL);
-                ButilExitOnRootFailure(hr, "Failed to copy the property value to the output buffer.");
+                *psczBundleId = sczUninstallSubKey;
+                sczUninstallSubKey = NULL;
             }
 
             break;
@@ -235,6 +264,34 @@ LExit:
 }
 
 
+DAPI_(HRESULT) BundleEnumRelatedBundleFixed(
+    __in_z LPCWSTR wzUpgradeCode,
+    __in BUNDLE_INSTALL_CONTEXT context,
+    __inout PDWORD pdwStartIndex,
+    __out_ecount(MAX_GUID_CHARS+1) LPWSTR wzBundleId
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczValue = NULL;
+    size_t cchValue = 0;
+
+    hr = BundleEnumRelatedBundle(wzUpgradeCode, context, pdwStartIndex, &sczValue);
+    if (SUCCEEDED(hr) && wzBundleId)
+    {
+        hr = ::StringCchLengthW(sczValue, STRSAFE_MAX_CCH, &cchValue);
+        ButilExitOnRootFailure(hr, "Failed to calculate length of string.");
+
+        hr = ::StringCchCopyNExW(wzBundleId, MAX_GUID_CHARS + 1, sczValue, cchValue, NULL, NULL, STRSAFE_FILL_BEHIND_NULL);
+        ButilExitOnRootFailure(hr, "Failed to copy the property value to the output buffer.");
+    }
+
+LExit:
+    ReleaseStr(sczValue);
+
+    return hr;
+}
+
+
 DAPI_(HRESULT) BundleGetBundleVariable(
     __in_z LPCWSTR wzBundleId,
     __in_z LPCWSTR wzVariable,
@@ -282,6 +339,34 @@ LExit:
     return hr;
 }
 
+
+DAPI_(HRESULT) BundleGetBundleVariableFixed(
+    __in_z LPCWSTR wzBundleId,
+    __in_z LPCWSTR wzVariable,
+    __out_ecount_opt(*pcchValue) LPWSTR wzValue,
+    __inout SIZE_T* pcchValue
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczValue = NULL;
+
+    if (!pcchValue)
+    {
+        ButilExitWithRootFailure(hr, E_INVALIDARG, "An invalid parameter was passed to the function.");
+    }
+
+    hr = BundleGetBundleVariable(wzBundleId, wzVariable, &sczValue);
+    if (SUCCEEDED(hr))
+    {
+        hr = CopyStringToBuffer(sczValue, wzValue, pcchValue);
+    }
+
+LExit:
+    ReleaseStr(sczValue);
+
+    return hr;
+}
+
 static HRESULT LocateAndQueryBundleValue(
     __in_z LPCWSTR wzBundleId,
     __in_opt LPCWSTR wzSubKey,
@@ -356,3 +441,34 @@ LExit:
 
     return hr;
 }
+
+static HRESULT CopyStringToBuffer(
+    __in_z LPWSTR wzValue,
+    __in_z_opt LPWSTR wzBuffer,
+    __inout SIZE_T* pcchBuffer
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fTooSmall = !wzBuffer;
+
+    if (!fTooSmall)
+    {
+        hr = ::StringCchCopyExW(wzBuffer, *pcchBuffer, wzValue, NULL, NULL, STRSAFE_FILL_BEHIND_NULL);
+        if (STRSAFE_E_INSUFFICIENT_BUFFER == hr)
+        {
+            fTooSmall = TRUE;
+        }
+    }
+
+    if (fTooSmall)
+    {
+        hr = ::StringCchLengthW(wzValue, STRSAFE_MAX_LENGTH, reinterpret_cast<size_t*>(pcchBuffer));
+        if (SUCCEEDED(hr))
+        {
+            hr = E_MOREDATA;
+            *pcchBuffer += 1; // null terminator.
+        }
+    }
+
+    return hr;
+}
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/butil.h b/src/libs/dutil/WixToolset.DUtil/inc/butil.h
index 989d4237..0405be8b 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/butil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/butil.h
@@ -30,15 +30,54 @@ RETURNS:
     All other returns are unexpected returns from other dutil methods.
 ********************************************************************/
 HRESULT DAPI BundleGetBundleInfo(
-    __in_z LPCWSTR szBundleId,
-    __in_z LPCWSTR szAttribute,
+    __in_z LPCWSTR wzBundleId,
+    __in_z LPCWSTR wzAttribute,
     __deref_out_z LPWSTR* psczValue
     );
 
 /********************************************************************
+BundleGetBundleInfoFixed - Queries the bundle installation metadata for a given property
+
+RETURNS:
+    E_INVALIDARG
+        An invalid parameter was passed to the function.
+    HRESULT_FROM_WIN32(ERROR_UNKNOWN_PRODUCT)
+        The bundle is not installed
+    HRESULT_FROM_WIN32(ERROR_UNKNOWN_PROPERTY)
+        The property is unrecognized
+    HRESULT_FROM_WIN32(ERROR_MORE_DATA)
+        A buffer is too small to hold the requested data.
+    E_NOTIMPL:
+        Tried to read a bundle attribute for a type which has not been implemented
+
+    All other returns are unexpected returns from other dutil methods.
+********************************************************************/
+HRESULT DAPI BundleGetBundleInfoFixed(
+    __in_z LPCWSTR wzBundleId,
+    __in_z LPCWSTR wzAttribute,
+    __out_ecount_opt(*pcchValue) LPWSTR wzValue,
+    __inout SIZE_T* pcchValue
+    );
+
+/********************************************************************
 BundleEnumRelatedBundle - Queries the bundle installation metadata for installs with the given upgrade code
+RETURNS:
+    E_INVALIDARG
+        An invalid parameter was passed to the function.
 
-NOTE: lpBundleIdBuff is a buffer to receive the bundle GUID. This buffer must be 39 characters long. 
+    All other returns are unexpected returns from other dutil methods.
+********************************************************************/
+HRESULT DAPI BundleEnumRelatedBundle(
+    __in_z LPCWSTR wzUpgradeCode,
+    __in BUNDLE_INSTALL_CONTEXT context,
+    __inout PDWORD pdwStartIndex,
+    __deref_out_z LPWSTR* psczBundleId
+    );
+
+/********************************************************************
+BundleEnumRelatedBundleFixed - Queries the bundle installation metadata for installs with the given upgrade code
+
+NOTE: lpBundleIdBuff is a buffer to receive the bundle GUID. This buffer must be 39 characters long.
         The first 38 characters are for the GUID, and the last character is for the terminating null character.
 RETURNS:
     E_INVALIDARG
@@ -46,11 +85,11 @@ RETURNS:
 
     All other returns are unexpected returns from other dutil methods.
 ********************************************************************/
-HRESULT DAPI BundleEnumRelatedBundle(
-  __in_z   LPCWSTR lpUpgradeCode,
-  __in     BUNDLE_INSTALL_CONTEXT context,
-  __inout  PDWORD pdwStartIndex,
-  __out_ecount(MAX_GUID_CHARS+1)  LPWSTR lpBundleIdBuf
+HRESULT DAPI BundleEnumRelatedBundleFixed(
+    __in_z LPCWSTR wzUpgradeCode,
+    __in BUNDLE_INSTALL_CONTEXT context,
+    __inout PDWORD pdwStartIndex,
+    __out_ecount(MAX_GUID_CHARS+1) LPWSTR wzBundleId
     );
 
 /********************************************************************
@@ -77,6 +116,32 @@ HRESULT DAPI BundleGetBundleVariable(
     __deref_out_z LPWSTR* psczValue
     );
 
+/********************************************************************
+BundleGetBundleVariableFixed - Queries the bundle installation metadata for a given variable
+
+RETURNS:
+    S_OK
+        Success, if the variable had a value, it's returned in psczValue
+    E_INVALIDARG
+        An invalid parameter was passed to the function.
+    HRESULT_FROM_WIN32(ERROR_UNKNOWN_PRODUCT)
+        The bundle is not installed
+    HRESULT_FROM_WIN32(ERROR_UNKNOWN_PROPERTY)
+        The variable is unrecognized
+    HRESULT_FROM_WIN32(ERROR_MORE_DATA)
+        A buffer is too small to hold the requested data.
+    E_NOTIMPL:
+        Tried to read a bundle variable for a type which has not been implemented
+
+    All other returns are unexpected returns from other dutil methods.
+********************************************************************/
+HRESULT DAPI BundleGetBundleVariableFixed(
+    __in_z LPCWSTR wzBundleId,
+    __in_z LPCWSTR wzVariable,
+    __out_ecount_opt(*pcchValue) LPWSTR wzValue,
+    __inout SIZE_T* pcchValue
+    );
+
 
 #ifdef __cplusplus
 }