From 34ceb5eeb4c8cf1cb42ed943798fd92c643ab7c4 Mon Sep 17 00:00:00 2001
From: Jacques Eloff <joeloff@users.noreply.github.com>
Date: Thu, 8 Feb 2024 11:19:10 -0800
Subject: [PATCH] Add install state cleanup to finalizer (#18266)

---
 src/finalizer/CMakeLists.txt |  1 +
 src/finalizer/finalizer.cpp  | 64 ++++++++++++++++++++++++++++++++++--
 src/finalizer/precomp.h      |  6 ++++
 3 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/src/finalizer/CMakeLists.txt b/src/finalizer/CMakeLists.txt
index aac5c20b3..80c5be98d 100644
--- a/src/finalizer/CMakeLists.txt
+++ b/src/finalizer/CMakeLists.txt
@@ -55,6 +55,7 @@ target_link_libraries(Finalizer shell32.lib)
 target_link_libraries(Finalizer advapi32.lib)
 target_link_libraries(Finalizer version.lib)
 target_link_libraries(Finalizer msi.lib)
+target_link_libraries(Finalizer shlwapi.lib)
 
 # Add WiX libraries
 target_link_libraries(Finalizer wcautil.lib)
diff --git a/src/finalizer/finalizer.cpp b/src/finalizer/finalizer.cpp
index d68418b9b..03a1213ba 100644
--- a/src/finalizer/finalizer.cpp
+++ b/src/finalizer/finalizer.cpp
@@ -492,12 +492,64 @@ LExit:
     return hr;
 }
 
+void RemoveInstallStateFile(LPWSTR sczSdkFeatureBandVersion, LPWSTR sczPlatform)
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczProgramData = NULL;
+    LPWSTR sczInstallStatePath = NULL;
+    LPWSTR sczPath = NULL;
+
+    hr = ShelGetFolder(&sczProgramData, CSIDL_COMMON_APPDATA);
+    ExitOnFailure(hr, "Failed to get shell folder.");
+
+    hr = PathConcat(sczProgramData, L"dotnet", &sczInstallStatePath);
+    ExitOnFailure(hr, "Failed to concat dotnet to install state path.");
+
+    hr = PathConcat(sczInstallStatePath, L"workloads", &sczInstallStatePath);
+    ExitOnFailure(hr, "Failed to concat workloads to install state path.");
+
+    hr = PathConcat(sczInstallStatePath, sczPlatform, &sczInstallStatePath);
+    ExitOnFailure(hr, "Failed to concat platform (%ls) to install state path.", sczPlatform);
+
+    hr = PathConcat(sczInstallStatePath, sczSdkFeatureBandVersion, &sczInstallStatePath);
+    ExitOnFailure(hr, "Failed to concat feature band (%ls) to install state path.", sczSdkFeatureBandVersion);
+
+    hr = PathConcat(sczInstallStatePath, L"installstate", &sczInstallStatePath);
+    ExitOnFailure(hr, "Failed to concat installstate to install state path.");
+
+    hr = PathConcat(sczInstallStatePath, L"default.json", &sczInstallStatePath);
+    ExitOnFailure(hr, "Failed to concat default.json to install state path.");
+
+    if (FileExistsEx(sczInstallStatePath, NULL))
+    {
+        LogStringLine(REPORT_STANDARD, "Deleting install state file: %ls", sczInstallStatePath);
+        hr = FileEnsureDelete(sczInstallStatePath);
+        ExitOnFailure(hr, "Failed to delete install state file: %ls", sczInstallStatePath);
+
+        hr = PathGetParentPath(sczInstallStatePath, &sczPath);
+        ExitOnFailure(hr, "Failed to get parent path of install state file.");
+
+        LogStringLine(REPORT_STANDARD, "Cleaning up empty workload folders.");
+        DirDeleteEmptyDirectoriesToRoot(sczPath, 0);
+    }
+    else
+    {
+        LogStringLine(REPORT_STANDARD, "Install state file does not exist: %ls", sczInstallStatePath);
+    }
+
+LExit:
+    ReleaseStr(sczPath);
+    ReleaseStr(sczInstallStatePath)
+    ReleaseStr(sczProgramData);
+}
+
 int wmain(int argc, wchar_t* argv[])
 {
     HRESULT hr = S_OK;
     DWORD dwExitCode = 0;
     LPWSTR sczDependent = NULL;
     LPWSTR sczFeatureBandVersion = NULL;
+    LPWSTR sczPlatform = NULL;
     BOOL bRestartRequired = FALSE;
     BOOL bSdkFeatureBandInstalled = FALSE;
     int iMajor = 0;
@@ -507,16 +559,19 @@ int wmain(int argc, wchar_t* argv[])
     hr = ::Initialize(argc, argv);
     ExitOnFailure(hr, "Failed to initialize.");
 
+    hr = StrAllocString(&sczPlatform, argv[3], 0);
+    ExitOnFailure(hr, "Failed to copy platform argument.");
+
     // Convert the full SDK version to a feature band version
     hr = ParseSdkVersion(argv[2], &sczFeatureBandVersion);
     ExitOnFailure(hr, "Failed to parse version, %ls.", argv[2]);
 
     // Create the dependent value, e.g., Microsoft.NET.Sdk,6.0.300,arm64
-    hr = StrAllocFormatted(&sczDependent, L"Microsoft.NET.Sdk,%ls,%ls", sczFeatureBandVersion, argv[3]);
+    hr = StrAllocFormatted(&sczDependent, L"Microsoft.NET.Sdk,%ls,%ls", sczFeatureBandVersion, sczPlatform);
     ExitOnFailure(hr, "Failed to create dependent.");
     LogStringLine(REPORT_STANDARD, "Setting target dependent to %ls.", sczDependent);
 
-    hr = ::DetectSdk(sczFeatureBandVersion, argv[3], &bSdkFeatureBandInstalled);
+    hr = ::DetectSdk(sczFeatureBandVersion, sczPlatform, &bSdkFeatureBandInstalled);
     ExitOnFailure(hr, "Failed to detect installed SDKs.");
 
     // If the feature band is still present, do not remove workloads.
@@ -529,7 +584,7 @@ int wmain(int argc, wchar_t* argv[])
     hr = ::RemoveDependent(sczDependent, &bRestartRequired);
     ExitOnFailure(hr, "Failed to remove dependent \"%ls\".", sczDependent);
 
-    hr = ::DeleteWorkloadRecords(sczFeatureBandVersion, argv[3]);
+    hr = ::DeleteWorkloadRecords(sczFeatureBandVersion, sczPlatform);
     ExitOnFailure(hr, "Failed to remove workload records.");
 
     if (bRestartRequired)
@@ -537,9 +592,12 @@ int wmain(int argc, wchar_t* argv[])
         dwExitCode = ERROR_SUCCESS_REBOOT_REQUIRED;
     }
 
+    RemoveInstallStateFile(sczFeatureBandVersion, sczPlatform);
+
 LExit:
     ReleaseStr(sczDependent);
     ReleaseStr(sczFeatureBandVersion);
+    ReleaseStr(sczPlatform);
     LogUninitialize(TRUE);
     RegUninitialize();
     WiuUninitialize();
diff --git a/src/finalizer/precomp.h b/src/finalizer/precomp.h
index 1a6cd5df0..97bc60b13 100644
--- a/src/finalizer/precomp.h
+++ b/src/finalizer/precomp.h
@@ -12,6 +12,8 @@
 #include <winreg.h>
 #include <msi.h>
 #include <pathcch.h>
+#include <shlobj.h>
+#include <shlwapi.h>
 
 // Configure some logging parameters for WiX
 #define ExitTrace LogErrorString
@@ -26,3 +28,7 @@
 #include "pathutil.h"
 #include "strutil.h"
 #include "wiutil.h"
+#include "dirutil.h"
+#include "fileutil.h"
+#include "shelutil.h"
+