Windows: code refactoring for handling of ESP files (DcsProp and PlatformInfo).

author: Mounir IDRASSI <mounir.idrassi@idrix.fr> 2017-06-11 01:27:38 +0200
committer: Mounir IDRASSI <mounir.idrassi@idrix.fr> 2017-06-11 01:28:42 +0200
commit: 4208b4358136d1e52f9b6be17491f1daabe0c72d (patch)
tree: 4ac633f47cba49c7bf42d409c58089aebb0c5f26
parent: 374ba4c831af904d4284e628a48c9ca7a39aa8e6 (diff)
download: VeraCrypt-4208b4358136d1e52f9b6be17491f1daabe0c72d.tar.gz
VeraCrypt-4208b4358136d1e52f9b6be17491f1daabe0c72d.zip
3 files changed, 89 insertions, 67 deletions
diff --git a/src/Common/BootEncryption.cpp b/src/Common/BootEncryption.cpp
index 79f62a16..1c68a2fe 100644
--- a/src/Common/BootEncryption.cpp
+++ b/src/Common/BootEncryption.cpp
@@ -1565,9 +1565,6 @@ namespace VeraCrypt
 		{
 			if (GetSystemDriveConfiguration().SystemPartition.IsGPT)
 			{
-				byte confContent[4096*8];
-				DWORD dwSize;
-
 				// for now, we don't support any boot config flags, like hidden OS one
 				if (config)
 					memset (config, 0, bufLength);
@@ -1575,12 +1572,10 @@ namespace VeraCrypt
 				// call ReadEfiConfig only when needed since it requires elevation
 				if (userConfig || customUserMessage || bootLoaderVersion)
 				{
-					ReadEfiConfig (L"\\EFI\\VeraCrypt\\DcsProp", confContent, sizeof (confContent) - 1, &dwSize);
-					
-					confContent[dwSize] = 0;
+					std::string confContent = ReadESPFile (L"\\EFI\\VeraCrypt\\DcsProp", true);
 
 					EfiBootConf conf;
-					conf.Load ((char*) confContent);
+					conf.Load ((char*) confContent.c_str());
 
 					if (userConfig)
 					{
@@ -2095,29 +2090,96 @@ namespace VeraCrypt
 	void 
 	GetVolumeESP(wstring& path) 
 	{
-		ULONG    len;
-		NTSTATUS res;
-		WCHAR tempBuf[1024];
-		memset(tempBuf, 0, sizeof(tempBuf));
+		static wstring g_EspPath;
+		static bool g_EspPathInitialized = false;
 
-		// Load NtQuerySystemInformation function point
-		if (!NtQuerySystemInformationPtr)
+		if (!g_EspPathInitialized)
 		{
-			NtQuerySystemInformationPtr = (NtQuerySystemInformationFn) GetProcAddress (GetModuleHandle (L"ntdll.dll"), "NtQuerySystemInformation");
+			ULONG    len;
+			NTSTATUS res;
+			WCHAR tempBuf[1024];
+			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;
+			g_EspPath = L"\\\\?";
+			g_EspPath += &pStr->Buffer[7];
+			g_EspPathInitialized = true;
 		}
 
-		res = NtQuerySystemInformationPtr((SYSTEM_INFORMATION_CLASS)SYSPARTITIONINFORMATION, tempBuf, sizeof(tempBuf), &len);
-		if (res != S_OK)
+		path = g_EspPath;
+	}
+
+	std::string ReadESPFile (LPCWSTR szFilePath, bool bSkipUTF8BOM)
+	{
+		if (!szFilePath || !szFilePath[0])
+			throw ParameterIncorrect (SRC_POS);
+
+		ByteArray fileContent;
+		DWORD dwSize = 0, dwOffset = 0;
+		std::wstring pathESP;
+
+		GetVolumeESP(pathESP);
+		if (szFilePath[0] != L'\\')
+			pathESP += L"\\";
+		File f(pathESP + szFilePath, true);
+		f.GetFileSize(dwSize);
+		fileContent.resize(dwSize + 1);
+		fileContent[dwSize] = 0;
+		f.Read(fileContent.data(), dwSize);
+		f.Close();
+
+		if (bSkipUTF8BOM)
 		{
-			SetLastError (res);
-			throw SystemException (SRC_POS);
-		}		
+			// remove UTF-8 BOM if any
+			if (0 == memcmp (fileContent.data(), "\xEF\xBB\xBF", 3))
+			{
+				dwOffset = 3;
+			}
+		}
 
-		PUNICODE_STRING pStr = (PUNICODE_STRING) tempBuf;
-		path = L"\\\\?";
-		path += &pStr->Buffer[7];
+		return (const char*) &fileContent[dwOffset];
+	}
+
+	void WriteESPFile (LPCWSTR szFilePath, LPBYTE pbData, DWORD dwDataLen, bool bAddUTF8BOM)
+	{
+		if (!szFilePath || !szFilePath[0] || !pbData || !dwDataLen)
+			throw ParameterIncorrect (SRC_POS);
+
+		ByteArray fileContent;
+		DWORD dwSize = dwDataLen, dwOffset = 0;
+		std::wstring pathESP;
+
+		if (bAddUTF8BOM)
+		{
+			dwSize += 3;
+			dwOffset = 3;
+		}
+
+		GetVolumeESP(pathESP);
+		if (szFilePath[0] != L'\\')
+			pathESP += L"\\";
+		File f(pathESP + szFilePath, false, true);
+		fileContent.resize(dwSize);
+		if (bAddUTF8BOM)
+			memcpy (fileContent.data(), "\xEF\xBB\xBF", 3);
+		memcpy (&fileContent[dwOffset], pbData, dwDataLen);
+		f.Write(fileContent.data(), dwSize);
+		f.Close();
 	}
 
 	EfiBoot::EfiBoot() {
diff --git a/src/Common/BootEncryption.h b/src/Common/BootEncryption.h
index 80ca3b52..0c4094b4 100644
--- a/src/Common/BootEncryption.h
+++ b/src/Common/BootEncryption.h
@@ -187,6 +187,8 @@ namespace VeraCrypt
 	};
 
 	void GetVolumeESP(wstring& path);
+	std::string ReadESPFile (LPCWSTR szFilePath, bool bSkipUTF8BOM);
+	void WriteESPFile (LPCWSTR szFilePath, LPBYTE pbData, DWORD dwDataLen, bool bAddUTF8BOM);
 
 	class EfiBoot {
 	public:
diff --git a/src/Mount/Mount.c b/src/Mount/Mount.c
index f678123c..6b4d92dc 100644
--- a/src/Mount/Mount.c
+++ b/src/Mount/Mount.c
@@ -11063,24 +11063,7 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 					// read PlatformInfo file if it exists
 					try
 					{
-						ByteArray fileContent;
-						DWORD sz, offset;
-						std::wstring path;
-						GetVolumeESP(path);
-						path += L"\\EFI\\VeraCrypt\\PlatformInfo";
-						File fPlatformInfo(path);
-						fPlatformInfo.GetFileSize(sz);
-						fileContent.resize(sz + 1);
-						fileContent[sz] = 0;
-						fPlatformInfo.Read((byte*)&fileContent[0], sz);
-						// remove UTF-8 BOM if any
-						if (0 == memcmp (fileContent.data(), "\xEF\xBB\xBF", 3))
-						{
-							offset = 3;
-						}
-						else
-							offset = 0;
-						platforminfo = (const char*) &fileContent[offset];						
+						platforminfo = ReadESPFile (L"\\EFI\\VeraCrypt\\PlatformInfo", true);						
 					}
 					catch (Exception &e)	{}
 
@@ -11130,38 +11113,13 @@ static BOOL CALLBACK BootLoaderPreferencesDlgProc (HWND hwndDlg, UINT msg, WPARA
 			{
 				try
 				{
-					std::string dcsprop;
-					ByteArray fileContent;
-					DWORD sz, offset;
-					std::wstring path;
-					GetVolumeESP(path);
-					path += L"\\EFI\\VeraCrypt\\DcsProp";
-					File f1(path);
-					f1.GetFileSize(sz);
-					fileContent.resize(sz + 1);
-					fileContent[sz] = 0;
-					f1.Read((byte*)&fileContent[0], sz);
-					f1.Close();
-					// remove UTF-8 BOM if any
-					if (0 == memcmp (fileContent.data(), "\xEF\xBB\xBF", 3))
-					{
-						offset = 3;
-					}
-					else
-						offset = 0;
+					std::string dcsprop = ReadESPFile (L"\\EFI\\VeraCrypt\\DcsProp", true);
 
-					dcsprop = (const char*) &fileContent[offset];
 					while (TextEditDialogBox(FALSE, hwndDlg, GetString ("BOOT_LOADER_CONFIGURATION_FILE"), dcsprop) == IDOK)
 					{
 						if (validateDcsPropXml (dcsprop.c_str()))
 						{
-							// Add UTF-8 BOM
-							fileContent.resize (dcsprop.length() + 3);
-							memcpy (fileContent.data(), "\xEF\xBB\xBF", 3);
-							memcpy (&fileContent[3], &dcsprop[0], dcsprop.length());
-							File f2(path,false,true);
-							f2.Write(fileContent.data(), fileContent.size());
-							f2.Close();
+							WriteESPFile (L"\\EFI\\VeraCrypt\\DcsProp", (LPBYTE) dcsprop.c_str(), (DWORD) dcsprop.size(), true);
 							break;
 						}
 						else
author	Mounir IDRASSI <mounir.idrassi@idrix.fr>	2017-06-11 01:27:38 +0200
committer	Mounir IDRASSI <mounir.idrassi@idrix.fr>	2017-06-11 01:28:42 +0200
commit	4208b4358136d1e52f9b6be17491f1daabe0c72d (patch)
tree	4ac633f47cba49c7bf42d409c58089aebb0c5f26
parent	374ba4c831af904d4284e628a48c9ca7a39aa8e6 (diff)
download	VeraCrypt-4208b4358136d1e52f9b6be17491f1daabe0c72d.tar.gz VeraCrypt-4208b4358136d1e52f9b6be17491f1daabe0c72d.zip