6 files changed, 43 insertions, 33 deletions

diff --git a/src/burn/engine/ba.h b/src/burn/engine/ba.h
index c092fedf..3561805a 100644
--- a/src/burn/engine/ba.h
+++ b/src/burn/engine/ba.h
@@ -111,6 +111,7 @@ HRESULT BootstrapperApplicationInterpretExecuteResult(
 );
 
 HRESULT BootstrapperApplicationEnsureWorkingFolder(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __deref_out_z LPWSTR* psczUserExperienceWorkingFolder
 );
diff --git a/src/burn/engine/bootstrapperapplication.cpp b/src/burn/engine/bootstrapperapplication.cpp
index 402f7015..947b3720 100644
--- a/src/burn/engine/bootstrapperapplication.cpp
+++ b/src/burn/engine/bootstrapperapplication.cpp
@@ -276,6 +276,7 @@ EXTERN_C HRESULT BootstrapperApplicationInterpretExecuteResult(
 }
 
 EXTERN_C HRESULT BootstrapperApplicationEnsureWorkingFolder(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __deref_out_z LPWSTR* psczUserExperienceWorkingFolder
     )
@@ -283,7 +284,7 @@ EXTERN_C HRESULT BootstrapperApplicationEnsureWorkingFolder(
     HRESULT hr = S_OK;
     LPWSTR sczWorkingFolder = NULL;
 
-    hr = CacheEnsureBaseWorkingFolder(pCache, &sczWorkingFolder);
+    hr = CacheEnsureBaseWorkingFolder(fElevated, pCache, &sczWorkingFolder);
     ExitOnFailure(hr, "Failed to create working folder.");
 
     hr = StrAllocFormatted(psczUserExperienceWorkingFolder, L"%ls%ls\\", sczWorkingFolder, L".ba");
diff --git a/src/burn/engine/cache.cpp b/src/burn/engine/cache.cpp
index c0ac3ecd..358327a2 100644
--- a/src/burn/engine/cache.cpp
+++ b/src/burn/engine/cache.cpp
@@ -106,6 +106,7 @@ static HRESULT SecurePath(
     __in LPCWSTR wzPath
     );
 static HRESULT CopyEngineToWorkingFolder(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __in_z LPCWSTR wzSourcePath,
     __in_z LPCWSTR wzWorkingFolderName,
@@ -330,6 +331,7 @@ LExit:
 }
 
 extern "C" HRESULT CacheEnsureBaseWorkingFolder(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __deref_out_z_opt LPWSTR* psczBaseWorkingFolder
     )
@@ -338,15 +340,32 @@ extern "C" HRESULT CacheEnsureBaseWorkingFolder(
 
     HRESULT hr = S_OK;
     LPWSTR sczPotential = NULL;
+    PSECURITY_DESCRIPTOR psd = NULL;
+    LPSECURITY_ATTRIBUTES pWorkingFolderAcl = NULL;
 
     if (!pCache->fInitializedBaseWorkingFolder)
     {
+        // If elevated, allocate the pWorkingFolderAcl to protect the working folder to only SYSTEM and Admins.
+        if (fElevated)
+        {
+            LPCWSTR wzSddl = L"D:PAI(A;;FA;;;BA)(A;OICIIO;GA;;;BA)(A;;FA;;;SY)(A;OICIIO;GA;;;SY)";
+            if (!::ConvertStringSecurityDescriptorToSecurityDescriptorW(wzSddl, SDDL_REVISION_1, &psd, NULL))
+            {
+                ExitWithLastError(hr, "Failed to create the security descriptor for the working folder.");
+            }
+
+            pWorkingFolderAcl = reinterpret_cast<LPSECURITY_ATTRIBUTES>(MemAlloc(sizeof(SECURITY_ATTRIBUTES), TRUE));
+            pWorkingFolderAcl->nLength = sizeof(SECURITY_ATTRIBUTES);
+            pWorkingFolderAcl->lpSecurityDescriptor = psd;
+            pWorkingFolderAcl->bInheritHandle = FALSE;
+        }
+
         for (DWORD i = 0; i < pCache->cPotentialBaseWorkingFolders; ++i)
         {
             hr = PathConcatRelativeToFullyQualifiedBase(pCache->rgsczPotentialBaseWorkingFolders[i], pCache->wzGuid, &sczPotential);
             if (SUCCEEDED(hr))
             {
-                hr = DirEnsureExists(sczPotential, NULL);
+                hr = DirEnsureExists(sczPotential, pWorkingFolderAcl);
                 if (SUCCEEDED(hr))
                 {
                     pCache->sczBaseWorkingFolder = sczPotential;
@@ -373,6 +392,11 @@ extern "C" HRESULT CacheEnsureBaseWorkingFolder(
     }
 
 LExit:
+    ReleaseMem(pWorkingFolderAcl);
+    if (psd)
+    {
+        ::LocalFree(psd);
+    }
     ReleaseStr(sczPotential);
 
     return hr;
@@ -888,6 +912,7 @@ extern "C" HRESULT CachePreparePackage(
 }
 
 extern "C" HRESULT CacheBundleToWorkingDirectory(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __in_z LPCWSTR wzExecutableName,
     __in BURN_SECTION* pSection,
@@ -912,7 +937,7 @@ extern "C" HRESULT CacheBundleToWorkingDirectory(
     }
     else // otherwise, carry on putting the bundle in the working folder.
     {
-        hr = CopyEngineToWorkingFolder(pCache, sczSourcePath, BUNDLE_WORKING_FOLDER_NAME, wzExecutableName, pSection, psczEngineWorkingPath);
+        hr = CopyEngineToWorkingFolder(fElevated, pCache, sczSourcePath, BUNDLE_WORKING_FOLDER_NAME, wzExecutableName, pSection, psczEngineWorkingPath);
         ExitOnFailure(hr, "Failed to copy engine to working folder.");
     }
 
@@ -2063,6 +2088,7 @@ LExit:
 
 
 static HRESULT CopyEngineToWorkingFolder(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __in_z LPCWSTR wzSourcePath,
     __in_z LPCWSTR wzWorkingFolderName,
@@ -2079,7 +2105,7 @@ static HRESULT CopyEngineToWorkingFolder(
     LPWSTR sczPayloadSourcePath = NULL;
     LPWSTR sczPayloadTargetPath = NULL;
 
-    hr = CacheEnsureBaseWorkingFolder(pCache, &sczWorkingFolder);
+    hr = CacheEnsureBaseWorkingFolder(fElevated, pCache, &sczWorkingFolder);
     ExitOnFailure(hr, "Failed to create working path to copy engine.");
 
     hr = PathConcatRelativeToFullyQualifiedBase(sczWorkingFolder, wzWorkingFolderName, &sczTargetDirectory);
diff --git a/src/burn/engine/cache.h b/src/burn/engine/cache.h
index 3f0ba749..7c4dfaa1 100644
--- a/src/burn/engine/cache.h
+++ b/src/burn/engine/cache.h
@@ -96,6 +96,7 @@ HRESULT CacheEnsureAcquisitionFolder(
     __in BURN_CACHE* pCache
     );
 HRESULT CacheEnsureBaseWorkingFolder(
+    __in BOOL fElevated,
     __in BURN_CACHE* pCache,
     __deref_out_z_opt LPWSTR* psczBaseWorkingFolder
     );
@@ -171,6 +172,7 @@ HRESULT CachePreparePackage(
     __in BURN_PACKAGE* pPackage
     );
 HRESULT CacheBundleToWorkingDirectory(
+    __in BOOL fElvated,
     __in BURN_CACHE* pCache,
     __in_z LPCWSTR wzExecutableName,
     __in BURN_SECTION* pSection,
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index ae74fdfd..a85e6f18 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -165,7 +165,7 @@ extern "C" HRESULT CoreInitialize(
     if (BURN_MODE_NORMAL == pEngineState->internalCommand.mode || BURN_MODE_EMBEDDED == pEngineState->internalCommand.mode)
     {
         // Extract all UX payloads to working folder.
-        hr = BootstrapperApplicationEnsureWorkingFolder(&pEngineState->cache, &pEngineState->userExperience.sczTempDirectory);
+        hr = BootstrapperApplicationEnsureWorkingFolder(pEngineState->internalCommand.fInitiallyElevated, &pEngineState->cache, &pEngineState->userExperience.sczTempDirectory);
         ExitOnFailure(hr, "Failed to get unique temporary folder for bootstrapper application.");
 
         hr = PayloadExtractUXContainer(&pEngineState->userExperience.payloads, &containerContext, pEngineState->userExperience.sczTempDirectory);
@@ -588,7 +588,7 @@ extern "C" HRESULT CoreElevate(
         // If the elevated companion pipe isn't created yet, let's make that happen.
         if (!pEngineState->sczBundleEngineWorkingPath)
         {
-            hr = CacheBundleToWorkingDirectory(&pEngineState->cache, pEngineState->registration.sczExecutableName, &pEngineState->section, &pEngineState->sczBundleEngineWorkingPath);
+            hr = CacheBundleToWorkingDirectory(pEngineState->internalCommand.fInitiallyElevated, &pEngineState->cache, pEngineState->registration.sczExecutableName, &pEngineState->section, &pEngineState->sczBundleEngineWorkingPath);
             ExitOnFailure(hr, "Failed to cache engine to working directory.");
         }
 
@@ -697,7 +697,7 @@ extern "C" HRESULT CoreApply(
     // Ensure the engine is cached to the working path.
     if (!pEngineState->sczBundleEngineWorkingPath)
     {
-        hr = CacheBundleToWorkingDirectory(&pEngineState->cache, pEngineState->registration.sczExecutableName, &pEngineState->section, &pEngineState->sczBundleEngineWorkingPath);
+        hr = CacheBundleToWorkingDirectory(pEngineState->internalCommand.fInitiallyElevated, &pEngineState->cache, pEngineState->registration.sczExecutableName, &pEngineState->section, &pEngineState->sczBundleEngineWorkingPath);
         ExitOnFailure(hr, "Failed to cache engine to working directory.");
     }
 
diff --git a/src/dtf/SfxCA/SfxUtil.cpp b/src/dtf/SfxCA/SfxUtil.cpp
index 2e6b0555..32dc6e04 100644
--- a/src/dtf/SfxCA/SfxUtil.cpp
+++ b/src/dtf/SfxCA/SfxUtil.cpp
@@ -164,38 +164,18 @@ bool ExtractToTempDirectory(__in MSIHANDLE hSession, __in HMODULE hModule,
         StringCchCopy(szTempDir, cchTempDirBuf, szModule);
         StringCchCat(szTempDir, cchTempDirBuf, L"-");
 
+        BOOL fCreatedDirectory = FALSE;
         DWORD cchTempDir = (DWORD) wcslen(szTempDir);
-        for (int i = 0; DirectoryExists(szTempDir); i++)
+        for (int i = 0; i < 10000 && !fCreatedDirectory; i++)
         {
                 swprintf_s(szTempDir + cchTempDir, cchTempDirBuf - cchTempDir, L"%d", i);
+                fCreatedDirectory = ::CreateDirectory(szTempDir, NULL);
         }
 
-        if (!CreateDirectory(szTempDir, NULL))
+        if (!fCreatedDirectory)
         {
-                cchCopied = GetTempPath(cchTempDirBuf, szTempDir);
-                if (cchCopied == 0 || cchCopied >= cchTempDirBuf)
-                {
-                        Log(hSession, L"Failed to get temp directory. Error code %d", GetLastError());
-                        return false;
-                }
-
-                wchar_t* szModuleName = wcsrchr(szModule, L'\\');
-                if (szModuleName == NULL) szModuleName = szModule;
-                else szModuleName = szModuleName + 1;
-                StringCchCat(szTempDir, cchTempDirBuf, szModuleName);
-                StringCchCat(szTempDir, cchTempDirBuf, L"-");
-
-                cchTempDir = (DWORD) wcslen(szTempDir);
-                for (int i = 0; DirectoryExists(szTempDir); i++)
-                {
-                        swprintf_s(szTempDir + cchTempDir, cchTempDirBuf - cchTempDir, L"%d", i);
-                }
-
-                if (!CreateDirectory(szTempDir, NULL))
-                {
-                        Log(hSession, L"Failed to create temp directory. Error code %d", GetLastError());
-                        return false;
-                }
+                Log(hSession, L"Failed to create temp directory. Error code %d", ::GetLastError());
+                return false;
         }
 
         Log(hSession, L"Extracting custom action to temporary directory: %s\\", szTempDir);