path: root/main/lua-ldbus/0001-fix-lua_State-usage-for-callbacks.patch

From 09fd539b42fd74e13b46970b8718f11bebbb23c2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20Ter=C3=A4s?= <timo.teras@iki.fi>
Date: Thu, 15 Oct 2015 13:44:00 +0300
Subject: [PATCH] fix lua_State usage for callbacks

We cannot store the lua_State at callback registration time, as that
lua coroutine might've been GC'ed at the callback invocation time.

This patch solves it by storing the invoking lua_State before calling
any dbus function that might call the callbacks. The callbacks are
then executed in the lua_State which called the dbus method invoking
the callback.

As side effect, we need to only track the lua reference for callbacks,
and avoid one malloc/free pair.
---
 src/connection.c   | 127 ++++++++++++++++++++++-------------------------------
 src/ldbus.c        |  20 +++++++++
 src/ldbus.h        |  23 +++++++---
 src/pending_call.c |  26 ++++-------
 src/timeout.c      |  26 +++--------
 src/timeout.h      |   4 +-
 src/watch.c        |  28 +++---------
 src/watch.h        |   4 +-
 8 files changed, 115 insertions(+), 143 deletions(-)

diff --git a/src/connection.c b/src/connection.c
index 8bbe74b..3c71034 100644
--- a/src/connection.c
+++ b/src/connection.c
@@ -46,10 +46,10 @@ static int ldbus_connection_gc(lua_State *L) {
 		   If you dbus_connection_unref() for the last time without closing the connection,
 		   the results are undefined
 		*/
-		dbus_connection_close(udata->connection);
+		ldbus_call(L, dbus_connection_close(udata->connection));
 	}
 
-	dbus_connection_unref(udata->connection);
+	ldbus_call(L, dbus_connection_unref(udata->connection));
 
 	return 0;
 }
@@ -96,7 +96,7 @@ static int ldbus_connection_send(lua_State *L) {
 	DBusMessage *message = check_DBusMessage(L, 2);
 
 	unsigned int serial;
-	lua_pushboolean(L, dbus_connection_send(connection, message, &serial));
+	lua_pushboolean(L, ldbus_wrap(L, dbus_connection_send(connection, message, &serial)));
 	lua_pushinteger(L, serial);
 
 	return 2;
@@ -108,7 +108,7 @@ static int ldbus_connection_send_with_reply(lua_State *L) {
 	int timeout_milliseconds = lua_isnoneornil(L, 3) ? (-1) : (luaL_checknumber(L, 3) * 1000);
 
 	DBusPendingCall *pending;
-	if (!dbus_connection_send_with_reply(connection, message, &pending, timeout_milliseconds)) {
+	if (!ldbus_wrap(L, dbus_connection_send_with_reply(connection, message, &pending, timeout_milliseconds))) {
 		return luaL_error(L, LDBUS_NO_MEMORY);
 	}
 
@@ -123,7 +123,7 @@ static int ldbus_connection_send_with_reply_and_block(lua_State *L) {
 	DBusMessage *reply;
 	DBusError error;
 	dbus_error_init(&error);
-	reply = dbus_connection_send_with_reply_and_block(connection, message, timeout_milliseconds, &error);
+	reply = ldbus_wrap(L, dbus_connection_send_with_reply_and_block(connection, message, timeout_milliseconds, &error));
 	if (dbus_error_is_set(&error)) {
 		lua_pushnil(L);
 		lua_pushstring(L, error.message);
@@ -147,7 +147,7 @@ static int ldbus_connection_read_write_dispatch(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
 	int timeout_milliseconds = luaL_optinteger(L, 2, -1);
 
-	lua_pushboolean(L, dbus_connection_read_write_dispatch(connection, timeout_milliseconds));
+	lua_pushboolean(L, ldbus_wrap(L, dbus_connection_read_write_dispatch(connection, timeout_milliseconds)));
 
 	return 1;
 }
@@ -156,7 +156,7 @@ static int ldbus_connection_read_write(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
 	int timeout_milliseconds = luaL_optinteger(L, 2, -1);
 
-	lua_pushboolean(L, dbus_connection_read_write(connection, timeout_milliseconds));
+	lua_pushboolean(L, ldbus_wrap(L, dbus_connection_read_write(connection, timeout_milliseconds)));
 
 	return 1;
 }
@@ -164,7 +164,7 @@ static int ldbus_connection_read_write(lua_State *L) {
 static int ldbus_connection_pop_message(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
 
-	DBusMessage *message = dbus_connection_pop_message(connection);
+	DBusMessage *message = ldbus_wrap(L, dbus_connection_pop_message(connection));
 
 	if (message == NULL) {
 		lua_pushnil(L);
@@ -183,16 +183,15 @@ static int ldbus_connection_get_dispatch_status(lua_State *L) {
 
 static int ldbus_connection_dispatch(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
-
-	lua_pushstring(L, DispatchStatus_lst [ dbus_connection_dispatch(connection) ]);
+	lua_pushstring(L, DispatchStatus_lst [ ldbus_wrap(L, dbus_connection_dispatch(connection)) ]);
 	return 1;
 }
 
 static int ldbus_connection_set_watch_functions(lua_State *L) {
-	ldbus_watch_udata *data;
-
 	DBusConnection *connection = check_DBusConnection(L, 1);
 	int has_toggle = lua_isnil(L, 4);
+	intptr_t ref;
+
 	lua_settop(L, 4);
 	/* Place a table below the 3 callback argument */
 	lua_createtable(L, 0, 3);
@@ -206,25 +205,22 @@ static int ldbus_connection_set_watch_functions(lua_State *L) {
 	luaL_requiref(L, "ldbus.watch", luaopen_ldbus_watch, FALSE);
 	lua_pop(L, 1);
 
-	if ((data = malloc(sizeof(ldbus_watch_udata))) == NULL) return luaL_error(L, LDBUS_NO_MEMORY);
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-
-	if (!dbus_connection_set_watch_functions(connection,
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	if (!ldbus_wrap(L, dbus_connection_set_watch_functions(connection,
 			ldbus_watch_add_function, ldbus_watch_remove_function,
 			has_toggle ? NULL : ldbus_watch_toggled_function,
-			(void *)data, ldbus_watch_free_data_function)) {
-		free(data);
+			(void *) ref, ldbus_watch_free_data_function))) {
+		luaL_unref(L, LUA_REGISTRYINDEX, ref);
 		return luaL_error(L, LDBUS_NO_MEMORY);
-	};
+	}
 	lua_pushboolean(L, TRUE);
 	return 1;
 }
 
 static int ldbus_connection_set_timeout_functions(lua_State *L) {
-	ldbus_timeout_udata *data;
-
 	DBusConnection *connection = check_DBusConnection(L, 1);
+	intptr_t ref;
+
 	lua_settop(L, 4);
 	/* Place a table below the 3 callback argument */
 	lua_createtable(L, 0, 3);
@@ -238,52 +234,41 @@ static int ldbus_connection_set_timeout_functions(lua_State *L) {
 	luaL_requiref(L, "ldbus.timeout", lua_open_ldbus_timeout, FALSE);
 	lua_pop(L, 1);
 
-	if ((data = malloc(sizeof(ldbus_timeout_udata))) == NULL) return luaL_error(L, LDBUS_NO_MEMORY);
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-
-	if (!dbus_connection_set_timeout_functions(connection,
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	if (!ldbus_wrap(L, dbus_connection_set_timeout_functions(connection,
 			ldbus_timeout_add_function, ldbus_timeout_remove_function, ldbus_timeout_toggled_function,
-			(void *)data, ldbus_timeout_free_data_function)) {
-		free(data);
+			(void *)ref, ldbus_timeout_free_data_function))) {
+		luaL_unref(L, LUA_REGISTRYINDEX, ref);
 		return luaL_error(L, LDBUS_NO_MEMORY);
 	};
 	lua_pushboolean(L, TRUE);
 	return 1;
 }
 
-static void free_data_function(void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
-	luaL_unref(L, LUA_REGISTRYINDEX, ref);
-	free(data);
-}
-
 static void wakeup_main_function(void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
+
 	lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
 	lua_pcall(L, 0, 0, 0);
 }
 
 static int ldbus_connection_set_wakeup_main_function(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
-	ldbus_callback_udata *data;
+	intptr_t ref;
+
 	luaL_checktype(L, 2, LUA_TFUNCTION);
 	lua_settop(L, 2);
-	if ((data = malloc(sizeof(ldbus_callback_udata))) == NULL) {
-		return luaL_error(L, LDBUS_NO_MEMORY);
-	}
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-	dbus_connection_set_wakeup_main_function(connection, wakeup_main_function, data, free_data_function);
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	dbus_connection_set_wakeup_main_function(connection, wakeup_main_function, (void*) ref, ldbus_free_ref);
 	lua_pushboolean(L, 1);
 	return 1;
 }
 
 static void dispatch_status_function(DBusConnection *connection, DBusDispatchStatus new_status, void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
+
 	UNUSED(connection);
 	lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
 	lua_pushstring(L, DispatchStatus_lst[new_status]);
@@ -292,15 +277,12 @@ static void dispatch_status_function(DBusConnection *connection, DBusDispatchSta
 
 static int ldbus_connection_set_dispatch_status_function(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
-	ldbus_callback_udata *data;
+	intptr_t ref;
+
 	luaL_checktype(L, 2, LUA_TFUNCTION);
 	lua_settop(L, 2);
-	if ((data = malloc(sizeof(ldbus_callback_udata))) == NULL) {
-		return luaL_error(L, LDBUS_NO_MEMORY);
-	}
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-	dbus_connection_set_dispatch_status_function(connection, dispatch_status_function, data, free_data_function);
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	dbus_connection_set_dispatch_status_function(connection, dispatch_status_function, (void *) ref, ldbus_free_ref);
 	lua_pushboolean(L, 1);
 	return 1;
 }
@@ -356,16 +338,17 @@ static int ldbus_connection_has_messages_to_send(lua_State *L) {
 }
 
 static void unregister_function(DBusConnection *connection, void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
+
 	UNUSED(connection);
 	luaL_unref(L, LUA_REGISTRYINDEX, ref);
-	free(data);
 }
 
 static DBusHandlerResult message_function(DBusConnection *connection, DBusMessage *message, void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
+
 	UNUSED(connection);
 	if (!lua_checkstack(L, 2)) {
 		return DBUS_HANDLER_RESULT_NEED_MEMORY;
@@ -396,16 +379,13 @@ static const DBusObjectPathVTable VTable = {
 static int ldbus_connection_register_object_path(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
 	const char *path = luaL_checkstring(L, 2);
-	ldbus_callback_udata *data;
+	intptr_t ref;
+
 	luaL_checktype(L, 3, LUA_TFUNCTION);
 	lua_settop(L, 3);
-	if ((data = malloc(sizeof(ldbus_callback_udata))) == NULL) {
-		return luaL_error(L, LDBUS_NO_MEMORY);
-	}
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-	if (!dbus_connection_register_object_path(connection, path, &VTable, data)) {
-		free(data);
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	if (!dbus_connection_register_object_path(connection, path, &VTable, (void*) ref)) {
+		luaL_unref(L, LUA_REGISTRYINDEX, ref);
 		return luaL_error(L, "unknown error");
 	}
 	lua_pushboolean(L, 1);
@@ -415,16 +395,13 @@ static int ldbus_connection_register_object_path(lua_State *L) {
 static int ldbus_connection_register_fallback(lua_State *L) {
 	DBusConnection *connection = check_DBusConnection(L, 1);
 	const char *path = luaL_checkstring(L, 2);
-	ldbus_callback_udata *data;
+	intptr_t ref;
+
 	luaL_checktype(L, 3, LUA_TFUNCTION);
 	lua_settop(L, 3);
-	if ((data = malloc(sizeof(ldbus_callback_udata))) == NULL) {
-		return luaL_error(L, LDBUS_NO_MEMORY);
-	}
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-	if (!dbus_connection_register_fallback(connection, path, &VTable, data)) {
-		free(data);
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	if (!dbus_connection_register_fallback(connection, path, &VTable, (void*) ref)) {
+		luaL_unref(L, LUA_REGISTRYINDEX, ref);
 		return luaL_error(L, "unknown error");
 	}
 	lua_pushboolean(L, 1);
diff --git a/src/ldbus.c b/src/ldbus.c
index 0ce3333..ab32257 100644
--- a/src/ldbus.c
+++ b/src/ldbus.c
@@ -12,6 +12,26 @@
 #include "watch.h"
 #include "ldbus.h"
 
+static __thread lua_State *__ldbus_state;
+
+lua_State *ldbus_get_state()
+{
+	return __ldbus_state;
+}
+
+lua_State *ldbus_set_state(lua_State *L)
+{
+	lua_State *old = __ldbus_state;
+	__ldbus_state = L;
+	return old;
+}
+
+void ldbus_free_ref(void *data) {
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
+	luaL_unref(L, LUA_REGISTRYINDEX, ref);
+}
+
 int tostring(lua_State *L) {
 	if (!luaL_getmetafield(L, 1, "__udtype")) {
 		lua_pushstring(L, "object with a generic __tostring metamethod but no __type metafield");
diff --git a/src/ldbus.h b/src/ldbus.h
index a33422f..9c01263 100644
--- a/src/ldbus.h
+++ b/src/ldbus.h
@@ -1,23 +1,36 @@
 #ifndef LDBUS_H
 #define LDBUS_H
 
+#include <stdint.h>
 #include <lua.h>
 
 #define UNUSED(x) (void)(x)
 #define LDBUS_INTERNAL __attribute__ ((visibility ("internal")))
 #define LDBUS_NO_MEMORY "no memory"
 
-typedef struct {
-	lua_State* L;
-	int ref;
-} ldbus_callback_udata;
-
 enum ldbus_callback_indexes {
 	DBUS_LUA_FUNC_ADD,
 	DBUS_LUA_FUNC_REMOVE,
 	DBUS_LUA_FUNC_TOGGLE
 };
 
+#define ldbus_wrap(L, func) __extension__ ({	\
+	lua_State *old = ldbus_set_state(L);	\
+	__auto_type result = (func);		\
+	ldbus_set_state(old);			\
+	result;					\
+})
+
+#define ldbus_call(L, func) __extension__ ({	\
+	lua_State *old = ldbus_set_state(L);	\
+	(func);					\
+	ldbus_set_state(old);			\
+})
+
+LDBUS_INTERNAL lua_State *ldbus_get_state();
+LDBUS_INTERNAL lua_State *ldbus_set_state(lua_State *);
+LDBUS_INTERNAL void ldbus_free_ref(void *ref);
+
 LDBUS_INTERNAL int tostring(lua_State *);
 
 int luaopen_ldbus(lua_State *);
diff --git a/src/pending_call.c b/src/pending_call.c
index 0dbc6ef..98b5d11 100644
--- a/src/pending_call.c
+++ b/src/pending_call.c
@@ -17,36 +17,28 @@ static DBusPendingCall* checkPendingCall(lua_State *L, int arg) {
 
 static int ldbus_pending_call_unref(lua_State *L) {
 	DBusPendingCall* pending = checkPendingCall(L, 1);
-	dbus_pending_call_unref(pending);
+	ldbus_call(L, dbus_pending_call_unref(pending));
 	return 0;
 }
 
 static void pending_notify_function(DBusPendingCall *pending, void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
+
 	UNUSED(pending);
 	lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
 	lua_pcall(L, 0, 0, 0);
 }
 
-static void free_data_function(void *data) {
-	lua_State *L = ((ldbus_callback_udata*)data)->L;
-	int ref = ((ldbus_callback_udata*)data)->ref;
-	luaL_unref(L, LUA_REGISTRYINDEX, ref);
-	free(data);
-}
-
 static int ldbus_pending_call_set_notify(lua_State *L) {
 	DBusPendingCall* pending = checkPendingCall(L, 1);
-	ldbus_callback_udata *data;
+	intptr_t ref;
+
 	luaL_checktype(L, 2, LUA_TFUNCTION);
 	lua_settop(L, 2);
-	if ((data = malloc(sizeof(ldbus_callback_udata))) == NULL) {
-		return luaL_error(L, LDBUS_NO_MEMORY);
-	}
-	data->L = L;
-	data->ref = luaL_ref(L, LUA_REGISTRYINDEX);
-	if (!dbus_pending_call_set_notify(pending, pending_notify_function, data, free_data_function)) {
+	ref = luaL_ref(L, LUA_REGISTRYINDEX);
+	if (!ldbus_wrap(L, dbus_pending_call_set_notify(pending, pending_notify_function, (void*) ref, ldbus_free_ref))) {
+		luaL_unref(L, LUA_REGISTRYINDEX, ref);
 		return luaL_error(L, LDBUS_NO_MEMORY);
 	}
 	lua_pushboolean(L, 1);
diff --git a/src/timeout.c b/src/timeout.c
index 573700d..b4cdfe5 100644
--- a/src/timeout.c
+++ b/src/timeout.c
@@ -27,7 +27,7 @@ static int ldbus_timeout_handle(lua_State *L) {
 	if (timeout == NULL) {
 		lua_pushnil(L);
 	} else {
-		lua_pushboolean(L, dbus_timeout_handle(timeout));
+		lua_pushboolean(L, ldbus_wrap(L, dbus_timeout_handle(timeout)));
 	}
 	return 1;
 }
@@ -46,8 +46,8 @@ LDBUS_INTERNAL void push_DBusTimeout(lua_State *L, DBusTimeout *timeout) {
 }
 
 LDBUS_INTERNAL dbus_bool_t ldbus_timeout_add_function(DBusTimeout *timeout, void *data) {
-	lua_State *L = ((ldbus_timeout_udata*)data)->L;
-	int ref = ((ldbus_timeout_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
 
 	if (!lua_checkstack(L, 4)) return FALSE;
 
@@ -63,9 +63,8 @@ LDBUS_INTERNAL dbus_bool_t ldbus_timeout_add_function(DBusTimeout *timeout, void
 }
 
 LDBUS_INTERNAL void ldbus_timeout_remove_function(DBusTimeout *timeout, void *data) {
-	lua_State *L = ((ldbus_timeout_udata*)data)->L;
-	int ref = ((ldbus_timeout_udata*)data)->ref;
-
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
 	DBusTimeout **udata;
 
 	/* Lookup remove callback from ref in `data` */
@@ -94,13 +93,11 @@ LDBUS_INTERNAL void ldbus_timeout_remove_function(DBusTimeout *timeout, void *da
 	if (udata != NULL) {
 		*udata = NULL;
 	}
-
-	return;
 }
 
 LDBUS_INTERNAL void ldbus_timeout_toggled_function(DBusTimeout *timeout, void *data) {
-	lua_State *L = ((ldbus_timeout_udata*)data)->L;
-	int ref = ((ldbus_timeout_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
 
 	/* Lookup remove callback from ref in `data` */
 	lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
@@ -114,15 +111,6 @@ LDBUS_INTERNAL void ldbus_timeout_toggled_function(DBusTimeout *timeout, void *d
 
 	/* Call the callback, with looked-up timeout as argument */
 	lua_pcall(L, 1, 0, 0);
-	return;
-}
-
-LDBUS_INTERNAL void ldbus_timeout_free_data_function(void *data) {
-	lua_State *L = ((ldbus_timeout_udata*)data)->L;
-	int ref = ((ldbus_timeout_udata*)data)->ref;
-
-	luaL_unref(L, LUA_REGISTRYINDEX, ref);
-	free(data);
 }
 
 int lua_open_ldbus_timeout(lua_State *L) {
diff --git a/src/timeout.h b/src/timeout.h
index 5794563..4eab6e6 100644
--- a/src/timeout.h
+++ b/src/timeout.h
@@ -12,13 +12,11 @@
 
 #define check_DBusTimeout(L, arg) (*(DBusTimeout **)luaL_checkudata((L), (arg), DBUS_TIMEOUT_METATABLE))
 
-typedef ldbus_callback_udata ldbus_timeout_udata;
-
 LDBUS_INTERNAL void push_DBusTimeout(lua_State *L, DBusTimeout *timeout);
 LDBUS_INTERNAL dbus_bool_t ldbus_timeout_add_function(DBusTimeout *timeout, void *data);
 LDBUS_INTERNAL void ldbus_timeout_remove_function(DBusTimeout *timeout, void *data);
 LDBUS_INTERNAL void ldbus_timeout_toggled_function(DBusTimeout *timeout, void *data);
-LDBUS_INTERNAL void ldbus_timeout_free_data_function(void *data);
+#define ldbus_timeout_free_data_function ldbus_free_ref
 
 int lua_open_ldbus_timeout(lua_State *L);
 
diff --git a/src/watch.c b/src/watch.c
index 2773212..173bc96 100644
--- a/src/watch.c
+++ b/src/watch.c
@@ -42,11 +42,9 @@ static int ldbus_watch_get_flags(lua_State *L) {
 static int ldbus_watch_handle(lua_State *L) {
 	DBusWatch *watch = check_DBusWatch(L, 1);
 	int flags;
-	dbus_bool_t ok;
 	luaL_argcheck(L, watch != NULL, 1, "watch invalid");
 	flags = luaL_checkinteger(L, 2);
-	ok = dbus_watch_handle(watch, flags);
-	lua_pushboolean(L, ok);
+	lua_pushboolean(L, ldbus_wrap(L, dbus_watch_handle(watch, flags)));
 	return 1;
 }
 
@@ -64,8 +62,8 @@ LDBUS_INTERNAL void push_DBusWatch(lua_State *L, DBusWatch *watch) {
 }
 
 LDBUS_INTERNAL dbus_bool_t ldbus_watch_add_function(DBusWatch *watch, void *data) {
-	lua_State *L = ((ldbus_watch_udata*)data)->L;
-	int ref = ((ldbus_watch_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
 
 	if (!lua_checkstack(L, 4)) return FALSE;
 
@@ -81,9 +79,8 @@ LDBUS_INTERNAL dbus_bool_t ldbus_watch_add_function(DBusWatch *watch, void *data
 }
 
 LDBUS_INTERNAL void ldbus_watch_remove_function(DBusWatch *watch, void *data) {
-	lua_State *L = ((ldbus_watch_udata*)data)->L;
-	int ref = ((ldbus_watch_udata*)data)->ref;
-
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
 	DBusWatch **udata;
 
 	/* Lookup remove callback from ref in `data` */
@@ -112,13 +109,11 @@ LDBUS_INTERNAL void ldbus_watch_remove_function(DBusWatch *watch, void *data) {
 	if (udata != NULL) {
 		*udata = NULL;
 	}
-
-	return;
 }
 
 LDBUS_INTERNAL void ldbus_watch_toggled_function(DBusWatch *watch, void *data) {
-	lua_State *L = ((ldbus_watch_udata*)data)->L;
-	int ref = ((ldbus_watch_udata*)data)->ref;
+	lua_State *L = ldbus_get_state();
+	intptr_t ref = (intptr_t) data;
 
 	/* Lookup remove callback from ref in `data` */
 	lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
@@ -132,15 +127,6 @@ LDBUS_INTERNAL void ldbus_watch_toggled_function(DBusWatch *watch, void *data) {
 
 	/* Call the callback, with looked-up watch as argument */
 	lua_pcall(L, 1, 0, 0);
-	return;
-}
-
-LDBUS_INTERNAL void ldbus_watch_free_data_function(void *data) {
-	lua_State *L = ((ldbus_watch_udata*)data)->L;
-	int ref = ((ldbus_watch_udata*)data)->ref;
-
-	luaL_unref(L, LUA_REGISTRYINDEX, ref);
-	free(data);
 }
 
 int luaopen_ldbus_watch(lua_State *L) {
diff --git a/src/watch.h b/src/watch.h
index f37f877..d3abf08 100644
--- a/src/watch.h
+++ b/src/watch.h
@@ -12,13 +12,11 @@
 
 #define check_DBusWatch(L, arg) (*(DBusWatch **)luaL_checkudata((L), (arg), DBUS_WATCH_METATABLE))
 
-typedef ldbus_callback_udata ldbus_watch_udata;
-
 LDBUS_INTERNAL void push_DBusWatch(lua_State *L, DBusWatch *watch);
 LDBUS_INTERNAL dbus_bool_t ldbus_watch_add_function(DBusWatch *watch, void *data);
 LDBUS_INTERNAL void ldbus_watch_remove_function(DBusWatch *watch, void *data);
 LDBUS_INTERNAL void ldbus_watch_toggled_function(DBusWatch *watch, void *data);
-LDBUS_INTERNAL void ldbus_watch_free_data_function(void *data);
+#define ldbus_watch_free_data_function ldbus_free_ref
 
 int luaopen_ldbus_watch(lua_State *L);
 
-- 
2.6.1