From 3659ddd70ce8f1616a44e5bcce793eda076b835d Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Sun, 22 Apr 2018 16:14:53 +0200
Subject: Windows: enhance ReflectDrivers mechanism by persisting it across
 major Windows upgrades.

---
 src/Common/BootEncryption.cpp | 35 ++++++++++++++++++++++++++++++++---
 src/Common/BootEncryption.h   |  4 +++-
 src/Mount/Mount.c             |  7 +++++++
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/Common/BootEncryption.cpp b/src/Common/BootEncryption.cpp
index c62df958..9c69ef04 100644
--- a/src/Common/BootEncryption.cpp
+++ b/src/Common/BootEncryption.cpp
@@ -2722,8 +2722,37 @@ namespace VeraCrypt
 
 				StringCchCatW (szSetupconfigLocation, ARRAYSIZE (szSetupconfigLocation), L"SetupConfig.ini");
 
-				if (bForInstall || FileExists (szSetupconfigLocation))
-					WritePrivateProfileStringW (L"SetupConfig", L"ReflectDrivers", bForInstall? szInstallPath : NULL, szSetupconfigLocation);
+				if (bForInstall)
+				{
+					wstring szPathParam = L"\"";
+					szPathParam += szInstallPath;
+					szPathParam += L"\"";
+					WritePrivateProfileStringW (L"SetupConfig", L"ReflectDrivers", szPathParam.c_str(), szSetupconfigLocation);
+
+					szPathParam = GetProgramConfigPath (L"SetupComplete.cmd");
+					FILE* scriptFile = _wfopen (szPathParam.c_str(), L"w");
+					if (scriptFile)
+					{
+						fwprintf (scriptFile, L"\"%s\\VeraCrypt.exe\" /PostOOBE\n", szInstallPath);
+						fclose (scriptFile);
+
+						WritePrivateProfileStringW (L"SetupConfig", L"PostOOBE", szPathParam.c_str(), szSetupconfigLocation);
+					}
+				}
+				else
+				{
+					if (FileExists (szSetupconfigLocation))
+					{
+						WritePrivateProfileStringW (L"SetupConfig", L"ReflectDrivers", NULL, szSetupconfigLocation);
+						WritePrivateProfileStringW (L"SetupConfig", L"PostOOBE", NULL, szSetupconfigLocation);
+					}
+
+					wstring scriptFilePath = GetProgramConfigPath (L"SetupComplete.cmd");
+					if (FileExists (scriptFilePath.c_str()))
+					{
+						::DeleteFileW (scriptFilePath.c_str());
+					}
+				}
 			}
 		}
 	}
@@ -3717,7 +3746,7 @@ namespace VeraCrypt
 
 			finally_do ({ EfiBootInst.DismountBootPartition(); });
 
-			EfiBootInst.MountBootPartition(0);			
+			EfiBootInst.MountBootPartition(0);		
 
 			EfiBootInst.GetFileSize(Is64BitOs()? L"\\EFI\\Boot\\bootx64.efi" : L"\\EFI\\Boot\\bootia32.efi", loaderSize);
 
diff --git a/src/Common/BootEncryption.h b/src/Common/BootEncryption.h
index 6eb42b50..63ebe353 100644
--- a/src/Common/BootEncryption.h
+++ b/src/Common/BootEncryption.h
@@ -310,7 +310,7 @@ namespace VeraCrypt
 		void GetEfiBootDeviceNumber (PSTORAGE_DEVICE_NUMBER pSdn);
 		void BackupSystemLoader ();
 		void RestoreSystemLoader ();
-		void UpdateSetupConfigFile (bool bForInstall);
+		static void UpdateSetupConfigFile (bool bForInstall);
 	protected:
 		static const uint32 RescueIsoImageSize = 1835008; // Size of ISO9660 image with bootable emulated 1.44MB floppy disk image
 
@@ -354,4 +354,6 @@ namespace VeraCrypt
 #define	TC_SYSTEM_FAVORITES_SERVICE_LOAD_ORDER_GROUP	L"Event Log"
 #define TC_SYSTEM_FAVORITES_SERVICE_CMDLINE_OPTION		L"/systemFavoritesService"
 
+#define VC_WINDOWS_UPGRADE_POSTOOBE_CMDLINE_OPTION		L"/PostOOBE"
+
 #endif // TC_HEADER_Common_BootEncryption
diff --git a/src/Mount/Mount.c b/src/Mount/Mount.c
index bbf956f0..8965984f 100644
--- a/src/Mount/Mount.c
+++ b/src/Mount/Mount.c
@@ -9383,6 +9383,13 @@ int WINAPI wWinMain (HINSTANCE hInstance, HINSTANCE hPrevInstance, wchar_t *lpsz
 	if (argv && argc == 2 && wstring (TC_SYSTEM_FAVORITES_SERVICE_CMDLINE_OPTION) == argv[1])
 		return StartSystemFavoritesService() ? 0 : 1;
 
+	if (argv && argc == 2 && wstring (VC_WINDOWS_UPGRADE_POSTOOBE_CMDLINE_OPTION) == argv[1])
+	{
+		InitOSVersionInfo();
+		BootEncryption::UpdateSetupConfigFile (true);
+		return 0;
+	}
+
 	int status;
 	atexit (localcleanup);
 	SetProcessShutdownParameters (0x100, 0);
-- 
cgit v1.2.3