Make rz_flag_set_next() try the addr before incremental numbers as su…

…ffix

In binaries with a large number of flags of the same name, like in macho
with many relocs to the same symbol, using an incremental number as
suffix to generate unique flag names leads to quadratic runtime. We can
avoid this by first trying to append the address, which as a nice side
effect also results in more meaningful names.

librz/core/cbin.c

@@ -1149,7 +1149,9 @@ static void reloc_set_flag(RzCore *core, RzBinReloc *reloc, const char *prefix,
 		return;
 	}
 	RzFlagItem *fi = rz_flag_set_next(core->flags, flag_name, flag_addr, bin_reloc_size(reloc));
-	rz_flag_item_set_realname(fi, reloc_name);
+	if (fi) {
+		rz_flag_item_set_realname(fi, reloc_name);
+	}
 
 	free(reloc_name);
 	free(flag_name);

librz/flag/flag.c

@@ -489,30 +489,57 @@ RZ_API char *rz_flag_get_liststr(RzFlag *f, ut64 off) {
 	return p;
 }
 
-// Set a new flag named `name` at offset `off`. If there's already a flag with
-// the same name, slightly change the name by appending ".%d" as suffix
+/**
+ * Set a flag if there is not already a flag with the same name that does
+ * not match the given \p off and \p size.
+ * \return whether to stop searching for another name, using bool instead of returning
+ *         a pointer to distinguish between existing name and failed malloc.
+ */
+bool try_set_flag(RzFlag *f, const char *name, ut64 off, ut32 size, RzFlagItem **r) {
+	RzFlagItem *fi = rz_flag_get(f, name);
+	if (fi) {
+		if (fi->offset == off && fi->size == size) {
+			*r = fi;
+			return true;
+		}
+		return false;
+	}
+	*r = rz_flag_set(f, name, off, size);
+	return true;
+}
+
+/**
+ * Set a new flag named \p name at \p off. If there's already a flag with
+ * the same name, slightly change the name by appending the address or ".%d" as suffix.
+ * If there is a flag at \p off of size \p size and a matching name, that flag is returned
+ * instead of creating a new one.
+ */
 RZ_API RzFlagItem *rz_flag_set_next(RzFlag *f, const char *name, ut64 off, ut32 size) {
 	rz_return_val_if_fail(f && name, NULL);
-	if (!rz_flag_get(f, name)) {
-		return rz_flag_set(f, name, off, size);
-	}
-	int i, newNameSize = strlen(name);
-	char *newName = malloc(newNameSize + 16);
-	if (!newName) {
+	RzFlagItem *r = NULL;
+	if (try_set_flag(f, name, off, size, &r)) {
+		return r;
+	}
+	size_t name_len = strlen(name);
+	static const size_t suffix_size = 16 + 2; // max size of a 64bit addr + '.' + '\0'
+	char *new_name = malloc(name_len + suffix_size);
+	if (!new_name) {
 		return NULL;
 	}
-	strcpy(newName, name);
-	for (i = 0;; i++) {
-		snprintf(newName + newNameSize, 15, ".%d", i);
-		if (!rz_flag_get(f, newName)) {
-			RzFlagItem *fi = rz_flag_set(f, newName, off, size);
-			if (fi) {
-				free(newName);
-				return fi;
+	memcpy(new_name, name, name_len);
+	snprintf(new_name + name_len, suffix_size, ".%" PFMT64x, off);
+	int i;
+	if (!try_set_flag(f, name, off, size, &r)) {
+		r = rz_flag_set(f, new_name, off, size);
+		for (i = 0; i < 1024 /* some upper bound to prevent unreasonable looping */; i++) {
+			snprintf(new_name + name_len, 17, ".%d", i);
+			if (try_set_flag(f, name, off, size, &r)) {
+				break;
 			}
 		}
 	}
-	return NULL;
+	free(new_name);
+	return r;
 }
 
 /* create or modify an existing flag item with the given name and parameters.

test/db/archos/darwin-arm64/dbg

@@ -42,14 +42,14 @@ NAME=maps
 FILE=bins/mach0/hello-macos-arm64
 ARGS=-d
 CMDS=<<EOF
-dm~hello~[3-]
+dm~hello~[3-9]
 EOF
-REGEXP_FILTER_OUT=([a-zA-Z0-9_\.-]+\s+)
+REGEXP_FILTER_OUT=([a-zA-Z0-9_-]+[\.rwx_-]*\s+)
 EXPECT=<<EOF
-- usr 16K u r-x hello-macos-arm64 hello-macos-arm64 hello_macos_arm64.r_x
-- usr 16K u rw- hello-macos-arm64 hello-macos-arm64 hello_macos_arm64.rw
-- usr 16K u rw- hello-macos-arm64 hello-macos-arm64 hello_macos_arm64.rw.0
-- usr 16K u r-- hello-macos-arm64 hello-macos-arm64 hello_macos_arm64.r
+- usr 16K u r-x hello-macos-arm64 hello-macos-arm64
+- usr 16K u rw- hello-macos-arm64 hello-macos-arm64
+- usr 16K u rw- hello-macos-arm64 hello-macos-arm64
+- usr 16K u r-- hello-macos-arm64 hello-macos-arm64
 EOF
 RUN
 
@@ -61,10 +61,11 @@ fl@F:maps~hello~[1-]
 echo --
 dm*~hello~[0-2]
 EOF
+REGEXP_FILTER_OUT=([0-9]+ [a-zA-Z0-9_-]+\.[\.rwx_]*|[f-].+|\n)
 EXPECT=<<EOF
 16384 hello_macos_arm64.r_x
 16384 hello_macos_arm64.rw
-16384 hello_macos_arm64.rw.0
+16384 hello_macos_arm64.rw.
 16384 hello_macos_arm64.r
 --
 f+ map.hello_macos_arm64.r_x 0x00004000

test/db/formats/elf/helloworld-gcc-elf

@@ -488,7 +488,7 @@ EXPECT=<<EOF
 0x08049694 0 sym..got.plt
 0x08049694 0 obj._GLOBAL_OFFSET_TABLE
 0x080496a0 4 reloc.puts
-0x080496a4 4 reloc.__gmon_start.0
+0x080496a4 4 reloc.__gmon_start.80496a4
 0x080496a8 4 reloc.__libc_start_main
 0x080496ac 8 section..data
 0x080496ac 0 sym..data

test/db/formats/elf/reloc

@@ -154,38 +154,38 @@ EXPECT=<<EOF
 0x08000040 4 reloc.target..text
 0x08000073 4 reloc..rodata.str1.1
 0x0800007a 4 reloc.strcmp
-0x0800009c 4 reloc..rodata.str1.1.0
-0x080000a3 4 reloc.strcmp.0
-0x080000be 4 reloc..rodata.str1.1.1
-0x080000c5 4 reloc.strcmp.1
-0x080000f8 8 reloc..rodata.str1.1.2
+0x0800009c 4 reloc..rodata.str1.1.800009c
+0x080000a3 4 reloc.strcmp.80000a3
+0x080000be 4 reloc..rodata.str1.1.80000be
+0x080000c5 4 reloc.strcmp.80000c5
+0x080000f8 8 reloc..rodata.str1.1.80000f8
 0x08000125 4 reloc.stat
 0x08000144 4 reloc.strlen
 0x08000199 4 reloc.memcpy
-0x080001be 8 reloc..rodata.str1.1.3
-0x080001e0 4 reloc.strlen.0
-0x08000231 4 reloc.memcpy.0
-0x08000267 4 reloc.memcpy.1
+0x080001be 8 reloc..rodata.str1.1.80001be
+0x080001e0 4 reloc.strlen.80001e0
+0x08000231 4 reloc.memcpy.8000231
+0x08000267 4 reloc.memcpy.8000267
 0x080002a2 4 reloc.perror
 0x080002c9 4 reloc.time
-0x080003a9 4 reloc..rodata.str1.1.4
+0x080003a9 4 reloc..rodata.str1.1.80003a9
 0x080003b8 4 reloc..L.str.5
-0x0800050f 8 reloc..rodata.str1.1.5
-0x0800053e 8 reloc..rodata.str1.1.6
+0x0800050f 8 reloc..rodata.str1.1.800050f
+0x0800053e 8 reloc..rodata.str1.1.800053e
 0x08000555 4 reloc.exit
 0x0800057f 4 reloc.write
 0x080005be 4 reloc.utilvserver_fmt_ulong_base
-0x080005f5 4 reloc.strlen.1
-0x0800061b 8 reloc..rodata.str1.1.7
-0x08000632 4 reloc.exit.0
+0x080005f5 4 reloc.strlen.80005f5
+0x0800061b 8 reloc..rodata.str1.1.800061b
+0x08000632 4 reloc.exit.8000632
 0x08000636 4 reloc.target..rodata.str1.1
 0x08000689 4 reloc.target..L.str.5
 0x08000880 4 reloc..text
-0x0800089c 4 reloc..text.0
-0x080008b8 4 reloc..text.1
-0x080008d4 4 reloc..text.2
-0x080008f0 4 reloc..text.3
-0x0800090c 4 reloc..text.4
+0x0800089c 4 reloc..text.800089c
+0x080008b8 4 reloc..text.80008b8
+0x080008d4 4 reloc..text.80008d4
+0x080008f0 4 reloc..text.80008f0
+0x0800090c 4 reloc..text.800090c
 0x08000f00 4 reloc.target.strcmp
 0x08000f08 4 reloc.target.stat
 0x08000f10 4 reloc.target.strlen

test/db/formats/mach0/relocs

@@ -1459,7 +1459,7 @@ vaddr      paddr      target     type    name
  4 fd: 5 +0x00003000 0x00004000 - 0x00005fff r-- vmap.__LINKEDIT
 --
 0860000010600000
-0x00004074 0x00006008  .`.. @ reloc.NSAutoreleasePool.0 24584 NSAutoreleasePool R 0x0
+0x00004074 0x00006008  .`.. @ reloc.NSAutoreleasePool.4074 24584 NSAutoreleasePool R 0x0
 0x00004078 0x00006010  .`.. @ reloc.NSFileManager 24592 NSFileManager R 0x0
 EOF
 RUN

test/db/formats/mach0/thumb

@@ -18,7 +18,7 @@ NAME=mach0 non-thumb symbol on thumb bin
 FILE=bins/mach0/thumb-test
 CMDS=pi 1@ sym.imp.strcmp
 EXPECT=<<EOF
-ldr pc, [reloc.strcmp.0]
+ldr pc, [reloc.strcmp.c000]
 EOF
 RUN
 

test/unit/test_flags.c

@@ -131,10 +131,47 @@ bool test_rz_flag_get_at() {
 	mu_end;
 }
 
+bool test_rz_flag_set_next() {
+	RzFlag *flag = rz_flag_new();
+
+	RzFlagItem *fi = rz_flag_set_next(flag, "catastasis", 0x100, 0);
+	mu_assert_notnull(fi, "set flag");
+	mu_assert_streq(fi->name, "catastasis", "flag_name");
+	fi = rz_flag_set_next(flag, "catastasis", 0xfa1afe1, 0);
+	mu_assert_notnull(fi, "set flag");
+	mu_assert_streq(fi->name, "catastasis.fa1afe1", "flag_name");
+	fi = rz_flag_set_next(flag, "catastasis", (ut64)-1, 0);
+	mu_assert_notnull(fi, "set flag");
+	mu_assert_streq(fi->name, "catastasis.ffffffffffffffff", "flag_name");
+	fi = rz_flag_set_next(flag, "catastasis", 0xfa1afe1, 0);
+	mu_assert_notnull(fi, "set flag");
+	mu_assert_streq(fi->name, "catastasis.0", "flag_name");
+	fi = rz_flag_set_next(flag, "catastasis", 0xfa1afe1, 0);
+	mu_assert_notnull(fi, "set flag");
+	mu_assert_streq(fi->name, "catastasis.1", "flag_name");
+	fi = rz_flag_set_next(flag, "catastasis.fa1afe1", 0x200, 0);
+	mu_assert_notnull(fi, "set flag");
+	mu_assert_streq(fi->name, "catastasis.fa1afe1.200", "flag_name");
+
+	RzList *all = rz_flag_all_list(flag, false);
+	mu_assert_eq(rz_list_length(all), 6, "all flags count");
+	mu_assert_streq(((RzFlagItem *)rz_list_get_n(all, 0))->name, "catastasis", "flag name");
+	mu_assert_streq(((RzFlagItem *)rz_list_get_n(all, 1))->name, "catastasis.fa1afe1.200", "flag name");
+	mu_assert_streq(((RzFlagItem *)rz_list_get_n(all, 2))->name, "catastasis.fa1afe1", "flag name");
+	mu_assert_streq(((RzFlagItem *)rz_list_get_n(all, 3))->name, "catastasis.0", "flag name");
+	mu_assert_streq(((RzFlagItem *)rz_list_get_n(all, 4))->name, "catastasis.1", "flag name");
+	mu_assert_streq(((RzFlagItem *)rz_list_get_n(all, 5))->name, "catastasis.ffffffffffffffff", "flag name");
+	rz_list_free(all);
+
+	rz_flag_free(flag);
+	mu_end;
+}
+
 int all_tests(void) {
 	mu_run_test(test_rz_flag_get_set);
 	mu_run_test(test_rz_flag_by_spaces);
 	mu_run_test(test_rz_flag_get_at);
+	mu_run_test(test_rz_flag_set_next);
 	return tests_passed != tests_run;
 }