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

7 files changed, 171 insertions, 193 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);