From cf8f661b736a8dd9bae2a4ef06b1aeff2041fc6f Mon Sep 17 00:00:00 2001
From: Khairul Azhar Kasmiran <kazarmy@gmail.com>
Date: Tue, 16 Feb 2021 16:38:29 +0800
Subject: [PATCH] aae: Realign on fcn start if not in bb (#622)

---
 librz/core/canalysis.c  | 28 ++++++++++++++++++++--------
 test/db/analysis/x86_32 | 19 ++++++++++++++++++-
 2 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/librz/core/canalysis.c b/librz/core/canalysis.c
index 3571690cce..20d454bf87 100644
--- a/librz/core/canalysis.c
+++ b/librz/core/canalysis.c
@@ -5490,6 +5490,7 @@ RZ_API void rz_core_analysis_esil(RzCore *core, const char *str, const char *tar
 
 		rz_analysis_op_fini(&op);
 		rz_asm_set_pc(core->rasm, cur);
+		size_t i_old = i;
 		if (!rz_analysis_op(core->analysis, &op, cur, buf + i, iend - i, RZ_ANALYSIS_OP_MASK_ESIL | RZ_ANALYSIS_OP_MASK_VAL | RZ_ANALYSIS_OP_MASK_HINT)) {
 			i += minopsize - 1; //   XXX dupe in op.size below
 		}
@@ -5497,21 +5498,21 @@ RZ_API void rz_core_analysis_esil(RzCore *core, const char *str, const char *tar
 		if (op.type == RZ_ANALYSIS_OP_TYPE_ILL || op.type == RZ_ANALYSIS_OP_TYPE_UNK) {
 			// i += 2
 			rz_analysis_op_fini(&op);
-			continue;
+			goto repeat;
 		}
 		//we need to check again i because buf+i may goes beyond its boundaries
 		//because of i+= minopsize - 1
 		if (i > iend) {
-			break;
+			goto repeat;
 		}
 		if (op.size < 1) {
 			i += minopsize - 1;
-			continue;
+			goto repeat;
 		}
 		if (emu_lazy) {
 			if (op.type & RZ_ANALYSIS_OP_TYPE_REP) {
 				i += op.size - 1;
-				continue;
+				goto repeat;
 			}
 			switch (op.type & RZ_ANALYSIS_OP_TYPE_MASK) {
 			case RZ_ANALYSIS_OP_TYPE_JMP:
@@ -5533,12 +5534,12 @@ RZ_API void rz_core_analysis_esil(RzCore *core, const char *str, const char *tar
 			case RZ_ANALYSIS_OP_TYPE_CSWI:
 			case RZ_ANALYSIS_OP_TYPE_TRAP:
 				i += op.size - 1;
-				continue;
+				goto repeat;
 			//  those require write support
 			case RZ_ANALYSIS_OP_TYPE_PUSH:
 			case RZ_ANALYSIS_OP_TYPE_POP:
 				i += op.size - 1;
-				continue;
+				goto repeat;
 			}
 		}
 		if (sn && op.type == RZ_ANALYSIS_OP_TYPE_SWI) {
@@ -5559,7 +5560,7 @@ RZ_API void rz_core_analysis_esil(RzCore *core, const char *str, const char *tar
 		const char *esilstr = RZ_STRBUF_SAFEGET(&op.esil);
 		i += op.size - 1;
 		if (!esilstr || !*esilstr) {
-			continue;
+			goto repeat;
 		}
 		rz_analysis_esil_set_pc(ESIL, cur);
 		rz_reg_setv(core->analysis->reg, pcname, cur + op.size);
@@ -5708,7 +5709,18 @@ RZ_API void rz_core_analysis_esil(RzCore *core, const char *str, const char *tar
 			break;
 		}
 		rz_analysis_esil_stack_free(ESIL);
-	repeat:;
+	repeat:
+		if (!rz_analysis_get_block_at(core->analysis, cur)) {
+			for (size_t fcn_i = i_old + 1; fcn_i <= i; fcn_i++) {
+				if (rz_analysis_get_function_at(core->analysis, start + fcn_i)) {
+					i = fcn_i - 1;
+					break;
+				}
+			}
+		}
+		if (i > iend) {
+			break;
+		}
 	} while (get_next_i(&ictx, &i));
 	free(buf);
 	ESIL->cb.hook_mem_read = NULL;
diff --git a/test/db/analysis/x86_32 b/test/db/analysis/x86_32
index 40a9e523f8..2b5f296e17 100644
--- a/test/db/analysis/x86_32
+++ b/test/db/analysis/x86_32
@@ -3411,7 +3411,7 @@ NAME=tiny1 aaef
 FILE=bins/elf/analysis/tiny1
 CMDS=<<EOF
 aF
-aaef  # TODO Fix if aae here
+aaef
 pdf
 # TODO eip flag?
 EOF
@@ -3425,3 +3425,20 @@ EXPECT=<<EOF
 \           0x08048059      int   0x80
 EOF
 RUN
+
+NAME=tiny1 aae
+FILE=bins/elf/analysis/tiny1
+CMDS=<<EOF
+aF
+aae
+pdf
+EOF
+EXPECT=<<EOF
+/ 7: entry0 ();
+|           0x08048054      mov   bl, 0x2a                             ; '*' ; 42
+|           0x08048056      xor   eax, eax
+|           0x08048058      inc   eax
+|           ;-- syscall.exit:
+\           0x08048059      int   0x80
+EOF
+RUN