Move test_fault_handler.h into test_runner.c

misc/test_runner/include/ia2_test_runner.h

@@ -1,5 +1,6 @@
 #pragma once
 #include <stdio.h>
+#include <stdbool.h>
 
 struct fake_criterion_test {
   void (*test)(void);
@@ -25,3 +26,30 @@ struct fake_criterion_test {
     fprintf(stderr, s "\n"); \
     exit(1);                 \
   } while (0)
+
+/*
+ * This header defines a test framework for detecting MPK violations using
+ * signal handlers. This file must be included exactly once from a source file
+ * in the main binary with IA2_DEFINE_TEST_HANDLER defined by the preprocessor.
+ * This will define the functions and variables used by the test handler, ensure
+ * it is initialized before main and provide access to the LOG and
+ * CHECK_VIOLATION macros. Other files which need CHECK_VIOLATION or LOG may
+ * include the header without defining IA2_DEFINE_TEST_HANDLER. Using
+ * CHECK_VIOLATION without defining the test handler will trigger a linker error
+ * when building the shared object.
+ */
+
+// Configure the signal handler to expect an mpk violation when `expr` is
+// evaluated. If `expr` doesn't trigger a fault, this macro manually raises a
+// fault with a different message.
+#define CHECK_VIOLATION(expr)                                                  \
+  ({                                                                           \
+    expect_fault = true;                                                       \
+    asm volatile("" : : : "memory");                                           \
+    volatile typeof(expr) _tmp = expr;                                         \
+    printf("CHECK_VIOLATION: did not seg fault as expected\n");                \
+    _exit(1);                                                                  \
+    _tmp;                                                                      \
+  })
+
+extern bool expect_fault;

misc/test_runner/test_runner.c

@@ -1,6 +1,7 @@
 #include "include/ia2_test_runner.h"
 #include <stdbool.h>
 #include <stdio.h>
+#include <string.h>
 #include <sys/wait.h>
 #include <unistd.h>
 
@@ -41,3 +42,72 @@ int main() {
   }
   return 0;
 }
+
+// This is shared data to allow checking for violations in multiple
+// compartments. We avoid using IA2_SHARED_DATA here to avoid including ia2.h
+// since that would pull in libia2 as a dependency (the libia2 build generates a
+// header included in ia2.h).
+bool expect_fault __attribute__((section("ia2_shared_data"))) = false;
+
+// Create a stack for the signal handler to use
+char sighandler_stack[4 * 1024] __attribute__((section("ia2_shared_data")))
+__attribute__((aligned(16))) = {0};
+char *sighandler_sp __attribute__((section("ia2_shared_data"))) =
+    &sighandler_stack[(4 * 1024) - 8];
+
+// This function must be declared naked because it's not necessarily safe for it
+// to write to the stack in its prelude (the stack isn't written to when the
+// function itself is called because it's only invoked as a signal handler).
+#if defined(__x86_64__)
+__attribute__((naked)) void handle_segfault(int sig) {
+  // This asm must preserve %rdi which contains the argument since
+  // print_mpk_message reads it
+  __asm__(
+      // Signal handlers are defined in the main binary, but they don't run with
+      // the same pkru state as the interrupted context. This means we have to
+      // remove all MPK restrictions to ensure can run it correctly.
+      "xorl %ecx, %ecx\n"
+      "xorl %edx, %edx\n"
+      "xorl %eax, %eax\n"
+      "wrpkru\n"
+      // Switch the stack to a shared buffer. There's only one u32 argument and
+      // no returns so we don't need a full wrapper here.
+      "movq sighandler_sp@GOTPCREL(%rip), %rsp\n"
+      "movq (%rsp), %rsp\n"
+      "callq print_mpk_message");
+}
+#elif defined(__aarch64__)
+#warning "Review test_fault_handler implementation after enabling x18 switching"
+void print_mpk_message(int sig);
+void handle_segfault(int sig) {
+    print_mpk_message(sig);
+}
+#endif
+
+// The test output should be checked to see that the segfault occurred at the
+// expected place.
+void print_mpk_message(int sig) {
+  if (sig == SIGSEGV) {
+    // Write directly to stdout since printf is not async-signal-safe
+    const char *ok_msg = "CHECK_VIOLATION: seg faulted as expected\n";
+    const char *early_fault_msg = "CHECK_VIOLATION: unexpected seg fault\n";
+    const char *msg;
+    if (expect_fault) {
+      msg = ok_msg;
+    } else {
+      msg = early_fault_msg;
+    }
+    write(1, msg, strlen(msg));
+    if (!expect_fault) {
+      _exit(-1);
+    }
+  }
+  _exit(0);
+}
+
+// Installs the previously defined signal handler and disables buffering on
+// stdout to allow using printf prior to the sighandler
+__attribute__((constructor)) void install_segfault_handler(void) {
+  setbuf(stdout, NULL);
+  signal(SIGSEGV, handle_segfault);
+}

tests/destructors/main.c

@@ -6,7 +6,7 @@
 #include <ia2.h>
 #include "plugin.h"
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 // This test uses two protection keys
 INIT_RUNTIME(2);

tests/destructors/plugin.c

@@ -3,7 +3,7 @@
 #include <stdio.h>
 #include <ia2.h>
 #include "exported_fn.h"
-#include "test_fault_handler.h"
+
 
 #define IA2_COMPARTMENT 2
 #include <ia2_compartment_init.inc>

tests/heap_two_keys/main.c

@@ -10,7 +10,7 @@ RUN: sh -c 'if [ ! -s "heap_two_keys_call_gates_0.ld" ]; then echo "No link args
 #include <ia2_allocator.h>
 #include "plugin.h"
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 // This test uses two protection keys
 INIT_RUNTIME(2);

tests/heap_two_keys/plugin.c

@@ -4,7 +4,7 @@ RUN: cat heap_two_keys_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
 #include <stdio.h>
 #include <ia2.h>
 #include "exported_fn.h"
-#include "test_fault_handler.h"
+
 
 #define IA2_COMPARTMENT 2
 #include <ia2_compartment_init.inc>

tests/mmap_loop/main.c

@@ -9,7 +9,7 @@ RUN: sh -c 'if [ ! -s "mmap_loop_call_gates_0.ld" ]; then echo "No link args as
 #include <ia2_test_runner.h>
 
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 /*
     This program tests that mmap and heap allocations are handled properly.

tests/protected_threads/main.c

@@ -9,7 +9,7 @@ RUN: sh -c 'if [ ! -s "protected_threads_call_gates_0.ld" ]; then echo "No link
 #include <sys/wait.h>
 #define IA2_DEFINE_TEST_HANDLER
 #include <ia2_test_runner.h>
-#include <test_fault_handler.h>
+
 #include <unistd.h>
 
 INIT_RUNTIME(2);

tests/read_config/main.c

@@ -15,7 +15,7 @@ RUN: cat main.c | FileCheck --match-full-lines --check-prefix=REWRITER %s
 // including plugin.h (which does include the output header) before core.h.
 #include "core.h"
 #define IA2_DEFINE_TEST_HANDLER
-#include <test_fault_handler.h>
+
 
 /*
     This test is modeled after nginx's function pointer usage. In this test,

tests/recursion/main.c

@@ -8,7 +8,7 @@ RUN: cat recursion_call_gates_2.ld | FileCheck --check-prefix=LINKARGS %s
 #include <ia2.h>
 #include <stdio.h>
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 INIT_RUNTIME(2);
 #define IA2_COMPARTMENT 1

tests/ro_sharing/main.c

@@ -7,7 +7,7 @@ RUN: sh -c 'if [ ! -s "ro_sharing_call_gates_0.ld" ]; then echo "No link args as
 #include <ia2.h>
 #include <stdio.h>
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 // This test checks that all RO data mapped in from executable files is shared.
 // This is needed so that the dynamic linker can read ELF metadata. Read-only

tests/ro_sharing/plugin.c

@@ -1,7 +1,7 @@
 /*
 RUN: cat ro_sharing_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
 */
-#include "test_fault_handler.h"
+
 #include <ia2_test_runner.h>
 
 #include <ia2.h>

tests/should_segfault/main.c

@@ -8,7 +8,7 @@ RUN: sh -c 'if [ ! -s "should_segfault_call_gates_0.ld" ]; then echo "No link ar
 #include <print_secret.h>
 #include <ia2.h>
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 INIT_RUNTIME(1);
 #define IA2_COMPARTMENT 1

tests/should_segfault/print_secret.c

@@ -6,7 +6,7 @@ RUN: cat should_segfault_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
 #include <stdlib.h>
 #include <stdbool.h>
 #include "print_secret.h"
-#include "test_fault_handler.h"
+
 
 static bool early_fault = false;
 

tests/sighandler/lib.c

@@ -8,7 +8,7 @@ RUN: cat sighandler_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
 #include <ia2_test_runner.h>
 
 #define IA2_DEFINE_TEST_HANDLER
-#include <test_fault_handler.h>
+
 
 #define IA2_COMPARTMENT 2
 #include <ia2_compartment_init.inc>

tests/sighandler/main.c

@@ -8,7 +8,7 @@ We need this because lib.c uses LINKARGS checks but not this file.
 #include <signal.h>
 #include <ia2.h>
 #define IA2_DEFINE_TEST_HANDLER
-#include <test_fault_handler.h>
+
 #include <ia2_test_runner.h>
 
 INIT_RUNTIME(2);

tests/threads/main.c

@@ -10,7 +10,7 @@ RUN: sh -c 'if [ ! -s "threads_call_gates_0.ld" ]; then echo "No link args as ex
 #include <string.h>
 #include <sys/wait.h>
 #define IA2_DEFINE_TEST_HANDLER
-#include <test_fault_handler.h>
+
 #include <unistd.h>
 
 INIT_RUNTIME(1);

tests/three_keys_minimal/main.c

@@ -11,7 +11,7 @@ INIT_RUNTIME(3);
 #include <threads.h>
 
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 void main_noop(void) {
 }

tests/tls_protected/library.c

@@ -3,7 +3,7 @@
 
 // Check that readelf shows exactly one executable segment
 #include "library.h"
-#include "test_fault_handler.h"
+
 #include <ia2_test_runner.h>
 
 #include <ia2.h>

tests/tls_protected/main.c

@@ -6,7 +6,7 @@
 #include <stdint.h>
 #include <unistd.h>
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 #include <threads.h>
 
 INIT_RUNTIME(2);

tests/trusted_direct/main.c

@@ -10,7 +10,7 @@ RUN: cat trusted_direct_call_gates_0.ld | FileCheck --check-prefix=LINKARGS %s
 #include <ia2.h>
 #include "plugin.h"
 #define IA2_DEFINE_TEST_HANDLER
-#include "test_fault_handler.h"
+
 
 // This test checks that an untrusted library can call a trusted main binary. An
 // MPK violation is triggered from the untrusted library if no arguments are