From 2dea40bfaba6b6e24b398200658141839b7b012b Mon Sep 17 00:00:00 2001
From: Dmitry Stogov <dmitry@zend.com>
Date: Wed, 26 Oct 2022 13:44:44 +0300
Subject: [PATCH] Add API to patch native code

---
 ir.h         |   4 +
 ir_gdb.c     |   6 --
 ir_patch.c   | 262 +++++++++++++++++++++++++++++++++++++++++++++++++++
 ir_private.h |   8 ++
 ir_x86.dasc  |   8 --
 5 files changed, 274 insertions(+), 14 deletions(-)
 create mode 100644 ir_patch.c

diff --git a/ir.h b/ir.h
index 6e90c44..2696797 100644
--- a/ir.h
+++ b/ir.h
@@ -610,6 +610,10 @@ int ir_emit_c(ir_ctx *ctx, FILE *f);
 bool ir_check(ir_ctx *ctx);
 void ir_consistency_check(void);
 
+/* Code patching (implementation in ir_patch.c) */
+int ir_patch(const void *code, size_t size, uint32_t jmp_table_size, const void *from_addr, const void *to_addr);
+
+
 /* IR Memmory Allocation */
 #ifndef ir_mem_malloc
 # define ir_mem_malloc   malloc
diff --git a/ir_gdb.c b/ir_gdb.c
index 2c7f7f6..38a199c 100644
--- a/ir_gdb.c
+++ b/ir_gdb.c
@@ -11,12 +11,6 @@
 #include "ir_private.h"
 #include "ir_elf.h"
 
-#if defined(HAVE_ATTRIBUTE_ALIGNED)
-# define IR_SET_ALIGNED(alignment, decl) decl __attribute__ ((__aligned__ (alignment)))
-#else
-# define IR_SET_ALIGNED(alignment, decl) decl
-#endif
-
 /* DWARF definitions. */
 #define DW_CIE_VERSION  1
 
diff --git a/ir_patch.c b/ir_patch.c
new file mode 100644
index 0000000..f3944a7
--- /dev/null
+++ b/ir_patch.c
@@ -0,0 +1,262 @@
+#include "ir.h"
+#include "ir_private.h"
+
+#if defined(IR_TARGET_X86) || defined(IR_TARGET_X64)
+/* This taken from LuaJIT. Thanks to Mike Pall. */
+static uint32_t _asm_x86_inslen(const uint8_t* p)
+{
+	static const uint8_t map_op1[256] = {
+		0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x20,
+		0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x51,0x51,
+		0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,
+		0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,0x92,0x92,0x92,0x92,0x52,0x45,0x10,0x51,
+#if defined(__x86_64__) || defined(_M_X64)
+		0x10,0x10,0x10,0x10,0x10,0x10,0x10,0x10,0x14,0x14,0x14,0x14,0x14,0x14,0x14,0x14,
+#else
+		0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,
+#endif
+		0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,
+		0x51,0x51,0x92,0x92,0x10,0x10,0x12,0x11,0x45,0x86,0x52,0x93,0x51,0x51,0x51,0x51,
+		0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,
+		0x93,0x86,0x93,0x93,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,
+		0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x51,0x47,0x51,0x51,0x51,0x51,0x51,
+#if defined(__x86_64__) || defined(_M_X64)
+		0x59,0x59,0x59,0x59,0x51,0x51,0x51,0x51,0x52,0x45,0x51,0x51,0x51,0x51,0x51,0x51,
+#else
+		0x55,0x55,0x55,0x55,0x51,0x51,0x51,0x51,0x52,0x45,0x51,0x51,0x51,0x51,0x51,0x51,
+#endif
+		0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x05,0x05,0x05,0x05,0x05,0x05,0x05,0x05,
+		0x93,0x93,0x53,0x51,0x70,0x71,0x93,0x86,0x54,0x51,0x53,0x51,0x51,0x52,0x51,0x51,
+		0x92,0x92,0x92,0x92,0x52,0x52,0x51,0x51,0x92,0x92,0x92,0x92,0x92,0x92,0x92,0x92,
+		0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x45,0x45,0x47,0x52,0x51,0x51,0x51,0x51,
+		0x10,0x51,0x10,0x10,0x51,0x51,0x63,0x66,0x51,0x51,0x51,0x51,0x51,0x51,0x92,0x92
+	};
+	static const uint8_t map_op2[256] = {
+		0x93,0x93,0x93,0x93,0x52,0x52,0x52,0x52,0x52,0x52,0x51,0x52,0x51,0x93,0x52,0x94,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x53,0x53,0x53,0x53,0x53,0x53,0x53,0x53,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x34,0x51,0x35,0x51,0x51,0x51,0x51,0x51,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x53,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x94,0x54,0x54,0x54,0x93,0x93,0x93,0x52,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,0x46,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x52,0x52,0x52,0x93,0x94,0x93,0x51,0x51,0x52,0x52,0x52,0x93,0x94,0x93,0x93,0x93,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x94,0x93,0x93,0x93,0x93,0x93,
+		0x93,0x93,0x94,0x93,0x94,0x94,0x94,0x93,0x52,0x52,0x52,0x52,0x52,0x52,0x52,0x52,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,
+		0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x93,0x52
+	};
+	uint32_t result = 0;
+	uint32_t prefixes = 0;
+	uint32_t x = map_op1[*p];
+
+	for (;;) {
+		switch (x >> 4) {
+			case 0:
+				return result + x + (prefixes & 4);
+			case 1:
+				prefixes |= x;
+				x = map_op1[*++p];
+				result++;
+				break;
+			case 2:
+				x = map_op2[*++p];
+				break;
+			case 3:
+				p++;
+				goto mrm;
+			case 4:
+				result -= (prefixes & 2);
+				/* fallthrough */
+			case 5:
+				return result + (x & 15);
+			case 6: /* Group 3. */
+				if (p[1] & 0x38) {
+					x = 2;
+				} else if ((prefixes & 2) && (x == 0x66)) {
+					x = 4;
+				}
+				goto mrm;
+			case 7: /* VEX c4/c5. */
+#if !defined(__x86_64__) && !defined(_M_X64)
+				if (p[1] < 0xc0) {
+					x = 2;
+					goto mrm;
+				}
+#endif
+				if (x == 0x70) {
+					x = *++p & 0x1f;
+					result++;
+					if (x >= 2) {
+						p += 2;
+						result += 2;
+						goto mrm;
+					}
+				}
+				p++;
+				result++;
+				x = map_op2[*++p];
+				break;
+			case 8:
+				result -= (prefixes & 2);
+				/* fallthrough */
+			case 9:
+mrm:
+				/* ModR/M and possibly SIB. */
+				result += (x & 15);
+				x = *++p;
+				switch (x >> 6) {
+					case 0:
+						if ((x & 7) == 5) {
+							return result + 4;
+						}
+						break;
+					case 1:
+						result++;
+						break;
+					case 2:
+						result += 4;
+						break;
+					case 3:
+						return result;
+				}
+				if ((x & 7) == 4) {
+					result++;
+					if (x < 0x40 && (p[1] & 7) == 5) {
+						result += 4;
+					}
+				}
+				return result;
+		}
+	}
+}
+
+typedef IR_SET_ALIGNED(1, uint16_t unaligned_uint16_t);
+typedef IR_SET_ALIGNED(1, int32_t unaligned_int32_t);
+
+static int ir_patch_code(const void *code, size_t size, const void *from_addr, const void *to_addr)
+{
+	int ret = 0;
+	uint8_t *p, *end;
+
+	p = (uint8_t*)code;
+	end = p + size - 5;
+	while (p < end) {
+		if ((*(unaligned_uint16_t*)p & 0xf0ff) == 0x800f && p + *(unaligned_int32_t*)(p+2) == (uint8_t*)from_addr - 6) {
+			*(unaligned_int32_t*)(p+2) = ((uint8_t*)to_addr - (p + 6));
+			ret++;
+		} else if (*p == 0xe9 && p + *(unaligned_int32_t*)(p+1) == (uint8_t*)from_addr - 5) {
+			*(unaligned_int32_t*)(p+1) = ((uint8_t*)to_addr - (p + 5));
+			ret++;
+		}
+		p += _asm_x86_inslen(p);
+	}
+	if (ret) {
+		ir_mem_flush((void*)code, size);
+	}
+	return ret;
+}
+
+#elif defined(IR_TARGET_AARCH64)
+
+static int ir_patch_code(const void *code, size_t size, const void *from_addr, const void *to_addr)
+{
+	int ret = 0;
+	uint8_t *p, *end;
+	const void *veneer = NULL;
+	ptrdiff_t delta;
+
+	end = (uint8_t*)code;
+	p = end + size;
+	while (p > end) {
+		uint32_t *ins_ptr;
+		uint32_t ins;
+
+		p -= 4;
+		ins_ptr = (uint32_t*)p;
+		ins = *ins_ptr;
+		if ((ins & 0xfc000000u) == 0x14000000u) {
+			// B (imm26:0..25)
+			delta = (uint32_t*)from_addr - ins_ptr;
+			if (((ins ^ (uint32_t)delta) & 0x01ffffffu) == 0) {
+				delta = (uint32_t*)to_addr - ins_ptr;
+				if (((delta + 0x02000000) >> 26) != 0) {
+					abort(); // branch target out of range
+				}
+				*ins_ptr = (ins & 0xfc000000u) | ((uint32_t)delta & 0x03ffffffu);
+				ret++;
+				if (!veneer) {
+					veneer = p;
+				}
+			}
+		} else if ((ins & 0xff000000u) == 0x54000000u ||
+		           (ins & 0x7e000000u) == 0x34000000u) {
+			// B.cond, CBZ, CBNZ (imm19:5..23)
+			delta = (uint32_t*)from_addr - ins_ptr;
+			if (((ins ^ ((uint32_t)delta << 5)) & 0x00ffffe0u) == 0) {
+				delta = (uint32_t*)to_addr - ins_ptr;
+				if (((delta + 0x40000) >> 19) != 0) {
+					if (veneer) {
+						delta = (uint32_t*)veneer - ins_ptr;
+						if (((delta + 0x40000) >> 19) != 0) {
+							abort(); // branch target out of range
+						}
+					} else {
+						abort(); // branch target out of range
+					}
+				}
+				*ins_ptr = (ins & 0xff00001fu) | (((uint32_t)delta & 0x7ffffu) << 5);
+				ret++;
+			}
+		} else if ((ins & 0x7e000000u) == 0x36000000u) {
+			// TBZ, TBNZ (imm14:5..18)
+			delta = (uint32_t*)from_addr - ins_ptr;
+			if (((ins ^ ((uint32_t)delta << 5)) & 0x0007ffe0u) == 0) {
+				delta = (uint32_t*)to_addr - ins_ptr;
+				if (((delta + 0x2000) >> 14) != 0) {
+					if (veneer) {
+						delta = (uint32_t*)veneer - ins_ptr;
+						if (((delta + 0x2000) >> 14) != 0) {
+							abort(); // branch target out of range
+						}
+					} else {
+						abort(); // branch target out of range
+					}
+				}
+				*ins_ptr = (ins & 0xfff8001fu) | (((uint32_t)delta & 0x3fffu) << 5);
+				ret++;
+			}
+		}
+	}
+
+	if (ret) {
+		ir_mem_flush((void*)code, size);
+	}
+
+	return ret;
+}
+#endif
+
+int ir_patch(const void *code, size_t size, uint32_t jmp_table_size, const void *from_addr, const void *to_addr)
+{
+	int ret = 0;
+
+	if (jmp_table_size) {
+		const void **jmp_slot = (const void **)((char*)code + IR_ALIGNED_SIZE(size, sizeof(void*)));
+
+		do {
+			if (*jmp_slot == from_addr) {
+				*jmp_slot = to_addr;
+				ret++;
+			}
+			jmp_slot++;
+		} while (--jmp_table_size);
+	}
+
+	ret += ir_patch_code(code, size, from_addr, to_addr);
+
+	return ret;
+}
diff --git a/ir_private.h b/ir_private.h
index d7736da..0a90f14 100644
--- a/ir_private.h
+++ b/ir_private.h
@@ -21,6 +21,11 @@
 # if __has_attribute(noinline)
 #  define IR_NEVER_INLINE __attribute__((noinline))
 # endif
+# if __has_attribute(__aligned__)
+#  define IR_SET_ALIGNED(alignment, decl) decl __attribute__ ((__aligned__ (alignment)))
+# endif
+#elif defined(_WIN32)
+# define IR_SET_ALIGNED(alignment, decl) __declspec(align(alignment)) decl
 #endif
 #ifndef EXPECTED
 # define EXPECTED(condition)   (condition)
@@ -32,6 +37,9 @@
 #ifndef IR_NEVER_INLINE
 # define IR_NEVER_INLINE
 #endif
+#ifndef IR_SET_ALIGNED
+#  define IR_SET_ALIGNED(alignment, decl) decl
+#endif
 
 /*** Helper routines ***/
 
diff --git a/ir_x86.dasc b/ir_x86.dasc
index 7304530..4f53117 100644
--- a/ir_x86.dasc
+++ b/ir_x86.dasc
@@ -2,14 +2,6 @@
 #include "ir_x86.h"
 #include "ir_private.h"
 
-#ifdef _WIN32
-# define IR_SET_ALIGNED(alignment, decl) __declspec(align(alignment)) decl
-#elif defined(HAVE_ATTRIBUTE_ALIGNED)
-# define IR_SET_ALIGNED(alignment, decl) decl __attribute__ ((__aligned__ (alignment)))
-#else
-# define IR_SET_ALIGNED(alignment, decl) decl
-#endif
-
 #define DASM_M_GROW(ctx, t, p, sz, need) \
   do { \
     size_t _sz = (sz), _need = (need); \