From 8810aa8908ed7887616d86dd5fb821fcfa92f444 Mon Sep 17 00:00:00 2001
From: Sean Hall <r.sean.hall@gmail.com>
Date: Fri, 3 Jun 2022 17:50:22 -0500
Subject: Update Burn algorithm for picking elevated temp path to use
 SystemTemp.

---
 src/burn/engine/cache.cpp                       | 19 ++++++++++++++++---
 src/burn/engine/precomp.h                       |  1 +
 src/burn/stub/stub.vcxproj                      |  4 ++--
 src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj |  2 +-
 src/burn/test/BurnUnitTest/CacheTest.cpp        |  2 ++
 5 files changed, 22 insertions(+), 6 deletions(-)

(limited to 'src/burn')

diff --git a/src/burn/engine/cache.cpp b/src/burn/engine/cache.cpp
index 7fca8cc7..5eacb20d 100644
--- a/src/burn/engine/cache.cpp
+++ b/src/burn/engine/cache.cpp
@@ -1389,6 +1389,8 @@ static HRESULT CalculatePotentialBaseWorkingFolders(
     Assert(!pCache->rgsczPotentialBaseWorkingFolders && !pCache->cPotentialBaseWorkingFolders);
     HRESULT hr = S_OK;
     LPWSTR sczTemp = NULL;
+    LPWSTR sczPolicy = NULL;
+    BOOL fNeedsExpansion = FALSE;
 
     hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&pCache->rgsczPotentialBaseWorkingFolders), 6, sizeof(LPWSTR), 6);
     ExitOnFailure(hr, "Failed to initialize array.");
@@ -1408,12 +1410,22 @@ static HRESULT CalculatePotentialBaseWorkingFolders(
     // but only use it if elevated because it should be secured against non-admin users.
     if (pInternalCommand->fInitiallyElevated)
     {
-        hr = PolcReadString(POLICY_BURN_REGISTRY_PATH, L"EngineWorkingDirectory", NULL, &sczTemp);
+        hr = PolcReadUnexpandedString(POLICY_BURN_REGISTRY_PATH, L"EngineWorkingDirectory", NULL, &fNeedsExpansion, &sczPolicy);
         ExitOnFailure(hr, "Failed to read EngineWorkingDirectory policy directory.");
 
-        if (sczTemp)
+        if (S_FALSE != hr)
         {
-            // PolcReadString is supposed to automatically expand REG_EXPAND_SZ values.
+            if (fNeedsExpansion)
+            {
+                hr = EnvExpandEnvironmentStringsForUser(NULL, sczPolicy, &sczTemp, NULL);
+                ExitOnFailure(hr, "Failed to expand EngineWorkingDirectory policy directory.");
+            }
+            else
+            {
+                sczTemp = sczPolicy;
+                sczPolicy = NULL;
+            }
+
             pCache->rgsczPotentialBaseWorkingFolders[pCache->cPotentialBaseWorkingFolders] = sczTemp;
             sczTemp = NULL;
             ++pCache->cPotentialBaseWorkingFolders;
@@ -1438,6 +1450,7 @@ static HRESULT CalculatePotentialBaseWorkingFolders(
 
 LExit:
     ReleaseStr(sczTemp);
+    ReleaseStr(sczPolicy);
 
     return hr;
 }
diff --git a/src/burn/engine/precomp.h b/src/burn/engine/precomp.h
index 19a6d48b..bc7046f6 100644
--- a/src/burn/engine/precomp.h
+++ b/src/burn/engine/precomp.h
@@ -31,6 +31,7 @@
 #include <certutil.h>
 #include <cryputil.h>
 #include <dirutil.h>
+#include <envutil.h>
 #include <fileutil.h>
 #include <guidutil.h>
 #include <logutil.h>
diff --git a/src/burn/stub/stub.vcxproj b/src/burn/stub/stub.vcxproj
index d73b9f49..0a2e1cae 100644
--- a/src/burn/stub/stub.vcxproj
+++ b/src/burn/stub/stub.vcxproj
@@ -55,14 +55,14 @@
 
   <PropertyGroup>
     <ProjectAdditionalIncludeDirectories>$(ProjectDir)..\engine\inc</ProjectAdditionalIncludeDirectories>
-    <ProjectAdditionalLinkLibraries>cabinet.lib;crypt32.lib;msi.lib;rpcrt4.lib;shlwapi.lib;wininet.lib;wintrust.lib;wuguid.lib;engine.res</ProjectAdditionalLinkLibraries>
+    <ProjectAdditionalLinkLibraries>cabinet.lib;crypt32.lib;msi.lib;rpcrt4.lib;shlwapi.lib;userenv.lib;wininet.lib;wintrust.lib;wuguid.lib;engine.res</ProjectAdditionalLinkLibraries>
   </PropertyGroup>
 
   <ItemDefinitionGroup>
     <Link>
       <SwapRunFromCD>true</SwapRunFromCD>
       <SwapRunFromNET>true</SwapRunFromNET>
-      <DelayLoadDLLs>cabinet.dll;crypt32.dll;msi.dll;shlwapi.dll;version.dll;wininet.dll;wintrust.dll</DelayLoadDLLs>
+      <DelayLoadDLLs>cabinet.dll;crypt32.dll;msi.dll;shlwapi.dll;userenv.dll;version.dll;wininet.dll;wintrust.dll</DelayLoadDLLs>
     </Link>
   </ItemDefinitionGroup>
 
diff --git a/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj b/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj
index b05666b2..e1a28712 100644
--- a/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj
+++ b/src/burn/test/BurnUnitTest/BurnUnitTest.vcxproj
@@ -40,7 +40,7 @@
 
   <PropertyGroup>
     <ProjectAdditionalIncludeDirectories>$(ProjectAdditionalIncludeDirectories);$(BurnGeneratedHeaderDirectory);..\..\engine;..\..\..\api\burn\WixToolset.BootstrapperCore.Native\inc;..\..\..\libs\dutil\WixToolset.Dutil\inc</ProjectAdditionalIncludeDirectories>
-    <ProjectAdditionalLinkLibraries>cabinet.lib;crypt32.lib;msi.lib;rpcrt4.lib;shlwapi.lib;wininet.lib;wintrust.lib;$(RootBuildFolder)libs\$(Configuration)\$(WixNativeSdkLibraryToolset)\$(PlatformTarget)\dutil.lib;engine.res</ProjectAdditionalLinkLibraries>
+    <ProjectAdditionalLinkLibraries>cabinet.lib;crypt32.lib;msi.lib;rpcrt4.lib;shlwapi.lib;userenv.lib;wininet.lib;wintrust.lib;$(RootBuildFolder)libs\$(Configuration)\$(WixNativeSdkLibraryToolset)\$(PlatformTarget)\dutil.lib;engine.res</ProjectAdditionalLinkLibraries>
   </PropertyGroup>
 
   <ItemGroup>
diff --git a/src/burn/test/BurnUnitTest/CacheTest.cpp b/src/burn/test/BurnUnitTest/CacheTest.cpp
index eb0b31ab..2bec05fe 100644
--- a/src/burn/test/BurnUnitTest/CacheTest.cpp
+++ b/src/burn/test/BurnUnitTest/CacheTest.cpp
@@ -92,6 +92,8 @@ namespace Bootstrapper
                 VerifyBaseWorkingFolder(L"B:\\TEST\\TMP\\", cache.rgsczPotentialBaseWorkingFolders[0]);
                 CacheUninitialize(&cache);
 
+                // First fallback is impractical to mock out - %windir%\SystemTemp on Win11 when running as SYSTEM.
+
                 hr = RegCreate(HKEY_LOCAL_MACHINE, L"SOFTWARE\\Policies\\WiX\\Burn", GENERIC_WRITE, &hkBurnPolicy);
                 NativeAssert::Succeeded(hr, "Failed to create Burn policy key.");
 
-- 
cgit v1.2.3-55-g6feb