Alexandre Julliard : ntdll: Validate the extended context before modifying the other registers.

[Alexandre Julliard]

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

Author: Alexandre Julliard <julliard at winehq.org>
Date:   Mon Jun  7 13:16:24 2021 +0200

ntdll: Validate the extended context before modifying the other registers.

Signed-off-by: Alexandre Julliard <julliard at winehq.org>

---

 dlls/ntdll/unix/signal_i386.c   | 27 ++++++++++++++++-----------
 dlls/ntdll/unix/signal_x86_64.c | 27 ++++++++++++++++-----------
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/dlls/ntdll/unix/signal_i386.c b/dlls/ntdll/unix/signal_i386.c
index c3f8c01d350..b66a7bd8995 100644
--- a/dlls/ntdll/unix/signal_i386.c
+++ b/dlls/ntdll/unix/signal_i386.c
@@ -994,7 +994,19 @@ NTSTATUS WINAPI NtSetContextThread( HANDLE handle, const CONTEXT *context )
     struct syscall_frame *frame = x86_thread_data()->syscall_frame;
     DWORD flags = context->ContextFlags & ~CONTEXT_i386;
     BOOL self = (handle == GetCurrentThread());
-    XSTATE *xs;
+
+    if ((flags & CONTEXT_XSTATE) && (cpu_info.ProcessorFeatureBits & CPU_FEATURE_AVX))
+    {
+        CONTEXT_EX *context_ex = (CONTEXT_EX *)(context + 1);
+        XSTATE *xs = (XSTATE *)((char *)context_ex + context_ex->XState.Offset);
+
+        if (context_ex->XState.Length < offsetof(XSTATE, YmmContext) ||
+            context_ex->XState.Length > sizeof(XSTATE))
+            return STATUS_INVALID_PARAMETER;
+        if ((xs->Mask & XSTATE_MASK_GSSE) && (context_ex->XState.Length < sizeof(XSTATE)))
+            return STATUS_BUFFER_OVERFLOW;
+    }
+    else flags &= ~CONTEXT_XSTATE;
 
     /* debug registers require a server call */
     if (self && (flags & CONTEXT_DEBUG_REGISTERS))
@@ -1066,25 +1078,18 @@ NTSTATUS WINAPI NtSetContextThread( HANDLE handle, const CONTEXT *context )
         }
         xsave->xstate.mask |= XSTATE_MASK_LEGACY_FLOATING_POINT;
     }
-    if ((cpu_info.ProcessorFeatureBits & CPU_FEATURE_AVX) && (xs = xstate_from_context( context )))
+    if (flags & CONTEXT_XSTATE)
     {
         struct syscall_xsave *xsave = get_syscall_xsave( frame );
         CONTEXT_EX *context_ex = (CONTEXT_EX *)(context + 1);
-
-        if (context_ex->XState.Length < offsetof(XSTATE, YmmContext)
-            || context_ex->XState.Length > sizeof(XSTATE))
-            return STATUS_INVALID_PARAMETER;
+        XSTATE *xs = (XSTATE *)((char *)context_ex + context_ex->XState.Offset);
 
         if (xs->Mask & XSTATE_MASK_GSSE)
         {
-            if (context_ex->XState.Length < sizeof(XSTATE))
-                return STATUS_BUFFER_OVERFLOW;
-
             xsave->xstate.mask |= XSTATE_MASK_GSSE;
             memcpy( &xsave->xstate.ymm_high, &xs->YmmContext, sizeof(xsave->xstate.ymm_high) );
         }
-        else
-            xsave->xstate.mask &= ~XSTATE_MASK_GSSE;
+        else xsave->xstate.mask &= ~XSTATE_MASK_GSSE;
     }
 
     return STATUS_SUCCESS;
diff --git a/dlls/ntdll/unix/signal_x86_64.c b/dlls/ntdll/unix/signal_x86_64.c
index 59c030ae6a4..086bfdfe066 100644
--- a/dlls/ntdll/unix/signal_x86_64.c
+++ b/dlls/ntdll/unix/signal_x86_64.c
@@ -1618,7 +1618,19 @@ NTSTATUS WINAPI NtSetContextThread( HANDLE handle, const CONTEXT *context )
     BOOL self = (handle == GetCurrentThread());
     struct syscall_frame *frame;
     struct syscall_xsave *xsave;
-    XSTATE *xs;
+
+    if ((flags & CONTEXT_XSTATE) && (cpu_info.ProcessorFeatureBits & CPU_FEATURE_AVX))
+    {
+        CONTEXT_EX *context_ex = (CONTEXT_EX *)(context + 1);
+        XSTATE *xs = (XSTATE *)((char *)context_ex + context_ex->XState.Offset);
+
+        if (context_ex->XState.Length < offsetof(XSTATE, YmmContext) ||
+            context_ex->XState.Length > sizeof(XSTATE))
+            return STATUS_INVALID_PARAMETER;
+        if ((xs->Mask & XSTATE_MASK_GSSE) && (context_ex->XState.Length < sizeof(XSTATE)))
+            return STATUS_BUFFER_OVERFLOW;
+    }
+    else flags &= ~CONTEXT_XSTATE;
 
     /* debug registers require a server call */
     if (self && (flags & CONTEXT_DEBUG_REGISTERS))
@@ -1684,24 +1696,17 @@ NTSTATUS WINAPI NtSetContextThread( HANDLE handle, const CONTEXT *context )
         xsave->xsave = context->u.FltSave;
         xsave->xstate.Mask |= XSTATE_MASK_LEGACY;
     }
-    if ((cpu_info.ProcessorFeatureBits & CPU_FEATURE_AVX) && (xs = xstate_from_context( context )))
+    if (flags & CONTEXT_XSTATE)
     {
         CONTEXT_EX *context_ex = (CONTEXT_EX *)(context + 1);
-
-        if (context_ex->XState.Length < offsetof(XSTATE, YmmContext)
-            || context_ex->XState.Length > sizeof(XSTATE))
-            return STATUS_INVALID_PARAMETER;
+        XSTATE *xs = (XSTATE *)((char *)context_ex + context_ex->XState.Offset);
 
         if (xs->Mask & XSTATE_MASK_GSSE)
         {
-            if (context_ex->XState.Length < sizeof(XSTATE))
-                return STATUS_BUFFER_OVERFLOW;
-
             xsave->xstate.Mask |= XSTATE_MASK_GSSE;
             memcpy( &xsave->xstate.YmmContext, &xs->YmmContext, sizeof(xs->YmmContext) );
         }
-        else
-            xsave->xstate.Mask &= ~XSTATE_MASK_GSSE;
+        else xsave->xstate.Mask &= ~XSTATE_MASK_GSSE;
     }
     return STATUS_SUCCESS;
 }