4 files changed, 112 insertions, 15 deletions

diff --git a/src/burn/engine/cache.cpp b/src/burn/engine/cache.cpp
index 61416aed..a35b83bc 100644
--- a/src/burn/engine/cache.cpp
+++ b/src/burn/engine/cache.cpp
@@ -1281,7 +1281,7 @@ static HRESULT CalculateWorkingFolder(
     HRESULT hr = S_OK;
     RPC_STATUS rs = RPC_S_OK;
     BOOL fElevated = FALSE;
-    WCHAR wzTempPath[MAX_PATH] = { };
+    LPWSTR sczTempPath = NULL;
     UUID guid = {};
     WCHAR wzGuid[39];
 
@@ -1291,20 +1291,13 @@ static HRESULT CalculateWorkingFolder(
 
         if (fElevated)
         {
-            if (!::GetWindowsDirectoryW(wzTempPath, countof(wzTempPath)))
-            {
-                ExitWithLastError(hr, "Failed to get windows path for working folder.");
-            }
-
-            hr = PathFixedBackslashTerminate(wzTempPath, countof(wzTempPath));
-            ExitOnFailure(hr, "Failed to ensure windows path for working folder ended in backslash.");
-
-            hr = ::StringCchCatW(wzTempPath, countof(wzTempPath), L"Temp\\");
-            ExitOnFailure(hr, "Failed to concat Temp directory on windows path for working folder.");
+            hr = PathGetSystemTempPath(&sczTempPath);
+            ExitOnFailure(hr, "Failed to get system temp folder path for working folder.");
         }
-        else if (0 == ::GetTempPathW(countof(wzTempPath), wzTempPath))
+        else
         {
-            ExitWithLastError(hr, "Failed to get temp path for working folder.");
+            hr = PathGetTempPath(&sczTempPath);
+            ExitOnFailure(hr, "Failed to get temp folder path for working folder.");
         }
 
         rs = ::UuidCreate(&guid);
@@ -1317,7 +1310,7 @@ static HRESULT CalculateWorkingFolder(
             ExitOnRootFailure(hr, "Failed to convert working folder guid into string.");
         }
 
-        hr = StrAllocFormatted(&vsczWorkingFolder, L"%ls%ls\\", wzTempPath, wzGuid);
+        hr = StrAllocFormatted(&vsczWorkingFolder, L"%ls%ls\\", sczTempPath, wzGuid);
         ExitOnFailure(hr, "Failed to append bundle id on to temp path for working folder.");
     }
 
@@ -1325,6 +1318,8 @@ static HRESULT CalculateWorkingFolder(
     ExitOnFailure(hr, "Failed to copy working folder path.");
 
 LExit:
+    ReleaseStr(sczTempPath);
+
     return hr;
 }
 
diff --git a/src/burn/engine/variable.cpp b/src/burn/engine/variable.cpp
index e2b1f1f2..e7391a9f 100644
--- a/src/burn/engine/variable.cpp
+++ b/src/burn/engine/variable.cpp
@@ -1916,7 +1916,6 @@ static HRESULT InitializeVariableTempFolder(
     HRESULT hr = S_OK;
     WCHAR wzPath[MAX_PATH] = { };
 
-    // get volume path name
     if (!::GetTempPathW(MAX_PATH, wzPath))
     {
         ExitWithLastError(hr, "Failed to get temp path.");
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h b/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h
index 579b8454..0ae9f437 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/pathutil.h
@@ -151,6 +151,22 @@ DAPI_(HRESULT) PathCreateTempDirectory(
     );
 
 /*******************************************************************
+ PathGetTempPath - returns the path to the temp folder
+    that is backslash terminated.
+*******************************************************************/
+DAPI_(HRESULT) PathGetTempPath(
+    __out_z LPWSTR* psczTempPath
+    );
+
+/*******************************************************************
+ PathGetSystemTempPath - returns the path to the system temp folder
+    that is backslash terminated.
+*******************************************************************/
+DAPI_(HRESULT) PathGetSystemTempPath(
+    __out_z LPWSTR* psczSystemTempPath
+    );
+
+/*******************************************************************
  PathGetKnownFolder - returns the path to a well-known shell folder
 
 *******************************************************************/
diff --git a/src/libs/dutil/WixToolset.DUtil/pathutil.cpp b/src/libs/dutil/WixToolset.DUtil/pathutil.cpp
index 7c3cfe06..5fad519b 100644
--- a/src/libs/dutil/WixToolset.DUtil/pathutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/pathutil.cpp
@@ -9,6 +9,7 @@
 #define PathExitWithLastError(x, s, ...) ExitWithLastErrorSource(DUTIL_SOURCE_PATHUTIL, x, s, __VA_ARGS__)
 #define PathExitOnFailure(x, s, ...) ExitOnFailureSource(DUTIL_SOURCE_PATHUTIL, x, s, __VA_ARGS__)
 #define PathExitOnRootFailure(x, s, ...) ExitOnRootFailureSource(DUTIL_SOURCE_PATHUTIL, x, s, __VA_ARGS__)
+#define PathExitWithRootFailure(x, e, s, ...) ExitWithRootFailureSource(DUTIL_SOURCE_PATHUTIL, x, e, s, __VA_ARGS__)
 #define PathExitOnFailureDebugTrace(x, s, ...) ExitOnFailureDebugTraceSource(DUTIL_SOURCE_PATHUTIL, x, s, __VA_ARGS__)
 #define PathExitOnNull(p, x, e, s, ...) ExitOnNullSource(DUTIL_SOURCE_PATHUTIL, p, x, e, s, __VA_ARGS__)
 #define PathExitOnNullWithLastError(p, x, s, ...) ExitOnNullWithLastErrorSource(DUTIL_SOURCE_PATHUTIL, p, x, s, __VA_ARGS__)
@@ -814,6 +815,92 @@ LExit:
 }
 
 
+DAPI_(HRESULT) PathGetTempPath(
+    __out_z LPWSTR* psczTempPath
+    )
+{
+    HRESULT hr = S_OK;
+    WCHAR wzTempPath[MAX_PATH + 1] = { };
+    DWORD cch = 0;
+
+    cch = ::GetTempPathW(countof(wzTempPath), wzTempPath);
+    if (!cch)
+    {
+        PathExitWithLastError(hr, "Failed to GetTempPath.");
+    }
+    else if (cch >= countof(wzTempPath))
+    {
+        PathExitWithRootFailure(hr, E_INSUFFICIENT_BUFFER, "TEMP directory path too long.");
+    }
+
+    hr = StrAllocString(psczTempPath, wzTempPath, cch);
+    PathExitOnFailure(hr, "Failed to copy TEMP directory path.");
+
+LExit:
+    return hr;
+}
+
+
+DAPI_(HRESULT) PathGetSystemTempPath(
+    __out_z LPWSTR* psczSystemTempPath
+    )
+{
+    HRESULT hr = S_OK;
+    HKEY hKey = NULL;
+    WCHAR wzTempPath[MAX_PATH + 1] = { };
+    DWORD cch = 0;
+
+    // There is no documented API to get system environment variables, so read them from the registry.
+    hr = RegOpen(HKEY_LOCAL_MACHINE, L"System\\CurrentControlSet\\Control\\Session Manager\\Environment", KEY_READ, &hKey);
+    if (E_FILENOTFOUND != hr)
+    {
+        PathExitOnFailure(hr, "Failed to open system environment registry key.");
+
+        // Follow documented precedence rules for TMP/TEMP from ::GetTempPath.
+        // TODO: values will be expanded with the current environment variables instead of the system environment variables.
+        hr = RegReadString(hKey, L"TMP", psczSystemTempPath);
+        if (E_FILENOTFOUND != hr)
+        {
+            PathExitOnFailure(hr, "Failed to get system TMP value.");
+
+            hr = PathBackslashTerminate(psczSystemTempPath);
+            PathExitOnFailure(hr, "Failed to backslash terminate system TMP value.");
+
+            ExitFunction();
+        }
+
+        hr = RegReadString(hKey, L"TEMP", psczSystemTempPath);
+        if (E_FILENOTFOUND != hr)
+        {
+            PathExitOnFailure(hr, "Failed to get system TEMP value.");
+
+            hr = PathBackslashTerminate(psczSystemTempPath);
+            PathExitOnFailure(hr, "Failed to backslash terminate system TEMP value.");
+
+            ExitFunction();
+        }
+    }
+
+    cch = ::GetSystemWindowsDirectoryW(wzTempPath, countof(wzTempPath));
+    if (!cch)
+    {
+        PathExitWithLastError(hr, "Failed to get Windows directory path.");
+    }
+    else if (cch >= countof(wzTempPath))
+    {
+        PathExitWithRootFailure(hr, E_INSUFFICIENT_BUFFER, "Windows directory path too long.");
+    }
+
+    hr = PathConcat(wzTempPath, L"TEMP\\", psczSystemTempPath);
+    PathExitOnFailure(hr, "Failed to concat Temp directory on Windows directory path.");
+
+LExit:
+    ReleaseRegKey(hKey);
+
+    return hr;
+}
+
+
 DAPI_(HRESULT) PathGetKnownFolder(
     __in int csidl,
     __out LPWSTR* psczKnownFolder