11 files changed, 248 insertions, 33 deletions

diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index 716e5af1..90bbc8f6 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -12,6 +12,7 @@ struct BURN_CACHE_THREAD_CONTEXT
 };
 
 
+static PFN_CREATEPROCESSW vpfnCreateProcessW = ::CreateProcessW;
 static PFN_PROCWAITFORCOMPLETION vpfnProcWaitForCompletion = ProcWaitForCompletion;
 
 
@@ -1947,12 +1948,46 @@ LExit:
 }
 
 extern "C" void CoreFunctionOverride(
+    __in_opt PFN_CREATEPROCESSW pfnCreateProcessW,
     __in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
     )
 {
+    vpfnCreateProcessW = pfnCreateProcessW;
     vpfnProcWaitForCompletion = pfnProcWaitForCompletion;
 }
 
+extern "C" HRESULT CoreCreateProcess(
+    __in_opt LPCWSTR wzApplicationName,
+    __inout_opt LPWSTR sczCommandLine,
+    __in BOOL fInheritHandles,
+    __in DWORD dwCreationFlags,
+    __in_opt LPCWSTR wzCurrentDirectory,
+    __in WORD wShowWindow,
+    __out LPPROCESS_INFORMATION pProcessInformation
+    )
+{
+    HRESULT hr = S_OK;
+    STARTUPINFOW si = { };
+    size_t cchCurrentDirectory = 0;
+
+    // CreateProcessW has undocumented MAX_PATH restriction for lpCurrentDirectory even when long path support is enabled.
+    if (wzCurrentDirectory && FAILED(::StringCchLengthW(wzCurrentDirectory, MAX_PATH - 1, &cchCurrentDirectory)))
+    {
+        wzCurrentDirectory = NULL;
+    }
+
+    si.cb = sizeof(si);
+    si.wShowWindow = wShowWindow;
+
+    if (!vpfnCreateProcessW(wzApplicationName, sczCommandLine, NULL, NULL, fInheritHandles, dwCreationFlags, NULL, wzCurrentDirectory, &si, pProcessInformation))
+    {
+        ExitWithLastError(hr, "CreateProcessW failed with return code: %d", Dutil_er);
+    }
+
+LExit:
+    return hr;
+}
+
 extern "C" HRESULT DAPI CoreWaitForProcCompletion(
     __in HANDLE hProcess,
     __in DWORD dwTimeout,
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 14dbabcc..31f63ed4 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -174,6 +174,19 @@ typedef struct _BURN_APPLY_CONTEXT
     DWORD dwCacheCheckpoint;
 } BURN_APPLY_CONTEXT;
 
+typedef BOOL (STDAPICALLTYPE *PFN_CREATEPROCESSW)(
+    __in_opt LPCWSTR lpApplicationName,
+    __inout_opt LPWSTR lpCommandLine,
+    __in_opt LPSECURITY_ATTRIBUTES lpProcessAttributes,
+    __in_opt LPSECURITY_ATTRIBUTES lpThreadAttributes,
+    __in BOOL bInheritHandles,
+    __in DWORD dwCreationFlags,
+    __in_opt LPVOID lpEnvironment,
+    __in_opt LPCWSTR lpCurrentDirectory,
+    __in LPSTARTUPINFOW lpStartupInfo,
+    __out LPPROCESS_INFORMATION lpProcessInformation
+    );
+
 typedef HRESULT (DAPI *PFN_PROCWAITFORCOMPLETION)(
     __in HANDLE hProcess,
     __in DWORD dwTimeout,
@@ -287,8 +300,18 @@ HRESULT CoreParseCommandLine(
     __inout HANDLE* phSourceEngineFile
     );
 void CoreFunctionOverride(
+    __in_opt PFN_CREATEPROCESSW pfnCreateProcessW,
     __in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
     );
+HRESULT CoreCreateProcess(
+    __in_opt LPCWSTR wzApplicationName,
+    __inout_opt LPWSTR sczCommandLine,
+    __in BOOL fInheritHandles,
+    __in DWORD dwCreationFlags,
+    __in_opt LPCWSTR wzCurrentDirectory,
+    __in WORD wShowWindow,
+    __out LPPROCESS_INFORMATION pProcessInformation
+    );
 HRESULT DAPI CoreWaitForProcCompletion(
     __in HANDLE hProcess,
     __in DWORD dwTimeout,
diff --git a/src/burn/engine/embedded.cpp b/src/burn/engine/embedded.cpp
index 58af5574..4c3de98e 100644
--- a/src/burn/engine/embedded.cpp
+++ b/src/burn/engine/embedded.cpp
@@ -52,7 +52,6 @@ extern "C" HRESULT EmbeddedRunBundle(
     DWORD dwCurrentProcessId = ::GetCurrentProcessId();
     HANDLE hCreatedPipesEvent = NULL;
     LPWSTR sczCommand = NULL;
-    STARTUPINFOW si = { };
     PROCESS_INFORMATION pi = { };
     BURN_PIPE_RESULT result = { };
 
@@ -79,10 +78,8 @@ extern "C" HRESULT EmbeddedRunBundle(
         ExitOnFailure(hr, "Failed to append user args.");
     }
 
-    if (!::CreateProcessW(wzExecutablePath, sczCommand, NULL, NULL, TRUE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi))
-    {
-        ExitWithLastError(hr, "Failed to create embedded process at path: %ls", wzExecutablePath);
-    }
+    hr = CoreCreateProcess(wzExecutablePath, sczCommand, TRUE, CREATE_NO_WINDOW, NULL, 0, &pi);
+    ExitOnFailure(hr, "Failed to create embedded process at path: %ls", wzExecutablePath);
 
     connection.dwProcessId = ::GetProcessId(pi.hProcess);
     connection.hProcess = pi.hProcess;
@@ -95,7 +92,7 @@ extern "C" HRESULT EmbeddedRunBundle(
     ExitOnFailure(hr, "Failed to process messages from embedded message.");
 
     // Get the return code from the embedded process.
-    hr = ProcWaitForCompletion(connection.hProcess, INFINITE, pdwExitCode);
+    hr = CoreWaitForProcCompletion(connection.hProcess, INFINITE, pdwExitCode);
     ExitOnFailure(hr, "Failed to wait for embedded executable: %ls", wzExecutablePath);
 
 LExit:
diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
index 13075497..413db69b 100644
--- a/src/burn/engine/engine.cpp
+++ b/src/burn/engine/engine.cpp
@@ -445,7 +445,6 @@ static HRESULT RunUntrusted(
     LPWSTR sczCachedCleanRoomBundlePath = NULL;
     LPWSTR sczParameters = NULL;
     LPWSTR sczFullCommandLine = NULL;
-    STARTUPINFOW si = { };
     PROCESS_INFORMATION pi = { };
     HANDLE hFileAttached = NULL;
     HANDLE hFileSelf = NULL;
@@ -484,12 +483,8 @@ static HRESULT RunUntrusted(
     hr = StrAllocFormattedSecure(&sczFullCommandLine, L"\"%ls\" %ls", wzCleanRoomBundlePath, sczParameters);
     ExitOnFailure(hr, "Failed to allocate full command-line.");
 
-    si.cb = sizeof(si);
-    si.wShowWindow = static_cast<WORD>(pEngineState->command.nCmdShow);
-    if (!::CreateProcessW(wzCleanRoomBundlePath, sczFullCommandLine, NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi))
-    {
-        ExitWithLastError(hr, "Failed to launch clean room process: %ls", sczFullCommandLine);
-    }
+    hr = CoreCreateProcess(wzCleanRoomBundlePath, sczFullCommandLine, TRUE, 0, NULL, static_cast<WORD>(pEngineState->command.nCmdShow), &pi);
+    ExitOnFailure(hr, "Failed to launch clean room process: %ls", sczFullCommandLine);
 
     hProcess = pi.hProcess;
     pi.hProcess = NULL;
diff --git a/src/burn/engine/exeengine.cpp b/src/burn/engine/exeengine.cpp
index 701adb74..1168f5ea 100644
--- a/src/burn/engine/exeengine.cpp
+++ b/src/burn/engine/exeengine.cpp
@@ -687,7 +687,6 @@ extern "C" HRESULT ExeEngineRunProcess(
 {
     HRESULT hr = S_OK;
     LPWSTR sczCommand = NULL;
-    STARTUPINFOW si = { };
     PROCESS_INFORMATION pi = { };
     GENERIC_EXECUTE_MESSAGE message = { };
     int nResult = IDNOACTION;
@@ -695,7 +694,6 @@ extern "C" HRESULT ExeEngineRunProcess(
     BOOL fDelayedCancel = FALSE;
     BOOL fFireAndForget = BURN_PACKAGE_TYPE_EXE == pPackage->type && pPackage->Exe.fFireAndForget;
     BOOL fInheritHandles = BURN_PACKAGE_TYPE_BUNDLE == pPackage->type;
-    size_t cchCachedDirectory = 0;
 
     // Always add user supplied arguments last.
     if (wzUserArgs)
@@ -704,19 +702,10 @@ extern "C" HRESULT ExeEngineRunProcess(
         ExitOnFailure(hr, "Failed to append user args.");
     }
 
-    // CreateProcessW has undocumented MAX_PATH restriction for lpCurrentDirectory even when long path support is enabled.
-    if (wzCachedDirectory && FAILED(::StringCchLengthW(wzCachedDirectory, MAX_PATH - 1, &cchCachedDirectory)))
-    {
-        wzCachedDirectory = NULL;
-    }
-
     // Make the cache location of the executable the current directory to help those executables
     // that expect stuff to be relative to them.
-    si.cb = sizeof(si);
-    if (!::CreateProcessW(wzExecutablePath, sczCommand ? sczCommand : sczBaseCommand, NULL, NULL, fInheritHandles, CREATE_NO_WINDOW, NULL, wzCachedDirectory, &si, &pi))
-    {
-        ExitWithLastError(hr, "Failed to CreateProcess on path: %ls", wzExecutablePath);
-    }
+    hr = CoreCreateProcess(wzExecutablePath, sczCommand ? sczCommand : sczBaseCommand, fInheritHandles, CREATE_NO_WINDOW, wzCachedDirectory, 0, &pi);
+    ExitOnFailure(hr, "Failed to CreateProcess on path: %ls", wzExecutablePath);
 
     message.type = GENERIC_EXECUTE_MESSAGE_PROCESS_STARTED;
     message.dwUIHint = MB_OK;
diff --git a/src/burn/engine/netfxchainer.cpp b/src/burn/engine/netfxchainer.cpp
index 68ddb093..86a06373 100644
--- a/src/burn/engine/netfxchainer.cpp
+++ b/src/burn/engine/netfxchainer.cpp
@@ -345,7 +345,6 @@ extern "C" HRESULT NetFxRunChainer(
     LPWSTR sczSectionName = NULL;
     LPWSTR sczCommand = NULL;
     NetFxChainer* pNetfxChainer = NULL;
-    STARTUPINFOW si = { };
     PROCESS_INFORMATION pi = { };
     HRESULT hrInternalError = 0;
 
@@ -372,11 +371,8 @@ extern "C" HRESULT NetFxRunChainer(
         ExitOnFailure(hr, "Failed to append user args.");
     }
 
-    si.cb = sizeof(si);
-    if (!::CreateProcessW(wzExecutablePath, sczCommand, NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi))
-    {
-        ExitWithLastError(hr, "Failed to CreateProcess on path: %ls", wzExecutablePath);
-    }
+    hr = CoreCreateProcess(wzExecutablePath, sczCommand, FALSE, CREATE_NO_WINDOW, NULL, 0, &pi);
+    ExitOnFailure(hr, "Failed to CreateProcess on path: %ls", wzExecutablePath);
 
     HANDLE handles[2] = { pi.hProcess, pNetfxChainer->hEventChaineeSend };
 
diff --git a/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj b/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj
index f07418a7..5f8004b4 100644
--- a/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj
+++ b/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj
@@ -47,6 +47,7 @@
     <ClCompile Include="AssemblyInfo.cpp" />
     <ClCompile Include="CacheTest.cpp" />
     <ClCompile Include="ElevationTest.cpp" />
+    <ClCompile Include="EmbeddedTest.cpp" />
     <ClCompile Include="ManifestHelpers.cpp" />
     <ClCompile Include="ManifestTest.cpp" />
     <ClCompile Include="PlanTest.cpp" />
diff --git a/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj.filters b/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj.filters
index 90290f52..9cfe9ee6 100644
--- a/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj.filters
+++ b/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj.filters
@@ -24,6 +24,9 @@
     <ClCompile Include="ElevationTest.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="EmbeddedTest.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="ManifestHelpers.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
diff --git a/src/burn/test/BurnUnitTest/ElevationTest.cpp b/src/burn/test/BurnUnitTest/ElevationTest.cpp
index ce34b4b1..8d4cc7ff 100644
--- a/src/burn/test/BurnUnitTest/ElevationTest.cpp
+++ b/src/burn/test/BurnUnitTest/ElevationTest.cpp
@@ -61,7 +61,7 @@ namespace Bootstrapper
             try
             {
                 ShelFunctionOverride(ElevateTest_ShellExecuteExW);
-                CoreFunctionOverride(ThrdWaitForCompletion);
+                CoreFunctionOverride(NULL, ThrdWaitForCompletion);
 
                 PipeConnectionInitialize(pConnection);
 
diff --git a/src/burn/test/BurnUnitTest/EmbeddedTest.cpp b/src/burn/test/BurnUnitTest/EmbeddedTest.cpp
new file mode 100644
index 00000000..cac1ba81
--- /dev/null
+++ b/src/burn/test/BurnUnitTest/EmbeddedTest.cpp
@@ -0,0 +1,175 @@
+// 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.
+
+#include "precomp.h"
+
+
+const HRESULT S_TEST_SUCCEEDED = 0x3133;
+const DWORD TEST_EXIT_CODE = 666;
+
+struct BUNDLE_RUNNER_CONTEXT
+{
+    DWORD dwResult;
+};
+
+
+static BOOL STDAPICALLTYPE EmbeddedTest_CreateProcessW(
+    __in_opt LPCWSTR lpApplicationName,
+    __inout_opt LPWSTR lpCommandLine,
+    __in_opt LPSECURITY_ATTRIBUTES lpProcessAttributes,
+    __in_opt LPSECURITY_ATTRIBUTES lpThreadAttributes,
+    __in BOOL bInheritHandles,
+    __in DWORD dwCreationFlags,
+    __in_opt LPVOID lpEnvironment,
+    __in_opt LPCWSTR lpCurrentDirectory,
+    __in LPSTARTUPINFOW lpStartupInfo,
+    __out LPPROCESS_INFORMATION lpProcessInformation
+    );
+static DWORD CALLBACK EmbeddedTest_ThreadProc(
+    __in LPVOID lpThreadParameter
+    );
+static int EmbeddedTest_GenericMessageHandler(
+    __in GENERIC_EXECUTE_MESSAGE* pMessage,
+    __in LPVOID pvContext
+    );
+
+namespace Microsoft
+{
+namespace Tools
+{
+namespace WindowsInstallerXml
+{
+namespace Test
+{
+namespace Bootstrapper
+{
+    using namespace System;
+    using namespace System::IO;
+    using namespace System::Threading;
+    using namespace Xunit;
+
+    public ref class EmbeddedTest : BurnUnitTest
+    {
+    public:
+        EmbeddedTest(BurnTestFixture^ fixture) : BurnUnitTest(fixture)
+        {
+        }
+
+        [Fact]
+        void EmbeddedProtocolTest()
+        {
+            HRESULT hr = S_OK;
+            BUNDLE_RUNNER_CONTEXT bundleRunnerContext = { };
+            DWORD dwExitCode = 0;
+
+            try
+            {
+                CoreFunctionOverride(EmbeddedTest_CreateProcessW, ThrdWaitForCompletion);
+
+                //
+                // bundle runner setup
+                //
+                hr = EmbeddedRunBundle(L"C:\\ignored\\target.exe", L"\"C:\\ignored\\target.exe\"", NULL, EmbeddedTest_GenericMessageHandler, &bundleRunnerContext, &dwExitCode);
+                TestThrowOnFailure(hr, L"Failed to run embedded bundle.");
+
+                // check results
+                Assert::Equal<HRESULT>(S_TEST_SUCCEEDED, (HRESULT)bundleRunnerContext.dwResult);
+                Assert::Equal<DWORD>(TEST_EXIT_CODE, dwExitCode);
+            }
+            finally
+            {
+            }
+        }
+    };
+}
+}
+}
+}
+}
+
+
+static BOOL STDAPICALLTYPE EmbeddedTest_CreateProcessW(
+    __in_opt LPCWSTR /*lpApplicationName*/,
+    __inout_opt LPWSTR lpCommandLine,
+    __in_opt LPSECURITY_ATTRIBUTES /*lpProcessAttributes*/,
+    __in_opt LPSECURITY_ATTRIBUTES /*lpThreadAttributes*/,
+    __in BOOL /*bInheritHandles*/,
+    __in DWORD /*dwCreationFlags*/,
+    __in_opt LPVOID /*lpEnvironment*/,
+    __in_opt LPCWSTR /*lpCurrentDirectory*/,
+    __in LPSTARTUPINFOW /*lpStartupInfo*/,
+    __out LPPROCESS_INFORMATION lpProcessInformation
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR scz = NULL;
+    LPCWSTR wzArgs = lpCommandLine + 24; //skip '"C:\ignored\target.exe" '
+
+    hr = StrAllocString(&scz, wzArgs, 0);
+    ExitOnFailure(hr, "Failed to copy arguments.");
+
+    // Pretend this thread is the embedded process.
+    lpProcessInformation->hProcess = ::CreateThread(NULL, 0, EmbeddedTest_ThreadProc, scz, 0, NULL);
+    ExitOnNullWithLastError(lpProcessInformation->hProcess, hr, "Failed to create thread.");
+
+    scz = NULL;
+
+LExit:
+    ReleaseStr(scz);
+
+    return SUCCEEDED(hr);
+}
+
+static DWORD CALLBACK EmbeddedTest_ThreadProc(
+    __in LPVOID lpThreadParameter
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczArguments = (LPWSTR)lpThreadParameter;
+    BURN_ENGINE_STATE engineState = { };
+    BURN_PIPE_CONNECTION* pConnection = &engineState.embeddedConnection;
+    DWORD dwResult = 0;
+
+    engineState.internalCommand.mode = BURN_MODE_EMBEDDED;
+
+    PipeConnectionInitialize(pConnection);
+
+    StrAlloc(&pConnection->sczName, MAX_PATH);
+    StrAlloc(&pConnection->sczSecret, MAX_PATH);
+
+    // parse command line arguments
+    if (3 != swscanf_s(sczArguments, L"-burn.embedded %s %s %u", pConnection->sczName, MAX_PATH, pConnection->sczSecret, MAX_PATH, &pConnection->dwProcessId))
+    {
+        ExitWithRootFailure(hr, E_INVALIDARG, "Failed to parse argument string.");
+    }
+
+    // set up connection with parent bundle runner
+    hr = PipeChildConnect(pConnection, FALSE);
+    ExitOnFailure(hr, "Failed to connect to parent bundle runner.");
+
+    // post known message
+    hr = ExternalEngineSendEmbeddedError(&engineState, S_TEST_SUCCEEDED, NULL, 0, reinterpret_cast<int*>(&dwResult));
+    ExitOnFailure(hr, "Failed to post known message to parent bundle runner.");
+
+LExit:
+    PipeConnectionUninitialize(pConnection);
+    ReleaseStr(sczArguments);
+
+    return FAILED(hr) ? (DWORD)hr : dwResult;
+}
+
+static int EmbeddedTest_GenericMessageHandler(
+    __in GENERIC_EXECUTE_MESSAGE* pMessage,
+    __in LPVOID pvContext
+    )
+{
+    BUNDLE_RUNNER_CONTEXT* pContext = reinterpret_cast<BUNDLE_RUNNER_CONTEXT*>(pvContext);
+    DWORD dwResult = 0;
+
+    if (GENERIC_EXECUTE_MESSAGE_ERROR == pMessage->type)
+    {
+        pContext->dwResult = pMessage->error.dwErrorCode;
+        dwResult = TEST_EXIT_CODE;
+    }
+
+    return dwResult;
+}
diff --git a/src/burn/test/BurnUnitTest/precomp.h b/src/burn/test/BurnUnitTest/precomp.h
index 69c62cb2..84e989a4 100644
--- a/src/burn/test/BurnUnitTest/precomp.h
+++ b/src/burn/test/BurnUnitTest/precomp.h
@@ -73,6 +73,7 @@
 #include "manifest.h"
 #include "splashscreen.h"
 #include "detect.h"
+#include "externalengine.h"
 
 #include "engine.version.h"