3 files changed, 191 insertions, 26 deletions
diff --git a/src/Driver/DriveFilter.c b/src/Driver/DriveFilter.c
index 66f49816..b8aad224 100644
--- a/src/Driver/DriveFilter.c
+++ b/src/Driver/DriveFilter.c
@@ -33,6 +33,10 @@ static BOOL DeviceFilterActive = FALSE;
 
 BOOL BootArgsValid = FALSE;
 BootArguments BootArgs;
+byte*  BootSecRegionData = NULL;
+uint32 BootSecRegionSize = 0;
+uint32 BootPkcs5;
+
 static uint64 BootLoaderArgsPtr;
 static BOOL BootDriveSignatureValid = FALSE;
 
@@ -76,10 +80,11 @@ NTSTATUS LoadBootArguments ()
 	NTSTATUS status = STATUS_UNSUCCESSFUL;
 	PHYSICAL_ADDRESS bootArgsAddr;
 	byte *mappedBootArgs;
+	byte *mappedCryptoInfo = NULL;
 	uint16 bootLoaderArgsIndex;
 
 	KeInitializeMutex (&MountMutex, 0);
-
+//	__debugbreak();
 	for (bootLoaderArgsIndex = 0;
 		bootLoaderArgsIndex < sizeof(BootArgsRegions)/ sizeof(BootArgsRegions[1]) && status != STATUS_SUCCESS;
 		++bootLoaderArgsIndex)
@@ -126,22 +131,61 @@ NTSTATUS LoadBootArguments ()
 				Dump ("BootDriveSignature = %x\n", BootArgs.BootDriveSignature);
 				Dump ("BootArgumentsCrc32 = %x\n", BootArgs.BootArgumentsCrc32);
 
-				if (CacheBootPassword && BootArgs.BootPassword.Length > 0)
-				{
-					int pim = CacheBootPim? (int) (BootArgs.Flags >> 16) : 0;
-					AddPasswordToCache (&BootArgs.BootPassword, pim);
-				}
-
 				// clear fingerprint
 				burn (BootLoaderFingerprint, sizeof (BootLoaderFingerprint));
+				MmUnmapIoSpace (mappedBootArgs, sizeof (BootArguments));
 
+				if (BootArgs.CryptoInfoLength > 0)
+				{
+					PHYSICAL_ADDRESS cryptoInfoAddress;
+					cryptoInfoAddress.QuadPart = BootLoaderArgsPtr + BootArgs.CryptoInfoOffset;
+					Dump ("CryptoInfo memory %x %d\n", cryptoInfoAddress.LowPart, BootArgs.CryptoInfoLength);
+
+					mappedCryptoInfo = MmMapIoSpace (cryptoInfoAddress, BootArgs.CryptoInfoLength, MmCached);
+					if (mappedCryptoInfo)
+					{
+						/* Get the parameters used for booting to speed up driver startup and avoid testing irrelevant PRFs */
+						BOOT_CRYPTO_HEADER* pBootCryptoInfo = (BOOT_CRYPTO_HEADER*) mappedCryptoInfo;
+						BootPkcs5 = pBootCryptoInfo->pkcs5;
+
+						BootSecRegionData = NULL;
+						BootSecRegionSize = 0;
+
+						if(BootArgs.CryptoInfoLength > (sizeof(BOOT_CRYPTO_HEADER) + sizeof(SECREGION_BOOT_PARAMS)) ) {
+							uint32   crc;
+							PHYSICAL_ADDRESS SecRegionAddress;
+							SECREGION_BOOT_PARAMS* SecRegionParams = (SECREGION_BOOT_PARAMS*) (mappedCryptoInfo + sizeof(BOOT_CRYPTO_HEADER) + 2);
+							byte *secRegionData = NULL;
+
+							SecRegionAddress.QuadPart = SecRegionParams->Ptr;
+							Dump ("SecRegion memory 0x%x %d\n", SecRegionAddress.LowPart, SecRegionParams->Size);
+
+							if( (SecRegionParams->Ptr != 0) && (SecRegionParams->Size > 0)) {
+								crc = GetCrc32((byte*)SecRegionParams, 12);
+								if(crc == SecRegionParams->Crc) {
+									Dump ("SecRegion crc ok\n");
+									secRegionData = MmMapIoSpace (SecRegionAddress, SecRegionParams->Size, MmCached);
+									BootSecRegionData = TCalloc (SecRegionParams->Size);
+									if(BootSecRegionData != NULL) {
+										BootSecRegionSize = SecRegionParams->Size;
+										memcpy(BootSecRegionData, secRegionData, SecRegionParams->Size);
+									}
+									burn (secRegionData, SecRegionParams->Size);
+									MmUnmapIoSpace (secRegionData,  SecRegionParams->Size);
+								}
+								// Erase boot loader scheduled keys
+								burn (mappedCryptoInfo, BootArgs.CryptoInfoLength);
+								MmUnmapIoSpace (mappedCryptoInfo, BootArgs.CryptoInfoLength);
+							}
+						}
+					}
+				}
 				status = STATUS_SUCCESS;
 			}
+		} else {
+			MmUnmapIoSpace (mappedBootArgs, sizeof (BootArguments));
 		}
-
-		MmUnmapIoSpace (mappedBootArgs, sizeof (BootArguments));
 	}
-
 	return status;
 }
 
@@ -327,30 +371,55 @@ static NTSTATUS MountDrive (DriveFilterExtension *Extension, Password *password,
 	Dump ("MountDrive pdo=%p\n", Extension->Pdo);
 	ASSERT (KeGetCurrentIrql() == PASSIVE_LEVEL);
 
-	// Check boot drive signature first (header CRC search could fail if a user restored the header to a non-boot drive)
-	if (BootDriveSignatureValid)
-	{
+	if (BootSecRegionData != NULL && BootSecRegionSize >= 1024) {
 		byte mbr[TC_SECTOR_SIZE_BIOS];
-
+		DCS_DISK_ENTRY_LIST* DeList = (DCS_DISK_ENTRY_LIST*)(BootSecRegionData + 512);
 		offset.QuadPart = 0;
 		status = TCReadDevice (Extension->LowerDeviceObject, mbr, offset, TC_SECTOR_SIZE_BIOS);
 
-		if (NT_SUCCESS (status) && BootArgs.BootDriveSignature != *(uint32 *) (mbr + 0x1b8))
+		if (NT_SUCCESS (status) && DeList->DE[DE_IDX_DISKID].DiskId.MbrID != *(uint32 *) (mbr + 0x1b8))
 			return STATUS_UNSUCCESSFUL;
-	}
 
-	header = TCalloc (TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE);
-	if (!header)
-		return STATUS_INSUFFICIENT_RESOURCES;
+		offset.QuadPart = 512;
+		status = TCReadDevice (Extension->LowerDeviceObject, mbr, offset, TC_SECTOR_SIZE_BIOS);
+		if (NT_SUCCESS (status) && memcmp(&DeList->DE[DE_IDX_DISKID].DiskId.GptID, mbr + 0x38, sizeof(DCS_GUID)) != 0)
+			return STATUS_UNSUCCESSFUL;
 
-	offset.QuadPart = hiddenVolume ? hiddenHeaderOffset : TC_BOOT_VOLUME_HEADER_SECTOR_OFFSET;
-	Dump ("Reading volume header at %I64u\n", offset.QuadPart);
+		header = TCalloc (TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE);
+		if (!header)
+			return STATUS_INSUFFICIENT_RESOURCES;
+		memcpy(header, BootSecRegionData, 512);
+		// Set extra data for the disk
+		Extension->Queue.SecRegionData = BootSecRegionData;
+		Extension->Queue.SecRegionSize = BootSecRegionSize;
+	} else {
+		// Check boot drive signature first (header CRC search could fail if a user restored the header to a non-boot drive)
+		if (BootDriveSignatureValid)
+		{
+			byte mbr[TC_SECTOR_SIZE_BIOS];
 
-	status = TCReadDevice (Extension->LowerDeviceObject, header, offset, TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE);
-	if (!NT_SUCCESS (status))
-	{
-		Dump ("TCReadDevice error %x\n", status);
-		goto ret;
+			offset.QuadPart = 0;
+			status = TCReadDevice (Extension->LowerDeviceObject, mbr, offset, TC_SECTOR_SIZE_BIOS);
+
+			if (NT_SUCCESS (status) && BootArgs.BootDriveSignature != *(uint32 *) (mbr + 0x1b8))
+				return STATUS_UNSUCCESSFUL;
+		}
+
+		header = TCalloc (TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE);
+		if (!header)
+			return STATUS_INSUFFICIENT_RESOURCES;
+
+		offset.QuadPart = hiddenVolume ? hiddenHeaderOffset : TC_BOOT_VOLUME_HEADER_SECTOR_OFFSET;
+		Dump ("Reading volume header at %I64u\n", offset.QuadPart);
+
+		status = TCReadDevice (Extension->LowerDeviceObject, header, offset, TC_BOOT_ENCRYPTION_VOLUME_HEADER_SIZE);
+		if (!NT_SUCCESS (status))
+		{
+			Dump ("TCReadDevice error %x\n", status);
+			goto ret;
+		}
+		Extension->Queue.SecRegionData = NULL;
+		Extension->Queue.SecRegionSize = 0;
 	}
 
 	if (headerSaltCrc32)
@@ -469,6 +538,41 @@ static NTSTATUS MountDrive (DriveFilterExtension *Extension, Password *password,
 		BootDriveFound = Extension->BootDrive = Extension->DriveMounted = Extension->VolumeHeaderPresent = TRUE;
 		BootDriveFilterExtension->MagicNumber = TC_BOOT_DRIVE_FILTER_EXTENSION_MAGIC_NUMBER;
 
+		if (BootSecRegionData != NULL && BootSecRegionSize > 1024) {
+			DCS_DISK_ENTRY_LIST* DeList = (DCS_DISK_ENTRY_LIST*)(BootSecRegionData + 512);
+			uint32 crc;
+			uint32 crcSaved;
+			crcSaved = DeList->CRC32;
+			DeList->CRC32 = 0;
+			crc = GetCrc32((byte*)DeList, 512);
+			if(crc == crcSaved){
+				if(DeList->DE[DE_IDX_PWDCACHE].Type == DE_PwdCache) {
+					uint64 sector = 0;
+					DCS_DEP_PWD_CACHE* pwdCache = (DCS_DEP_PWD_CACHE*)(BootSecRegionData + DeList->DE[DE_IDX_PWDCACHE].Sectors.Offset);
+					DecryptDataUnits((unsigned char*)pwdCache, (UINT64_STRUCT*)&sector, 1, Extension->Queue.CryptoInfo);
+					crcSaved = pwdCache->CRC;
+					pwdCache->CRC = 0;
+					crc = GetCrc32((unsigned char*)pwdCache, 512);
+					if(crcSaved == crc && pwdCache->Count < CACHE_SIZE){
+						uint32 i;
+						for(i = 0; i<pwdCache->Count; ++i){
+							if (CacheBootPassword && pwdCache->Pwd[i].Length > 0)	{
+								int pim = CacheBootPim? (int) (pwdCache->Pim[i]) : 0;
+								AddPasswordToCache (&pwdCache->Pwd[i], pim);
+							}
+						}
+						burn(pwdCache, sizeof(*pwdCache));
+					}
+				}
+			}
+		}
+
+		if (CacheBootPassword && BootArgs.BootPassword.Length > 0)
+		{
+			int pim = CacheBootPim? (int) (BootArgs.Flags >> 16) : 0;
+			AddPasswordToCache (&BootArgs.BootPassword, pim);
+		}
+
 		burn (&BootArgs.BootPassword, sizeof (BootArgs.BootPassword));
 
 		{
diff --git a/src/Driver/EncryptedIoQueue.c b/src/Driver/EncryptedIoQueue.c
index 029a42f5..7f50ec30 100644
--- a/src/Driver/EncryptedIoQueue.c
+++ b/src/Driver/EncryptedIoQueue.c
@@ -225,6 +225,47 @@ static void ReleaseFragmentBuffer (EncryptedIoQueue *queue, byte *buffer)
 	}
 }
 
+BOOL 
+UpdateBuffer(
+	byte*     buffer,
+	byte*     secRegion,
+	uint64    bufferDiskOffset,
+	uint32    bufferLength,
+	BOOL      doUpadte
+	) 
+{
+	uint64       intersectStart;
+	uint32       intersectLength;
+	uint32       i;
+	DCS_DISK_ENTRY_LIST *DeList = (DCS_DISK_ENTRY_LIST*)(secRegion + 512);
+	BOOL         updated = FALSE;
+
+	if (secRegion == NULL) return FALSE;
+	for (i = 0; i < DeList->Count; ++i) {
+		if (DeList->DE[i].Type == DE_Sectors) {
+			GetIntersection(
+				bufferDiskOffset, bufferLength,
+				DeList->DE[i].Sectors.Start, DeList->DE[i].Sectors.Start + DeList->DE[i].Sectors.Length - 1,
+				&intersectStart, &intersectLength
+				);
+			if (intersectLength != 0) {
+				updated = TRUE;
+				if(doUpadte && buffer != NULL) {
+//					Dump("Subst data\n");
+					memcpy(
+						buffer + (intersectStart - bufferDiskOffset),
+						secRegion + DeList->DE[i].Sectors.Offset + (intersectStart - DeList->DE[i].Sectors.Start),
+						intersectLength
+						);
+				} else {
+					return TRUE;
+				}
+			}
+		}
+	}
+	return updated;
+}
+
 
 static VOID CompletionThreadProc (PVOID threadArg)
 {
@@ -260,6 +301,11 @@ static VOID CompletionThreadProc (PVOID threadArg)
 
 				DecryptDataUnits (request->Data + request->EncryptedOffset, &dataUnit, request->EncryptedLength / ENCRYPTION_DATA_UNIT_SIZE, queue->CryptoInfo);
 			}
+//			Dump("Read sector %lld count %d\n", request->Offset.QuadPart >> 9, request->Length >> 9);
+			// Update subst sectors
+			if((queue->SecRegionData != NULL) && (queue->SecRegionSize > 512)) {
+				UpdateBuffer(request->Data, queue->SecRegionData, request->Offset.QuadPart, request->Length, TRUE);
+			}
 
 			if (request->CompleteOriginalIrp)
 			{
@@ -609,6 +655,10 @@ static VOID MainThreadProc (PVOID threadArg)
 							DecryptDataUnits (buffer + (intersectStart - alignedOffset.QuadPart), &dataUnit, intersectLength / ENCRYPTION_DATA_UNIT_SIZE, queue->CryptoInfo);
 						}
 					}
+					// Update subst sectors
+ 					if((queue->SecRegionData != NULL) && (queue->SecRegionSize > 512)) {
+ 						UpdateBuffer(buffer, queue->SecRegionData, alignedOffset.QuadPart, alignedLength, TRUE);
+ 					}
 
 					memcpy (dataBuffer, buffer + (item->OriginalOffset.LowPart & (ENCRYPTION_DATA_UNIT_SIZE - 1)), item->OriginalLength);
 				}
@@ -697,6 +747,15 @@ static VOID MainThreadProc (PVOID threadArg)
 				Dump ("Preventing write to boot loader or host protected area\n");
 				CompleteOriginalIrp (item, STATUS_MEDIA_WRITE_PROTECTED, 0);
 				continue;
+			} 
+			else if (item->Write
+				&& (queue->SecRegionData != NULL) && (queue->SecRegionSize > 512)
+				&& UpdateBuffer (NULL, queue->SecRegionData, item->OriginalOffset.QuadPart, (uint32)(item->OriginalOffset.QuadPart + item->OriginalLength - 1), FALSE))
+			{
+				// Prevent inappropriately designed software from damaging important data
+				Dump ("Preventing write to the system GPT area\n");
+				CompleteOriginalIrp (item, STATUS_MEDIA_WRITE_PROTECTED, 0);
+				continue;
 			}
 
 			dataBuffer = (PUCHAR) MmGetSystemAddressForMdlSafe (irp->MdlAddress, HighPagePriority);
diff --git a/src/Driver/EncryptedIoQueue.h b/src/Driver/EncryptedIoQueue.h
index ffe71471..fe9365ed 100644
--- a/src/Driver/EncryptedIoQueue.h
+++ b/src/Driver/EncryptedIoQueue.h
@@ -117,6 +117,8 @@ typedef struct
 	LARGE_INTEGER LastPerformanceCounter;
 #endif
 
+ 	byte*  SecRegionData;
+ 	SIZE_T SecRegionSize;
 }  EncryptedIoQueue;