From 35682d30795498632c0406011e6c56f97a73f1d6 Mon Sep 17 00:00:00 2001
From: George Marques <george@gmarqu.es>
Date: Fri, 2 Apr 2021 10:34:44 -0300
Subject: [PATCH] GDScript: Properly validate return type

When the type cannot be validated at compile time, the runtime must do a
check to ensure type safety is kept, as the code might be assuming the
return type is correct in another place, leading to crashes if the
contract is broken.
---
 modules/gdscript/gdscript_byte_codegen.cpp |  81 ++++++++-
 modules/gdscript/gdscript_disassembler.cpp |  33 ++++
 modules/gdscript/gdscript_function.h       |   4 +
 modules/gdscript/gdscript_vm.cpp           | 192 ++++++++++++++++++++-
 4 files changed, 304 insertions(+), 6 deletions(-)

diff --git a/modules/gdscript/gdscript_byte_codegen.cpp b/modules/gdscript/gdscript_byte_codegen.cpp
index e27329d18bce..a31d46bbba6f 100644
--- a/modules/gdscript/gdscript_byte_codegen.cpp
+++ b/modules/gdscript/gdscript_byte_codegen.cpp
@@ -1286,8 +1286,85 @@ void GDScriptByteCodeGenerator::write_newline(int p_line) {
 }
 
 void GDScriptByteCodeGenerator::write_return(const Address &p_return_value) {
-	append(GDScriptFunction::OPCODE_RETURN, 1);
-	append(p_return_value);
+	if (!function->return_type.has_type || p_return_value.type.has_type) {
+		// Either the function is untyped or the return value is also typed.
+
+		// If this is a typed function, then we need to check for potential conversions.
+		if (function->return_type.has_type) {
+			if (function->return_type.kind == GDScriptDataType::BUILTIN && function->return_type.builtin_type == Variant::ARRAY && function->return_type.has_container_element_type()) {
+				// Typed array.
+				const GDScriptDataType &element_type = function->return_type.get_container_element_type();
+
+				Variant script = function->return_type.script_type;
+				int script_idx = get_constant_pos(script);
+				script_idx |= (GDScriptFunction::ADDR_TYPE_LOCAL_CONSTANT << GDScriptFunction::ADDR_BITS);
+
+				append(GDScriptFunction::OPCODE_RETURN_TYPED_ARRAY, 2);
+				append(p_return_value);
+				append(script_idx);
+				append(element_type.kind == GDScriptDataType::BUILTIN ? element_type.builtin_type : Variant::OBJECT);
+				append(element_type.native_type);
+			} else if (function->return_type.kind == GDScriptDataType::BUILTIN && p_return_value.type.kind == GDScriptDataType::BUILTIN && function->return_type.builtin_type != p_return_value.type.builtin_type) {
+				// Add conversion.
+				append(GDScriptFunction::OPCODE_RETURN_TYPED_BUILTIN, 1);
+				append(p_return_value);
+				append(function->return_type.builtin_type);
+			} else {
+				// Just assign.
+				append(GDScriptFunction::OPCODE_RETURN, 1);
+				append(p_return_value);
+			}
+		} else {
+			append(GDScriptFunction::OPCODE_RETURN, 1);
+			append(p_return_value);
+		}
+	} else {
+		switch (function->return_type.kind) {
+			case GDScriptDataType::BUILTIN: {
+				if (function->return_type.builtin_type == Variant::ARRAY && function->return_type.has_container_element_type()) {
+					const GDScriptDataType &element_type = function->return_type.get_container_element_type();
+
+					Variant script = function->return_type.script_type;
+					int script_idx = get_constant_pos(script);
+					script_idx |= (GDScriptFunction::ADDR_TYPE_LOCAL_CONSTANT << GDScriptFunction::ADDR_BITS);
+
+					append(GDScriptFunction::OPCODE_RETURN_TYPED_ARRAY, 2);
+					append(p_return_value);
+					append(script_idx);
+					append(element_type.kind == GDScriptDataType::BUILTIN ? element_type.builtin_type : Variant::OBJECT);
+					append(element_type.native_type);
+				} else {
+					append(GDScriptFunction::OPCODE_RETURN_TYPED_BUILTIN, 1);
+					append(p_return_value);
+					append(function->return_type.builtin_type);
+				}
+			} break;
+			case GDScriptDataType::NATIVE: {
+				append(GDScriptFunction::OPCODE_RETURN_TYPED_NATIVE, 2);
+				append(p_return_value);
+				int class_idx = GDScriptLanguage::get_singleton()->get_global_map()[function->return_type.native_type];
+				class_idx |= (GDScriptFunction::ADDR_TYPE_GLOBAL << GDScriptFunction::ADDR_BITS);
+				append(class_idx);
+			} break;
+			case GDScriptDataType::GDSCRIPT:
+			case GDScriptDataType::SCRIPT: {
+				Variant script = function->return_type.script_type;
+				int script_idx = get_constant_pos(script);
+				script_idx |= (GDScriptFunction::ADDR_TYPE_LOCAL_CONSTANT << GDScriptFunction::ADDR_BITS);
+
+				append(GDScriptFunction::OPCODE_RETURN_TYPED_SCRIPT, 2);
+				append(p_return_value);
+				append(script_idx);
+			} break;
+			default: {
+				ERR_PRINT("Compiler bug: unresolved return.");
+
+				// Shouldn't get here, but fail-safe to a regular return;
+				append(GDScriptFunction::OPCODE_RETURN, 1);
+				append(p_return_value);
+			} break;
+		}
+	}
 }
 
 void GDScriptByteCodeGenerator::write_assert(const Address &p_test, const Address &p_message) {
diff --git a/modules/gdscript/gdscript_disassembler.cpp b/modules/gdscript/gdscript_disassembler.cpp
index 32adca29ed7b..76ebdc8bed7f 100644
--- a/modules/gdscript/gdscript_disassembler.cpp
+++ b/modules/gdscript/gdscript_disassembler.cpp
@@ -761,6 +761,39 @@ void GDScriptFunction::disassemble(const Vector<String> &p_code_lines) const {
 
 				incr = 2;
 			} break;
+			case OPCODE_RETURN_TYPED_BUILTIN: {
+				text += "return typed builtin (";
+				text += Variant::get_type_name((Variant::Type)_code_ptr[ip + 2]);
+				text += ") ";
+				text += DADDR(1);
+
+				incr += 3;
+			} break;
+			case OPCODE_RETURN_TYPED_ARRAY: {
+				text += "return typed array ";
+				text += DADDR(1);
+
+				incr += 5;
+			} break;
+			case OPCODE_RETURN_TYPED_NATIVE: {
+				text += "return typed native (";
+				text += DADDR(2);
+				text += ") ";
+				text += DADDR(1);
+
+				incr += 3;
+			} break;
+			case OPCODE_RETURN_TYPED_SCRIPT: {
+				Variant script = _constants_ptr[_code_ptr[ip + 2]];
+				Script *sc = Object::cast_to<Script>(script.operator Object *());
+
+				text += "return typed script (";
+				text += sc->get_path();
+				text += ") ";
+				text += DADDR(1);
+
+				incr += 3;
+			} break;
 
 #define DISASSEMBLE_ITERATE(m_type)      \
 	case OPCODE_ITERATE_##m_type: {      \
diff --git a/modules/gdscript/gdscript_function.h b/modules/gdscript/gdscript_function.h
index 41080ef37199..9fc75b66ce4a 100644
--- a/modules/gdscript/gdscript_function.h
+++ b/modules/gdscript/gdscript_function.h
@@ -306,6 +306,10 @@ public:
 		OPCODE_JUMP_IF_NOT,
 		OPCODE_JUMP_TO_DEF_ARGUMENT,
 		OPCODE_RETURN,
+		OPCODE_RETURN_TYPED_BUILTIN,
+		OPCODE_RETURN_TYPED_ARRAY,
+		OPCODE_RETURN_TYPED_NATIVE,
+		OPCODE_RETURN_TYPED_SCRIPT,
 		OPCODE_ITERATE_BEGIN,
 		OPCODE_ITERATE_BEGIN_INT,
 		OPCODE_ITERATE_BEGIN_FLOAT,
diff --git a/modules/gdscript/gdscript_vm.cpp b/modules/gdscript/gdscript_vm.cpp
index ffc871bdcbc3..6b7da4a4670a 100644
--- a/modules/gdscript/gdscript_vm.cpp
+++ b/modules/gdscript/gdscript_vm.cpp
@@ -128,7 +128,10 @@ Variant *GDScriptFunction::_get_variant(int p_address, GDScriptInstance *p_insta
 
 #ifdef DEBUG_ENABLED
 static String _get_script_name(const Ref<Script> p_script) {
-	if (p_script->get_name().is_empty()) {
+	Ref<GDScript> gdscript = p_script;
+	if (gdscript.is_valid()) {
+		return gdscript->get_script_class_name();
+	} else if (p_script->get_name().is_empty()) {
 		return p_script->get_path().get_file();
 	} else {
 		return p_script->get_name();
@@ -293,6 +296,10 @@ String GDScriptFunction::_get_call_error(const Callable::CallError &p_err, const
 		&&OPCODE_JUMP_IF_NOT,                        \
 		&&OPCODE_JUMP_TO_DEF_ARGUMENT,               \
 		&&OPCODE_RETURN,                             \
+		&&OPCODE_RETURN_TYPED_BUILTIN,               \
+		&&OPCODE_RETURN_TYPED_ARRAY,                 \
+		&&OPCODE_RETURN_TYPED_NATIVE,                \
+		&&OPCODE_RETURN_TYPED_SCRIPT,                \
 		&&OPCODE_ITERATE_BEGIN,                      \
 		&&OPCODE_ITERATE_BEGIN_INT,                  \
 		&&OPCODE_ITERATE_BEGIN_FLOAT,                \
@@ -569,7 +576,7 @@ Variant GDScriptFunction::call(GDScriptInstance *p_instance, const Variant **p_a
 
 #ifdef DEBUG_ENABLED
 	OPCODE_WHILE(ip < _code_size) {
-		int last_opcode = _code_ptr[ip];
+		int last_opcode = _code_ptr[ip] & INSTR_MASK;
 #else
 	OPCODE_WHILE(true) {
 #endif
@@ -1113,14 +1120,14 @@ Variant GDScriptFunction::call(GDScriptInstance *p_instance, const Variant **p_a
 #ifdef DEBUG_ENABLED
 					err_text = "Trying to assign value of type '" + Variant::get_type_name(src->get_type()) +
 							   "' to a variable of type '" + +"'.";
-					OPCODE_BREAK;
 #endif
+					OPCODE_BREAK;
 				}
 				if (!dst_arr->typed_assign(*src)) {
 #ifdef DEBUG_ENABLED
 					err_text = "Trying to assign a typed array with an array of different type.'";
-					OPCODE_BREAK;
 #endif
+					OPCODE_BREAK;
 				}
 
 				ip += 3;
@@ -2143,6 +2150,183 @@ Variant GDScriptFunction::call(GDScriptInstance *p_instance, const Variant **p_a
 				OPCODE_BREAK;
 			}
 
+			OPCODE(OPCODE_RETURN_TYPED_BUILTIN) {
+				CHECK_SPACE(3);
+				GET_INSTRUCTION_ARG(r, 0);
+
+				Variant::Type ret_type = (Variant::Type)_code_ptr[ip + 2];
+				GD_ERR_BREAK(ret_type < 0 || ret_type >= Variant::VARIANT_MAX);
+
+				if (r->get_type() != ret_type) {
+					if (Variant::can_convert_strict(r->get_type(), ret_type)) {
+						Callable::CallError ce;
+						Variant::construct(ret_type, retvalue, const_cast<const Variant **>(&r), 1, ce);
+					} else {
+#ifdef DEBUG_ENABLED
+						err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "%s".)",
+								Variant::get_type_name(r->get_type()), Variant::get_type_name(ret_type));
+#endif // DEBUG_ENABLED
+
+						// Construct a base type anyway so type constraints are met.
+						Callable::CallError ce;
+						Variant::construct(ret_type, retvalue, nullptr, 0, ce);
+						OPCODE_BREAK;
+					}
+				} else {
+					retvalue = *r;
+				}
+#ifdef DEBUG_ENABLED
+				exit_ok = true;
+#endif // DEBUG_ENABLED
+				OPCODE_BREAK;
+			}
+
+			OPCODE(OPCODE_RETURN_TYPED_ARRAY) {
+				CHECK_SPACE(5);
+				GET_INSTRUCTION_ARG(r, 0);
+
+				GET_INSTRUCTION_ARG(script_type, 1);
+				Variant::Type builtin_type = (Variant::Type)_code_ptr[ip + 3];
+				int native_type_idx = _code_ptr[ip + 4];
+				GD_ERR_BREAK(native_type_idx < 0 || native_type_idx >= _global_names_count);
+				const StringName native_type = _global_names_ptr[native_type_idx];
+
+				if (r->get_type() != Variant::ARRAY) {
+#ifdef DEBUG_ENABLED
+					err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "Array[%s]".)",
+							Variant::get_type_name(r->get_type()), Variant::get_type_name(builtin_type));
+#endif
+					OPCODE_BREAK;
+				}
+
+				Array array;
+				array.set_typed(builtin_type, native_type, script_type);
+
+#ifdef DEBUG_ENABLED
+				bool valid = array.typed_assign(*VariantInternal::get_array(r));
+#else
+				array.typed_assign(*VariantInternal::get_array(r));
+#endif // DEBUG_ENABLED
+
+				// Assign the return value anyway since we want it to be the valid type.
+				retvalue = array;
+
+#ifdef DEBUG_ENABLED
+				if (!valid) {
+					err_text = "Trying to return a typed array with an array of different type.'";
+					OPCODE_BREAK;
+				}
+
+				exit_ok = true;
+#endif // DEBUG_ENABLED
+				OPCODE_BREAK;
+			}
+
+			OPCODE(OPCODE_RETURN_TYPED_NATIVE) {
+				CHECK_SPACE(3);
+				GET_INSTRUCTION_ARG(r, 0);
+
+				GET_INSTRUCTION_ARG(type, 1);
+				GDScriptNativeClass *nc = Object::cast_to<GDScriptNativeClass>(type->operator Object *());
+				GD_ERR_BREAK(!nc);
+
+				if (r->get_type() != Variant::OBJECT && r->get_type() != Variant::NIL) {
+					err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "%s".)",
+							Variant::get_type_name(r->get_type()), nc->get_name());
+					OPCODE_BREAK;
+				}
+
+#ifdef DEBUG_ENABLED
+				bool freed = false;
+				Object *ret_obj = r->get_validated_object_with_check(freed);
+
+				if (freed) {
+					err_text = "Trying to return a previously freed instance.";
+					OPCODE_BREAK;
+				}
+#else
+				Object *ret_obj = r->operator Object *();
+#endif // DEBUG_ENABLED
+				if (ret_obj && !ClassDB::is_parent_class(ret_obj->get_class_name(), nc->get_name())) {
+#ifdef DEBUG_ENABLED
+					err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "%s".)",
+							ret_obj->get_class_name(), nc->get_name());
+#endif // DEBUG_ENABLED
+					OPCODE_BREAK;
+				}
+				retvalue = *r;
+
+#ifdef DEBUG_ENABLED
+				exit_ok = true;
+#endif // DEBUG_ENABLED
+				OPCODE_BREAK;
+			}
+
+			OPCODE(OPCODE_RETURN_TYPED_SCRIPT) {
+				CHECK_SPACE(3);
+				GET_INSTRUCTION_ARG(r, 0);
+
+				GET_INSTRUCTION_ARG(type, 1);
+				Script *base_type = Object::cast_to<Script>(type->operator Object *());
+				GD_ERR_BREAK(!base_type);
+
+				if (r->get_type() != Variant::OBJECT && r->get_type() != Variant::NIL) {
+#ifdef DEBUG_ENABLED
+					err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "%s".)",
+							Variant::get_type_name(r->get_type()), _get_script_name(Ref<Script>(base_type)));
+#endif // DEBUG_ENABLED
+					OPCODE_BREAK;
+				}
+
+#ifdef DEBUG_ENABLED
+				bool freed = false;
+				Object *ret_obj = r->get_validated_object_with_check(freed);
+
+				if (freed) {
+					err_text = "Trying to return a previously freed instance.";
+					OPCODE_BREAK;
+				}
+#else
+				Object *ret_obj = r->operator Object *();
+#endif // DEBUG_ENABLED
+
+				if (ret_obj) {
+					ScriptInstance *ret_inst = ret_obj->get_script_instance();
+					if (!ret_inst) {
+#ifdef DEBUG_ENABLED
+						err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "%s".)",
+								ret_obj->get_class_name(), _get_script_name(Ref<GDScript>(base_type)));
+#endif // DEBUG_ENABLED
+						OPCODE_BREAK;
+					}
+
+					Script *ret_type = ret_obj->get_script_instance()->get_script().ptr();
+					bool valid = false;
+
+					while (ret_type) {
+						if (ret_type == base_type) {
+							valid = true;
+							break;
+						}
+						ret_type = ret_type->get_base_script().ptr();
+					}
+
+					if (!valid) {
+#ifdef DEBUG_ENABLED
+						err_text = vformat(R"(Trying to return value of type "%s" from a function which the return type is "%s".)",
+								_get_script_name(ret_obj->get_script_instance()->get_script()), _get_script_name(Ref<GDScript>(base_type)));
+#endif // DEBUG_ENABLED
+						OPCODE_BREAK;
+					}
+				}
+				retvalue = *r;
+
+#ifdef DEBUG_ENABLED
+				exit_ok = true;
+#endif // DEBUG_ENABLED
+				OPCODE_BREAK;
+			}
+
 			OPCODE(OPCODE_ITERATE_BEGIN) {
 				CHECK_SPACE(8); // Space for this and a regular iterate.