[PATCH 5/9] ntdll: Add process-local handle tracking (part 2)

[Daniel Santos]

Adds a facility for ntdll to track handles locally and attach
information to them, allowing for some operations to be performed
directly by the client process.

Signed-off-by: Daniel Santos 
---
 dlls/ntdll/om.c | 186 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 186 insertions(+)

diff --git a/dlls/ntdll/om.c b/dlls/ntdll/om.c
index 6233633..dfb0f76 100644
--- a/dlls/ntdll/om.c
+++ b/dlls/ntdll/om.c
@@ -302,6 +302,192 @@ static struct {
     PTHREAD_MUTEX_INITIALIZER                   /* non-recursive fast lock */
 };

+/* Adds the object to the handle database, increasing its reference count to account for
+ * it being referenced by the handle tree its self (don't ntdll_object_release it to
+ * resolve this increment).
+ *
+ * Locking:
+ *      Locks handles.mutex */
+__attribute__((noinline, flatten))
+int ntdll_handle_add(struct ntdll_object *obj)
+{
+    int ret = 0;
+
+    TRACE_(ntdll_obj)("%p\n", obj);
+
+try_again:
+    pthread_mutex_lock(&handles.mutex);
+    ret = wine_rb_put(&handles.tree, &obj->h, &obj->tree_entry);
+
+    /* increase the ref count */
+    if (!ret)
+        ntdll_object_grab(obj);
+    pthread_mutex_unlock(&handles.mutex);
+
+    if (ret)
+    {
+        struct ntdll_object *existing = ntdll_handle_find(obj->h);
+
+        /* bug if we get here */
+        ERR("Failed to insert handle %p.\n", obj->h);
+        if (existing)
+            ERR("Existing object: %s\n", ntdll_object_dump(existing));
+        ERR("New object: %s\n", ntdll_object_dump(obj));
+        assert(0);
+
+        /* should probably just exit(1) here */
+        ntdll_handle_remove(obj->h);
+        goto try_again;
+    }
+
+    TRACE("Added obj %s\n", ntdll_object_dump(obj));
+
+    return ret;
+}
+
+/* Remove an object from the handles tree. If no object with the specified handle is found
+ * then nothing is done. If an object is found but it's refcount is 0 then a failed assertion
+ * will result.
+ *
+ * Locking:
+ *      Locks handles.mutex */
+__attribute__((noinline))
+void ntdll_handle_remove(const HANDLE h)
+{
+    struct ntdll_object *obj = ntdll_handle_find(h);
+
+    if (!obj)
+        return;
+
+    TRACE_(ntdll_obj)("h = %p) obj = %s\n", h, ntdll_object_dump(obj));
+    assert(obj->refcount >= 1);
+
+    pthread_mutex_lock(&handles.mutex);
+    /* FIXME: a more efficient wine_rb_remove_by_entry would be nice here */
+    wine_rb_remove(&handles.tree, &h);
+    pthread_mutex_unlock(&handles.mutex);
+
+    ntdll_object_release(obj);
+}
+
+
+
+/* Searches for an object in the process-local database with the specified handle.
+ * If found, the object's refcount is incremented and a pointer to the object is
+ * returned -- it will then be the responsibility of the caller to release the
+ * object when done.
+ *
+ * Locking:
+ *      Locks handles.mutex */
+struct ntdll_object *ntdll_handle_find(const HANDLE h)
+{
+    struct wine_rb_entry *entry;
+    struct ntdll_object *ret = NULL;
+
+    pthread_mutex_lock(&handles.mutex);
+    entry = wine_rb_get(&handles.tree, &h);
+
+    if (entry)
+    {
+        ret = WINE_RB_ENTRY_VALUE(entry, struct ntdll_object, tree_entry);
+        ntdll_object_grab(ret);
+    }
+    pthread_mutex_unlock(&handles.mutex);
+
+    //TRACE_(ntdll_obj)("(%p) result: %p\n", h, ret);
+
+    return ret;
+}
+
+
+/* callback for ntdll_objects_cleanup()
+ *
+ * Locking:
+ *      Locks objects.mutex
+ *      Called when handles.mutex already locked
+ */
+
+static void ntdll_objects_cleanup_cb(struct wine_rb_entry *entry, void *context)
+{
+    struct ntdll_object *obj = WINE_RB_ENTRY_VALUE(entry, struct ntdll_object, tree_entry);
+    size_t *leaked_handles = (size_t *)context;
+
+    ++(*leaked_handles);
+
+    ERR_(ntdll_obj)("Leaked object handle: %s\n", ntdll_object_dump(obj));
+
+    /* handles.mutex already locked */
+    wine_rb_remove(&handles.tree, &obj->h);
+
+    ntdll_object_release(obj);
+}
+
+/* atexit() cleanup
+ *
+ * Locking:
+ *      Locks handles.mutex --> then objects.mutex
+ */
+static void ntdll_objects_cleanup(void)
+{
+    size_t leaked_handles = 0;
+    size_t leaked_objects = 0;
+
+    TRACE_(ntdll_obj)("\n");
+    pthread_mutex_lock(&handles.mutex);
+    wine_rb_for_each_entry(&handles.tree, ntdll_objects_cleanup_cb, &leaked_handles);
+    pthread_mutex_unlock(&handles.mutex);
+
+    pthread_mutex_lock(&objects.mutex);
+    leaked_objects = list_count(&objects.list);
+    if (leaked_objects)
+    {
+        struct ntdll_object *obj;
+        //struct list *i;
+        LIST_FOR_EACH_ENTRY( obj, &objects.list, struct ntdll_object , list_entry )
+        {
+            ERR_(ntdll_obj)("Leaked object: %s\n", ntdll_object_dump(obj));
+        }
+    }
+    pthread_mutex_unlock(&objects.mutex);
+
+    if (leaked_handles || leaked_objects)
+        ERR_(ntdll_obj)("*** %zu leaked handles found, %zu leaked objects remain.\n",
+                        leaked_handles, leaked_objects);
+
+}
+
+/* initialize objects list and handles tree */
+NTSTATUS ntdll_object_db_init(void)
+{
+    NTSTATUS ret = 0;
+
+    TRACE_(ntdll_obj)("\n");
+
+    /* init objects list if not already inited */
+    pthread_mutex_lock(&objects.mutex);
+    if (!objects.list.next)
+        list_init(&objects.list);
+    pthread_mutex_unlock(&objects.mutex);
+
+    /* init red-black handle-to-object tree if not already inited */
+    pthread_mutex_lock(&handles.mutex);
+    if (!handles.tree.stack.entries)
+    {
+        /* passing handles.tree.functions instead of obj_handles_rb_ops to aid -findirect-inline
+         * (it might not matter) */
+        if (wine_rb_init(&handles.tree, handles.tree.functions) == -1)
+        {
+            ERR("Failed to initialize ntdll object handle rbtree.\n");
+            ret = ERROR_OUTOFMEMORY;
+        }
+        else
+            atexit(ntdll_objects_cleanup);
+    }
+    pthread_mutex_unlock(&handles.mutex);
+
+    return ret;
+}
+

/*
  *	Generic object functions