3 files changed, 68 insertions, 6 deletions
diff --git a/src/Driver/DriveFilter.c b/src/Driver/DriveFilter.c
index d4d5e122..08bebe18 100644
--- a/src/Driver/DriveFilter.c
+++ b/src/Driver/DriveFilter.c
@@ -327,10 +327,14 @@ static void ComputeBootLoaderFingerprint(PDEVICE_OBJECT LowerDeviceObject, byte*
 	status = TCReadDevice (LowerDeviceObject, ioBuffer, offset, TC_SECTOR_SIZE_BIOS);
 	if (NT_SUCCESS (status))
 	{
-#if !defined (_WIN64)
-		KFLOATING_SAVE floatingPointState;
-		NTSTATUS saveStatus = STATUS_SUCCESS;
-		if (HasISSE()|| (HasSSE2() && HasMMX()))
+		NTSTATUS saveStatus = STATUS_INVALID_PARAMETER;
+#ifdef _WIN64
+		XSTATE_SAVE SaveState;
+		if (g_isIntel && HasSAVX())
+			saveStatus = KeSaveExtendedProcessorState(XSTATE_MASK_GSSE, &SaveState);
+#else
+		KFLOATING_SAVE floatingPointState;		
+		if (HasISSE() || (HasSSSE3() && HasMMX()))
 			saveStatus = KeSaveFloatingPointState (&floatingPointState);
 #endif
 		WHIRLPOOL_add (ioBuffer, TC_BOOT_SECTOR_PIM_VALUE_OFFSET, &whirlpool);
@@ -367,8 +371,10 @@ static void ComputeBootLoaderFingerprint(PDEVICE_OBJECT LowerDeviceObject, byte*
 			sha512_end (&BootLoaderFingerprint [WHIRLPOOL_DIGESTSIZE], &sha2);
 		}
 
-#if !defined (_WIN64)
-		if (NT_SUCCESS (saveStatus) && (HasISSE() || (HasSSE2() && HasMMX())))
+		if (NT_SUCCESS (saveStatus))
+#ifdef _WIN64
+			KeRestoreExtendedProcessorState(&SaveState);
+#else
 			KeRestoreFloatingPointState (&floatingPointState);
 #endif
 	}
diff --git a/src/Driver/DumpFilter.c b/src/Driver/DumpFilter.c
index 1b57bdbf..18feca06 100644
--- a/src/Driver/DumpFilter.c
+++ b/src/Driver/DumpFilter.c
@@ -14,6 +14,7 @@
 #include "DriveFilter.h"
 #include "Ntdriver.h"
 #include "Tests.h"
+#include "cpu.h"
 
 static DriveFilterExtension *BootDriveFilterExtension = NULL;
 static LARGE_INTEGER DumpPartitionOffset;
@@ -63,7 +64,21 @@ NTSTATUS DumpFilterEntry (PFILTER_EXTENSION filterExtension, PFILTER_INITIALIZAT
 	// KeSaveFloatingPointState() may generate a bug check during crash dump
 #if !defined (_WIN64)
 	if (filterExtension->DumpType == DumpTypeCrashdump)
+	{
 		dumpConfig.HwEncryptionEnabled = FALSE;
+		// disable also SSE optimizations
+		HasMMX() = 0;
+		HasISSE() = 0;
+		HasSSE2() = 0;
+		HasSSSE3() = 0;
+		HasSSE41() = 0;
+		HasSSE42() = 0;
+		HasAESNI() = 0;
+		HasCLMUL() = 0;
+		HasSAVX() = 0;
+		HasSAVX2() = 0;
+		HasSBMI2() = 0;
+	}
 #endif
 
 	EnableHwEncryption (dumpConfig.HwEncryptionEnabled);
diff --git a/src/Driver/Ntdriver.c b/src/Driver/Ntdriver.c
index ab555904..8f6f151f 100644
--- a/src/Driver/Ntdriver.c
+++ b/src/Driver/Ntdriver.c
@@ -73,6 +73,11 @@
 #pragma alloc_text(INIT,DriverEntry)
 #pragma alloc_text(INIT,TCCreateRootDeviceObject)
 
+/* We need to silence 'type cast' warning in order to use MmGetSystemRoutineAddress.
+ * MmGetSystemRoutineAddress() should have been declare FARPROC instead of PVOID.
+ */
+#pragma warning(disable:4055)
+
 PDRIVER_OBJECT TCDriverObject;
 PDEVICE_OBJECT RootDeviceObject = NULL;
 static KMUTEX RootDeviceControlMutex;
@@ -91,6 +96,8 @@ static size_t EncryptionThreadPoolFreeCpuCountLimit = 0;
 static BOOL SystemFavoriteVolumeDirty = FALSE;
 static BOOL PagingFileCreationPrevented = FALSE;
 static BOOL EnableExtendedIoctlSupport = FALSE;
+static KeSaveExtendedProcessorStateFn KeSaveExtendedProcessorStatePtr = NULL;
+static KeRestoreExtendedProcessorStateFn KeRestoreExtendedProcessorStatePtr = NULL;
 
 POOL_TYPE ExDefaultNonPagedPoolType = NonPagedPool;
 ULONG ExDefaultMdlProtection = 0;
@@ -119,6 +126,15 @@ NTSTATUS DriverEntry (PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
 		ExDefaultMdlProtection = MdlMappingNoExecute;
 	}
 
+	// KeSaveExtendedProcessorState/KeRestoreExtendedProcessorState are available starting from Windows 7
+	if ((OsMajorVersion > 6) || (OsMajorVersion == 6 && OsMinorVersion >= 1))
+	{
+		UNICODE_STRING funcName;
+		RtlInitUnicodeString(&funcName, L"KeSaveExtendedProcessorState");
+		KeSaveExtendedProcessorStatePtr = (KeSaveExtendedProcessorStateFn) MmGetSystemRoutineAddress(&funcName);
+		KeRestoreExtendedProcessorStatePtr = (KeRestoreExtendedProcessorStateFn) 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))))
 		return DumpFilterEntry ((PFILTER_EXTENSION) DriverObject, (PFILTER_INITIALIZATION_DATA) RegistryPath);
@@ -3960,3 +3976,28 @@ BOOL IsOSAtLeast (OSVersionEnum reqMinOS)
 	return ((OsMajorVersion << 16 | OsMinorVersion << 8)
 		>= (major << 16 | minor << 8));
 }
+
+NTSTATUS NTAPI KeSaveExtendedProcessorState (
+    __in ULONG64 Mask,
+    PXSTATE_SAVE XStateSave
+    )
+{
+	if (KeSaveExtendedProcessorStatePtr)
+	{
+		return (KeSaveExtendedProcessorStatePtr) (Mask, XStateSave);
+	}
+	else
+	{
+		return STATUS_SUCCESS;
+	}
+}
+
+VOID NTAPI KeRestoreExtendedProcessorState (
+	PXSTATE_SAVE XStateSave
+	)
+{
+	if (KeRestoreExtendedProcessorStatePtr)
+	{
+		(KeRestoreExtendedProcessorStatePtr) (XStateSave);
+	}
+}
+\ No newline at end of file