From 2722b465302626c2231f2d76e927bce3a9ccda4f Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Sun, 27 Oct 2019 01:39:37 +0200
Subject: Windows: code refactoring and convert NTSTATUS error code to WIN32
 equivalent before displaying error message.

---
 src/Common/BootEncryption.cpp | 59 +++++++++++++++----------------------------
 src/Common/BootEncryption.h   |  8 +++---
 2 files changed, 26 insertions(+), 41 deletions(-)

diff --git a/src/Common/BootEncryption.cpp b/src/Common/BootEncryption.cpp
index b047089c..5c72f9f5 100644
--- a/src/Common/BootEncryption.cpp
+++ b/src/Common/BootEncryption.cpp
@@ -2478,9 +2478,10 @@ namespace VeraCrypt
 	static const wchar_t*	EfiVarGuid = L"{8BE4DF61-93CA-11D2-AA0D-00E098032B8C}";
 
 	void 
-	GetVolumeESP(wstring& path) 
+	GetVolumeESP(wstring& path, wstring& bootVolumePath) 
 	{
 		static wstring g_EspPath;
+		static wstring g_BootVolumePath;
 		static bool g_EspPathInitialized = false;
 
 		if (!g_EspPathInitialized)
@@ -2501,17 +2502,29 @@ namespace VeraCrypt
 			res = NtQuerySystemInformationPtr((SYSTEM_INFORMATION_CLASS)SYSPARTITIONINFORMATION, tempBuf, sizeof(tempBuf), &len);
 			if (res != S_OK)
 			{
+				/* try to convert the returned NTSTATUS to a WIN32 system error using RtlNtStatusToDosError */
+				RtlNtStatusToDosErrorFn RtlNtStatusToDosErrorPtr = (RtlNtStatusToDosErrorFn) GetProcAddress (GetModuleHandle (L"ntdll.dll"), "RtlNtStatusToDosError");
+				if (RtlNtStatusToDosErrorPtr)
+				{
+					ULONG win32err = RtlNtStatusToDosErrorPtr (res);
+					if (win32err != ERROR_MR_MID_NOT_FOUND)
+						res = (NTSTATUS) win32err;
+				}
+
 				SetLastError (res);
 				throw SystemException (SRC_POS);
 			}		
 
 			PUNICODE_STRING pStr = (PUNICODE_STRING) tempBuf;
+
+			g_BootVolumePath = pStr->Buffer;
 			g_EspPath = L"\\\\?";
 			g_EspPath += &pStr->Buffer[7];
 			g_EspPathInitialized = true;
 		}
 
 		path = g_EspPath;
+		bootVolumePath = g_BootVolumePath;
 	}
 
 	std::string ReadESPFile (LPCWSTR szFilePath, bool bSkipUTF8BOM)
@@ -2521,9 +2534,9 @@ namespace VeraCrypt
 
 		ByteArray fileContent;
 		DWORD dwSize = 0, dwOffset = 0;
-		std::wstring pathESP;
+		std::wstring pathESP, bootVolumePath;
 
-		GetVolumeESP(pathESP);
+		GetVolumeESP(pathESP, bootVolumePath);
 		if (szFilePath[0] != L'\\')
 			pathESP += L"\\";
 		File f(pathESP + szFilePath, true);
@@ -2552,7 +2565,7 @@ namespace VeraCrypt
 
 		ByteArray fileContent;
 		DWORD dwSize = dwDataLen, dwOffset = 0;
-		std::wstring pathESP;
+		std::wstring pathESP, bootVolumePath;
 
 		if (bAddUTF8BOM)
 		{
@@ -2560,7 +2573,7 @@ namespace VeraCrypt
 			dwOffset = 3;
 		}
 
-		GetVolumeESP(pathESP);
+		GetVolumeESP(pathESP, bootVolumePath);
 		if (szFilePath[0] != L'\\')
 			pathESP += L"\\";
 
@@ -2580,42 +2593,12 @@ namespace VeraCrypt
 		ZeroMemory (&partInfo, sizeof (partInfo));
 		m_bMounted = false;
 		bDeviceInfoValid = false;
-		bBootVolumePathSelected = false;
-	}
-
-	void EfiBoot::SelectBootVolumeESP() {
-		NTSTATUS res;
-		ULONG    len;
-		memset(tempBuf, 0, sizeof(tempBuf));
-
-		// Load NtQuerySystemInformation function point
-		if (!NtQuerySystemInformationPtr)
-		{
-			NtQuerySystemInformationPtr = (NtQuerySystemInformationFn) GetProcAddress (GetModuleHandle (L"ntdll.dll"), "NtQuerySystemInformation");
-			if (!NtQuerySystemInformationPtr)
-				throw SystemException (SRC_POS);
-		}
-
-		res = NtQuerySystemInformationPtr((SYSTEM_INFORMATION_CLASS)SYSPARTITIONINFORMATION, tempBuf, sizeof(tempBuf), &len);
-		if (res != S_OK)
-		{
-			SetLastError (res);
-			throw SystemException (SRC_POS);
-		}		
-
-		PUNICODE_STRING pStr = (PUNICODE_STRING) tempBuf;
-		BootVolumePath = pStr->Buffer;
-
-		EfiBootPartPath = L"\\\\?";
-		EfiBootPartPath += &pStr->Buffer[7];
-
-		bBootVolumePathSelected = true;
 	}
 
 	void EfiBoot::PrepareBootPartition(bool bDisableException) {
-		if (!bBootVolumePathSelected) {
-			SelectBootVolumeESP();
-		}
+		
+		GetVolumeESP (EfiBootPartPath, BootVolumePath);
+
 		std::wstring devicePath = L"\\\\?\\GLOBALROOT";
 		devicePath += BootVolumePath;
 		Device  dev(devicePath.c_str(), TRUE);
diff --git a/src/Common/BootEncryption.h b/src/Common/BootEncryption.h
index 19ea372e..a94f1918 100644
--- a/src/Common/BootEncryption.h
+++ b/src/Common/BootEncryption.h
@@ -29,6 +29,10 @@ typedef NTSTATUS (WINAPI *NtQuerySystemInformationFn)(
 		PULONG                   ReturnLength
 );
 
+typedef ULONG (WINAPI *RtlNtStatusToDosErrorFn)(
+  NTSTATUS Status
+);
+
 using namespace std;
 
 namespace VeraCrypt
@@ -193,7 +197,7 @@ namespace VeraCrypt
 		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);
 
@@ -221,7 +225,6 @@ 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 () { if (bDeviceInfoValid) return &sdn; else { SetLastError (ERROR_INVALID_DRIVE); throw SystemException(SRC_POS);}}
 
 	protected:
@@ -231,7 +234,6 @@ namespace VeraCrypt
 		PARTITION_INFORMATION_EX partInfo;
 		bool bDeviceInfoValid;
 		WCHAR     tempBuf[1024];
-		bool  bBootVolumePathSelected;
 		std::wstring BootVolumePath;
 	};
 
-- 
cgit v1.2.3