Make afs respect the cc (Fix #3291)

Applying a function signature with afs to a function was using a
combination of the existing args and wrong stack addresses to apply the
args to function variables. This generally resulted in incorrect
locations. Instead, we now use the info from the cc.

librz/analysis/fcn.c

@@ -1883,67 +1883,47 @@ RZ_API RZ_OWN char *rz_analysis_function_get_signature(RZ_NONNULL RzAnalysisFunc
 /**
  * \brief Sets the RzCallable type for the given function
  *
- * Checks if the type is defined already for this function, if yes -
- * it removes the existing one and sets the one defined by the RzCallable.
- * If there is a mismatch between existing arguments - it overwrites
- * their types and names, removes arguments if necessary.
+ * Overwrites all arguments, the return type, calling convention and noreturn property of \p f to
+ * match the contents of \p callable. This is done according to the calling convention in
+ * \p callable, or \p f if it is not defined in \p callable.
  *
  * \param a RzAnalysis instance
  * \param f Function to update
- * \param callable A function type
+ * \param callable A function type to apply to \p f
  */
-RZ_API bool rz_analysis_function_set_type(RzAnalysis *a, RZ_NONNULL RzAnalysisFunction *f, RZ_NONNULL RzCallable *callable) {
-	rz_return_val_if_fail(a && f && callable, false);
-	// At first, we check if the arguments match, and rename/retype them
-	void **it;
-	size_t index = 0;
-	if (rz_pvector_empty(callable->args)) {
-		rz_analysis_function_delete_all_vars(f);
+RZ_API void rz_analysis_function_set_type(RzAnalysis *a, RZ_NONNULL RzAnalysisFunction *f, RZ_NONNULL RzCallable *callable) {
+	rz_return_if_fail(a && f && callable);
+	// Set the cc first, it will be used further down.
+	if (callable->cc) {
+		f->cc = rz_str_constpool_get(&a->constpool, callable->cc);
+	}
+	// All args will be overwritten
+	rz_analysis_function_delete_vars_by_arg_property(f, true);
+	RzStackAddr stack_off = rz_type_db_pointer_size(a->typedb) / 8; // return val
+	if (f->cc) {
+		stack_off += rz_analysis_cc_shadow_store(a, f->cc);
 	}
 	size_t args_count = rz_pvector_len(callable->args);
-	RzPVector *cloned_vars = rz_pvector_clone(&f->vars);
-	rz_pvector_foreach (cloned_vars, it) {
-		RzAnalysisVar *var = *it;
-		if (!rz_analysis_var_is_arg(var)) {
+	for (size_t index = 0; index < args_count; index++) {
+		RzCallableArg *arg = *rz_pvector_index_ptr(callable->args, index);
+		if (!arg || !arg->type) {
 			continue;
 		}
-		if (index < args_count) {
-			RzCallableArg *arg = *rz_pvector_index_ptr(callable->args, index);
-			if (arg) {
-				free(var->name);
-				if (arg->name) {
-					var->name = strdup(arg->name);
-				}
-				rz_type_free(var->type);
-				var->type = rz_type_clone(arg->type);
-			}
-			index++;
-		} else {
-			// There is no match for this argument in the RzCallable type,
-			// thus we remove it from the function
-			rz_analysis_function_delete_var(f, var);
-		}
-	}
-	// For f->vars is already empty, add args into it
-	for (; index < args_count; index++) {
-		RzCallableArg *arg = *rz_pvector_index_ptr(callable->args, index);
-		if (arg && arg->type) {
-			size_t size = rz_type_db_get_bitsize(a->typedb, arg->type);
-			// For user defined args, we set its delta and kind to its index and stack var by default
-			RzAnalysisVarStorage stor = { 0 };
+		RzAnalysisVarStorage stor = { 0 };
+		const char *loc = f->cc ? rz_analysis_cc_arg(a, f->cc, index) : "stack";
+		if (!loc || rz_str_startswith(loc, "stack")) {
 			stor.type = RZ_ANALYSIS_VAR_STORAGE_STACK;
-			stor.stack_off = index;
-			rz_analysis_function_set_var(f, &stor, arg->type, size, arg->name);
+			stor.stack_off = stack_off;
+			stack_off += (rz_type_db_get_bitsize(a->typedb, arg->type) + 7) / 8;
+		} else {
+			stor.type = RZ_ANALYSIS_VAR_STORAGE_REG;
+			stor.reg = rz_str_constpool_get(&a->constpool, loc);
 		}
+		rz_analysis_function_set_var(f, &stor, arg->type, 0, arg->name);
 	}
-
-	if (callable->noret) {
-		f->is_noreturn = true;
-	} else {
-		f->ret_type = callable->ret;
-	}
-	rz_pvector_free(cloned_vars);
-	return true;
+	f->is_noreturn = callable->noret;
+	rz_type_free(f->ret_type);
+	f->ret_type = callable->ret ? rz_type_clone(callable->ret) : NULL;
 }
 
 /**
@@ -1985,7 +1965,8 @@ RZ_API bool rz_analysis_function_set_type_str(RzAnalysis *a, RZ_NONNULL RzAnalys
 		RZ_LOG_ERROR("Parsed function signature should not be NULL\n");
 		return false;
 	}
-	return rz_analysis_function_set_type(a, f, result->callable);
+	rz_analysis_function_set_type(a, f, result->callable);
+	return true;
 }
 
 RZ_API RzAnalysisFunction *rz_analysis_fcn_next(RzAnalysis *analysis, ut64 addr) {

librz/analysis/function.c

@@ -115,7 +115,7 @@ RZ_API void rz_analysis_function_free(void *_fcn) {
 	ht_up_free(fcn->inst_vars);
 	ht_up_free(fcn->labels);
 	ht_pp_free(fcn->label_addrs);
-
+	rz_type_free(fcn->ret_type);
 	free(fcn->name);
 	rz_list_free(fcn->imports);
 	free(fcn);

librz/analysis/var.c

@@ -203,6 +203,24 @@ RZ_API void rz_analysis_function_delete_vars_by_storage_type(RzAnalysisFunction
 	}
 }
 
+/**
+ * Delete all variables from \p fcn that are either arguments or locals, depending on \p args
+ * \param args if true, all args will be deleted, otherwise all non-arg vars will be deleted
+ */
+RZ_API void rz_analysis_function_delete_vars_by_arg_property(RzAnalysisFunction *fcn, bool args) {
+	rz_return_if_fail(fcn);
+	size_t i;
+	for (i = 0; i < rz_pvector_len(&fcn->vars);) {
+		RzAnalysisVar *var = rz_pvector_at(&fcn->vars, i);
+		if (rz_analysis_var_is_arg(var)) {
+			rz_pvector_remove_at(&fcn->vars, i);
+			var_free(var);
+			continue;
+		}
+		i++;
+	}
+}
+
 RZ_API void rz_analysis_function_delete_all_vars(RzAnalysisFunction *fcn) {
 	void **it;
 	rz_pvector_foreach (&fcn->vars, it) {

librz/core/cmd/cmd_analysis.c

@@ -2344,8 +2344,10 @@ RZ_IPI RzCmdStatus rz_analysis_function_signature_type_handler(RzCore *core, int
 	}
 	if (!rz_type_func_ret_set(core->analysis->typedb, fcn->name, ret_type)) {
 		RZ_LOG_ERROR("core: Cannot find type %s\n", argv[1]);
+		rz_type_free(ret_type);
 		return RZ_CMD_STATUS_ERROR;
 	}
+	rz_type_free(fcn->ret_type);
 	fcn->ret_type = ret_type;
 	return RZ_CMD_STATUS_OK;
 }

librz/include/rz_analysis.h

@@ -1608,7 +1608,7 @@ RZ_API void rz_analysis_del_jmprefs(RzAnalysis *analysis, RzAnalysisFunction *fc
 RZ_API char *rz_analysis_function_get_json(RzAnalysisFunction *function);
 RZ_API RzAnalysisFunction *rz_analysis_fcn_next(RzAnalysis *analysis, ut64 addr);
 RZ_API RZ_OWN char *rz_analysis_function_get_signature(RZ_NONNULL RzAnalysisFunction *function);
-RZ_API bool rz_analysis_function_set_type(RzAnalysis *a, RZ_NONNULL RzAnalysisFunction *f, RZ_NONNULL RzCallable *callable);
+RZ_API void rz_analysis_function_set_type(RzAnalysis *a, RZ_NONNULL RzAnalysisFunction *f, RZ_NONNULL RzCallable *callable);
 RZ_API bool rz_analysis_function_set_type_str(RzAnalysis *a, RZ_NONNULL RzAnalysisFunction *f, RZ_NONNULL const char *sig);
 RZ_API int rz_analysis_fcn_count(RzAnalysis *a, ut64 from, ut64 to);
 RZ_API RzAnalysisBlock *rz_analysis_fcn_bbget_in(const RzAnalysis *analysis, RzAnalysisFunction *fcn, ut64 addr);
@@ -1640,6 +1640,7 @@ RZ_API RZ_BORROW RzAnalysisVar *rz_analysis_function_get_reg_var_at(RzAnalysisFu
 RZ_API RZ_BORROW RzAnalysisVar *rz_analysis_function_get_stack_var_in(RzAnalysisFunction *fcn, RzStackAddr stack_off);
 RZ_API RZ_BORROW RzAnalysisVar *rz_analysis_function_get_var_byname(RzAnalysisFunction *fcn, const char *name);
 RZ_API void rz_analysis_function_delete_vars_by_storage_type(RzAnalysisFunction *fcn, RzAnalysisVarStorageType stor);
+RZ_API void rz_analysis_function_delete_vars_by_arg_property(RzAnalysisFunction *fcn, bool args);
 RZ_API void rz_analysis_function_delete_all_vars(RzAnalysisFunction *fcn);
 RZ_API void rz_analysis_function_delete_unused_vars(RzAnalysisFunction *fcn);
 RZ_API void rz_analysis_function_delete_var(RzAnalysisFunction *fcn, RzAnalysisVar *var);

librz/include/rz_type.h

@@ -413,7 +413,7 @@ RZ_API void rz_type_func_delete_all(RzTypeDB *typedb);
 RZ_API bool rz_type_func_exist(RzTypeDB *typedb, RZ_NONNULL const char *func_name);
 
 RZ_API RZ_BORROW RzType *rz_type_func_ret(RzTypeDB *typedb, RZ_NONNULL const char *func_name);
-RZ_API bool rz_type_func_ret_set(RzTypeDB *typedb, const char *func_name, RZ_OWN RZ_NONNULL RzType *type);
+RZ_API bool rz_type_func_ret_set(RzTypeDB *typedb, const char *func_name, RZ_BORROW RZ_NONNULL RzType *type);
 
 RZ_API RZ_BORROW const char *rz_type_func_cc(RzTypeDB *typedb, RZ_NONNULL const char *func_name);
 RZ_API bool rz_type_func_cc_set(RzTypeDB *typedb, const char *name, const char *cc);

librz/type/function.c

@@ -345,13 +345,14 @@ RZ_API bool rz_type_func_arg_add(RzTypeDB *typedb, RZ_NONNULL const char *func_n
  * \param name Name of the callable to search
  * \param type RzType return type
  */
-RZ_API bool rz_type_func_ret_set(RzTypeDB *typedb, const char *name, RZ_OWN RZ_NONNULL RzType *type) {
+RZ_API bool rz_type_func_ret_set(RzTypeDB *typedb, const char *name, RZ_BORROW RZ_NONNULL RzType *type) {
 	rz_return_val_if_fail(typedb && name && type, false);
 	RzCallable *callable = rz_type_func_get(typedb, name);
 	if (!callable) {
 		return false;
 	}
-	callable->ret = type;
+	rz_type_free(callable->ret);
+	callable->ret = rz_type_clone(type);
 	return true;
 }
 

test/db/cmd/types

@@ -2028,6 +2028,56 @@ size_t strlen(const char *s);
 EOF
 RUN
 
+NAME=afs fastcall
+FILE=bins/elf/analysis/fast
+CMDS=<<EOF
+s sym.fastcaslled
+af
+afc fastcall
+afs "int fastcalled(int sarg0, int sarg1, int sarg2, int sarg3)"
+afs
+afv
+afi~name,convention
+EOF
+EXPECT=<<EOF
+int fastcalled(int sarg0, int sarg1, int sarg2, int sarg3);
+var int32_t var_24h @ stack - 0x24
+var int32_t var_20h @ stack - 0x20
+var int32_t var_14h @ stack - 0x14
+var int32_t var_10h @ stack - 0x10
+arg int sarg2 @ stack + 0x4
+arg int sarg3 @ stack + 0x8
+arg int sarg0 @ ecx
+arg int sarg1 @ edx
+name: fastcalled
+call-convention: fastcall
+EOF
+RUN
+
+NAME=afsr
+FILE=bins/elf/analysis/fast
+CMDS=<<EOF
+s sym.fastcaslled
+af
+afsr uint32_t
+afs
+EOF
+EXPECT=<<EOF
+int fastcalled(int sarg0, int sarg1, int sarg2, int sarg3);
+var int32_t var_24h @ stack - 0x24
+var int32_t var_20h @ stack - 0x20
+var int32_t var_14h @ stack - 0x14
+var int32_t var_10h @ stack - 0x10
+arg int sarg2 @ stack + 0x4
+arg int sarg3 @ stack + 0x8
+arg int sarg0 @ ecx
+arg int sarg1 @ edx
+name: fastcalled
+call-convention: fastcall
+EOF
+RUN
+
+
 NAME=td crash
 FILE==
 CMDS=<<EOF

test/unit/test_analysis_cc.c

@@ -109,4 +109,4 @@ bool all_tests() {
 	return tests_passed != tests_run;
 }
 
-mu_main(all_tests)
+mu_main(all_tests)

test/unit/test_analysis_function.c

@@ -3,6 +3,7 @@
 
 #include <rz_analysis.h>
 #include "minunit.h"
+#include "test_config.h"
 
 #include "test_analysis_block_invars.inl"
 
@@ -322,6 +323,116 @@ bool test_initial_underscore(void) {
 	mu_end;
 }
 
+bool test_rz_analysis_function_set_type() {
+	RzAnalysis *analysis = rz_analysis_new();
+	rz_analysis_use(analysis, "x86");
+	rz_analysis_set_bits(analysis, 32);
+	rz_analysis_cc_set(analysis, "eax sectarian(ecx, edx, stack)");
+	rz_type_db_purge(analysis->typedb);
+	const char *types_dir = TEST_BUILD_TYPES_DIR;
+	rz_type_db_init(analysis->typedb, types_dir, "x86", 32, "linux");
+
+	// Only setup here
+	RzAnalysisFunction *f = rz_analysis_create_function(analysis, "postcard", 0x100, RZ_ANALYSIS_FCN_TYPE_NULL);
+	RzAnalysisVarStorage stor = { 0 };
+	stor.type = RZ_ANALYSIS_VAR_STORAGE_REG;
+	stor.reg = "edi";
+	rz_analysis_function_set_var(f, &stor, NULL, 4, "oldarg0");
+	stor.type = RZ_ANALYSIS_VAR_STORAGE_REG;
+	stor.reg = "ecx";
+	rz_analysis_function_set_var(f, &stor, NULL, 4, "oldarg1");
+	stor.type = RZ_ANALYSIS_VAR_STORAGE_STACK;
+	stor.stack_off = 1000;
+	rz_analysis_function_set_var(f, &stor, NULL, 4, "oldarg2");
+	stor.type = RZ_ANALYSIS_VAR_STORAGE_STACK;
+	stor.stack_off = -8;
+	rz_analysis_function_set_var(f, &stor, NULL, 4, "oldvar");
+	mu_assert_eq(rz_pvector_len(&f->vars), 4, "initial vars");
+	// The order in the vars vector is allowed to be different. It is only assumed here because
+	// it is currently deterministic and this simplifies the test code.
+	RzAnalysisVar *var = rz_pvector_at(&f->vars, 0);
+	mu_assert_streq(var->name, "oldarg0", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_REG, "var storage type");
+	mu_assert_streq(var->storage.reg, "edi", "var storage reg");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "int32_t", "var type");
+	var = rz_pvector_at(&f->vars, 1);
+	mu_assert_streq(var->name, "oldarg1", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_REG, "var storage type");
+	mu_assert_streq(var->storage.reg, "ecx", "var storage reg");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "int32_t", "var type");
+	var = rz_pvector_at(&f->vars, 2);
+	mu_assert_streq(var->name, "oldarg2", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_STACK, "var storage type");
+	mu_assert_eq(var->storage.stack_off, 1000, "var storage stack");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "int32_t", "var type");
+	var = rz_pvector_at(&f->vars, 3);
+	mu_assert_streq(var->name, "oldvar", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_STACK, "var storage type");
+	mu_assert_eq(var->storage.stack_off, -8, "var storage stack");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "int32_t", "var type");
+	f->is_noreturn = true;
+
+	RzCallable *c = rz_type_callable_new("nopartofme");
+	rz_type_callable_arg_add(c, rz_type_callable_arg_new(analysis->typedb, "arg0", rz_type_identifier_of_base_type_str(analysis->typedb, "uint8_t")));
+	rz_type_callable_arg_add(c, rz_type_callable_arg_new(analysis->typedb, "arg1", rz_type_identifier_of_base_type_str(analysis->typedb, "uint32_t")));
+	rz_type_callable_arg_add(c, rz_type_callable_arg_new(analysis->typedb, "arg2", rz_type_identifier_of_base_type_str(analysis->typedb, "int64_t")));
+	rz_type_callable_arg_add(c, rz_type_callable_arg_new(analysis->typedb, "arg3", rz_type_identifier_of_base_type_str(analysis->typedb, "uint32_t")));
+	c->noret = false;
+	c->ret = rz_type_identifier_of_base_type_str(analysis->typedb, "uint16_t");
+	c->cc = rz_str_constpool_get(&analysis->constpool, "sectarian");
+
+	// Actual testing
+
+	rz_analysis_function_set_type(analysis, f, c);
+	rz_type_callable_free(c);
+	mu_assert_streq(f->cc, "sectarian", "cc");
+	mu_assert_eq(rz_pvector_len(&f->vars), 5, "initial vars");
+	// Expected: only the var that was not an arg from before still exists, all
+	// args have been replaced by the ones from the callable.
+	// See note about the ordering above
+	var = rz_pvector_at(&f->vars, 0);
+	mu_assert_streq(var->name, "oldvar", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_STACK, "var storage type");
+	mu_assert_eq(var->storage.stack_off, -8, "var storage stack");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "int32_t", "var type");
+	var = rz_pvector_at(&f->vars, 1);
+	mu_assert_streq(var->name, "arg0", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_REG, "var storage type");
+	mu_assert_streq(var->storage.reg, "ecx", "var storage reg");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "uint8_t", "var type");
+	var = rz_pvector_at(&f->vars, 2);
+	mu_assert_streq(var->name, "arg1", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_REG, "var storage type");
+	mu_assert_streq(var->storage.reg, "edx", "var storage reg");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "uint32_t", "var type");
+	var = rz_pvector_at(&f->vars, 3);
+	mu_assert_streq(var->name, "arg2", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_STACK, "var storage type");
+	mu_assert_eq(var->storage.stack_off, 4, "var storage stack");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "int64_t", "var type");
+	var = rz_pvector_at(&f->vars, 4);
+	mu_assert_streq(var->name, "arg3", "var name");
+	mu_assert_eq(var->storage.type, RZ_ANALYSIS_VAR_STORAGE_STACK, "var storage type");
+	mu_assert_eq(var->storage.stack_off, 12, "var storage stack");
+	mu_assert_eq(var->type->kind, RZ_TYPE_KIND_IDENTIFIER, "var type kind");
+	mu_assert_streq(var->type->identifier.name, "uint32_t", "var type");
+
+	mu_assert_notnull(f->ret_type, "ret type");
+	mu_assert_eq(f->ret_type->kind, RZ_TYPE_KIND_IDENTIFIER, "ret type kind");
+	mu_assert_streq(f->ret_type->identifier.name, "uint16_t", "ret type");
+
+	rz_analysis_free(analysis);
+	mu_end;
+}
+
 int all_tests() {
 	mu_run_test(test_rz_analysis_function_relocate);
 	mu_run_test(test_rz_analysis_function_labels);
@@ -330,6 +441,7 @@ int all_tests() {
 	mu_run_test(test_dll_names);
 	mu_run_test(test_autonames);
 	mu_run_test(test_initial_underscore);
+	mu_run_test(test_rz_analysis_function_set_type);
 	return tests_passed != tests_run;
 }