Add ASan wrapper for jit callbacks from external code (AMD64)sandbox/ches01/asan_amd64

Signed-off-by: Vyacheslav Cherkashin <v.cherkashin@samsung.com>

3 files changed, 289 insertions, 49 deletions

diff --git a/src/vm/amd64/asan_wrapper.S b/src/vm/amd64/asan_wrapper.S
index a58828b2d3..70eb6e8acb 100644
--- a/src/vm/amd64/asan_wrapper.S
+++ b/src/vm/amd64/asan_wrapper.S
@@ -28,9 +28,13 @@ RETADDR_OFFSET=8*15
 	push %r14
 	push %r13
 	push %r12
+
+#	push %rax	# tmp for align stack
 .endm
 
 .macro POP_REGS
+#	pop %rax
+
 	pop %r12
 	pop %r13
 	pop %r14
@@ -54,30 +58,22 @@ RETADDR_OFFSET=8*15
 	.global hex_wrapper_size
 	.global hex_wrapper_entry_offset
 
-	# Extern functions
-	.global asan_enable_addr_offset
-	.global asan_disable_addr_offset
-	.global target_addr_offset
-	.global get_return_addr_offset
-
-
         .text
 	.code64
 
 hex_wrapper:
-asan_enable_addr:	.quad 0x0badc0de0badc0de
-asan_disable_addr:	.quad 0x0badc0de0badc0de
 target_addr:		.quad 0x0badc0de0badc0de
-get_return_addr:	.quad 0x0badc0de0badc0de
+push_addr:		.quad 0xdeadfeefdeadbeef
+pop_addr:		.quad 0xdeadc0dedeadc0de
 
 hex_wrapper_entry:
+	#int3
 	# Save context
 	PUSH_REGS
 
 	# Save the return adderss
-	call *get_return_addr(%rip)	# rax: get space for saving the return address
 	mov RETADDR_OFFSET(%rsp), %rdi	# rdi: get return address
-	mov %rdi, (%rax)		# save the return address to the received space
+	call *push_addr(%rip)		# save 'ret addr' and call pre_code
 
 	# Change the return address on the 'wrapper_second'
 	call next
@@ -85,43 +81,91 @@ next:	pop %rax				# rax: get current rip
 	add $(wrapper_second - next), %rax	# rax: add offset to 'wrapper_second'
 	mov %rax, RETADDR_OFFSET(%rsp)		# change the return address
 
-	# Enable ASan
-	call *asan_enable_addr(%rip)
-
 	# Restore context
 	POP_REGS
 
+	#int3
 	# Call original function
 	jmp *target_addr(%rip)
 wrapper_second:
+	#int3
 	# Save context
 	PUSH_REGS
 
-	# Align stack
-	add $8, %rsp
-
-	# Disable ASan
-	call *asan_disable_addr(%rip)
-
+	sub $8, %rsp
 	# Restore the return address
-	call *get_return_addr(%rip)	# rax: get pointer to saved the return address
-	sub $8, %rsp			# 'restore' aligned stack
-	mov (%rax), %rdi		# rdi: get the return address
-	mov %rdi, RETADDR_OFFSET(%rsp)	# restore the return address
+	call *pop_addr(%rip)		# rax: call post_code and get 'ret addr'
+	add $8, %rsp
+	mov %rax, RETADDR_OFFSET(%rsp)	# restore the return address
 
 	# Restore context
 	POP_REGS
 
-	# Return
+	#int3
+	# return
 	ret
 
 hex_wrapper_size:		.quad . - hex_wrapper
-
 hex_wrapper_entry_offset:	.quad hex_wrapper_entry - hex_wrapper
-asan_enable_addr_offset:	.quad asan_enable_addr - hex_wrapper
-asan_disable_addr_offset:	.quad asan_disable_addr - hex_wrapper
-target_addr_offset:		.quad target_addr - hex_wrapper
-get_return_addr_offset:		.quad get_return_addr - hex_wrapper
 
 
 # set environment LD_PRELOAD /root/asan/libasan.so
+# b TheUMEntryPrestub
+# b UMThunkStub
+
+
+        .global hex_wrapper2
+	.global hex_wrapper2_size
+	.global hex_wrapper2_entry_offset
+
+        .text
+	.code64
+
+hex_wrapper2:
+target_addr2:		.quad 0x0badc0de0badc0de
+push_addr2:		.quad 0xdeadfeefdeadbeef
+pop_addr2:		.quad 0xdeadc0dedeadc0de
+
+hex_wrapper2_entry:
+	#int3
+	# Save context
+	PUSH_REGS
+
+	# Save the return adderss
+	mov RETADDR_OFFSET(%rsp), %rdi	# rdi: get return address
+	call *push_addr2(%rip)		# save 'ret addr' and call pre_code
+
+	# Change the return address on the 'wrapper_second'
+	call next2
+next2:	pop %rax				# rax: get current rip
+	add $(wrapper_second2 - next2), %rax	# rax: add offset to 'wrapper_second'
+	mov %rax, RETADDR_OFFSET(%rsp)		# change the return address
+
+	# Restore context
+	POP_REGS
+
+	#int3
+	# Call original function
+	jmp *target_addr2(%rip)
+wrapper_second2:
+	sub $8, %rsp		# add space for 'ret addr'
+
+	#int3
+	# Save context
+	PUSH_REGS
+
+//	sub $8, %rsp
+	# Restore the return address
+	call *pop_addr2(%rip)		# rax: call post_code and get 'ret addr'
+//	add $8, %rsp
+	mov %rax, RETADDR_OFFSET(%rsp)	# restore the return address
+
+	# Restore context
+	POP_REGS
+
+	#int3
+	# return
+	ret
+
+hex_wrapper2_size:		.quad . - hex_wrapper2
+hex_wrapper2_entry_offset:	.quad hex_wrapper2_entry - hex_wrapper2
diff --git a/src/vm/amd64/cgenamd64.cpp b/src/vm/amd64/cgenamd64.cpp
index 814f04dd9f..2ea4971148 100644
--- a/src/vm/amd64/cgenamd64.cpp
+++ b/src/vm/amd64/cgenamd64.cpp
@@ -26,6 +26,8 @@
 #include "clrtocomcall.h"
 #endif // FEATURE_COMINTEROP
 
+extern "C" int dprintf(int fd, const char *format, ...);
+
 void UpdateRegDisplayFromCalleeSavedRegisters(REGDISPLAY * pRD, CalleeSavedRegisters * pRegs)
 {
     LIMITED_METHOD_CONTRACT;
@@ -553,8 +555,26 @@ void emitJump(LPBYTE pBuffer, LPVOID target)
     _ASSERTE(DbgIsExecutable(pBuffer, 12));
 }
 
+
+
+extern "C" LPVOID ASanCreateWraper2(LPVOID target, void (*push_addr)(LPVOID addr), LPVOID (*pop_addr)());
+extern "C" void my_push_addr_j2n(LPVOID addr);
+extern "C" LPVOID my_pop_addr_j2n();
+
+void change_my_target(void *old_target, void *new_target)
+{
+	dprintf(1, "#### change_my_target!!: old_target=%p new_target=%p\n", old_target, new_target);
+}
+
+
 void UMEntryThunkCode::Encode(BYTE* pTargetCode, void* pvSecretParam)
 {
+    dprintf(1, "#### UMEntryThunkCode::Encode: entry_pointer=%p pTargetCode=%p pvSecretParam=%p\n", &m_movR10, pTargetCode, pvSecretParam);
+
+    dprintf(1, "#### UMEntryThunkCode::Encode: target: TheUMEntryPrestub: %d UMThunkStub: %d\n",
+		    pTargetCode == (BYTE *)TheUMEntryPrestub,
+		    pTargetCode == (BYTE *)UMThunkStub
+	   );
     CONTRACTL
     {
         NOTHROW;
@@ -563,6 +583,13 @@ void UMEntryThunkCode::Encode(BYTE* pTargetCode, void* pvSecretParam)
     }
     CONTRACTL_END;
 
+    {
+	BYTE *old = pTargetCode;
+	pTargetCode = (BYTE *)ASanCreateWraper2((LPVOID)pTargetCode, my_push_addr_j2n, my_pop_addr_j2n);
+	dprintf(1, "#### UMEntryThunkCode::Encode: CHANGE!!! new_pTargetCode=%p\n", pTargetCode);
+	change_my_target(old, pTargetCode);
+    }
+
     // padding                  // CC CC CC CC
     // mov r10, pUMEntryThunk   // 49 ba xx xx xx xx xx xx xx xx    // METHODDESC_REGISTER
     // mov rax, pJmpDest        // 48 b8 xx xx xx xx xx xx xx xx    // need to ensure this imm64 is qword aligned
diff --git a/src/vm/dllimport.cpp b/src/vm/dllimport.cpp
index 76ca654622..6d3285c1f8 100644
--- a/src/vm/dllimport.cpp
+++ b/src/vm/dllimport.cpp
@@ -6857,31 +6857,78 @@ extern "C" void __sanitizer_disable_calls_forwarding();
 
 extern "C" long write(int fd, const void *buf, int count);
 
-void my_enable()
+
+static __thread bool is_enable = false;
+
+extern "C" void assert_my_enable()
+{
+	assert(is_enable == false);
+}
+
+static __thread int asan_cnt = 0;
+
+extern "C" void my_enable_j2n()
+{
+	if (asan_cnt == 0) {
+		char msg[] = "#### MY_ENABLE  J->N\n";
+		write(1, msg, sizeof(msg) - 1);
+		__sanitizer_disable_calls_forwarding();
+	} else {
+		char msg[] = "#### MY_ENABLE  J->N [SKIP]\n";
+		write(1, msg, sizeof(msg) - 1);
+	}
+	++asan_cnt;
+}
+
+extern "C" void my_disable_j2n()
 {
-	char msg[] = "#### MY_ENABLE\n";
-	write(1, msg, sizeof(msg) - 1);
-	__sanitizer_disable_calls_forwarding();
+	if (asan_cnt == 1) {
+		__sanitizer_enable_calls_forwarding();
+		char msg[] = "#### MY_DISABLE J->N\n";
+		write(1, msg, sizeof(msg) - 1);
+	} else {
+		char msg[] = "#### MY_DISABLE J->N [SKIP]\n";
+		write(1, msg, sizeof(msg) - 1);
+	}
+	--asan_cnt;
+}
+
+extern "C" void my_enable_n2j()
+{
+	if (asan_cnt == 0) {
+		char msg[] = "#### MY_ENABLE  N->J\n";
+		write(1, msg, sizeof(msg) - 1);
+		__sanitizer_disable_calls_forwarding();
+	} else {
+		char msg[] = "#### MY_ENABLE  N->J [SKIP]\n";
+		write(1, msg, sizeof(msg) - 1);
+	}
+	++asan_cnt;
 }
 
-void my_disable()
+extern "C" void my_disable_n2j()
 {
-	__sanitizer_enable_calls_forwarding();
-	char msg[] = "#### MY_DISABLE\n";
-	write(1, msg, sizeof(msg) - 1);
+	if (asan_cnt == 1) {
+		__sanitizer_enable_calls_forwarding();
+		char msg[] = "#### MY_DISABLE N->J\n";
+		write(1, msg, sizeof(msg) - 1);
+	} else {
+		char msg[] = "#### MY_DISABLE N->J [SKIP]\n";
+		write(1, msg, sizeof(msg) - 1);
+	}
+	--asan_cnt;
 }
 
 #include <pshpack1.h>
 struct my_struct {
-	void (*asan_enable)();
-	void (*asan_disable)();
 	UINT64 real_func;
-	UINT64 *(*get_addr_space)();
+	void (*push_addr)(LPVOID addr);
+	LPVOID (*pop_addr)();
 };
 
 
-static __thread UINT64 g_addr_space;
-UINT64 *get_return_addr_space()
+static __thread LPVOID g_addr_space;
+LPVOID *get_return_addr_space()
 {
 	return &g_addr_space;
 }
@@ -6890,29 +6937,151 @@ extern UINT64 hex_wrapper;
 extern UINT64 hex_wrapper_size;
 extern UINT64 hex_wrapper_entry_offset;
 
+extern UINT64 hex_wrapper2;
+extern UINT64 hex_wrapper2_size;
+extern UINT64 hex_wrapper2_entry_offset;
+
+
+// 2
+static __thread LPVOID g_addr2_space;
+extern "C" void my_space_push_addr(LPVOID addr)
+{
+	g_addr2_space = addr;
+}
 
-LPVOID ASanCreateWraper(LPVOID target)
+extern "C" LPVOID *my_space_top()
 {
+	return &g_addr2_space;
+}
+
+/*
+class MyStack {
+public:
+	MyStack() {
+	}
+
+
+
+private:
+	LPVOID m_stack[50];
+};
+*/
+
+template <typename Type, int STACK_SIZE>
+class StaticStack {	// Don't init because static TLS variable initialize only zeros
+public:
+	void push(Type addr) {
+		_ASSERTE(++m_pos <= STACK_SIZE);
+		m_data[m_pos - 1] = addr;
+	}
+
+	void pop() {
+		_ASSERTE(m_pos-- > 0);
+	}
+
+	Type top() {
+		_ASSERTE(m_pos > 0);
+		return m_data[m_pos - 1];
+	}
+
+	bool empty() {
+		return m_pos == 0;
+	}
+
+private:
+	int m_pos;
+	Type m_data[STACK_SIZE];
+};
+
+
+//#include <clr/stack.h>
+
+//using my_stack = clr::Stack<LPVOID>;
+//static __thread my_stack *g_addr_stack;
+static __thread StaticStack<LPVOID, 128> g_retaddr_stack;
+
+extern "C" void my_push_addr_j2n(LPVOID addr)
+{
+	my_disable_j2n();
+	g_retaddr_stack.push(addr);
+}
+extern "C" LPVOID my_pop_addr_j2n()
+{
+	LPVOID addr = g_retaddr_stack.top();
+	g_retaddr_stack.pop();
+
+	my_enable_j2n();
+
+	return addr;
+}
+
+extern "C" void my_push_addr_n2j(LPVOID addr)
+{
+	g_retaddr_stack.push(addr);
+	my_enable_n2j();
+}
+
+extern "C" LPVOID my_pop_addr_n2j()
+{
+	my_disable_n2j();
+
+	LPVOID addr = g_retaddr_stack.top();
+	g_retaddr_stack.pop();
+
+	return addr;
+}
+
+extern "C" LPVOID ASanCreateWraper(LPVOID target, void (*push_addr)(LPVOID addr), LPVOID (*pop_addr)())
+{
+//	AsanWrapperCode *wraper = (AsanWrapperCode *)(void *)SystemDomain::GetGlobalLoaderAllocator()->GetExecutableHeap()->AllocMem(S_SIZE_T(sizeof(AsanWrapperCode)));
 	void *wrapper_space = (void *)SystemDomain::GetGlobalLoaderAllocator()->GetExecutableHeap()->AllocMem(S_SIZE_T(0x200));
 
 	// copy asan_wrapper
 	memcpy(wrapper_space, &hex_wrapper, hex_wrapper_size);
 
 	struct my_struct pp = {
-		.asan_enable = my_enable,
-		.asan_disable = my_disable,
 		.real_func = (UINT64)target,
-		.get_addr_space = get_return_addr_space,
+		.push_addr = push_addr,// my_enable_n2j,
+		.pop_addr = pop_addr,// my_disable_n2j,
 	};
 
 	// copy my_struct
 	memcpy(wrapper_space, &pp, sizeof(pp));
 
+	dprintf(1, "#### ASanCreateWraper: wrapper_space[%p..%p] exec_space[%p..%p]\n",
+			wrapper_space, (UINT64)wrapper_space + 0x200,
+			(UINT64)wrapper_space + hex_wrapper_entry_offset,
+			(UINT64)wrapper_space + hex_wrapper_size);
+
 	return (LPVOID)((UINT64)wrapper_space + hex_wrapper_entry_offset);
 }
 
 
-extern "C" int dprintf(int fd, const char *format, ...);
+extern "C" LPVOID ASanCreateWraper2(LPVOID target, void (*push_addr)(LPVOID addr), LPVOID (*pop_addr)())
+{
+//	AsanWrapperCode *wraper = (AsanWrapperCode *)(void *)SystemDomain::GetGlobalLoaderAllocator()->GetExecutableHeap()->AllocMem(S_SIZE_T(sizeof(AsanWrapperCode)));
+	void *wrapper_space = (void *)SystemDomain::GetGlobalLoaderAllocator()->GetExecutableHeap()->AllocMem(S_SIZE_T(0x200));
+
+	// copy asan_wrapper
+	memcpy(wrapper_space, &hex_wrapper2, hex_wrapper2_size);
+
+	struct my_struct pp = {
+		.push_addr = push_addr,// my_enable_n2j,
+		.pop_addr = pop_addr,// my_disable_n2j,
+		.real_func = (UINT64)target,
+	};
+
+	// copy my_struct
+	memcpy(wrapper_space, &pp, sizeof(pp));
+
+	dprintf(1, "#### ASanCreateWraper2: wrapper_space[%p..%p] exec_space[%p..%p]\n",
+			wrapper_space, (UINT64)wrapper_space + 0x200,
+			(UINT64)wrapper_space + hex_wrapper2_entry_offset,
+			(UINT64)wrapper_space + hex_wrapper2_size);
+
+	return (LPVOID)((UINT64)wrapper_space + hex_wrapper2_entry_offset);
+}
+
 extern "C" void *dlopen(const char *filename, int flags);
 extern "C" int dlclose(void *handle);
 extern "C" void *dlsym(void *handle, const char *symbol);