From 0bcc7bb5c77fffba05476d38e16ea34c06184f29 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pedro=20J=2E=20Est=C3=A9banez?= <pedrojrulez@gmail.com>
Date: Fri, 3 Feb 2023 11:42:51 +0100
Subject: Avoid losing references to objects in the native-scripting boundary

---
 core/object/class_db.h           | 29 +++++++++++++++++++++++++++++
 core/object/method_bind.h        |  4 ++++
 modules/gdscript/gdscript_vm.cpp |  8 +++++++-
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/core/object/class_db.h b/core/object/class_db.h
index 0e408e8845..0b62cf40f7 100644
--- a/core/object/class_db.h
+++ b/core/object/class_db.h
@@ -40,6 +40,8 @@
 #include "core/object/callable_method_pointer.h"
 #include "core/templates/hash_set.h"
 
+#include <type_traits>
+
 #define DEFVAL(m_defval) (m_defval)
 
 #ifdef DEBUG_METHODS_ENABLED
@@ -241,6 +243,24 @@ public:
 
 	static uint64_t get_api_hash(APIType p_api);
 
+	template <typename>
+	struct member_function_traits;
+
+	template <typename R, typename T, typename... Args>
+	struct member_function_traits<R (T::*)(Args...)> {
+		using return_type = R;
+	};
+
+	template <typename R, typename T, typename... Args>
+	struct member_function_traits<R (T::*)(Args...) const> {
+		using return_type = R;
+	};
+
+	template <typename R, typename... Args>
+	struct member_function_traits<R (*)(Args...)> {
+		using return_type = R;
+	};
+
 	template <class N, class M, typename... VarArgs>
 	static MethodBind *bind_method(N p_method_name, M p_method, VarArgs... p_args) {
 		Variant args[sizeof...(p_args) + 1] = { p_args..., Variant() }; // +1 makes sure zero sized arrays are also supported.
@@ -249,6 +269,9 @@ public:
 			argptrs[i] = &args[i];
 		}
 		MethodBind *bind = create_method_bind(p_method);
+		if constexpr (std::is_same<typename member_function_traits<M>::return_type, Object *>::value) {
+			bind->set_return_type_is_raw_object_ptr(true);
+		}
 		return bind_methodfi(METHOD_FLAGS_DEFAULT, bind, p_method_name, sizeof...(p_args) == 0 ? nullptr : (const Variant **)argptrs, sizeof...(p_args));
 	}
 
@@ -261,6 +284,9 @@ public:
 		}
 		MethodBind *bind = create_static_method_bind(p_method);
 		bind->set_instance_class(p_class);
+		if constexpr (std::is_same<typename member_function_traits<M>::return_type, Object *>::value) {
+			bind->set_return_type_is_raw_object_ptr(true);
+		}
 		return bind_methodfi(METHOD_FLAGS_DEFAULT, bind, p_method_name, sizeof...(p_args) == 0 ? nullptr : (const Variant **)argptrs, sizeof...(p_args));
 	}
 
@@ -273,6 +299,9 @@ public:
 
 		bind->set_name(p_name);
 		bind->set_default_arguments(p_default_args);
+		if constexpr (std::is_same<typename member_function_traits<M>::return_type, Object *>::value) {
+			bind->set_return_type_is_raw_object_ptr(true);
+		}
 
 		String instance_type = bind->get_instance_class();
 
diff --git a/core/object/method_bind.h b/core/object/method_bind.h
index 8334a7eef6..d37479f45b 100644
--- a/core/object/method_bind.h
+++ b/core/object/method_bind.h
@@ -49,6 +49,7 @@ class MethodBind {
 	bool _static = false;
 	bool _const = false;
 	bool _returns = false;
+	bool _returns_raw_obj_ptr = false;
 
 protected:
 	Variant::Type *argument_types = nullptr;
@@ -121,6 +122,9 @@ public:
 	_FORCE_INLINE_ bool has_return() const { return _returns; }
 	virtual bool is_vararg() const { return false; }
 
+	_FORCE_INLINE_ bool is_return_type_raw_object_ptr() { return _returns_raw_obj_ptr; }
+	_FORCE_INLINE_ void set_return_type_is_raw_object_ptr(bool p_returns_raw_obj) { _returns_raw_obj_ptr = p_returns_raw_obj; }
+
 	void set_default_arguments(const Vector<Variant> &p_defargs);
 
 	uint32_t get_hash() const;
diff --git a/modules/gdscript/gdscript_vm.cpp b/modules/gdscript/gdscript_vm.cpp
index b99f5d2685..1ab8109dd4 100644
--- a/modules/gdscript/gdscript_vm.cpp
+++ b/modules/gdscript/gdscript_vm.cpp
@@ -1969,7 +1969,13 @@ Variant GDScriptFunction::call(GDScriptInstance *p_instance, const Variant **p_a
 				VariantInternal::initialize(ret, Variant::OBJECT);
 				Object **ret_opaque = VariantInternal::get_object(ret);
 				method->ptrcall(base_obj, argptrs, ret_opaque);
-				VariantInternal::update_object_id(ret);
+				if (method->is_return_type_raw_object_ptr()) {
+					// The Variant has to participate in the ref count since the method returns a raw Object *.
+					VariantInternal::object_assign(ret, *ret_opaque);
+				} else {
+					// The method, in case it returns something, returns an already encapsulated object.
+					VariantInternal::update_object_id(ret);
+				}
 
 #ifdef DEBUG_ENABLED
 				if (GDScriptLanguage::get_singleton()->profiling) {
-- 
cgit v1.2.3