From 5a5f9d21ddfff1fc678fbd31e45d1cd47353f009 Mon Sep 17 00:00:00 2001
From: Mark Johnston <markj@FreeBSD.org>
Date: Sat, 19 Jul 2014 02:27:31 +0000
Subject: [PATCH] Use a C wrapper for trap() instead of checking and calling
 the DTrace trap hook in assembly.

Suggested by:	kib
Reviewed by:	kib (original version)
X-MFC-With:	r268600
---
 sys/amd64/amd64/exception.S             | 19 +------------------
 sys/amd64/amd64/trap.c                  | 16 +++++++++++++++-
 sys/cddl/dev/dtrace/amd64/dtrace_subr.c |  4 +---
 sys/cddl/dev/fbt/fbt.c                  |  9 +++++++--
 4 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/sys/amd64/amd64/exception.S b/sys/amd64/amd64/exception.S
index bb5fd5669e74..2a908a980f8d 100644
--- a/sys/amd64/amd64/exception.S
+++ b/sys/amd64/amd64/exception.S
@@ -228,24 +228,7 @@ alltraps_pushregs_no_rdi:
 	.type	calltrap,@function
 calltrap:
 	movq	%rsp,%rdi
-#ifdef KDTRACE_HOOKS
-	/*
-	 * Give DTrace a chance to vet this trap and skip the call to trap() if
-	 * it turns out that it was caused by a DTrace probe.
-	 */
-	movq	dtrace_trap_func,%rax
-	testq	%rax,%rax
-	je	skiphook
-	call	*%rax
-	testq	%rax,%rax
-	jne	skiptrap
-	movq	%rsp,%rdi
-skiphook:
-#endif
-	call	trap
-#ifdef KDTRACE_HOOKS
-skiptrap:
-#endif
+	call	trap_check
 	MEXITCOUNT
 	jmp	doreti			/* Handle any pending ASTs */
 
diff --git a/sys/amd64/amd64/trap.c b/sys/amd64/amd64/trap.c
index d9203bc5e7c8..e9ba62da316c 100644
--- a/sys/amd64/amd64/trap.c
+++ b/sys/amd64/amd64/trap.c
@@ -97,7 +97,8 @@ PMC_SOFT_DEFINE( , , page_fault, write);
 #include <sys/dtrace_bsd.h>
 #endif
 
-extern void trap(struct trapframe *frame);
+extern void __noinline trap(struct trapframe *frame);
+extern void trap_check(struct trapframe *frame);
 extern void syscall(struct trapframe *frame);
 void dblfault_handler(struct trapframe *frame);
 
@@ -604,6 +605,19 @@ trap(struct trapframe *frame)
 	return;
 }
 
+/*
+ * Ensure that we ignore any DTrace-induced faults. This function cannot
+ * be instrumented, so it cannot generate such faults itself.
+ */
+void
+trap_check(struct trapframe *frame)
+{
+
+	if (dtrace_trap_func != NULL && (*dtrace_trap_func)(frame))
+		return;
+	trap(frame);
+}
+
 static int
 trap_pfault(frame, usermode)
 	struct trapframe *frame;
diff --git a/sys/cddl/dev/dtrace/amd64/dtrace_subr.c b/sys/cddl/dev/dtrace/amd64/dtrace_subr.c
index e761f380fa43..c123cc6aa5bf 100644
--- a/sys/cddl/dev/dtrace/amd64/dtrace_subr.c
+++ b/sys/cddl/dev/dtrace/amd64/dtrace_subr.c
@@ -462,9 +462,7 @@ dtrace_gethrestime(void)
 	return (current_time.tv_sec * 1000000000ULL + current_time.tv_nsec);
 }
 
-/*
- * Function to handle DTrace traps during probes. See amd64/amd64/exception.S.
- */
+/* Function to handle DTrace traps during probes. See amd64/amd64/trap.c. */
 int
 dtrace_trap(struct trapframe *frame)
 {
diff --git a/sys/cddl/dev/fbt/fbt.c b/sys/cddl/dev/fbt/fbt.c
index 7e256e7ba878..a4bffa25372e 100644
--- a/sys/cddl/dev/fbt/fbt.c
+++ b/sys/cddl/dev/fbt/fbt.c
@@ -232,13 +232,18 @@ fbt_provide_module_function(linker_file_t lf, int symindx,
 	int size;
 	u_int8_t *instr, *limit;
 
-	if (strncmp(name, "dtrace_", 7) == 0 &&
-	    strncmp(name, "dtrace_safe_", 12) != 0) {
+	if ((strncmp(name, "dtrace_", 7) == 0 &&
+	    strncmp(name, "dtrace_safe_", 12) != 0) ||
+	    strcmp(name, "trap_check") == 0) {
 		/*
 		 * Anything beginning with "dtrace_" may be called
 		 * from probe context unless it explicitly indicates
 		 * that it won't be called from probe context by
 		 * using the prefix "dtrace_safe_".
+		 *
+		 * Additionally, we avoid instrumenting trap_check() to avoid
+		 * the possibility of generating a fault in probe context before
+		 * DTrace's fault handler is called.
 		 */
 		return (0);
 	}