Add argument and policy setting to set Burn's base working directory.

Fixes #5856

5 files changed, 268 insertions, 117 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;
+}
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/apputil.h b/src/libs/dutil/WixToolset.DUtil/inc/apputil.h
index 1a1e14f7..11280102 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/apputil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/apputil.h
@@ -34,12 +34,51 @@ AppParseCommandLine - parses the command line using CommandLineToArgvW.
                       by calling AppFreeCommandLineArgs.
 
 ********************************************************************/
-DAPI_(HRESULT) AppParseCommandLine(
+HRESULT DAPI AppParseCommandLine(
     __in LPCWSTR wzCommandLine,
     __in int* argc,
     __in LPWSTR** pArgv
     );
 
+/*******************************************************************
+ AppAppendCommandLineArgument - appends a command line argument on to a
+    string such that ::CommandLineToArgv() will shred them correctly
+    (i.e. quote arguments with spaces in them).
+********************************************************************/
+HRESULT DAPI AppAppendCommandLineArgument(
+    __deref_inout_z LPWSTR* psczCommandLine,
+    __in_z LPCWSTR wzArgument
+    );
+
+HRESULT DAPIV AppAppendCommandLineArgumentFormatted(
+    __deref_inout_z LPWSTR* psczCommandLine,
+    __in __format_string LPCWSTR wzFormat,
+    ...
+    );
+
+HRESULT DAPI AppAppendCommandLineArgumentFormattedArgs(
+    __deref_inout_z LPWSTR* psczCommandLine,
+    __in __format_string LPCWSTR wzFormat,
+    __in va_list args
+    );
+
+/********************************************************************
+AppEscapeCommandLineArgumentFormatted - formats a string and then
+    escapes it such that ::CommandLineToArgv() will parse it back unaltered.
+
+********************************************************************/
+HRESULT DAPIV AppEscapeCommandLineArgumentFormatted(
+    __deref_inout_z LPWSTR* psczEscapedArgument,
+    __in __format_string LPCWSTR wzFormat,
+    ...
+    );
+
+HRESULT DAPI AppEscapeCommandLineArgumentFormattedArgs(
+    __deref_inout_z LPWSTR* psczEscapedArgument,
+    __in __format_string LPCWSTR wzFormat,
+    __in va_list args
+    );
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h b/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h
index 0ae9f437..00a468ce 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h
@@ -14,16 +14,6 @@ typedef enum PATH_EXPAND
 
 
 /*******************************************************************
- PathCommandLineAppend - appends a command line argument on to a
-    string such that ::CommandLineToArgv() will shred them correctly
-    (i.e. quote arguments with spaces in them).
-********************************************************************/
-DAPI_(HRESULT) PathCommandLineAppend(
-    __deref_inout_z LPWSTR* psczCommandLine,
-    __in_z LPCWSTR wzArgument
-    );
-
-/*******************************************************************
  PathFile -  returns a pointer to the file part of the path.
 ********************************************************************/
 DAPI_(LPWSTR) PathFile(
diff --git a/src/libs/dutil/WixToolset.DUtil/pathutil.cpp b/src/libs/dutil/WixToolset.DUtil/pathutil.cpp
index 5fad519b..7bac8ac3 100644
--- a/src/libs/dutil/WixToolset.DUtil/pathutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/pathutil.cpp
@@ -21,108 +21,6 @@
 #define PATH_GOOD_ENOUGH 64
 
 
-DAPI_(HRESULT) PathCommandLineAppend(
-    __deref_inout_z LPWSTR* psczCommandLine,
-    __in_z LPCWSTR wzArgument
-    )
-{
-    HRESULT hr = S_OK;
-    LPWSTR sczQuotedArg = NULL;
-    BOOL fRequiresQuoting = FALSE;
-    DWORD dwMaxEscapedSize = 0;
-
-    // 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;
-            ++dwMaxEscapedSize;
-        }
-        else if (L'\\' == *pwz) // some backslashes need escaping, so we'll count them all to make sure there is room.
-        {
-            ++dwMaxEscapedSize;
-        }
-
-        ++dwMaxEscapedSize;
-    }
-
-    // If we found anything in the argument that requires our argument to be quoted
-    if (fRequiresQuoting)
-    {
-        hr = StrAlloc(&sczQuotedArg, dwMaxEscapedSize + 3); // plus three for the start and end quote plus null terminator.
-        PathExitOnFailure(hr, "Failed to allocate argument to be quoted.");
-
-        LPCWSTR pwz = wzArgument; 
-        LPWSTR pwzQuoted = sczQuotedArg;
-
-        *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.
-    }
-
-    // If there is already data in the command line, append a space before appending the
-    // argument.
-    if (*psczCommandLine && **psczCommandLine)
-    {
-        hr = StrAllocConcat(psczCommandLine, L" ", 0);
-        PathExitOnFailure(hr, "Failed to append space to command line with existing data.");
-    }
-
-    hr = StrAllocConcat(psczCommandLine, sczQuotedArg ? sczQuotedArg : wzArgument, 0);
-    PathExitOnFailure(hr, "Failed to copy command line argument.");
-
-LExit:
-    ReleaseStr(sczQuotedArg);
-
-    return hr;
-}
-
-
 DAPI_(LPWSTR) PathFile(
     __in_z LPCWSTR wzPath
     )
diff --git a/src/libs/dutil/WixToolset.DUtil/precomp.h b/src/libs/dutil/WixToolset.DUtil/precomp.h
index f8f3b944..46d29f21 100644
--- a/src/libs/dutil/WixToolset.DUtil/precomp.h
+++ b/src/libs/dutil/WixToolset.DUtil/precomp.h
@@ -45,6 +45,7 @@
 #include "dutil.h"
 #include "verutil.h"
 #include "aclutil.h"
+#include "apputil.h"
 #include "atomutil.h"
 #include "buffutil.h"
 #include "butil.h"