1 files changed, 38 insertions, 21 deletions
diff --git a/src/Common/BootEncryption.h b/src/Common/BootEncryption.h
index f63ec541..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;
@@ -40,6 +34,7 @@ namespace VeraCrypt
 		File (wstring path,bool readOnly = false, bool create = false);
 		virtual ~File () { Close(); }
 
+		bool IsOpened () const { return FileOpen;}
 		void CheckOpened (const char* srcPos) { if (!FileOpen) { SetLastError (LastError); throw SystemException (srcPos);} }
 		void Close ();
 		DWORD Read (byte *buffer, DWORD size);
@@ -176,6 +171,7 @@ namespace VeraCrypt
 		int authorizeRetry;
 		int bmlLockFlags;
 		int bmlDriverEnabled;
+		string actionSuccessValue;
 
 		EfiBootConf();
 
@@ -187,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);
 
@@ -197,42 +195,42 @@ namespace VeraCrypt
 	public:
 		EfiBoot();
 
-		void MountBootPartition(WCHAR letter);
-		void DismountBootPartition();
+		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, wchar_t* nameNew, BOOL bForce);
+		BOOL RenameFile(const wchar_t* name, const wchar_t* nameNew, BOOL bForce);
 		BOOL DelFile(const wchar_t* name);
 		BOOL MkDir(const wchar_t* name, bool& bAlreadyExists);
 		BOOL ReadConfig (const wchar_t* name, EfiBootConf& conf);
 		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();
-		void SelectBootVolume(WCHAR* bootVolumePath);
-		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;
-		WCHAR	EfiBootPartPath[3];
+		std::wstring	EfiBootPartPath;
 		STORAGE_DEVICE_NUMBER sdn;
 		PARTITION_INFORMATION_EX partInfo;
+		bool bDeviceInfoValid;
 		WCHAR     tempBuf[1024];
-		bool  bBootVolumePathSelected;
-		WCHAR BootVolumePath[MAX_PATH];
+		std::wstring BootVolumePath;
 	};
 
 	class BootEncryption
 	{
 	public:
-		BootEncryption (HWND parent);
+		BootEncryption (HWND parent, bool postOOBE = false, bool setBootEntry = true, bool forceFirstBootEntry = true, bool setBootNext = false);
 		~BootEncryption ();
 
 		enum FilterType
@@ -279,16 +277,19 @@ namespace VeraCrypt
 		void ProbeRealSystemDriveSize ();
 		bool ReadBootSectorConfig (byte *config, size_t bufLength, byte *userConfig = nullptr, string *customUserMessage = nullptr, uint16 *bootLoaderVersion = nullptr);
 		uint32 ReadDriverConfigurationFlags ();
+		uint32 ReadServiceConfigurationFlags ();
 		void RegisterBootDriver (bool hiddenSystem);
 		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);
 		void InitialSecurityChecksForHiddenOS ();
 		void RestrictPagingFilesToSystemPartition ();
 		void SetDriverConfigurationFlag (uint32 flag, bool state);
+		void SetServiceConfigurationFlag (uint32 flag, bool state);
 		void SetDriverServiceStartType (DWORD startType);
 		void SetHiddenOSCreationPhase (unsigned int newPhase);
 		void StartDecryption (BOOL discardUnreadableEncryptedSectors);
@@ -310,7 +311,10 @@ namespace VeraCrypt
 		void GetEfiBootDeviceNumber (PSTORAGE_DEVICE_NUMBER pSdn);
 		void BackupSystemLoader ();
 		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
 
@@ -339,6 +343,10 @@ namespace VeraCrypt
 		bool RealSystemDriveSizeValid;
 		bool RescueVolumeHeaderValid;
 		bool VolumeHeaderValid;
+		bool PostOOBEMode;
+		bool SetBootNext;
+		bool SetBootEntry;
+		bool ForceFirstBootEntry;
 	};
 }
 
@@ -353,5 +361,14 @@ namespace VeraCrypt
 #define TC_SYSTEM_FAVORITES_SERVICE_NAME				_T(TC_APP_NAME) L"SystemFavorites"
 #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"
 
 #endif // TC_HEADER_Common_BootEncryption