x86/test_nx: update testing module after recent kernel changes

index 3f92ce0..98020bc 100644
@@ -12,6 +12,7 @@
 #include <linux/module.h>
 #include <linux/sort.h>
 #include <linux/slab.h>
+#include <linux/kdebug.h>
 
 #include <asm/uaccess.h>
 #include <asm/asm.h>

@@ -27,48 +28,30 @@ extern int rodata_test_data;
  *
  * To do this, the test code tries to execute memory in stack/kmalloc/etc,
  * and then checks if the expected trap happens.
- *
- * Sadly, this implies having a dynamic exception handling table entry.
- * ... which can be done (and will make Rusty cry)... but it can only
- * be done in a stand-alone module with only 1 entry total.
- * (otherwise we'd have to sort and that's just too messy)
  */
 
+static void *exception_label;
 
+void fixup_label(void);
 
-/*
- * We want to set up an exception handling point on our stack,
- * which means a variable value. This function is rather dirty
- * and walks the exception table of the module, looking for a magic
- * marker and replaces it with a specific function.
- */
-static void fudze_exception_table(void *marker, void *new)
+static int die_notify(struct notifier_block *nb, unsigned long val, void *data)
 {
-	struct module *mod = THIS_MODULE;
-	struct exception_table_entry *extable;
+	struct die_args *args = data;
+	struct pt_regs *regs = args->regs;
 
-	/*
-	 * Note: This module has only 1 exception table entry,
-	 * so searching and sorting is not needed. If that changes,
-	 * this would be the place to search and re-sort the exception
-	 * table.
-	 */
-	if (mod->num_exentries > 1) {
-		printk(KERN_ERR "test_nx: too many exception table entries!\n");
-		printk(KERN_ERR "test_nx: test results are not reliable.\n");
-		return;
+	if (val == DIE_PAGE_FAULT &&
+	    regs->ip == (unsigned long)exception_label) {
+		regs->ip = (unsigned long)fixup_label;
+		return NOTIFY_STOP;
 	}
-	extable = (struct exception_table_entry *)mod->extable;
-	extable[0].insn = (unsigned long)new;
-}
 
+	return NOTIFY_DONE;
+}
 
-/*
- * exception tables get their symbols translated so we need
- * to use a fake function to put in there, which we can then
- * replace at runtime.
- */
-void foo_label(void);
+struct notifier_block die_notify_block = {
+	.notifier_call = die_notify,
+	.priority = INT_MAX, /* Highest priority */
+};
 
 /*
  * returns 0 for not-executable, negative for executable

@@ -81,23 +64,20 @@ static noinline int test_address(void *address)
 {
 	unsigned long result;
 
-	/* Set up an exception table entry for our address */
-	fudze_exception_table(&foo_label, address);
+	/* Set up an exception label for our address */
+	exception_label = address;
+
 	result = 1;
 	asm volatile(
-		"foo_label:\n"
-		"0:	call *%[fake_code]\n"
-		"1:\n"
+		"	call *%[fake_code]\n"
 		".section .fixup,\"ax\"\n"
-		"2:	mov %[zero], %[rslt]\n"
+		"fixup_label:\n"
+		"	mov %[zero], %[rslt]\n"
 		"	ret\n"
 		".previous\n"
-		_ASM_EXTABLE(0b,2b)
 		: [rslt] "=r" (result)
 		: [fake_code] "r" (address), [zero] "r" (0UL), "0" (result)
 	);
-	/* change the exception table back for the next round */
-	fudze_exception_table(address, &foo_label);
 
 	if (result)
 		return -ENODEV;

@@ -108,10 +88,14 @@ static unsigned char test_data = 0xC3; /* 0xC3 is the opcode for "ret" */
 
 static int test_NX(void)
 {
-	int ret = 0;
 	/* 0xC3 is the opcode for "ret" */
 	char stackcode[] = {0xC3, 0x90, 0 };
 	char *heap;
+	int ret;
+
+	ret = register_die_notifier(&die_notify_block);
+	if (ret)
+		return ret;
 
 	test_data = 0xC3;
 

@@ -126,8 +110,10 @@ static int test_NX(void)
 
 	/* Test 2: Check if the heap is executable */
 	heap = kmalloc(64, GFP_KERNEL);
-	if (!heap)
-		return -ENOMEM;
+	if (!heap) {
+		ret = -ENOMEM;
+		goto out;
+	}
 	heap[0] = 0xC3; /* opcode for "ret" */
 
 	if (test_address(heap)) {

@@ -153,14 +139,16 @@ static int test_NX(void)
 	}
 #endif
 
-#if 0
+#ifdef CONFIG_DEBUG_SET_MODULE_RONX
 	/* Test 4: Check if the .data section of a module is executable */
 	if (test_address(&test_data)) {
 		printk(KERN_ERR "test_nx: .data section is executable\n");
 		ret = -ENODEV;
 	}
-
 #endif
+
+out:
+	unregister_die_notifier(&die_notify_block);
 	return ret;
 }
 

index 027088f..0d2665e 100644
@@ -664,6 +664,10 @@ no_context(struct pt_regs *regs, unsigned long error_code,
 	if (is_errata93(regs, address))
 		return;
 
+	if (notify_die(DIE_PAGE_FAULT, "page fault", regs, error_code,
+				X86_TRAP_PF, signal) == NOTIFY_STOP)
+		return;
+
 	/*
 	 * Oops. The kernel tried to access some bad page. We'll have to
 	 * terminate things with extreme prejudice: