Paul Gofman : ntdll: Restore AVX registers in NtSetContextThread() on i386.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 601175822e9c61ce02e21f20c76daf1b760d6924
URL:    https://source.winehq.org/git/wine.git/?a=commit;h=601175822e9c61ce02e21f20c...

Author: Paul Gofman <pgofman(a)codeweavers.com>
Date:   Fri Aug 21 02:42:15 2020 +0300

ntdll: Restore AVX registers in NtSetContextThread() on i386.

Signed-off-by: Paul Gofman <pgofman(a)codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard(a)winehq.org>

---

 dlls/ntdll/tests/exception.c  |  2 --
 dlls/ntdll/unix/signal_i386.c | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/dlls/ntdll/tests/exception.c b/dlls/ntdll/tests/exception.c
index aff18e3439..e75336ed11 100644
--- a/dlls/ntdll/tests/exception.c
+++ b/dlls/ntdll/tests/exception.c
@@ -5562,7 +5562,6 @@ 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 && 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 +5572,6 @@ 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 && 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_i386.c b/dlls/ntdll/unix/signal_i386.c
index 791ffe9094..6a27277bc2 100644
--- a/dlls/ntdll/unix/signal_i386.c
+++ b/dlls/ntdll/unix/signal_i386.c
@@ -687,7 +687,7 @@ static inline void save_fpux( CONTEXT *context )
 /***********************************************************************
  *           restore_fpu
  *
- * Restore the FPU context to a sigcontext.
+ * Restore the x87 FPU context
  */
 static inline void restore_fpu( const CONTEXT *context )
 {
@@ -701,7 +701,7 @@ static inline void restore_fpu( const CONTEXT *context )
 /***********************************************************************
  *           restore_fpux
  *
- * Restore the FPU extended context to a sigcontext.
+ * Restore the FPU extended context
  */
 static inline void restore_fpux( const CONTEXT *context )
 {
@@ -715,6 +715,38 @@ static inline void restore_fpux( const CONTEXT *context )
     __asm__ __volatile__( "fxrstor %0" : : "m" (*state) );
 }
 
+/***********************************************************************
+ *           restore_xstate
+ *
+ * Restore the XState context
+ */
+static inline 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. FPUX state should be restored by now, so we can reuse some space in
+         * ExtendedRegisters. */
+        XSAVE_FORMAT *fpux = (XSAVE_FORMAT *)context->ExtendedRegisters;
+        DWORD mxcsr, mxcsr_mask;
+
+        mxcsr = fpux->MxCsr;
+        mxcsr_mask = fpux->MxCsr_Mask;
+
+        assert( (void *)&xrstor_base->MxCsr > (void *)context->ExtendedRegisters );
+        xrstor_base->MxCsr = mxcsr;
+        xrstor_base->MxCsr_Mask = mxcsr_mask;
+    }
+    __asm__ volatile( "xrstor %0" : : "m"(*xrstor_base), "a" (4), "d" (0) );
+}
 
 /***********************************************************************
  *           fpux_to_fpu
@@ -1129,6 +1161,8 @@ NTSTATUS WINAPI NtSetContextThread( HANDLE handle, const CONTEXT *context )
     if (flags & CONTEXT_EXTENDED_REGISTERS) restore_fpux( context );
     else if (flags & CONTEXT_FLOATING_POINT) restore_fpu( context );
 
+    restore_xstate( context );
+
     if (flags & CONTEXT_FULL)
     {
         if (!(flags & CONTEXT_CONTROL))