1 files changed, 227 insertions, 4 deletions

diff --git a/src/libs/dutil/WixToolset.DUtil/apputil.cpp b/src/libs/dutil/WixToolset.DUtil/apputil.cpp
index 589a09dd..7e0bbc7b 100644
--- a/src/libs/dutil/WixToolset.DUtil/apputil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/apputil.cpp
@@ -20,7 +20,18 @@ const DWORD PRIVATE_LOAD_LIBRARY_SEARCH_SYSTEM32 = 0x00000800;
 typedef BOOL(WINAPI *LPFN_SETDEFAULTDLLDIRECTORIES)(DWORD);
 typedef BOOL(WINAPI *LPFN_SETDLLDIRECTORYW)(LPCWSTR);
 
-extern "C" void DAPI AppFreeCommandLineArgs(
+/********************************************************************
+EscapeCommandLineArgument - encodes wzArgument such that
+    ::CommandLineToArgv() will parse it back unaltered. If no escaping
+    was required, *psczEscaped is NULL.
+
+********************************************************************/
+static HRESULT EscapeCommandLineArgument(
+    __in_z LPCWSTR wzArgument,
+    __out_z LPWSTR* psczEscaped
+    );
+
+DAPI_(void) AppFreeCommandLineArgs(
     __in LPWSTR* argv
     )
 {
@@ -34,7 +45,7 @@ AppInitialize - initializes the standard safety precautions for an
                 installation application.
 
 ********************************************************************/
-extern "C" void DAPI AppInitialize(
+DAPI_(void) AppInitialize(
     __in_ecount(cSafelyLoadSystemDlls) LPCWSTR rgsczSafelyLoadSystemDlls[],
     __in DWORD cSafelyLoadSystemDlls
     )
@@ -85,12 +96,12 @@ extern "C" void DAPI AppInitialize(
     }
 }
 
-extern "C" void DAPI AppInitializeUnsafe()
+DAPI_(void) AppInitializeUnsafe()
 {
     ::HeapSetInformation(NULL, HeapEnableTerminationOnCorruption, NULL, 0);
 }
 
-extern "C" DAPI_(HRESULT) AppParseCommandLine(
+DAPI_(HRESULT) AppParseCommandLine(
     __in LPCWSTR wzCommandLine,
     __in int* pArgc,
     __in LPWSTR** pArgv
@@ -122,3 +133,215 @@ LExit:
 
     return hr;
 }
+
+DAPI_(HRESULT) AppAppendCommandLineArgument(
+    __deref_inout_z LPWSTR* psczCommandLine,
+    __in_z LPCWSTR wzArgument
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczQuotedArg = NULL;
+
+    hr = EscapeCommandLineArgument(wzArgument, &sczQuotedArg);
+    AppExitOnFailure(hr, "Failed to escape command line argument.");
+
+    // If there is already data in the command line,
+    // append a space before appending the argument.
+    if (*psczCommandLine && **psczCommandLine)
+    {
+        hr = StrAllocConcatSecure(psczCommandLine, L" ", 0);
+        AppExitOnFailure(hr, "Failed to append space to command line with existing data.");
+    }
+
+    hr = StrAllocConcatSecure(psczCommandLine, sczQuotedArg ? sczQuotedArg : wzArgument, 0);
+    AppExitOnFailure(hr, "Failed to copy command line argument.");
+
+LExit:
+    ReleaseStr(sczQuotedArg);
+
+    return hr;
+}
+
+DAPIV_(HRESULT) AppAppendCommandLineArgumentFormatted(
+    __deref_inout_z LPWSTR* psczCommandLine,
+    __in __format_string LPCWSTR wzFormat,
+    ...
+    )
+{
+    HRESULT hr = S_OK;
+    va_list args;
+
+    va_start(args, wzFormat);
+    hr = AppAppendCommandLineArgumentFormattedArgs(psczCommandLine, wzFormat, args);
+    va_end(args);
+
+    return hr;
+}
+
+DAPI_(HRESULT) AppAppendCommandLineArgumentFormattedArgs(
+    __deref_inout_z LPWSTR* psczCommandLine,
+    __in __format_string LPCWSTR wzFormat,
+    __in va_list args
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczQuotedArg = NULL;
+
+    hr = AppEscapeCommandLineArgumentFormattedArgs(&sczQuotedArg, wzFormat, args);
+    AppExitOnFailure(hr, "Failed to escape command line argument.");
+
+    // If there is already data in the command line,
+    // append a space before appending the argument.
+    if (*psczCommandLine && **psczCommandLine)
+    {
+        hr = StrAllocConcatSecure(psczCommandLine, L" ", 0);
+        AppExitOnFailure(hr, "Failed to append space to command line with existing data.");
+    }
+
+    hr = StrAllocConcatSecure(psczCommandLine, sczQuotedArg, 0);
+    AppExitOnFailure(hr, "Failed to copy command line argument.");
+
+LExit:
+    ReleaseStr(sczQuotedArg);
+
+    return hr;
+}
+
+DAPIV_(HRESULT) AppEscapeCommandLineArgumentFormatted(
+    __deref_inout_z LPWSTR* psczEscapedArgument,
+    __in __format_string LPCWSTR wzFormat,
+    ...
+    )
+{
+    HRESULT hr = S_OK;
+    va_list args;
+
+    va_start(args, wzFormat);
+    hr = AppEscapeCommandLineArgumentFormattedArgs(psczEscapedArgument, wzFormat, args);
+    va_end(args);
+
+    return hr;
+}
+
+DAPI_(HRESULT) AppEscapeCommandLineArgumentFormattedArgs(
+    __deref_inout_z LPWSTR* psczEscapedArgument,
+    __in __format_string LPCWSTR wzFormat,
+    __in va_list args
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczFormattedArg = NULL;
+    LPWSTR sczQuotedArg = NULL;
+
+    hr = StrAllocFormattedArgsSecure(&sczFormattedArg, wzFormat, args);
+    AppExitOnFailure(hr, "Failed to format command line argument.");
+
+    hr = EscapeCommandLineArgument(sczFormattedArg, &sczQuotedArg);
+    AppExitOnFailure(hr, "Failed to escape command line argument.");
+
+    if (sczQuotedArg)
+    {
+        *psczEscapedArgument = sczQuotedArg;
+        sczQuotedArg = NULL;
+    }
+    else
+    {
+        *psczEscapedArgument = sczFormattedArg;
+        sczFormattedArg = NULL;
+    }
+
+LExit:
+    ReleaseStr(sczFormattedArg);
+    ReleaseStr(sczQuotedArg);
+
+    return hr;
+}
+
+static HRESULT EscapeCommandLineArgument(
+    __in_z LPCWSTR wzArgument,
+    __out_z LPWSTR* psczEscaped
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fRequiresQuoting = FALSE;
+    SIZE_T cMaxEscapedSize = 0;
+
+    *psczEscaped = NULL;
+
+    // Loop through the argument determining if it needs to be quoted and what the maximum
+    // size would be if there are escape characters required.
+    for (LPCWSTR pwz = wzArgument; *pwz; ++pwz)
+    {
+        // Arguments with whitespace need quoting.
+        if (L' ' == *pwz || L'\t' == *pwz || L'\n' == *pwz || L'\v' == *pwz)
+        {
+            fRequiresQuoting = TRUE;
+        }
+        else if (L'"' == *pwz) // quotes need quoting and sometimes escaping.
+        {
+            fRequiresQuoting = TRUE;
+            ++cMaxEscapedSize;
+        }
+        else if (L'\\' == *pwz) // some backslashes need escaping, so we'll count them all to make sure there is room.
+        {
+            ++cMaxEscapedSize;
+        }
+
+        ++cMaxEscapedSize;
+    }
+
+    // If we found anything in the argument that requires our argument to be quoted
+    if (fRequiresQuoting)
+    {
+        hr = StrAlloc(psczEscaped, cMaxEscapedSize + 3); // plus three for the start and end quote plus null terminator.
+        AppExitOnFailure(hr, "Failed to allocate argument to be quoted.");
+
+        LPCWSTR pwz = wzArgument;
+        LPWSTR pwzQuoted = *psczEscaped;
+
+        *pwzQuoted = L'"';
+        ++pwzQuoted;
+        while (*pwz)
+        {
+            DWORD dwBackslashes = 0;
+            while (L'\\' == *pwz)
+            {
+                ++dwBackslashes;
+                ++pwz;
+            }
+
+            // Escape all backslashes at the end of the string.
+            if (!*pwz)
+            {
+                dwBackslashes *= 2;
+            }
+            else if (L'"' == *pwz) // escape all backslashes before the quote and escape the quote itself.
+            {
+                dwBackslashes = dwBackslashes * 2 + 1;
+            }
+            // the backslashes don't have to be escaped.
+
+            // Add the appropriate number of backslashes
+            for (DWORD i = 0; i < dwBackslashes; ++i)
+            {
+                *pwzQuoted = L'\\';
+                ++pwzQuoted;
+            }
+
+            // If there is a character, add it after all the escaped backslashes
+            if (*pwz)
+            {
+                *pwzQuoted = *pwz;
+                ++pwz;
+                ++pwzQuoted;
+            }
+        }
+
+        *pwzQuoted = L'"';
+        ++pwzQuoted;
+        *pwzQuoted = L'\0'; // ensure the arg is null terminated.
+    }
+
+LExit:
+    return hr;
+}