3 files changed, 119 insertions, 22 deletions
diff --git a/src/Mount/Mount.c b/src/Mount/Mount.c
index 30714560..7d4fa81b 100644
--- a/src/Mount/Mount.c
+++ b/src/Mount/Mount.c
@@ -9603,7 +9603,7 @@ static DWORD WINAPI SystemFavoritesServiceCtrlHandler (	DWORD dwControl,
 	case SERVICE_CONTROL_STOP:
 		SystemFavoritesServiceSetStatus (SERVICE_STOP_PENDING);
 
-		if (!(BootEncObj->ReadServiceConfigurationFlags () & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER))
+ 		if (!(BootEncObj->ReadServiceConfigurationFlags () & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER))
 		{
 			try
 			{
@@ -11818,6 +11818,17 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 				BOOL bPimCacheEnabled = (driverConfig & TC_DRIVER_CONFIG_CACHE_BOOT_PIM)? TRUE : FALSE;
 				BOOL bBlockSysEncTrimEnabled = (driverConfig & VC_DRIVER_CONFIG_BLOCK_SYS_TRIM)? TRUE : FALSE;
 				BOOL bClearKeysEnabled = (driverConfig & VC_DRIVER_CONFIG_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION)? TRUE : FALSE;
+				BOOL bAutoFixBootloader = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER)? FALSE : TRUE;
+				BOOL bForceVeraCryptNextBoot = FALSE;
+				BOOL bForceSetVeraCryptBootEntry = TRUE;
+				BOOL bForceVeraCryptFirstEntry = TRUE;
+				if (bSystemIsGPT)
+				{
+					bForceVeraCryptNextBoot = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_FORCE_SET_BOOTNEXT)? TRUE : FALSE;
+					bForceSetVeraCryptBootEntry = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_SET_BOOTENTRY)? FALSE : TRUE;
+					bForceVeraCryptFirstEntry = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_FORCE_FIRST_BOOTENTRY)? FALSE : TRUE;
+				}
+
 				BOOL bIsHiddenOS = IsHiddenOSRunning ();
 
 				if (bClearKeysEnabled)
@@ -11882,6 +11893,24 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 				}
 				else
 					CheckDlgButton (hwndDlg, IDC_BLOCK_SYSENC_TRIM, bBlockSysEncTrimEnabled ? BST_CHECKED : BST_UNCHECKED);
+
+				CheckDlgButton (hwndDlg, IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN, bAutoFixBootloader? BST_CHECKED : BST_UNCHECKED);
+				if (bSystemIsGPT)
+				{
+					if (!bAutoFixBootloader || bIsHiddenOS)
+					{
+						// we disable other options if updating bootloader is not allowed or if hidden OS us running
+						EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT), FALSE);
+						EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY), FALSE);
+						EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY), FALSE);
+					}
+					else
+					{
+						CheckDlgButton (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT, bForceVeraCryptNextBoot? BST_CHECKED : BST_UNCHECKED);
+						CheckDlgButton (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY, bForceSetVeraCryptBootEntry? BST_CHECKED : BST_UNCHECKED);
+						CheckDlgButton (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY, bForceVeraCryptFirstEntry? BST_CHECKED : BST_UNCHECKED);
+					}
+				}
 			}
 			catch (Exception &e)
 			{
@@ -11992,6 +12021,17 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 					BOOL bBlockSysEncTrimEnabled = IsDlgButtonChecked (hwndDlg, IDC_BLOCK_SYSENC_TRIM);
 					BOOL bClearKeysEnabled = IsDlgButtonChecked (hwndDlg, IDC_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION);
 
+					BOOL bAutoFixBootloader = IsDlgButtonChecked (hwndDlg, IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN);
+					BOOL bForceVeraCryptNextBoot = FALSE;
+					BOOL bForceSetVeraCryptBootEntry = TRUE;
+					BOOL bForceVeraCryptFirstEntry = TRUE;
+					if (bSystemIsGPT)
+					{
+						bForceVeraCryptNextBoot = IsDlgButtonChecked (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT);
+						bForceSetVeraCryptBootEntry = IsDlgButtonChecked (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY);
+						bForceVeraCryptFirstEntry = IsDlgButtonChecked (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY);
+					}
+
 					if (bClearKeysEnabled && !BootEncObj->IsSystemFavoritesServiceRunning())
 					{
 						// the system favorite service service should be running
@@ -12010,8 +12050,17 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 					SetDriverConfigurationFlag (TC_DRIVER_CONFIG_CACHE_BOOT_PIM, (bPasswordCacheEnabled && bPimCacheEnabled)? TRUE : FALSE);
 					SetDriverConfigurationFlag (TC_DRIVER_CONFIG_DISABLE_EVIL_MAID_ATTACK_DETECTION, IsDlgButtonChecked (hwndDlg, IDC_DISABLE_EVIL_MAID_ATTACK_DETECTION));
 					SetDriverConfigurationFlag (VC_DRIVER_CONFIG_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION, bClearKeysEnabled);
-					if (!IsHiddenOSRunning ()) /* we don't need to update TRIM config for hidden OS since it's always blocked */
+					SetDriverConfigurationFlag (VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_UPDATE_LOADER, bAutoFixBootloader? FALSE : TRUE);
+					if (bSystemIsGPT && !IsHiddenOSRunning ())
+					{
+						/* we don't need to update TRIM config for hidden OS since it's always blocked */
 						SetDriverConfigurationFlag (VC_DRIVER_CONFIG_BLOCK_SYS_TRIM, bBlockSysEncTrimEnabled);
+
+						/* we don't update bootloader settings since we never update bootloader under Hidden OS */
+						SetDriverConfigurationFlag (VC_SYSTEM_FAVORITES_SERVICE_CONFIG_FORCE_SET_BOOTNEXT, bForceVeraCryptNextBoot);
+						SetDriverConfigurationFlag (VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_SET_BOOTENTRY, bForceSetVeraCryptBootEntry? FALSE : TRUE);
+						SetDriverConfigurationFlag (VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_FORCE_FIRST_BOOTENTRY, bForceVeraCryptFirstEntry? FALSE : TRUE);
+					}
 				}
 				catch (Exception &e)
 				{
@@ -12072,6 +12121,39 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 			}
 
 			break;
+
+		case IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN:
+			if (bSystemIsGPT)
+			{
+				if (IsDlgButtonChecked (hwndDlg, IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN))
+				{
+					if (!IsHiddenOSRunning ())
+					{
+						uint32 driverConfig = ReadDriverConfigurationFlags();
+						BOOL bForceVeraCryptNextBoot = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_FORCE_SET_BOOTNEXT)? TRUE : FALSE;
+						BOOL bForceSetVeraCryptBootEntry = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_SET_BOOTENTRY)? FALSE : TRUE;
+						BOOL bForceVeraCryptFirstEntry = (driverConfig & VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_FORCE_FIRST_BOOTENTRY)? FALSE : TRUE;
+
+						EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT), TRUE);
+						EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY), TRUE);
+						EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY), TRUE);
+
+						CheckDlgButton (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT, bForceVeraCryptNextBoot? BST_CHECKED : BST_UNCHECKED);
+						CheckDlgButton (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY, bForceSetVeraCryptBootEntry? BST_CHECKED : BST_UNCHECKED);
+						CheckDlgButton (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY, bForceVeraCryptFirstEntry? BST_CHECKED : BST_UNCHECKED);
+					}
+				}
+				else
+				{
+					CheckDlgButton (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT, BST_UNCHECKED);
+					CheckDlgButton (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY, BST_UNCHECKED);
+					CheckDlgButton (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY, BST_UNCHECKED);
+					EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_NEXT_BOOT_VERACRYPT), FALSE);
+					EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_VERACRYPT_BOOT_ENTRY), FALSE);
+					EnableWindow (GetDlgItem (hwndDlg, IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY), FALSE);
+				}
+			}
+			break;
 		}
 		return 0;
 	}
diff --git a/src/Mount/Mount.rc b/src/Mount/Mount.rc
index cd4a3e13..f63afe5c 100644
--- a/src/Mount/Mount.rc
+++ b/src/Mount/Mount.rc
@@ -287,7 +287,7 @@ BEGIN
     LTEXT           "",IDT_PKCS11_LIB_HELP,16,63,286,65
 END
 
-IDD_EFI_SYSENC_SETTINGS DIALOGEX 0, 0, 375, 194
+IDD_EFI_SYSENC_SETTINGS DIALOGEX 0, 0, 375, 250
 STYLE DS_SETFONT | DS_MODALFRAME | DS_FIXEDSYS | DS_CENTER | WS_POPUP | WS_CAPTION | WS_SYSMENU
 CAPTION "VeraCrypt - System Encryption Settings"
 FONT 8, "MS Shell Dlg", 400, 0, 0x1
@@ -304,13 +304,21 @@ BEGIN
                     "Button",BS_AUTOCHECKBOX | WS_DISABLED | WS_TABSTOP,16,83,340,10
     CONTROL         "Block TRIM command on system partition/drive",IDC_BLOCK_SYSENC_TRIM,
                     "Button",BS_AUTOCHECKBOX | WS_TABSTOP,16,98,340,10
-    GROUPBOX        "Advanced Options",IDT_ADVANCED_OPTIONS,7,131,355,36
-    PUSHBUTTON      "Edit Boot Loader Configuration",IDC_EDIT_DCSPROP,10,144,173,14
-    PUSHBUTTON      "Display EFI Platform Information",IDC_SHOW_PLATFORMINFO,187,144,173,14
-    PUSHBUTTON      "Cancel",IDCANCEL,313,170,50,14
-    DEFPUSHBUTTON   "OK",IDOK,255,170,50,14
     CONTROL         "Clear encryption keys from memory if a new device is inserted",IDC_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION,
                     "Button",BS_AUTOCHECKBOX | WS_TABSTOP,16,112,340,10
+    GROUPBOX        "Advanced Options",IDT_ADVANCED_OPTIONS,7,131,355,91
+    CONTROL         "Automatically fix boot configuration issues that may prevent Windows from starting",IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN,
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,16,144,340,10
+    CONTROL         "Force machine to boot on VeraCrypt in the next startup",IDC_FORCE_NEXT_BOOT_VERACRYPT,
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,16,158,340,10
+    CONTROL         "Force the presence of VeraCrypt entry in the EFI firmware boot menu",IDC_FORCE_VERACRYPT_BOOT_ENTRY,
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,16,172,340,10
+    CONTROL         "Force VeraCrypt entry to be the first in the EFI firmware boot menu",IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY,
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,16,186,340,10
+    PUSHBUTTON      "Edit Boot Loader Configuration",IDC_EDIT_DCSPROP,10,201,173,14
+    PUSHBUTTON      "Display EFI Platform Information",IDC_SHOW_PLATFORMINFO,187,201,173,14
+    PUSHBUTTON      "Cancel",IDCANCEL,313,226,50,14
+    DEFPUSHBUTTON   "OK",IDOK,255,226,50,14
 END
 
 IDD_PERFORMANCE_SETTINGS DIALOGEX 0, 0, 371, 293
@@ -401,7 +409,7 @@ BEGIN
     CONTROL         "TrueCrypt Mode",IDC_TRUECRYPT_MODE,"Button",BS_AUTOCHECKBOX | WS_TABSTOP,7,7,76,10
 END
 
-IDD_SYSENC_SETTINGS DIALOGEX 0, 0, 371, 310
+IDD_SYSENC_SETTINGS DIALOGEX 0, 0, 371, 344
 STYLE DS_SETFONT | DS_MODALFRAME | DS_FIXEDSYS | DS_CENTER | WS_POPUP | WS_CAPTION | WS_SYSMENU
 CAPTION "VeraCrypt - System Encryption Settings"
 FONT 8, "MS Shell Dlg", 400, 0, 0x1
@@ -414,21 +422,24 @@ BEGIN
     CONTROL         "&Cache pre-boot authentication password in driver memory (for mounting of non-system volumes)",IDC_BOOT_LOADER_CACHE_PASSWORD,
                     "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,192,339,10
     CONTROL         "Include PIM when caching pre-boot authentication password",IDC_BOOT_LOADER_CACHE_PIM,
-                    "Button",BS_AUTOCHECKBOX | WS_DISABLED | WS_TABSTOP,18,207,340,10
+                    "Button",BS_AUTOCHECKBOX | WS_DISABLED | WS_TABSTOP,18,207,339,10
     CONTROL         "Allow pre-boot &authentication to be bypassed by pressing the Esc key (enables boot manager)",IDC_ALLOW_ESC_PBA_BYPASS,
-                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,222,340,10
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,222,339,10
     CONTROL         "Disable ""Evil Maid"" attack detection",IDC_DISABLE_EVIL_MAID_ATTACK_DETECTION,
-                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,237,340,10
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,237,339,10
     CONTROL         "Block TRIM command on system partition/drive",IDC_BLOCK_SYSENC_TRIM,
-                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,251,340,10
-    PUSHBUTTON      "Cancel",IDCANCEL,314,286,50,14
-    DEFPUSHBUTTON   "OK",IDOK,257,286,50,14
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,251,339,10
+    CONTROL         "Clear encryption keys from memory if a new device is inserted",IDC_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION,
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,265,339,10
+    CONTROL         "Automatically fix boot configuration issues that may prevent Windows from starting",IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN,
+                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,298,339,10
+    PUSHBUTTON      "Cancel",IDCANCEL,314,320,50,14
+    DEFPUSHBUTTON   "OK",IDOK,257,320,50,14
     LTEXT           "Display this custom message in the pre-boot authentication screen (24 characters maximum):",IDT_CUSTOM_BOOT_LOADER_MESSAGE,18,39,337,8
-    GROUPBOX        "Boot Loader Screen Options",IDT_BOOT_LOADER_SCREEN_OPTIONS,9,7,355,165
-    GROUPBOX        "Security Options",IDT_SECURITY_OPTIONS,9,177,355,105
     LTEXT           "",IDC_CUSTOM_BOOT_LOADER_MESSAGE_HELP,18,72,337,73
-    CONTROL         "Clear encryption keys from memory if a new device is inserted",IDC_CLEAR_KEYS_ON_NEW_DEVICE_INSERTION,
-                    "Button",BS_AUTOCHECKBOX | WS_TABSTOP,18,265,340,10
+    GROUPBOX        "Security Options",IDT_SECURITY_OPTIONS,9,177,355,105
+    GROUPBOX        "Boot Loader Screen Options",IDT_BOOT_LOADER_SCREEN_OPTIONS,9,7,355,165
+    GROUPBOX        "Advanced Options",IDT_ADVANCED_OPTIONS,9,285,355,29
 END
 
 /////////////////////////////////////////////////////////////////////////////
@@ -504,7 +515,7 @@ BEGIN
         LEFTMARGIN, 7
         RIGHTMARGIN, 368
         TOPMARGIN, 7
-        BOTTOMMARGIN, 184
+        BOTTOMMARGIN, 240
     END
 
     IDD_PERFORMANCE_SETTINGS, DIALOG
@@ -536,7 +547,7 @@ BEGIN
         LEFTMARGIN, 7
         RIGHTMARGIN, 364
         TOPMARGIN, 7
-        BOTTOMMARGIN, 300
+        BOTTOMMARGIN, 334
     END
 END
 #endif    // APSTUDIO_INVOKED
diff --git a/src/Mount/Resource.h b/src/Mount/Resource.h
index d102faf4..dc73d180 100644
--- a/src/Mount/Resource.h
+++ b/src/Mount/Resource.h
@@ -195,6 +195,10 @@
 #define IDC_ENABLE_CPU_RNG              1172
 #define IDC_ENABLE_RAM_ENCRYPTION       1173
 #define IDC_USE_LEGACY_MAX_PASSWORD_LENGTH 1174
+#define IDC_UPDATE_BOOTLOADER_ON_SHUTDOWN 1175
+#define IDC_FORCE_NEXT_BOOT_VERACRYPT   1176
+#define IDC_FORCE_VERACRYPT_BOOT_ENTRY  1177
+#define IDC_FORCE_VERACRYPT_FIRST_BOOT_ENTRY 1178
 #define IDM_HELP                        40001
 #define IDM_ABOUT                       40002
 #define IDM_UNMOUNT_VOLUME              40003
@@ -271,7 +275,7 @@
 #define _APS_NO_MFC                     1
 #define _APS_NEXT_RESOURCE_VALUE        120
 #define _APS_NEXT_COMMAND_VALUE         40069
-#define _APS_NEXT_CONTROL_VALUE         1175
+#define _APS_NEXT_CONTROL_VALUE         1179
 #define _APS_NEXT_SYMED_VALUE           101
 #endif
 #endif