Use balutil methods in Engine.cs to avoid size_t ugliness.

11 files changed, 466 insertions, 201 deletions

diff --git a/src/api/burn/WixToolset.Mba.Core/BalUtil.cs b/src/api/burn/WixToolset.Mba.Core/BalUtil.cs
index f478eca4..4c949a5a 100644
--- a/src/api/burn/WixToolset.Mba.Core/BalUtil.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BalUtil.cs
@@ -7,16 +7,39 @@ namespace WixToolset.Mba.Core
 
     internal static class BalUtil
     {
-        [DllImport("mbanative.dll", ExactSpelling = true, PreserveSig = false)]
-        internal static extern IBootstrapperEngine InitializeFromCreateArgs(
-            IntPtr pArgs,
-            ref Command pCommand
+        [DllImport("mbanative.dll", ExactSpelling = true)]
+        internal static extern int BalEscapeStringFromEngine(
+            [MarshalAs(UnmanagedType.Interface)] IBootstrapperEngine pEngine,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzIn,
+            ref StrUtil.StrHandle psczOut
+            );
+
+        [DllImport("mbanative.dll", ExactSpelling = true)]
+        internal static extern int BalFormatStringFromEngine(
+            [MarshalAs(UnmanagedType.Interface)] IBootstrapperEngine pEngine,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzFormat,
+            ref StrUtil.StrHandle psczOut
+            );
+
+        [DllImport("mbanative.dll", ExactSpelling = true)]
+        internal static extern int BalGetStringVariableFromEngine(
+            [MarshalAs(UnmanagedType.Interface)] IBootstrapperEngine pEngine,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzVariable,
+            ref StrUtil.StrHandle psczOut
+            );
+
+        [DllImport("mbanative.dll", ExactSpelling = true)]
+        internal static extern int BalGetVersionVariableFromEngine(
+            [MarshalAs(UnmanagedType.Interface)] IBootstrapperEngine pEngine,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzVariable,
+            ref StrUtil.StrHandle psczOut
             );
 
         [DllImport("mbanative.dll", ExactSpelling = true)]
-        internal static extern void StoreBAInCreateResults(
-            IntPtr pResults,
-            [MarshalAs(UnmanagedType.Interface)] IBootstrapperApplication pBA
+        [return: MarshalAs(UnmanagedType.Bool)]
+        internal static extern bool BalVariableExistsFromEngine(
+            [MarshalAs(UnmanagedType.Interface)] IBootstrapperEngine pEngine,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzVariable
             );
     }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/BaseBootstrapperApplicationFactory.cs b/src/api/burn/WixToolset.Mba.Core/BaseBootstrapperApplicationFactory.cs
index ad8a5dc0..b215681e 100644
--- a/src/api/burn/WixToolset.Mba.Core/BaseBootstrapperApplicationFactory.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BaseBootstrapperApplicationFactory.cs
@@ -44,7 +44,7 @@ namespace WixToolset.Mba.Core
             {
                 cbSize = Marshal.SizeOf(typeof(Command))
             };
-            var pEngine = BalUtil.InitializeFromCreateArgs(pArgs, ref pCommand);
+            var pEngine = MbaNative.InitializeFromCreateArgs(pArgs, ref pCommand);
             engine = new Engine(pEngine);
             bootstrapperCommand = pCommand.GetBootstrapperCommand();
         }
@@ -57,7 +57,7 @@ namespace WixToolset.Mba.Core
         /// <param name="ba">The <see cref="IBootstrapperApplication"/>.</param>
         public static void StoreBAInCreateResults(IntPtr pResults, IBootstrapperApplication ba)
         {
-            BalUtil.StoreBAInCreateResults(pResults, ba);
+            MbaNative.StoreBAInCreateResults(pResults, ba);
         }
     }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/Engine.cs b/src/api/burn/WixToolset.Mba.Core/Engine.cs
index aed420d5..5ebada36 100644
--- a/src/api/burn/WixToolset.Mba.Core/Engine.cs
+++ b/src/api/burn/WixToolset.Mba.Core/Engine.cs
@@ -13,8 +13,6 @@ namespace WixToolset.Mba.Core
     /// </summary>
     public sealed class Engine : IEngine
     {
-        // Burn errs on empty strings, so declare initial buffer size.
-        private const int InitialBufferSize = 80;
         private static readonly string normalizeVersionFormatString = "{0} must be less than or equal to " + UInt16.MaxValue;
 
         private IBootstrapperEngine engine;
@@ -62,9 +60,7 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public bool ContainsVariable(string name)
         {
-            IntPtr capacity = new IntPtr(0);
-            int ret = this.engine.GetVariableString(name, IntPtr.Zero, ref capacity);
-            return NativeMethods.E_NOTFOUND != ret;
+            return BalUtil.BalVariableExistsFromEngine(this.engine, name);
         }
 
         /// <inheritdoc/>
@@ -101,24 +97,21 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public string EscapeString(string input)
         {
-            IntPtr capacity = new IntPtr(InitialBufferSize);
-            StringBuilder sb = new StringBuilder(capacity.ToInt32());
-
-            // Get the size of the buffer.
-            int ret = this.engine.EscapeString(input, sb, ref capacity);
-            if (NativeMethods.E_INSUFFICIENT_BUFFER == ret || NativeMethods.E_MOREDATA == ret)
+            StrUtil.StrHandle handle = new StrUtil.StrHandle();
+            try
             {
-                capacity = new IntPtr(capacity.ToInt32() + 1); // Add one for the null terminator.
-                sb.Capacity = capacity.ToInt32();
-                ret = this.engine.EscapeString(input, sb, ref capacity);
-            }
+                int ret = BalUtil.BalEscapeStringFromEngine(this.engine, input, ref handle);
+                if (ret != NativeMethods.S_OK)
+                {
+                    throw new Win32Exception(ret);
+                }
 
-            if (NativeMethods.S_OK != ret)
+                return handle.ToUniString();
+            }
+            finally
             {
-                throw new Win32Exception(ret);
+                handle.Dispose();
             }
-
-            return sb.ToString();
         }
 
         /// <inheritdoc/>
@@ -133,24 +126,21 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public string FormatString(string format)
         {
-            IntPtr capacity = new IntPtr(InitialBufferSize);
-            StringBuilder sb = new StringBuilder(capacity.ToInt32());
-
-            // Get the size of the buffer.
-            int ret = this.engine.FormatString(format, sb, ref capacity);
-            if (NativeMethods.E_INSUFFICIENT_BUFFER == ret || NativeMethods.E_MOREDATA == ret)
+            StrUtil.StrHandle handle = new StrUtil.StrHandle();
+            try
             {
-                capacity = new IntPtr(capacity.ToInt32() + 1); // Add one for the null terminator.
-                sb.Capacity = capacity.ToInt32();
-                ret = this.engine.FormatString(format, sb, ref capacity);
-            }
+                int ret = BalUtil.BalFormatStringFromEngine(this.engine, format, ref handle);
+                if (ret != NativeMethods.S_OK)
+                {
+                    throw new Win32Exception(ret);
+                }
 
-            if (NativeMethods.S_OK != ret)
+                return handle.ToUniString();
+            }
+            finally
             {
-                throw new Win32Exception(ret);
+                handle.Dispose();
             }
-
-            return sb.ToString();
         }
 
         /// <inheritdoc/>
@@ -168,53 +158,60 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public SecureString GetVariableSecureString(string name)
         {
-            var pUniString = this.getStringVariable(name, out var length);
+            StrUtil.StrHandle handle = new StrUtil.StrHandle();
             try
             {
-                return this.convertToSecureString(pUniString, length);
+                int ret = BalUtil.BalGetStringVariableFromEngine(this.engine, name, ref handle);
+                if (ret != NativeMethods.S_OK)
+                {
+                    throw new Win32Exception(ret);
+                }
+
+                return handle.ToSecureString();
             }
             finally
             {
-                if (IntPtr.Zero != pUniString)
-                {
-                    Marshal.FreeCoTaskMem(pUniString);
-                }
+                handle.Dispose();
             }
         }
 
         /// <inheritdoc/>
         public string GetVariableString(string name)
         {
-            int length;
-            IntPtr pUniString = this.getStringVariable(name, out length);
+            StrUtil.StrHandle handle = new StrUtil.StrHandle();
             try
             {
-                return Marshal.PtrToStringUni(pUniString, length);
+                int ret = BalUtil.BalGetStringVariableFromEngine(this.engine, name, ref handle);
+                if (ret != NativeMethods.S_OK)
+                {
+                    throw new Win32Exception(ret);
+                }
+
+                return handle.ToUniString();
             }
             finally
             {
-                if (IntPtr.Zero != pUniString)
-                {
-                    Marshal.FreeCoTaskMem(pUniString);
-                }
+                handle.Dispose();
             }
         }
 
         /// <inheritdoc/>
         public string GetVariableVersion(string name)
         {
-            int length;
-            IntPtr pUniString = this.getVersionVariable(name, out length);
+            StrUtil.StrHandle handle = new StrUtil.StrHandle();
             try
             {
-                return Marshal.PtrToStringUni(pUniString, length);
+                int ret = BalUtil.BalGetVersionVariableFromEngine(this.engine, name, ref handle);
+                if (ret != NativeMethods.S_OK)
+                {
+                    throw new Win32Exception(ret);
+                }
+
+                return handle.ToUniString();
             }
             finally
             {
-                if (IntPtr.Zero != pUniString)
-                {
-                    Marshal.FreeCoTaskMem(pUniString);
-                }
+                handle.Dispose();
             }
         }
 
@@ -337,132 +334,6 @@ namespace WixToolset.Mba.Core
         }
 
         /// <summary>
-        /// Gets the variable given by <paramref name="name"/> as a string.
-        /// </summary>
-        /// <param name="name">The name of the variable to get.</param>
-        /// <param name="length">The length of the Unicode string.</param>
-        /// <returns>The value by a pointer to a Unicode string.  Must be freed by Marshal.FreeCoTaskMem.</returns>
-        /// <exception cref="Exception">An error occurred getting the variable.</exception>
-        internal IntPtr getStringVariable(string name, out int length)
-        {
-            IntPtr capacity = new IntPtr(InitialBufferSize);
-            bool success = false;
-            IntPtr pValue = Marshal.AllocCoTaskMem(capacity.ToInt32() * UnicodeEncoding.CharSize);
-            try
-            {
-                // Get the size of the buffer.
-                int ret = this.engine.GetVariableString(name, pValue, ref capacity);
-                if (NativeMethods.E_INSUFFICIENT_BUFFER == ret || NativeMethods.E_MOREDATA == ret)
-                {
-                    // Don't need to add 1 for the null terminator, the engine already includes that.
-                    pValue = Marshal.ReAllocCoTaskMem(pValue, capacity.ToInt32() * UnicodeEncoding.CharSize);
-                    ret = this.engine.GetVariableString(name, pValue, ref capacity);
-                }
-
-                if (NativeMethods.S_OK != ret)
-                {
-                    throw Marshal.GetExceptionForHR(ret);
-                }
-
-                // The engine only returns the exact length of the string if the buffer was too small, so calculate it ourselves.
-                int maxLength = capacity.ToInt32();
-                for (length = 0; length < maxLength; ++length)
-                {
-                    if (0 == Marshal.ReadInt16(pValue, length * UnicodeEncoding.CharSize))
-                    {
-                        break;
-                    }
-                }
-
-                success = true;
-                return pValue;
-            }
-            finally
-            {
-                if (!success && IntPtr.Zero != pValue)
-                {
-                    Marshal.FreeCoTaskMem(pValue);
-                }
-            }
-        }
-
-        /// <summary>
-        /// Gets the variable given by <paramref name="name"/> as a version string.
-        /// </summary>
-        /// <param name="name">The name of the variable to get.</param>
-        /// <param name="length">The length of the Unicode string.</param>
-        /// <returns>The value by a pointer to a Unicode string.  Must be freed by Marshal.FreeCoTaskMem.</returns>
-        /// <exception cref="Exception">An error occurred getting the variable.</exception>
-        internal IntPtr getVersionVariable(string name, out int length)
-        {
-            IntPtr capacity = new IntPtr(InitialBufferSize);
-            bool success = false;
-            IntPtr pValue = Marshal.AllocCoTaskMem(capacity.ToInt32() * UnicodeEncoding.CharSize);
-            try
-            {
-                // Get the size of the buffer.
-                int ret = this.engine.GetVariableVersion(name, pValue, ref capacity);
-                if (NativeMethods.E_INSUFFICIENT_BUFFER == ret || NativeMethods.E_MOREDATA == ret)
-                {
-                    // Don't need to add 1 for the null terminator, the engine already includes that.
-                    pValue = Marshal.ReAllocCoTaskMem(pValue, capacity.ToInt32() * UnicodeEncoding.CharSize);
-                    ret = this.engine.GetVariableVersion(name, pValue, ref capacity);
-                }
-
-                if (NativeMethods.S_OK != ret)
-                {
-                    throw Marshal.GetExceptionForHR(ret);
-                }
-
-                // The engine only returns the exact length of the string if the buffer was too small, so calculate it ourselves.
-                int maxLength = capacity.ToInt32();
-                for (length = 0; length < maxLength; ++length)
-                {
-                    if (0 == Marshal.ReadInt16(pValue, length * UnicodeEncoding.CharSize))
-                    {
-                        break;
-                    }
-                }
-
-                success = true;
-                return pValue;
-            }
-            finally
-            {
-                if (!success && IntPtr.Zero != pValue)
-                {
-                    Marshal.FreeCoTaskMem(pValue);
-                }
-            }
-        }
-
-        /// <summary>
-        /// Initialize a SecureString with the given Unicode string.
-        /// </summary>
-        /// <param name="pUniString">Pointer to Unicode string.</param>
-        /// <param name="length">The string's length.</param>
-        internal SecureString convertToSecureString(IntPtr pUniString, int length)
-        {
-            if (IntPtr.Zero == pUniString)
-            {
-                return null;
-            }
-
-            SecureString value = new SecureString();
-            short s;
-            char c;
-            for (int charIndex = 0; charIndex < length; charIndex++)
-            {
-                s = Marshal.ReadInt16(pUniString, charIndex * UnicodeEncoding.CharSize);
-                c = (char)s;
-                value.AppendChar(c);
-                s = 0;
-                c = (char)0;
-            }
-            return value;
-        }
-
-        /// <summary>
         /// Utility method for converting a <see cref="Version"/> into a <see cref="long"/>.
         /// </summary>
         /// <param name="version"></param>
diff --git a/src/api/burn/WixToolset.Mba.Core/MbaNative.cs b/src/api/burn/WixToolset.Mba.Core/MbaNative.cs
new file mode 100644
index 00000000..a68a3907
--- /dev/null
+++ b/src/api/burn/WixToolset.Mba.Core/MbaNative.cs
@@ -0,0 +1,22 @@
+// 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.
+
+namespace WixToolset.Mba.Core
+{
+    using System;
+    using System.Runtime.InteropServices;
+
+    internal static class MbaNative
+    {
+        [DllImport("mbanative.dll", ExactSpelling = true, PreserveSig = false)]
+        internal static extern IBootstrapperEngine InitializeFromCreateArgs(
+            IntPtr pArgs,
+            ref Command pCommand
+            );
+
+        [DllImport("mbanative.dll", ExactSpelling = true)]
+        internal static extern void StoreBAInCreateResults(
+            IntPtr pResults,
+            [MarshalAs(UnmanagedType.Interface)] IBootstrapperApplication pBA
+            );
+    }
+}
diff --git a/src/api/burn/WixToolset.Mba.Core/NativeMethods.cs b/src/api/burn/WixToolset.Mba.Core/NativeMethods.cs
index adb2256e..45a0bc4d 100644
--- a/src/api/burn/WixToolset.Mba.Core/NativeMethods.cs
+++ b/src/api/burn/WixToolset.Mba.Core/NativeMethods.cs
@@ -34,4 +34,16 @@ namespace WixToolset.Mba.Core
             );
         #endregion
     }
+
+    #region SafeHandles
+    internal abstract class SafeHandleZeroIsDefaultAndInvalid : SafeHandle
+    {
+        public SafeHandleZeroIsDefaultAndInvalid() : base(IntPtr.Zero, true) { }
+
+        public override bool IsInvalid
+        {
+            get { return this.handle == IntPtr.Zero; }
+        }
+    }
+    #endregion
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/StrUtil.cs b/src/api/burn/WixToolset.Mba.Core/StrUtil.cs
new file mode 100644
index 00000000..5daef8e3
--- /dev/null
+++ b/src/api/burn/WixToolset.Mba.Core/StrUtil.cs
@@ -0,0 +1,54 @@
+// 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.
+
+namespace WixToolset.Mba.Core
+{
+    using System;
+    using System.Runtime.InteropServices;
+    using System.Security;
+    using System.Text;
+
+    internal static class StrUtil
+    {
+        [DllImport("mbanative.dll", ExactSpelling = true)]
+        internal static extern void StrFree(
+            IntPtr scz
+            );
+
+        internal sealed class StrHandle : SafeHandleZeroIsDefaultAndInvalid
+        {
+            protected override bool ReleaseHandle()
+            {
+                StrFree(this.handle);
+                return true;
+            }
+
+            public string ToUniString()
+            {
+                return Marshal.PtrToStringUni(this.handle);
+            }
+
+            public SecureString ToSecureString()
+            {
+                if (this.handle == IntPtr.Zero)
+                {
+                    return null;
+                }
+
+                SecureString value = new SecureString();
+                char c;
+                for (int charIndex = 0; ; charIndex++)
+                {
+                    c = (char)Marshal.ReadInt16(this.handle, charIndex * UnicodeEncoding.CharSize);
+                    if (c == '\0')
+                    {
+                        break;
+                    }
+
+                    value.AppendChar(c);
+                }
+
+                return value;
+            }
+        }
+    }
+}
diff --git a/src/api/burn/WixToolset.Mba.Core/VerUtil.cs b/src/api/burn/WixToolset.Mba.Core/VerUtil.cs
index 81c5b716..a342f85c 100644
--- a/src/api/burn/WixToolset.Mba.Core/VerUtil.cs
+++ b/src/api/burn/WixToolset.Mba.Core/VerUtil.cs
@@ -83,12 +83,8 @@ namespace WixToolset.Mba.Core
             }
         }
 
-        internal sealed class VersionHandle : SafeHandle
+        internal sealed class VersionHandle : SafeHandleZeroIsDefaultAndInvalid
         {
-            public VersionHandle() : base(IntPtr.Zero, true) { }
-
-            public override bool IsInvalid => false;
-
             protected override bool ReleaseHandle()
             {
                 VerFreeVersion(this.handle);
diff --git a/src/api/burn/balutil/balutil.cpp b/src/api/burn/balutil/balutil.cpp
index 7a638219..5671da4e 100644
--- a/src/api/burn/balutil/balutil.cpp
+++ b/src/api/burn/balutil/balutil.cpp
@@ -89,6 +89,66 @@ LExit:
 }
 
 
+DAPI_(HRESULT) BalEscapeString(
+    __in_z LPCWSTR wzIn,
+    __inout LPWSTR* psczOut
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (!vpEngine)
+    {
+        hr = E_POINTER;
+        ExitOnRootFailure(hr, "BalInitialize() must be called first.");
+    }
+
+    hr = BalEscapeStringFromEngine(vpEngine, wzIn, psczOut);
+
+LExit:
+    return hr;
+}
+
+
+DAPI_(HRESULT) BalEscapeStringFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzIn,
+    __inout LPWSTR* psczOut
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T cch = 0;
+
+    if (*psczOut)
+    {
+        hr = StrMaxLength(*psczOut, &cch);
+        ExitOnFailure(hr, "Failed to determine length of value.");
+    }
+    else
+    {
+        hr = ::StringCchLengthW(wzIn, STRSAFE_MAX_LENGTH, reinterpret_cast<size_t*>(&cch));
+        ExitOnFailure(hr, "Failed to determine length of source.");
+
+        cch = min(STRSAFE_MAX_LENGTH, cch + VARIABLE_GROW_FACTOR);
+        hr = StrAlloc(psczOut, cch);
+        ExitOnFailure(hr, "Failed to pre-allocate value.");
+    }
+
+    hr = pEngine->EscapeString(wzIn, *psczOut, &cch);
+    if (E_MOREDATA == hr)
+    {
+        ++cch;
+
+        hr = StrAllocSecure(psczOut, cch);
+        ExitOnFailure(hr, "Failed to allocate value.");
+
+        hr = pEngine->EscapeString(wzIn, *psczOut, &cch);
+    }
+
+LExit:
+    return hr;
+}
+
+
 // The contents of psczOut may be sensitive, should keep encrypted and SecureZeroFree.
 DAPI_(HRESULT) BalFormatString(
     __in_z LPCWSTR wzFormat,
@@ -96,7 +156,6 @@ DAPI_(HRESULT) BalFormatString(
     )
 {
     HRESULT hr = S_OK;
-    SIZE_T cch = 0;
 
     if (!vpEngine)
     {
@@ -104,13 +163,39 @@ DAPI_(HRESULT) BalFormatString(
         ExitOnRootFailure(hr, "BalInitialize() must be called first.");
     }
 
+    hr = BalFormatStringFromEngine(vpEngine, wzFormat, psczOut);
+
+LExit:
+    return hr;
+}
+
+
+// The contents of psczOut may be sensitive, should keep encrypted and SecureZeroFree.
+DAPI_(HRESULT) BalFormatStringFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzFormat,
+    __inout LPWSTR* psczOut
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T cch = 0;
+
     if (*psczOut)
     {
         hr = StrMaxLength(*psczOut, &cch);
         ExitOnFailure(hr, "Failed to determine length of value.");
     }
+    else
+    {
+        hr = ::StringCchLengthW(wzFormat, STRSAFE_MAX_LENGTH, reinterpret_cast<size_t*>(&cch));
+        ExitOnFailure(hr, "Failed to determine length of source.");
+
+        cch = min(STRSAFE_MAX_LENGTH, cch + VARIABLE_GROW_FACTOR);
+        hr = StrAlloc(psczOut, cch);
+        ExitOnFailure(hr, "Failed to pre-allocate value.");
+    }
 
-    hr = vpEngine->FormatString(wzFormat, *psczOut, &cch);
+    hr = pEngine->FormatString(wzFormat, *psczOut, &cch);
     if (E_MOREDATA == hr)
     {
         ++cch;
@@ -118,7 +203,7 @@ DAPI_(HRESULT) BalFormatString(
         hr = StrAllocSecure(psczOut, cch);
         ExitOnFailure(hr, "Failed to allocate value.");
 
-        hr = vpEngine->FormatString(wzFormat, *psczOut, &cch);
+        hr = pEngine->FormatString(wzFormat, *psczOut, &cch);
     }
 
 LExit:
@@ -172,7 +257,7 @@ DAPI_(BOOL) BalVariableExists(
     )
 {
     HRESULT hr = S_OK;
-    SIZE_T cch = 0;
+    BOOL fExists = FALSE;
 
     if (!vpEngine)
     {
@@ -180,9 +265,23 @@ DAPI_(BOOL) BalVariableExists(
         ExitOnRootFailure(hr, "BalInitialize() must be called first.");
     }
 
-    hr = vpEngine->GetVariableString(wzVariable, NULL, &cch);
+    fExists = BalVariableExistsFromEngine(vpEngine, wzVariable);
 
 LExit:
+    return fExists;
+}
+
+
+DAPI_(BOOL) BalVariableExistsFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzVariable
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T cch = 0;
+
+    hr = pEngine->GetVariableString(wzVariable, NULL, &cch);
+
     return E_NOTFOUND != hr;
 }
 
@@ -194,7 +293,6 @@ DAPI_(HRESULT) BalGetStringVariable(
     )
 {
     HRESULT hr = S_OK;
-    SIZE_T cch = 0;
 
     if (!vpEngine)
     {
@@ -202,13 +300,36 @@ DAPI_(HRESULT) BalGetStringVariable(
         ExitOnRootFailure(hr, "BalInitialize() must be called first.");
     }
 
+    hr = BalGetStringVariableFromEngine(vpEngine, wzVariable, psczValue);
+
+LExit:
+    return hr;
+}
+
+
+// The contents of psczValue may be sensitive, if variable is hidden should keep value encrypted and SecureZeroFree.
+DAPI_(HRESULT) BalGetStringVariableFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzVariable,
+    __inout LPWSTR* psczValue
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T cch = 0;
+
     if (*psczValue)
     {
         hr = StrMaxLength(*psczValue, &cch);
         ExitOnFailure(hr, "Failed to determine length of value.");
     }
+    else
+    {
+        cch = VARIABLE_GROW_FACTOR;
+        hr = StrAlloc(psczValue, cch);
+        ExitOnFailure(hr, "Failed to pre-allocate value.");
+    }
 
-    hr = vpEngine->GetVariableString(wzVariable, *psczValue, &cch);
+    hr = pEngine->GetVariableString(wzVariable, *psczValue, &cch);
     if (E_MOREDATA == hr)
     {
         ++cch;
@@ -216,7 +337,7 @@ DAPI_(HRESULT) BalGetStringVariable(
         hr = StrAllocSecure(psczValue, cch);
         ExitOnFailure(hr, "Failed to allocate value.");
 
-        hr = vpEngine->GetVariableString(wzVariable, *psczValue, &cch);
+        hr = pEngine->GetVariableString(wzVariable, *psczValue, &cch);
     }
 
 LExit:
@@ -244,6 +365,82 @@ LExit:
 }
 
 
+DAPI_(HRESULT) BalGetVersionVariable(
+    __in_z LPCWSTR wzVariable,
+    __inout LPWSTR* psczValue
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (!vpEngine)
+    {
+        hr = E_POINTER;
+        ExitOnRootFailure(hr, "BalInitialize() must be called first.");
+    }
+
+    hr = BalGetVersionVariableFromEngine(vpEngine, wzVariable, psczValue);
+
+LExit:
+    return hr;
+}
+
+
+DAPI_(HRESULT) BalGetVersionVariableFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzVariable,
+    __inout LPWSTR* psczValue
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T cch = 0;
+
+    if (*psczValue)
+    {
+        hr = StrMaxLength(*psczValue, &cch);
+        ExitOnFailure(hr, "Failed to determine length of value.");
+    }
+    else
+    {
+        cch = VARIABLE_GROW_FACTOR;
+        hr = StrAlloc(psczValue, cch);
+        ExitOnFailure(hr, "Failed to pre-allocate value.");
+    }
+
+    hr = pEngine->GetVariableVersion(wzVariable, *psczValue, &cch);
+    if (E_MOREDATA == hr)
+    {
+        ++cch;
+
+        hr = StrAllocSecure(psczValue, cch);
+        ExitOnFailure(hr, "Failed to allocate value.");
+
+        hr = pEngine->GetVariableVersion(wzVariable, *psczValue, &cch);
+    }
+
+LExit:
+    return hr;
+}
+
+DAPI_(HRESULT) BalSetVersionVariable(
+    __in_z LPCWSTR wzVariable,
+    __in_z_opt LPCWSTR wzValue
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (!vpEngine)
+    {
+        hr = E_POINTER;
+        ExitOnRootFailure(hr, "BalInitialize() must be called first.");
+    }
+
+    hr = vpEngine->SetVariableVersion(wzVariable, wzValue);
+
+LExit:
+    return hr;
+}
+
+
 DAPIV_(HRESULT) BalLog(
     __in BOOTSTRAPPER_LOG_LEVEL level,
     __in_z __format_string LPCSTR szFormat,
diff --git a/src/api/burn/balutil/inc/balutil.h b/src/api/burn/balutil/inc/balutil.h
index fad8a471..82fd1fe5 100644
--- a/src/api/burn/balutil/inc/balutil.h
+++ b/src/api/burn/balutil/inc/balutil.h
@@ -79,6 +79,27 @@ DAPI_(HRESULT) BalEvaluateCondition(
     );
 
 /*******************************************************************
+BalEscapeString - escapes a string to use as part of a formatted string variable.
+
+ Note: Use StrFree() to release psczOut.
+********************************************************************/
+DAPI_(HRESULT) BalEscapeString(
+    __in_z LPCWSTR wzIn,
+    __inout LPWSTR* psczOut
+    );
+
+/*******************************************************************
+BalEscapeStringFromEngine - escapes a string to use as part of a formatted string variable.
+
+ Note: Use StrFree() to release psczOut.
+********************************************************************/
+DAPI_(HRESULT) BalEscapeStringFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzIn,
+    __inout LPWSTR* psczOut
+    );
+
+/*******************************************************************
 BalFormatString - formats a string using variables in the engine.
 
  Note: Use StrFree() to release psczOut.
@@ -89,6 +110,17 @@ DAPI_(HRESULT) BalFormatString(
     );
 
 /*******************************************************************
+BalFormatStringFromEngine - formats a string using variables in the engine.
+
+ Note: Use StrFree() to release psczOut.
+********************************************************************/
+DAPI_(HRESULT) BalFormatStringFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzFormat,
+    __inout LPWSTR* psczOut
+    );
+
+/*******************************************************************
 BalGetNumericVariable - gets a number from a variable in the engine.
 
  Note: Returns E_NOTFOUND if variable does not exist.
@@ -116,6 +148,15 @@ DAPI_(BOOL) BalVariableExists(
     );
 
 /*******************************************************************
+BalVariableExistsFromEngine - checks if a variable exists in the engine.
+
+********************************************************************/
+DAPI_(BOOL) BalVariableExistsFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzVariable
+    );
+
+/*******************************************************************
 BalGetStringVariable - gets a string from a variable in the engine.
 
  Note: Use StrFree() to release psczValue.
@@ -126,7 +167,19 @@ DAPI_(HRESULT) BalGetStringVariable(
     );
 
 /*******************************************************************
+BalGetStringVariableFromEngine - gets a string from a variable in the engine.
+
+ Note: Use StrFree() to release psczValue.
+********************************************************************/
+DAPI_(HRESULT) BalGetStringVariableFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzVariable,
+    __inout LPWSTR* psczValue
+    );
+
+/*******************************************************************
 BalSetStringVariable - sets a string variable in the engine.
+    If the value contains unexpanded variables, set fFormatted to true.
 
 ********************************************************************/
 DAPI_(HRESULT) BalSetStringVariable(
@@ -136,6 +189,36 @@ DAPI_(HRESULT) BalSetStringVariable(
     );
 
 /*******************************************************************
+BalGetVersionVariable - gets a version from a variable in the engine.
+
+ Note: Use StrFree() to release psczValue.
+********************************************************************/
+DAPI_(HRESULT) BalGetVersionVariable(
+    __in_z LPCWSTR wzVariable,
+    __inout LPWSTR* psczValue
+    );
+
+/*******************************************************************
+BalGetVersionVariableFromEngine - gets a version from a variable in the engine.
+
+ Note: Use StrFree() to release psczValue.
+********************************************************************/
+DAPI_(HRESULT) BalGetVersionVariableFromEngine(
+    __in IBootstrapperEngine* pEngine,
+    __in_z LPCWSTR wzVariable,
+    __inout LPWSTR* psczValue
+    );
+
+/*******************************************************************
+BalSetVersionVariable - sets a version variable in the engine.
+
+********************************************************************/
+DAPI_(HRESULT) BalSetVersionVariable(
+    __in_z LPCWSTR wzVariable,
+    __in_z_opt LPCWSTR wzValue
+    );
+
+/*******************************************************************
  BalLog - logs a message with the engine.
 
 ********************************************************************/
diff --git a/src/api/burn/balutil/precomp.h b/src/api/burn/balutil/precomp.h
index c500060a..207c8ca6 100644
--- a/src/api/burn/balutil/precomp.h
+++ b/src/api/burn/balutil/precomp.h
@@ -7,6 +7,7 @@
 #include <msi.h>
 #include <wininet.h>
 #include <CommCtrl.h>
+#include <strsafe.h>
 
 #include <dutil.h>
 #include <pathutil.h>
diff --git a/src/api/burn/mbanative/mbanative.def b/src/api/burn/mbanative/mbanative.def
index 28e923b6..4bad14d0 100644
--- a/src/api/burn/mbanative/mbanative.def
+++ b/src/api/burn/mbanative/mbanative.def
@@ -2,8 +2,14 @@
 
 
 EXPORTS
+    BalEscapeStringFromEngine
+    BalFormatStringFromEngine
+    BalGetStringVariableFromEngine
+    BalGetVersionVariableFromEngine
+    BalVariableExistsFromEngine
     InitializeFromCreateArgs
     StoreBAInCreateResults
+    StrFree
     VerCompareParsedVersions
     VerCompareStringVersions
     VerCopyVersion