From 2c3719c7a9b370461d18c4422dd965dda80cb859 Mon Sep 17 00:00:00 2001
From: Paul Gofman <pgofman@codeweavers.com>
Date: Fri, 21 Aug 2020 02:42:13 +0300
Subject: [PATCH] ntdll: Restore AVX registers in NtSetContextThread() on
 x86_64.

Signed-off-by: Paul Gofman <pgofman@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
---
 dlls/ntdll/tests/exception.c    |  4 ++--
 dlls/ntdll/unix/signal_x86_64.c | 30 ++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/dlls/ntdll/tests/exception.c b/dlls/ntdll/tests/exception.c
index 041a5f258d1..f92491bb341 100644
--- a/dlls/ntdll/tests/exception.c
+++ b/dlls/ntdll/tests/exception.c
@@ -5562,7 +5562,7 @@ static void test_extended_context(void)
     for (i = 0; i < 8; ++i)
     {
         /* Older Windows version do not reset AVX context to INIT_STATE on x86. */
-        todo_wine_if(i >= 4)
+        todo_wine_if(i >= 4 && sizeof(void *) == 4)
         ok(!data[i] || broken(i >= 4 && sizeof(void *) == 4 && data[i] == test_extended_context_spoil_data2[i]),
                 "Got unexpected data %#x, i %u.\n", data[i], i);
     }
@@ -5573,7 +5573,7 @@ static void test_extended_context(void)
             ARRAY_SIZE(except_code_set_ymm0), PAGE_EXECUTE_READ);
 
     for (i = 0; i < 8; ++i)
-        todo_wine_if(i >= 4)
+        todo_wine_if(i >= 4 && sizeof(void *) == 4)
         ok(data[i] == test_extended_context_data[i], "Got unexpected data %#x, i %u.\n", data[i], i);
 }
 #endif
diff --git a/dlls/ntdll/unix/signal_x86_64.c b/dlls/ntdll/unix/signal_x86_64.c
index a98dc8657d6..7c13b4f7d20 100644
--- a/dlls/ntdll/unix/signal_x86_64.c
+++ b/dlls/ntdll/unix/signal_x86_64.c
@@ -1535,6 +1535,34 @@ __ASM_GLOBAL_FUNC( set_full_cpu_context,
                    "iretq" );
 
 
+/***********************************************************************
+ *           restore_xstate
+ *
+ * Restore the XState context.
+ */
+static void restore_xstate( const CONTEXT *context )
+{
+    XSAVE_FORMAT *xrstor_base;
+    XSTATE *xs;
+
+    if (!(xs = xstate_from_context( context )))
+        return;
+
+    xrstor_base = (XSAVE_FORMAT *)xs - 1;
+
+    if (!(xs->CompactionMask & ((ULONG64)1 << 63)))
+    {
+        /* Non-compacted xrstor will load Mxcsr regardless of the specified mask. Loading garbage there
+         * may lead to fault. We have only padding, no more used EXCEPTION_RECORD or unused context fields
+         * at the MxCsr restore location, so just put it there. */
+        assert( (void *)&xrstor_base->MxCsr > (void *)context->VectorRegister );
+        xrstor_base->MxCsr = context->u.FltSave.MxCsr;
+        xrstor_base->MxCsr_Mask = context->u.FltSave.MxCsr_Mask;
+    }
+
+    __asm__ volatile( "xrstor64 %0" : : "m"(*xrstor_base), "a" (4), "d" (0) );
+}
+
 /***********************************************************************
  *           get_server_context_flags
  *
@@ -1725,6 +1753,8 @@ NTSTATUS WINAPI NtSetContextThread( HANDLE handle, const CONTEXT *context )
         }
     }
 
+    restore_xstate( context );
+
     if (flags & CONTEXT_FULL)
     {
         if (!(flags & CONTEXT_CONTROL))