From 57095a91b053a5570bd797782b51b84fdd42173c Mon Sep 17 00:00:00 2001
From: Paul Gofman <pgofman@codeweavers.com>
Date: Wed, 19 Jul 2023 10:02:06 -0600
Subject: [PATCH] server: Cancel socket asyncs when the last handle in process
 is closed.

---
 dlls/ws2_32/tests/afd.c | 11 +++++++----
 server/async.c          | 22 ++++++++++++++++++++++
 server/file.h           |  1 +
 server/handle.c         | 15 +++++++++++++++
 server/handle.h         |  1 +
 server/sock.c           |  3 +--
 6 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/dlls/ws2_32/tests/afd.c b/dlls/ws2_32/tests/afd.c
index ef3e6385c7f..97139605bf1 100644
--- a/dlls/ws2_32/tests/afd.c
+++ b/dlls/ws2_32/tests/afd.c
@@ -2858,17 +2858,20 @@ static void test_async_cancel_on_handle_close(void)
 
             closesocket(listener);
 
+            /* Canceled asyncs with completion port and no event do not update IOSB before removing completion. */
+            todo_wine_if(other_process && tests[i].apc_context && !tests[i].event)
             ok(io.Status == 0xcccccccc, "got %#lx\n", io.Status);
+
             memset(&io, 0xcc, sizeof(io));
             key = 0xcc;
             value = 0;
             ret = NtRemoveIoCompletion(port, &key, &value, &io, &zero);
             if (other_process && tests[i].apc_context && !tests[i].event)
             {
-                todo_wine ok(!ret, "got %#lx\n", ret);
-                todo_wine ok(!key, "got key %#Ix\n", key);
-                todo_wine ok(value == 0xdeadbeef, "got value %#Ix\n", value);
-                todo_wine ok(io.Status == STATUS_CANCELLED, "got %#lx\n", io.Status);
+                ok(!ret, "got %#lx\n", ret);
+                ok(!key, "got key %#Ix\n", key);
+                ok(value == 0xdeadbeef, "got value %#Ix\n", value);
+                ok(io.Status == STATUS_CANCELLED, "got %#lx\n", io.Status);
             }
             else
             {
diff --git a/server/async.c b/server/async.c
index 26946b5f5ce..9cb251df5ce 100644
--- a/server/async.c
+++ b/server/async.c
@@ -614,6 +614,28 @@ void cancel_process_asyncs( struct process *process )
     cancel_async( process, NULL, NULL, 0 );
 }
 
+int async_close_obj_handle( struct object *obj, struct process *process, obj_handle_t handle )
+{
+    /* Handle a special case when the last object handle in the given process is closed.
+     * If this is the last object handle overall that is handled in object's close_handle and
+     * destruction. */
+    struct async *async;
+
+    if (obj->handle_count == 1 || get_obj_handle_count( process, obj ) != 1) return 1;
+
+restart:
+    LIST_FOR_EACH_ENTRY( async, &process->asyncs, struct async, process_entry )
+    {
+        if (async->terminated || async->canceled || get_fd_user( async->fd ) != obj) continue;
+        if (!async->completion || !async->data.apc_context || async->event) continue;
+
+        async->canceled = 1;
+        fd_cancel_async( async->fd, async );
+        goto restart;
+    }
+    return 1;
+}
+
 void cancel_terminating_thread_asyncs( struct thread *thread )
 {
     struct async *async;
diff --git a/server/file.h b/server/file.h
index 210fcec5e78..a43b4afad08 100644
--- a/server/file.h
+++ b/server/file.h
@@ -246,6 +246,7 @@ extern struct thread *async_get_thread( struct async *async );
 extern struct async *find_pending_async( struct async_queue *queue );
 extern void cancel_process_asyncs( struct process *process );
 extern void cancel_terminating_thread_asyncs( struct thread *thread );
+extern int async_close_obj_handle( struct object *obj, struct process *process, obj_handle_t handle );
 
 static inline void init_async_queue( struct async_queue *queue )
 {
diff --git a/server/handle.c b/server/handle.c
index 38ad80da267..0595fdb403b 100644
--- a/server/handle.c
+++ b/server/handle.c
@@ -520,6 +520,21 @@ obj_handle_t find_inherited_handle( struct process *process, const struct object
     return 0;
 }
 
+/* return number of open handles to the object in the process */
+unsigned int get_obj_handle_count( struct process *process, const struct object *obj )
+{
+    struct handle_table *table = process->handles;
+    struct handle_entry *ptr;
+    unsigned int count = 0;
+    int i;
+
+    if (!table) return 0;
+
+    for (i = 0, ptr = table->entries; i <= table->last; i++, ptr++)
+        if (ptr->ptr == obj) ++count;
+    return count;
+}
+
 /* get/set the handle reserved flags */
 /* return the old flags (or -1 on error) */
 static int set_handle_flags( struct process *process, obj_handle_t handle, int mask, int flags )
diff --git a/server/handle.h b/server/handle.h
index ac3104dc003..1d02e040258 100644
--- a/server/handle.h
+++ b/server/handle.h
@@ -48,6 +48,7 @@ extern obj_handle_t open_object( struct process *process, obj_handle_t parent, u
                                  const struct object_ops *ops, const struct unicode_str *name,
                                  unsigned int attr );
 extern obj_handle_t find_inherited_handle( struct process *process, const struct object_ops *ops );
+extern unsigned int get_obj_handle_count( struct process *process, const struct object *obj );
 extern void close_process_handles( struct process *process );
 extern struct handle_table *alloc_handle_table( struct process *process, int count );
 extern struct handle_table *copy_handle_table( struct process *process, struct process *parent,
diff --git a/server/sock.c b/server/sock.c
index 34be6ea22ef..2aae9aefbf0 100644
--- a/server/sock.c
+++ b/server/sock.c
@@ -1655,8 +1655,7 @@ static int sock_close_handle( struct object *obj, struct process *process, obj_h
             if (signaled) complete_async_poll( poll_req, STATUS_SUCCESS );
         }
     }
-
-    return 1;
+    return async_close_obj_handle( obj, process, handle );
 }
 
 static void sock_destroy( struct object *obj )
-- 
2.11.4.GIT