1 files changed, 22 insertions, 15 deletions
diff --git a/src/Common/BootEncryption.h b/src/Common/BootEncryption.h
index 0b5fe4f0..ddf6f3e5 100644
--- a/src/Common/BootEncryption.h
+++ b/src/Common/BootEncryption.h
@@ -18,15 +18,9 @@
 #include "Exception.h"
 #include "Platform/PlatformBase.h"
 #include "Volumes.h"
-#include <Winternl.h>
 
-#define SYSPARTITIONINFORMATION 0x62
-
-typedef NTSTATUS (WINAPI *NtQuerySystemInformationFn)(
-		SYSTEM_INFORMATION_CLASS SystemInformationClass,
-		PVOID                    SystemInformation,
-      ULONG                    SystemInformationLength,
-		PULONG                   ReturnLength
+typedef ULONG (WINAPI *RtlNtStatusToDosErrorFn)(
+  NTSTATUS Status
 );
 
 using namespace std;
@@ -189,9 +183,11 @@ namespace VeraCrypt
 		BOOL Load (const wchar_t* fileName);
 		void Load (char* configContent);
 		BOOL Save (const wchar_t* fileName, HWND hwnd);
+		static BOOL IsPostExecFileField (const string& szFieldValue, string& filePath);
+		static BOOL IsPostExecFileField (const string& szFieldValue, wstring& filePath);
 	};
 
-	void GetVolumeESP(wstring& path);
+	void GetVolumeESP(wstring& path, wstring& bootVolumePath);
 	std::string ReadESPFile (LPCWSTR szFilePath, bool bSkipUTF8BOM);
 	void WriteESPFile (LPCWSTR szFilePath, LPBYTE pbData, DWORD dwDataLen, bool bAddUTF8BOM);
 
@@ -199,16 +195,18 @@ namespace VeraCrypt
 	public:
 		EfiBoot();
 
-		void PrepareBootPartition();
+		void PrepareBootPartition(bool bDisableException = false);
 		bool IsEfiBoot();
 
 		void DeleteStartExec(uint16 statrtOrderNum = 0xDC5B, wchar_t* type = NULL);
-		void SetStartExec(wstring description, wstring execPath, uint16 statrtOrderNum = 0xDC5B, wchar_t* type = NULL, uint32 attr = 1);
+		void SetStartExec(wstring description, wstring execPath, bool setBootEntry = true, bool forceFirstBootEntry = true, bool setBootNext = true, uint16 statrtOrderNum = 0xDC5B, wchar_t* type = NULL, uint32 attr = 1);
 		void SaveFile(const wchar_t* name, byte* data, DWORD size);
 		void GetFileSize(const wchar_t* name, unsigned __int64& size);
 		void ReadFile(const wchar_t* name, byte* data, DWORD size);
 		void CopyFile(const wchar_t* name, const wchar_t* targetName);
 		bool FileExists(const wchar_t* name);
+		static bool CompareFiles (const wchar_t* fileName1, const wchar_t* fileName2);
+		static bool CompareFileData (const wchar_t* fileName, const byte* data, DWORD size);
 
 		BOOL RenameFile(const wchar_t* name, const wchar_t* nameNew, BOOL bForce);
 		BOOL DelFile(const wchar_t* name);
@@ -217,23 +215,22 @@ namespace VeraCrypt
 		BOOL UpdateConfig (const wchar_t* name, int pim, int hashAlgo, HWND hwndDlg);
 		BOOL WriteConfig (const wchar_t* name, bool preserveUserConfig, int pim, int hashAlgo, const char* passPromptMsg, HWND hwndDlg);
 		BOOL DelDir(const wchar_t* name);
-		void SelectBootVolumeESP();
-		PSTORAGE_DEVICE_NUMBER GetStorageDeviceNumber () { return &sdn;}
+		PSTORAGE_DEVICE_NUMBER GetStorageDeviceNumber () { if (bDeviceInfoValid) return &sdn; else { SetLastError (ERROR_INVALID_DRIVE); throw SystemException(SRC_POS);}}
 
 	protected:
 		bool m_bMounted;
 		std::wstring	EfiBootPartPath;
 		STORAGE_DEVICE_NUMBER sdn;
 		PARTITION_INFORMATION_EX partInfo;
+		bool bDeviceInfoValid;
 		WCHAR     tempBuf[1024];
-		bool  bBootVolumePathSelected;
 		std::wstring BootVolumePath;
 	};
 
 	class BootEncryption
 	{
 	public:
-		BootEncryption (HWND parent, bool postOOBE = false);
+		BootEncryption (HWND parent, bool postOOBE = false, bool setBootEntry = true, bool forceFirstBootEntry = true, bool setBootNext = false);
 		~BootEncryption ();
 
 		enum FilterType
@@ -285,6 +282,7 @@ namespace VeraCrypt
 		void RegisterFilterDriver (bool registerDriver, FilterType filterType);
 		void RegisterSystemFavoritesService (BOOL registerService);
 		void RegisterSystemFavoritesService (BOOL registerService, BOOL noFileHandling);
+		bool IsSystemFavoritesServiceRunning ();
 		void UpdateSystemFavoritesService ();
 		void RenameDeprecatedSystemLoaderBackup ();
 		bool RestartComputer (BOOL bShutdown = FALSE);
@@ -315,6 +313,8 @@ namespace VeraCrypt
 		void RestoreSystemLoader ();
 		static void UpdateSetupConfigFile (bool bForInstall);
 		void GetSecureBootConfig (BOOL* pSecureBootEnabled, BOOL *pVeraCryptKeysLoaded);
+		bool IsUsingUnsupportedAlgorithm(LONG driverVersion);
+		void NotifyService (DWORD dwNotifyCmd);
 	protected:
 		static const uint32 RescueIsoImageSize = 1835008; // Size of ISO9660 image with bootable emulated 1.44MB floppy disk image
 
@@ -344,6 +344,9 @@ namespace VeraCrypt
 		bool RescueVolumeHeaderValid;
 		bool VolumeHeaderValid;
 		bool PostOOBEMode;
+		bool SetBootNext;
+		bool SetBootEntry;
+		bool ForceFirstBootEntry;
 	};
 }
 
@@ -359,8 +362,12 @@ 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
+#define VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_SET_BOOTENTRY			0x4
+#define VC_SYSTEM_FAVORITES_SERVICE_CONFIG_DONT_FORCE_FIRST_BOOTENTRY	0x8
 
 #define VC_WINDOWS_UPGRADE_POSTOOBE_CMDLINE_OPTION		L"/PostOOBE"