4 files changed, 80 insertions, 30 deletions
diff --git a/src/Common/BootEncryption.h b/src/Common/BootEncryption.h
index 7f5c3b16..e25427dd 100644
--- a/src/Common/BootEncryption.h
+++ b/src/Common/BootEncryption.h
@@ -360,6 +360,7 @@ namespace VeraCrypt
 #define	TC_SYSTEM_FAVORITES_SERVICE_LOAD_ORDER_GROUP	L"Event Log"
 #define TC_SYSTEM_FAVORITES_SERVICE_CMDLINE_OPTION		L"/systemFavoritesService"
 #define VC_SYSTEM_FAVORITES_SERVICE_ARG_SKIP_MOUNT		L"/SkipMount"
+#define VC_SYSTEM_FAVORITES_SERVICE_ARG_UPDATE_LOADER	L"/UpdateLoader"
 
 #define VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER			0x1
 #define VC_SYSTEM_FAVORITES_SERVICE_CONFIG_FORCE_SET_BOOTNEXT			0x2
diff --git a/src/Common/Dlgcode.c b/src/Common/Dlgcode.c
index 7f2218e4..683e5c43 100644
--- a/src/Common/Dlgcode.c
+++ b/src/Common/Dlgcode.c
@@ -175,6 +175,10 @@ BOOL bLanguageSetInSetup = FALSE;
 extern BOOL bMakePackage;
 #endif
 
+#ifdef TCMOUNT
+extern BOOL ServiceMode;
+#endif
+
 // Status of detection of hidden sectors (whole-system-drive encryption). 
 // 0 - Unknown/undetermined/completed, 1: Detection is or was in progress (but did not complete e.g. due to system crash).
 int HiddenSectorDetectionStatus = 0;	
@@ -4759,10 +4763,16 @@ load:
 			if (IsNonInstallMode () && CreateDriverSetupMutex () && DriverUnload () && nLoadRetryCount++ < 3)
 				goto load;
 
-			CloseDriverSetupMutex ();
-			CloseHandle (hDriver);
-			hDriver = INVALID_HANDLE_VALUE;
-			return ERR_DRIVER_VERSION;
+#ifdef TCMOUNT
+			// don't fail in case of service. This solves issues during upgrade when system encryption is enabled
+			if (!ServiceMode)
+#endif
+			{
+				CloseDriverSetupMutex ();
+				CloseHandle (hDriver);
+				hDriver = INVALID_HANDLE_VALUE;
+				return ERR_DRIVER_VERSION;
+			}
 		}
 #else
 		if (!bResult)
diff --git a/src/Mount/Mount.c b/src/Mount/Mount.c
index 8f7261b9..e3aca2ff 100644
--- a/src/Mount/Mount.c
+++ b/src/Mount/Mount.c
@@ -9663,28 +9663,39 @@ static void SystemFavoritesServiceSetStatus (DWORD status, DWORD waitHint = 0)
 	SetServiceStatus (SystemFavoritesServiceStatusHandle, &SystemFavoritesServiceStatus);
 }
 
-static void SystemFavoritesServiceUpdateLoaderProcessing ()
+static void SystemFavoritesServiceUpdateLoaderProcessing (BOOL bForce)
 {
-	if (!(BootEncObj->ReadServiceConfigurationFlags () & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER))
+	SystemFavoritesServiceLogInfo (L"SystemFavoritesServiceUpdateLoaderProcessing called");
+	if (bForce || !(BootEncObj->ReadServiceConfigurationFlags () & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER))
 	{
+		SystemFavoritesServiceLogInfo (L"SystemFavoritesServiceUpdateLoaderProcessing processing");
 		try
 		{
 			BootEncryption::UpdateSetupConfigFile (true);
+			SystemFavoritesServiceLogInfo (L"SystemFavoritesServiceUpdateLoaderProcessing: UpdateSetupConfigFile called");
 			if (!BootEncStatus.HiddenSystem)
 			{
 				// re-install our bootloader again in case the update process has removed it.
 				bool bForceSetNextBoot = false;
 				bool bSetBootentry = true;
 				bool bForceFirstBootEntry = true;
-				uint32 flags = BootEncObj->ReadServiceConfigurationFlags ();
-				if (flags & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_FORCE_SET_BOOTNEXT)
-					bForceSetNextBoot = true;
-				if (flags & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_SET_BOOTENTRY)
-					bSetBootentry = false;
-				if (flags & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_FORCE_FIRST_BOOTENTRY)
-					bForceFirstBootEntry = false;
-				BootEncryption bootEnc (NULL, true, bSetBootentry, bForceFirstBootEntry, bForceSetNextBoot);
+				bool bPostOOBE = true;
+				if (bForce)
+					bPostOOBE = false;
+				else
+				{
+					uint32 flags = BootEncObj->ReadServiceConfigurationFlags ();
+					if (flags & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_FORCE_SET_BOOTNEXT)
+						bForceSetNextBoot = true;
+					if (flags & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_SET_BOOTENTRY)
+						bSetBootentry = false;
+					if (flags & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_FORCE_FIRST_BOOTENTRY)
+						bForceFirstBootEntry = false;
+				}
+				BootEncryption bootEnc (NULL, bPostOOBE, bSetBootentry, bForceFirstBootEntry, bForceSetNextBoot);
+				SystemFavoritesServiceLogInfo (L"SystemFavoritesServiceUpdateLoaderProcessing: InstallBootLoader calling");
 				bootEnc.InstallBootLoader (true);
+				SystemFavoritesServiceLogInfo (L"SystemFavoritesServiceUpdateLoaderProcessing: InstallBootLoader called");
 			}
 		}
 		catch (...)
@@ -9704,7 +9715,7 @@ static DWORD WINAPI SystemFavoritesServiceCtrlHandler (	DWORD dwControl,
 	case SERVICE_CONTROL_STOP:
 		SystemFavoritesServiceSetStatus (SERVICE_STOP_PENDING);
 
-		SystemFavoritesServiceUpdateLoaderProcessing ();
+		SystemFavoritesServiceUpdateLoaderProcessing (FALSE);
 
 		/* clear VC_DRIVER_CONFIG_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION flag */
 		SetDriverConfigurationFlag (VC_DRIVER_CONFIG_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION, FALSE);
@@ -9717,7 +9728,7 @@ static DWORD WINAPI SystemFavoritesServiceCtrlHandler (	DWORD dwControl,
 			/* perform fixing of bootloader and SetupConfig.ini when the system resumes from sleep */
 			if (dwEventType == PBT_APMRESUMEAUTOMATIC)
 			{
-				SystemFavoritesServiceUpdateLoaderProcessing ();
+				SystemFavoritesServiceUpdateLoaderProcessing (FALSE);
 			}
 		}
 		break;
@@ -9726,7 +9737,7 @@ static DWORD WINAPI SystemFavoritesServiceCtrlHandler (	DWORD dwControl,
 			/* perform fixing of bootloader and SetupConfig.ini when the user logs in or when he unlocks his locked session */
 			if ((dwEventType == WTS_SESSION_UNLOCK) || (dwEventType == WTS_SESSION_LOGON))
 			{
-				SystemFavoritesServiceUpdateLoaderProcessing ();
+				SystemFavoritesServiceUpdateLoaderProcessing (FALSE);
 			}
 		}
 		break;
@@ -9791,14 +9802,21 @@ static VOID WINAPI SystemFavoritesServiceMain (DWORD argc, LPTSTR *argv)
 	BOOL status = FALSE;
 	DEV_BROADCAST_DEVICEINTERFACE hdr;
 	BOOL bSkipMount = FALSE;
+	BOOL bUpdateLoader = FALSE;
+	DWORD i;
 	memset (&SystemFavoritesServiceStatus, 0, sizeof (SystemFavoritesServiceStatus));
 	SystemFavoritesServiceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
 	SystemFavoritesServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
 	if (IsOSAtLeast (WIN_VISTA))
 		SystemFavoritesServiceStatus.dwControlsAccepted |= SERVICE_ACCEPT_PRESHUTDOWN | SERVICE_ACCEPT_SESSIONCHANGE | SERVICE_ACCEPT_POWEREVENT;
 
-	if ((argc >= 2) && (0 == _wcsicmp (argv[1], VC_SYSTEM_FAVORITES_SERVICE_ARG_SKIP_MOUNT)))
-		bSkipMount = TRUE;
+	for (i = 1; i < argc; i++)
+	{
+		if (0 == _wcsicmp (argv[i], VC_SYSTEM_FAVORITES_SERVICE_ARG_SKIP_MOUNT))
+			bSkipMount = TRUE;
+		else if (0 == _wcsicmp (argv[i], VC_SYSTEM_FAVORITES_SERVICE_ARG_UPDATE_LOADER))
+			bUpdateLoader = TRUE;
+	}
 
 	ZeroMemory (&hdr, sizeof(hdr));
 	hdr.dbcc_size = sizeof (hdr);
@@ -9817,6 +9835,12 @@ static VOID WINAPI SystemFavoritesServiceMain (DWORD argc, LPTSTR *argv)
 	SetUnhandledExceptionFilter (SystemFavoritesServiceExceptionHandler);
 	_set_invalid_parameter_handler (SystemFavoritesServiceInvalidParameterHandler);
 
+	if (bUpdateLoader)
+	{
+		SystemFavoritesServiceSetStatus (SERVICE_START_PENDING, 120000);
+		SystemFavoritesServiceUpdateLoaderProcessing (TRUE);
+	}
+
 	if (!bSkipMount)
 	{
 		InitGlobalLocks ();
@@ -9846,7 +9870,7 @@ static VOID WINAPI SystemFavoritesServiceMain (DWORD argc, LPTSTR *argv)
 		if (!(ReadDriverConfigurationFlags() & TC_DRIVER_CONFIG_CACHE_BOOT_PASSWORD))
 			WipeCache (NULL, TRUE);
 
-		SystemFavoritesServiceUpdateLoaderProcessing ();
+		SystemFavoritesServiceUpdateLoaderProcessing (FALSE);
 	}
 
 	SystemFavoritesServiceSetStatus (SERVICE_RUNNING);
diff --git a/src/SetupDLL/Setup.c b/src/SetupDLL/Setup.c
index 4aec04f9..f262a914 100644
--- a/src/SetupDLL/Setup.c
+++ b/src/SetupDLL/Setup.c
@@ -1360,7 +1360,10 @@ BOOL StartStopService_Dll (MSIHANDLE hInstaller, wchar_t *lpszService, BOOL bSta
 	if (bStart)
 	{
 		if (!StartService (hService, argc, argv) && (GetLastError () != ERROR_SERVICE_ALREADY_RUNNING))
+		{
+			MSILog(hInstaller, MSI_ERROR_LEVEL, L"Failed to start %s. Error 0x%.8X", lpszService, GetLastError ());
 			goto error;
+		}
 	}
 	else
 		ControlService (hService, SERVICE_CONTROL_STOP, &status);
@@ -1379,10 +1382,16 @@ BOOL StartStopService_Dll (MSIHANDLE hInstaller, wchar_t *lpszService, BOOL bSta
 
 	bRet = QueryServiceStatus (hService, &status);
 	if (bRet != TRUE)
+	{
+		MSILog(hInstaller, MSI_ERROR_LEVEL, L"Failed to query status of %s. Error 0x%.8X", lpszService, GetLastError ());
 		goto error;
+	}
 
 	if (status.dwCurrentState != dwExpectedState)
+	{
+		MSILog(hInstaller, MSI_ERROR_LEVEL, L"Current state of %s (0x%.8X) is different from expected one (0x%.8X).", lpszService, status.dwCurrentState, dwExpectedState);
 		goto error;
+	}
 
 	bOK = TRUE;
 
@@ -1511,11 +1520,6 @@ BOOL DoDriverUnload_Dll (MSIHANDLE hInstaller, HWND hwnd)
 			BootEncryption bootEnc (hwnd);
 			if (bootEnc.GetDriverServiceStartType() == SERVICE_BOOT_START)
 			{
-				MSILogAndShow(hInstaller, MSI_ERROR_LEVEL, L"VeraCrypt MSI installation is currently not compatible with Windows System Encryption");
-				bOK = FALSE;
-				goto end;
-
-				/************* TODO: find a workaround to the fact that MSI service cannot access EFI environment variable
 				try
 				{
 					// Check hidden OS update consistency
@@ -1559,7 +1563,6 @@ BOOL DoDriverUnload_Dll (MSIHANDLE hInstaller, HWND hwnd)
 					SystemEncryptionUpdate = TRUE;
 					PortableMode = FALSE;
 				}
-				*****************************************/
 			}
 		}
 		catch (...)	{ }
@@ -1612,6 +1615,12 @@ BOOL DoDriverUnload_Dll (MSIHANDLE hInstaller, HWND hwnd)
 
 			if (TCWindowClosed)
 				Sleep (2000);
+
+			// stop service
+			if (SystemEncryptionUpdate)
+			{
+				StartStopService_Dll (hInstaller, TC_SYSTEM_FAVORITES_SERVICE_NAME, FALSE, 0, NULL);
+			}
 		}
 
 		// Test for any applications attached to driver
@@ -1928,7 +1937,8 @@ BOOL UpgradeBootLoader_Dll (MSIHANDLE hInstaller, HWND hwndDlg)
 		{
 			MSILog (hInstaller, MSI_INFO_LEVEL, GetString("INSTALLER_UPDATING_BOOT_LOADER"));
 
-			bootEnc.InstallBootLoader (true);
+			// this is done by the service now
+			//bootEnc.InstallBootLoader (true);
 
 			if (bootEnc.GetInstalledBootLoaderVersion() <= TC_RESCUE_DISK_UPGRADE_NOTICE_MAX_VERSION)
 			{
@@ -2776,15 +2786,20 @@ EXTERN_C UINT STDAPICALLTYPE VC_CustomAction_PostInstall(MSIHANDLE hInstaller)
 
 					if (StartStopService_Dll (hInstaller, TC_SYSTEM_FAVORITES_SERVICE_NAME, FALSE, 0, NULL))
 					{
-						// we tell the service not to load system favorites on startup
-						LPCWSTR szArgs[2] = { TC_SYSTEM_FAVORITES_SERVICE_NAME, VC_SYSTEM_FAVORITES_SERVICE_ARG_SKIP_MOUNT};
+						// we tell the service not to load system favorites on startup and to update bootloader on startup
+						LPCWSTR szArgs[3] = { TC_SYSTEM_FAVORITES_SERVICE_NAME, VC_SYSTEM_FAVORITES_SERVICE_ARG_SKIP_MOUNT, VC_SYSTEM_FAVORITES_SERVICE_ARG_UPDATE_LOADER};
 						if (!CopyFile (szTmp, servicePath.c_str(), FALSE))
 							ForceCopyFile (szTmp, servicePath.c_str());
 
-						StartStopService_Dll (hInstaller, TC_SYSTEM_FAVORITES_SERVICE_NAME, TRUE, 2, szArgs);
+						MSILog(hInstaller, MSI_ERROR_LEVEL, L"VC_CustomAction_PostInstall: SystemEncryptionUpdate = %s", SystemEncryptionUpdate? L"TRUE" : L"FALSE");
+
+						StartStopService_Dll (hInstaller, TC_SYSTEM_FAVORITES_SERVICE_NAME, TRUE, SystemEncryptionUpdate? 3 : 2, szArgs);
 					}
 					else
+					{
+						MSILog(hInstaller, MSI_ERROR_LEVEL, L"VC_CustomAction_PostInstall: failed to stop %S", servicePath.c_str());
 						ForceCopyFile (szTmp, servicePath.c_str());
+					}
 
 					BootEncObj.SetDriverConfigurationFlag (driverFlags, true);