3 files changed, 53 insertions, 10 deletions
diff --git a/src/Driver/DriveFilter.c b/src/Driver/DriveFilter.c
index 5fbacac4..c9efd7fb 100644
--- a/src/Driver/DriveFilter.c
+++ b/src/Driver/DriveFilter.c
@@ -75,28 +75,31 @@ static int64 DecoySystemWipedAreaEnd;
 PKTHREAD DecoySystemWipeThread = NULL;
 static NTSTATUS DecoySystemWipeResult;
 
-uint64 BootArgsRegions[] = { EFI_BOOTARGS_REGIONS };
+static uint64 BootArgsRegionsDefault[] = { EFI_BOOTARGS_REGIONS_DEFAULT };
+static uint64 BootArgsRegionsEFI[] = { EFI_BOOTARGS_REGIONS_EFI };
 
-NTSTATUS LoadBootArguments ()
+NTSTATUS LoadBootArguments (BOOL bIsEfi)
 {
 	NTSTATUS status = STATUS_UNSUCCESSFUL;
 	PHYSICAL_ADDRESS bootArgsAddr;
 	byte *mappedBootArgs;
 	byte *mappedCryptoInfo = NULL;
 	uint16 bootLoaderArgsIndex;
+	uint64* BootArgsRegionsPtr = bIsEfi? BootArgsRegionsEFI : BootArgsRegionsDefault;
+	size_t BootArgsRegionsCount = bIsEfi? sizeof(BootArgsRegionsEFI)/ sizeof(BootArgsRegionsEFI[0]) : sizeof(BootArgsRegionsDefault)/ sizeof(BootArgsRegionsDefault[0]);
 
 	KeInitializeMutex (&MountMutex, 0);
 //	__debugbreak();
 	for (bootLoaderArgsIndex = 0;
-		bootLoaderArgsIndex < sizeof(BootArgsRegions)/ sizeof(BootArgsRegions[1]) && status != STATUS_SUCCESS;
+		bootLoaderArgsIndex < BootArgsRegionsCount && status != STATUS_SUCCESS;
 		++bootLoaderArgsIndex)
 	{
-		bootArgsAddr.QuadPart = BootArgsRegions[bootLoaderArgsIndex] + TC_BOOT_LOADER_ARGS_OFFSET;
+		bootArgsAddr.QuadPart = BootArgsRegionsPtr[bootLoaderArgsIndex] + TC_BOOT_LOADER_ARGS_OFFSET;
 		Dump ("Checking BootArguments at 0x%x\n", bootArgsAddr.LowPart);
 
-			mappedBootArgs = MmMapIoSpace (bootArgsAddr, sizeof (BootArguments), MmCached);
-			if (!mappedBootArgs)
-				return STATUS_INSUFFICIENT_RESOURCES;
+		mappedBootArgs = MmMapIoSpace (bootArgsAddr, sizeof (BootArguments), MmCached);
+		if (!mappedBootArgs)
+			return STATUS_INSUFFICIENT_RESOURCES;
 
 		if (TC_IS_BOOT_ARGUMENTS_SIGNATURE (mappedBootArgs))
 		{
@@ -118,7 +121,7 @@ NTSTATUS LoadBootArguments ()
 			// Sanity check: for valid boot argument, the password is less than 64 bytes long
 			if (bootArguments->BootPassword.Length <= MAX_LEGACY_PASSWORD)
 			{
-				BootLoaderArgsPtr = BootArgsRegions[bootLoaderArgsIndex];
+				BootLoaderArgsPtr = BootArgsRegionsPtr[bootLoaderArgsIndex];
 
 				BootArgs = *bootArguments;
 				BootArgsValid = TRUE;
diff --git a/src/Driver/DriveFilter.h b/src/Driver/DriveFilter.h
index f19609b0..b164fa5b 100644
--- a/src/Driver/DriveFilter.h
+++ b/src/Driver/DriveFilter.h
@@ -70,7 +70,7 @@ CRYPTO_INFO *GetSystemDriveCryptoInfo ();
 BOOL IsBootDriveMounted ();
 BOOL IsBootEncryptionSetupInProgress ();
 BOOL IsHiddenSystemRunning ();
-NTSTATUS LoadBootArguments ();
+NTSTATUS LoadBootArguments (BOOL bIsEfi);
 static NTSTATUS SaveDriveVolumeHeader (DriveFilterExtension *Extension);
 NTSTATUS StartBootEncryptionSetup (PDEVICE_OBJECT DeviceObject, PIRP irp, PIO_STACK_LOCATION irpSp);
 void EmergencyClearAllKeys (PIRP irp, PIO_STACK_LOCATION irpSp);
diff --git a/src/Driver/Ntdriver.c b/src/Driver/Ntdriver.c
index 97fb1bf1..bf57fcdc 100644
--- a/src/Driver/Ntdriver.c
+++ b/src/Driver/Ntdriver.c
@@ -140,12 +140,44 @@ static BOOL EnableExtendedIoctlSupport = FALSE;
 static BOOL AllowTrimCommand = FALSE;
 static KeSaveExtendedProcessorStateFn KeSaveExtendedProcessorStatePtr = NULL;
 static KeRestoreExtendedProcessorStateFn KeRestoreExtendedProcessorStatePtr = NULL;
+static ExGetFirmwareEnvironmentVariableFn ExGetFirmwareEnvironmentVariablePtr = NULL;
 
 POOL_TYPE ExDefaultNonPagedPoolType = NonPagedPool;
 ULONG ExDefaultMdlProtection = 0;
 
 PDEVICE_OBJECT VirtualVolumeDeviceObjects[MAX_MOUNTED_VOLUME_DRIVE_NUMBER + 1];
 
+BOOL IsUefiBoot ()
+{
+	BOOL bStatus = FALSE;
+	NTSTATUS ntStatus = STATUS_NOT_IMPLEMENTED;
+	
+	Dump ("IsUefiBoot BEGIN\n");
+	ASSERT (KeGetCurrentIrql() == PASSIVE_LEVEL);
+
+	if (ExGetFirmwareEnvironmentVariablePtr)
+	{
+		ULONG valueLengh = 0;
+		UNICODE_STRING emptyName;
+		GUID guid;
+		RtlInitUnicodeString(&emptyName, L"");
+		memset (&guid, 0, sizeof(guid));
+		Dump ("IsUefiBoot calling ExGetFirmwareEnvironmentVariable\n");
+		ntStatus = ExGetFirmwareEnvironmentVariablePtr (&emptyName, &guid, NULL, &valueLengh, NULL);
+		Dump ("IsUefiBoot ExGetFirmwareEnvironmentVariable returned 0x%08x\n", ntStatus);
+	}
+	else
+	{
+		Dump ("IsUefiBoot ExGetFirmwareEnvironmentVariable not found on the system\n");
+	}
+	
+	if (STATUS_NOT_IMPLEMENTED != ntStatus)
+		bStatus = TRUE;
+
+	Dump ("IsUefiBoot bStatus = %s END\n", bStatus? "TRUE" : "FALSE");
+	return bStatus;
+}
+
 void GetDriverRandomSeed (unsigned char* pbRandSeed, size_t cbRandSeed)
 {
 	LARGE_INTEGER iSeed, iSeed2;
@@ -248,6 +280,14 @@ NTSTATUS DriverEntry (PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
 		KeSaveExtendedProcessorStatePtr = (KeSaveExtendedProcessorStateFn) MmGetSystemRoutineAddress(&saveFuncName);
 		KeRestoreExtendedProcessorStatePtr = (KeRestoreExtendedProcessorStateFn) MmGetSystemRoutineAddress(&restoreFuncName);
 	}
+	
+	// ExGetFirmwareEnvironmentVariable is available starting from Windows 8
+	if ((OsMajorVersion > 6) || (OsMajorVersion == 6 && OsMinorVersion >= 2))
+	{
+		UNICODE_STRING funcName;
+		RtlInitUnicodeString(&funcName, L"ExGetFirmwareEnvironmentVariable");
+		ExGetFirmwareEnvironmentVariablePtr = (ExGetFirmwareEnvironmentVariableFn) MmGetSystemRoutineAddress(&funcName);
+	}
 
 	// Load dump filter if the main driver is already loaded
 	if (NT_SUCCESS (TCDeviceIoControl (NT_ROOT_PREFIX, TC_IOCTL_GET_DRIVER_VERSION, NULL, 0, &version, sizeof (version))))
@@ -278,7 +318,7 @@ NTSTATUS DriverEntry (PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
 					TC_BUG_CHECK (STATUS_INVALID_PARAMETER);
 			}
 
-			LoadBootArguments();
+			LoadBootArguments(IsUefiBoot ());
 			VolumeClassFilterRegistered = IsVolumeClassFilterRegistered();
 
 			DriverObject->DriverExtension->AddDevice = DriverAddDevice;