1 files changed, 172 insertions, 66 deletions
diff --git a/src/Common/BootEncryption.cpp b/src/Common/BootEncryption.cpp
index 43cee062..c01a8b4b 100644
--- a/src/Common/BootEncryption.cpp
+++ b/src/Common/BootEncryption.cpp
@@ -75,8 +75,28 @@ namespace VeraCrypt
 		static void CopyFile (const string &sourceFile, const string &destinationFile)
 		{
 			Elevate();
+			DWORD result;
+			CComBSTR sourceFileBstr, destinationFileBstr;
+			BSTR bstr = A2WBSTR(sourceFile.c_str());
+			if (bstr)
+			{
+				sourceFileBstr.Attach (bstr);
 
-			DWORD result = ElevatedComInstance->CopyFile (CComBSTR (sourceFile.c_str()), CComBSTR (destinationFile.c_str()));
+				bstr = A2WBSTR(destinationFile.c_str());
+				if (bstr)
+				{
+					destinationFileBstr.Attach (bstr);
+					result = ElevatedComInstance->CopyFile (sourceFileBstr, destinationFileBstr);
+				}
+				else
+				{
+					result = ERROR_OUTOFMEMORY;
+				}
+			}
+			else
+			{
+				result = ERROR_OUTOFMEMORY;
+			}
 
 			if (result != ERROR_SUCCESS)
 			{
@@ -88,8 +108,18 @@ namespace VeraCrypt
 		static void DeleteFile (const string &file)
 		{
 			Elevate();
-
-			DWORD result = ElevatedComInstance->DeleteFile (CComBSTR (file.c_str()));
+			CComBSTR fileBstr;
+			DWORD result;
+			BSTR bstr = A2WBSTR(file.c_str());
+			if (bstr)
+			{
+				fileBstr.Attach (bstr);
+				result = ElevatedComInstance->DeleteFile (fileBstr);
+			}
+			else
+			{
+				result = ERROR_OUTOFMEMORY;
+			}
 
 			if (result != ERROR_SUCCESS)
 			{
@@ -102,10 +132,20 @@ namespace VeraCrypt
 		{
 			Elevate();
 
-			CComBSTR bufferBstr;
+			DWORD result;
+			CComBSTR bufferBstr, fileBstr;
 			if (bufferBstr.AppendBytes ((const char *) buffer, size) != S_OK)
 				throw ParameterIncorrect (SRC_POS);
-			DWORD result = ElevatedComInstance->ReadWriteFile (write, device, CComBSTR (filePath.c_str()), &bufferBstr, offset, size, sizeDone);
+			BSTR bstr = A2WBSTR(filePath.c_str());
+			if (bstr)
+			{
+				fileBstr.Attach (bstr);
+				result = ElevatedComInstance->ReadWriteFile (write, device, fileBstr, &bufferBstr, offset, size, sizeDone);
+			}
+			else
+			{
+				result = ERROR_OUTOFMEMORY;
+			}
 
 			if (result != ERROR_SUCCESS)
 			{
@@ -127,8 +167,29 @@ namespace VeraCrypt
 		static void WriteLocalMachineRegistryDwordValue (char *keyPath, char *valueName, DWORD value)
 		{
 			Elevate();
+			DWORD result;
+			CComBSTR keyPathBstr, valueNameBstr;
+			BSTR bstr = A2WBSTR(keyPath);
+			if (bstr)
+			{
+				keyPathBstr.Attach (bstr);
 
-			DWORD result = ElevatedComInstance->WriteLocalMachineRegistryDwordValue (CComBSTR (keyPath), CComBSTR (valueName), value);
+				bstr = A2WBSTR(valueName);
+				if (bstr)
+				{
+					valueNameBstr.Attach (bstr);
+
+					result = ElevatedComInstance->WriteLocalMachineRegistryDwordValue (keyPathBstr, valueNameBstr, value);
+				}
+				else
+				{
+					result = ERROR_OUTOFMEMORY;
+				}
+			}
+			else
+			{
+				result = ERROR_OUTOFMEMORY;
+			}
 		
 			if (result != ERROR_SUCCESS)
 			{
@@ -233,26 +294,27 @@ namespace VeraCrypt
 #endif // SETUP
 
 
-	File::File (string path, bool readOnly, bool create) : Elevated (false), FileOpen (false)
+	File::File (string path, bool readOnly, bool create) : Elevated (false), FileOpen (false), LastError(0)
 	{
 		Handle = CreateFile (path.c_str(),
 			readOnly ? GENERIC_READ : GENERIC_READ | GENERIC_WRITE,
 			FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, create ? CREATE_ALWAYS : OPEN_EXISTING,
 			FILE_FLAG_RANDOM_ACCESS | FILE_FLAG_WRITE_THROUGH, NULL);
 
-		try
+		if (Handle != INVALID_HANDLE_VALUE)
 		{
-			throw_sys_if (Handle == INVALID_HANDLE_VALUE);
+			FileOpen = true;
 		}
-		catch (SystemException &)
+		else 
 		{
-			if (GetLastError() == ERROR_ACCESS_DENIED && IsUacSupported())
+			LastError = GetLastError();
+			if (LastError == ERROR_ACCESS_DENIED && IsUacSupported())
+			{
 				Elevated = true;
-			else
-				throw;
+				FileOpen = true;
+			}
 		}
 
-		FileOpen = true;
 		FilePointerPosition = 0;
 		IsDevice = false;
 		Path = path;
@@ -260,19 +322,25 @@ namespace VeraCrypt
 
 	void File::Close ()
 	{
-		if (FileOpen)
+		if (Handle != INVALID_HANDLE_VALUE)
 		{
-			if (!Elevated)
-				CloseHandle (Handle);
-
-			FileOpen = false;
+			CloseHandle (Handle);
+			Handle = INVALID_HANDLE_VALUE;
 		}
+
+		FileOpen = false;
 	}
 
 	DWORD File::Read (byte *buffer, DWORD size)
 	{
 		DWORD bytesRead;
 
+		if (!FileOpen)
+		{
+			SetLastError (LastError);
+			throw SystemException ();
+		}
+
 		if (Elevated)
 		{
 			DWORD bytesRead;
@@ -288,6 +356,12 @@ namespace VeraCrypt
 
 	void File::SeekAt (int64 position)
 	{
+		if (!FileOpen)
+		{
+			SetLastError (LastError);
+			throw SystemException ();
+		}
+
 		FilePointerPosition = position;
 
 		if (!Elevated)
@@ -302,6 +376,12 @@ namespace VeraCrypt
 	{
 		DWORD bytesWritten;
 
+		if (!FileOpen)
+		{
+			SetLastError (LastError);
+			throw SystemException ();
+		}
+
 		try
 		{
 			if (Elevated)
@@ -349,19 +429,20 @@ namespace VeraCrypt
 			FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
 			FILE_FLAG_RANDOM_ACCESS | FILE_FLAG_WRITE_THROUGH, NULL);
 
-		try
+		if (Handle != INVALID_HANDLE_VALUE)
 		{
-			throw_sys_if (Handle == INVALID_HANDLE_VALUE);
+			FileOpen = true;
 		}
-		catch (SystemException &)
+		else 
 		{
-			if (GetLastError() == ERROR_ACCESS_DENIED && IsUacSupported())
+			LastError = GetLastError ();
+			if (LastError == ERROR_ACCESS_DENIED && IsUacSupported())
+			{
 				Elevated = true;
-			else
-				throw;
+				FileOpen = true;
+			}
 		}
 
-		FileOpen = true;
 		FilePointerPosition = 0;
 		IsDevice = true;
 		Path = path;
@@ -393,7 +474,7 @@ namespace VeraCrypt
 	{
 		if (RescueIsoImage)
 			delete[] RescueIsoImage;
-
+				
 		Elevator::Release();
 	}
 
@@ -635,7 +716,7 @@ namespace VeraCrypt
 			stringstream partPath;
 			partPath << "\\Device\\Harddisk" << driveNumber << "\\Partition" << partNumber;
 
-			DISK_PARTITION_INFO_STRUCT diskPartInfo;
+			DISK_PARTITION_INFO_STRUCT diskPartInfo = {0};
 			StringCbPrintfW (diskPartInfo.deviceName, sizeof (diskPartInfo.deviceName), L"%hs", partPath.str().c_str());
 
 			try
@@ -647,50 +728,60 @@ namespace VeraCrypt
 				continue;
 			}
 
-			Partition part;
-			part.DevicePath = partPath.str();
-			part.Number = partNumber;
-			part.Info = diskPartInfo.partInfo;
-			part.IsGPT = diskPartInfo.IsGPT;
-
-			// Mount point
-			wstringstream ws;
-			ws << partPath.str().c_str();
-			int driveNumber = GetDiskDeviceDriveLetter ((wchar_t *) ws.str().c_str());
-
-			if (driveNumber >= 0)
+			if (	(diskPartInfo.IsGPT == TRUE || diskPartInfo.IsGPT == FALSE)
+				&&	(diskPartInfo.IsDynamic == TRUE || diskPartInfo.IsDynamic == FALSE)
+				&&	(diskPartInfo.partInfo.BootIndicator == TRUE || diskPartInfo.partInfo.BootIndicator == FALSE)
+				&&	(diskPartInfo.partInfo.RecognizedPartition == TRUE || diskPartInfo.partInfo.RecognizedPartition == FALSE)
+				&&	(diskPartInfo.partInfo.RewritePartition == TRUE || diskPartInfo.partInfo.RewritePartition == FALSE)
+				&&	(diskPartInfo.partInfo.StartingOffset.QuadPart >= 0)
+				&&	(diskPartInfo.partInfo.PartitionLength.QuadPart >= 0)
+				)
 			{
-				part.MountPoint += (char) (driveNumber + 'A');
-				part.MountPoint += ":";
-			}
+				Partition part;
+				part.DevicePath = partPath.str();
+				part.Number = partNumber;
+				part.Info = diskPartInfo.partInfo;
+				part.IsGPT = diskPartInfo.IsGPT;
+
+				// Mount point
+				wstringstream ws;
+				ws << partPath.str().c_str();
+				int driveNumber = GetDiskDeviceDriveLetter ((wchar_t *) ws.str().c_str());
+
+				if (driveNumber >= 0)
+				{
+					part.MountPoint += (char) (driveNumber + 'A');
+					part.MountPoint += ":";
+				}
 
-			// Volume ID
-			wchar_t volumePath[TC_MAX_PATH];
-			if (ResolveSymbolicLink ((wchar_t *) ws.str().c_str(), volumePath, sizeof(volumePath)))
-			{
-				wchar_t volumeName[TC_MAX_PATH];
-				HANDLE fh = FindFirstVolumeW (volumeName, array_capacity (volumeName));
-				if (fh != INVALID_HANDLE_VALUE)
+				// Volume ID
+				wchar_t volumePath[TC_MAX_PATH];
+				if (ResolveSymbolicLink ((wchar_t *) ws.str().c_str(), volumePath, sizeof(volumePath)))
 				{
-					do
+					wchar_t volumeName[TC_MAX_PATH];
+					HANDLE fh = FindFirstVolumeW (volumeName, array_capacity (volumeName));
+					if (fh != INVALID_HANDLE_VALUE)
 					{
-						wstring volumeNameStr = volumeName;
-						wchar_t devicePath[TC_MAX_PATH];
-
-						if (QueryDosDeviceW (volumeNameStr.substr (4, volumeNameStr.size() - 1 - 4).c_str(), devicePath, array_capacity (devicePath)) != 0
-							&& wcscmp (volumePath, devicePath) == 0)
+						do
 						{
-							part.VolumeNameId = volumeName;
-							break;
-						}
+							wstring volumeNameStr = volumeName;
+							wchar_t devicePath[TC_MAX_PATH];
 
-					} while (FindNextVolumeW (fh, volumeName, array_capacity (volumeName)));
+							if (QueryDosDeviceW (volumeNameStr.substr (4, volumeNameStr.size() - 1 - 4).c_str(), devicePath, array_capacity (devicePath)) != 0
+								&& wcscmp (volumePath, devicePath) == 0)
+							{
+								part.VolumeNameId = volumeName;
+								break;
+							}
+
+						} while (FindNextVolumeW (fh, volumeName, array_capacity (volumeName)));
 
-					FindVolumeClose (fh);
+						FindVolumeClose (fh);
+					}
 				}
-			}
 
-			partList.push_back (part);
+				partList.push_back (part);
+			}
 		}
 
 		return partList;
@@ -793,6 +884,7 @@ namespace VeraCrypt
 	bool BootEncryption::SystemDriveContainsPartitionType (byte type)
 	{
 		Device device (GetSystemDriveConfiguration().DevicePath, true);
+		device.CheckOpened ();
 
 		byte mbrBuf[TC_SECTOR_SIZE_BIOS];
 		device.SeekAt (0);
@@ -1103,6 +1195,7 @@ namespace VeraCrypt
 		if (hiddenOSCreation)
 		{
 			Device device (GetSystemDriveConfiguration().DevicePath);
+			device.CheckOpened ();
 			byte headerSector[TC_SECTOR_SIZE_BIOS];
 
 			device.SeekAt (HiddenOSCandidatePartition.Info.StartingOffset.QuadPart + HiddenOSCandidatePartition.Info.PartitionLength.QuadPart - TC_VOLUME_HEADER_GROUP_SIZE + TC_VOLUME_HEADER_EFFECTIVE_SIZE);
@@ -1191,6 +1284,7 @@ namespace VeraCrypt
 	void BootEncryption::WriteBootSectorConfig (const byte newConfig[])
 	{
 		Device device (GetSystemDriveConfiguration().DevicePath);
+		device.CheckOpened ();
 		byte mbr[TC_SECTOR_SIZE_BIOS];
 
 		device.SeekAt (0);
@@ -1213,6 +1307,7 @@ namespace VeraCrypt
 	void BootEncryption::WriteBootSectorUserConfig (byte userConfig, const string &customUserMessage)
 	{
 		Device device (GetSystemDriveConfiguration().DevicePath);
+		device.CheckOpened ();
 		byte mbr[TC_SECTOR_SIZE_BIOS];
 
 		device.SeekAt (0);
@@ -1330,6 +1425,7 @@ namespace VeraCrypt
 			throw ParameterIncorrect (SRC_POS);
 
 		Device device (GetSystemDriveConfiguration().DevicePath);
+		device.CheckOpened();
 		byte mbr[TC_SECTOR_SIZE_BIOS];
 
 		device.SeekAt (0);
@@ -1388,6 +1484,7 @@ namespace VeraCrypt
 
 		// Write MBR
 		Device device (GetSystemDriveConfiguration().DevicePath);
+		device.CheckOpened ();
 		byte mbr[TC_SECTOR_SIZE_BIOS];
 
 		device.SeekAt (0);
@@ -1542,6 +1639,7 @@ namespace VeraCrypt
 		else
 		{
 			Device bootDevice (GetSystemDriveConfiguration().DevicePath, true);
+			bootDevice.CheckOpened ();
 			bootDevice.SeekAt (TC_BOOT_VOLUME_HEADER_SECTOR_OFFSET);
 			bootDevice.Read (image + TC_CD_BOOTSECTOR_OFFSET + TC_BOOT_VOLUME_HEADER_SECTOR_OFFSET, TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE);
 		}
@@ -1550,6 +1648,7 @@ namespace VeraCrypt
 		try
 		{
 			File sysBakFile (GetSystemLoaderBackupPath(), true);
+			sysBakFile.CheckOpened ();
 			sysBakFile.Read (image + TC_CD_BOOTSECTOR_OFFSET + TC_ORIG_BOOT_LOADER_BACKUP_SECTOR_OFFSET, TC_BOOT_LOADER_AREA_SIZE);
 			
 			image[TC_CD_BOOTSECTOR_OFFSET + TC_BOOT_SECTOR_CONFIG_OFFSET] |= TC_BOOT_CFG_FLAG_RESCUE_DISK_ORIG_SYS_LOADER;
@@ -1605,6 +1704,7 @@ namespace VeraCrypt
 				path[0] = drive;
 
 				Device driveDevice (path, true);
+				driveDevice.CheckOpened ();
 				size_t verifiedSectorCount = (TC_CD_BOOTSECTOR_OFFSET + TC_ORIG_BOOT_LOADER_BACKUP_SECTOR_OFFSET + TC_BOOT_LOADER_AREA_SIZE) / 2048;
 				Buffer buffer ((verifiedSectorCount + 1) * 2048);
 
@@ -1638,7 +1738,8 @@ namespace VeraCrypt
 
 		// Initial rescue disk assumes encryption of the drive has been completed (EncryptedAreaLength == volumeSize)
 		memcpy (RescueVolumeHeader, VolumeHeader, sizeof (RescueVolumeHeader));
-		ReadVolumeHeader (TRUE, (char *) RescueVolumeHeader, password, pkcs5, FALSE, NULL, cryptoInfo);
+		if (0 != ReadVolumeHeader (TRUE, (char *) RescueVolumeHeader, password, pkcs5, FALSE, NULL, cryptoInfo))
+			throw ParameterIncorrect (SRC_POS);
 
 		DecryptBuffer (RescueVolumeHeader + HEADER_ENCRYPTED_DATA_OFFSET, HEADER_ENCRYPTED_DATA_SIZE, cryptoInfo);
 
@@ -1666,6 +1767,7 @@ namespace VeraCrypt
 			throw ParameterIncorrect (SRC_POS);
 
 		Device device (GetSystemDriveConfiguration().DevicePath);
+		device.CheckOpened ();
 
 		device.SeekAt (TC_BOOT_VOLUME_HEADER_SECTOR_OFFSET);
 		device.Write ((byte *) VolumeHeader, sizeof (VolumeHeader));
@@ -1697,7 +1799,7 @@ namespace VeraCrypt
 	void BootEncryption::BackupSystemLoader ()
 	{
 		Device device (GetSystemDriveConfiguration().DevicePath, true);
-	
+		device.CheckOpened ();
 		byte bootLoaderBuf[TC_BOOT_LOADER_AREA_SECTOR_COUNT * TC_SECTOR_SIZE_BIOS];
 
 		device.SeekAt (0);
@@ -1724,11 +1826,12 @@ namespace VeraCrypt
 		byte bootLoaderBuf[TC_BOOT_LOADER_AREA_SECTOR_COUNT * TC_SECTOR_SIZE_BIOS];
 
 		File backupFile (GetSystemLoaderBackupPath(), true);
-		
+		backupFile.CheckOpened();
 		if (backupFile.Read (bootLoaderBuf, sizeof (bootLoaderBuf)) != sizeof (bootLoaderBuf))
 			throw ParameterIncorrect (SRC_POS);
 
 		Device device (GetSystemDriveConfiguration().DevicePath);
+		device.CheckOpened ();
 
 		// Preserve current partition table
 		byte mbr[TC_SECTOR_SIZE_BIOS];
@@ -2065,6 +2168,7 @@ namespace VeraCrypt
 		{
 			// Verify CRC of header salt
 			Device device (config.DevicePath, true);
+			device.CheckOpened ();
 			byte header[TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE];
 
 			device.SeekAt (TC_BOOT_VOLUME_HEADER_SECTOR_OFFSET);
@@ -2128,6 +2232,7 @@ namespace VeraCrypt
 
 		char header[TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE];
 		Device device (config.DevicePath);
+		device.CheckOpened ();
 
 		// Only one algorithm is currently supported
 		if (pkcs5 != 0)
@@ -2335,6 +2440,7 @@ namespace VeraCrypt
 			Buffer sector (geometry.BytesPerSector);
 
 			Device device (config.DevicePath);
+			device.CheckOpened ();
 
 			try
 			{