From 96cc6950f07350c1bd19ccea7877074b8fc0ebca Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek@codeweavers.com>
Date: Wed, 15 Apr 2020 14:55:22 +0200
Subject: [PATCH] ntdll: Use select request instead of get_suspend_context to
 transfer context back to client.

Signed-off-by: Jacek Caban <jacek@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
---
 dlls/ntdll/exception.c | 15 ---------------
 dlls/ntdll/server.c    |  7 +++++++
 server/thread.c        | 24 ++++++++++++++++++++++++
 server/thread.h        |  1 +
 4 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/dlls/ntdll/exception.c b/dlls/ntdll/exception.c
index fe3af73d34b..589ceab390b 100644
--- a/dlls/ntdll/exception.c
+++ b/dlls/ntdll/exception.c
@@ -112,25 +112,10 @@ static ULONG remove_vectored_handler( struct list *handler_list, VECTORED_HANDLE
 void wait_suspend( CONTEXT *context )
 {
     int saved_errno = errno;
-    context_t server_context;
-    DWORD flags = context->ContextFlags;
 
     /* wait with 0 timeout, will only return once the thread is no longer suspended */
     server_select( NULL, 0, SELECT_INTERRUPTIBLE, 0, context, NULL );
 
-    /* retrieve the new context */
-    SERVER_START_REQ( get_suspend_context )
-    {
-        wine_server_set_reply( req, &server_context, sizeof(server_context) );
-        wine_server_call( req );
-        if (wine_server_reply_size( reply ))
-        {
-            context_from_server( context, &server_context );
-            context->ContextFlags |= flags;  /* unchanged registers are still available */
-        }
-    }
-    SERVER_END_REQ;
-
     errno = saved_errno;
 }
 
diff --git a/dlls/ntdll/server.c b/dlls/ntdll/server.c
index 9103941e1bd..be080a9f73b 100644
--- a/dlls/ntdll/server.c
+++ b/dlls/ntdll/server.c
@@ -639,9 +639,16 @@ unsigned int server_select( const select_op_t *select_op, data_size_t size, UINT
                     wine_server_add_data( req, &server_context, sizeof(server_context) );
                     suspend_context = FALSE; /* server owns the context now */
                 }
+                if (context) wine_server_set_reply( req, &server_context, sizeof(server_context) );
                 ret = server_call_unlocked( req );
                 apc_handle  = reply->apc_handle;
                 call        = reply->call;
+                if (wine_server_reply_size( reply ))
+                {
+                    DWORD context_flags = context->ContextFlags; /* unchanged registers are still available */
+                    context_from_server( context, &server_context );
+                    context->ContextFlags |= context_flags;
+                }
             }
             SERVER_END_REQ;
 
diff --git a/server/thread.c b/server/thread.c
index d75dc238d51..b93e91cefd5 100644
--- a/server/thread.c
+++ b/server/thread.c
@@ -780,6 +780,19 @@ static int send_thread_wakeup( struct thread *thread, client_ptr_t cookie, int s
     struct wake_up_reply reply;
     int ret;
 
+    /* check if we're waking current suspend wait */
+    if (thread->suspend_context && thread->suspend_cookie == cookie
+        && signaled != STATUS_KERNEL_APC && signaled != STATUS_USER_APC)
+    {
+        if (!thread->suspend_context->flags)
+        {
+            if (current->context == current->suspend_context) thread->context = NULL;
+            free( thread->suspend_context );
+            thread->suspend_context = NULL;
+        }
+        else signaled = STATUS_KERNEL_APC; /* signal a fake APC so that client calls select to get a new context */
+    }
+
     memset( &reply, 0, sizeof(reply) );
     reply.cookie   = cookie;
     reply.signaled = signaled;
@@ -1557,6 +1570,7 @@ DECL_HANDLER(select)
         memcpy( current->suspend_context, context, sizeof(*context) );
         current->suspend_context->flags = 0;  /* to keep track of what is modified */
         current->context = current->suspend_context;
+        current->suspend_cookie = req->cookie;
     }
 
     if (!req->cookie)
@@ -1627,6 +1641,16 @@ DECL_HANDLER(select)
         }
         release_object( apc );
     }
+    else if(get_error() != STATUS_PENDING && get_reply_max_size() == sizeof(context_t)
+            && current->suspend_context && current->suspend_cookie == req->cookie)
+    {
+        if (current->suspend_context->flags)
+            set_reply_data_ptr( current->suspend_context, sizeof(context_t) );
+        else
+            free( current->suspend_context );
+        if (current->context == current->suspend_context) current->context = NULL;
+        current->suspend_context = NULL;
+    }
 }
 
 /* queue an APC for a thread or process */
diff --git a/server/thread.h b/server/thread.h
index 66e35603d36..9aabcaca996 100644
--- a/server/thread.h
+++ b/server/thread.h
@@ -77,6 +77,7 @@ struct thread
     int                    unix_tid;      /* Unix tid of client */
     context_t             *context;       /* current context if in an exception handler */
     context_t             *suspend_context; /* current context if suspended */
+    client_ptr_t           suspend_cookie;/* wait cookie of suspending select */
     client_ptr_t           teb;           /* TEB address (in client address space) */
     client_ptr_t           entry_point;   /* entry point (in client address space) */
     affinity_t             affinity;      /* affinity mask */