Replace criterion with our test runner on x86

cmake/define-test.cmake

@@ -137,7 +137,7 @@ function(define_test)
 
     set(ADD_SANDBOX_DEP FALSE)
     if (DEFINE_TEST_CRITERION_TEST)
-        list(APPEND DEFINE_TEST_UNWRAPPED_LIBS criterion)
+        list(APPEND DEFINE_TEST_UNWRAPPED_LIBS ia2_test_runner)
         if (NOT DEFINE_TEST_NOT_IN_CHECK_IA2)
             if (LIBIA2_AARCH64)
                 # unless natively AArch64, default to running tests with qemu-aarch64 and a custom LD_LIBRARY_PATH

cmake/ia2.cmake

@@ -43,13 +43,8 @@ function(add_ia2_compartment NAME TYPE)
     add_library(${NAME} SHARED)
   endif()
 
-  # The x86 version is missing a dependency here and libs rely on include path overlap to pick up
-  # criterion header. I want to keep spoofed criterion static for simplicity so we add the -I
-  # directly here.
-  if (LIBIA2_AARCH64)
-    target_include_directories(${NAME} PRIVATE
-        ${CMAKE_SOURCE_DIR}/misc/spoofed_criterion/include)
-  endif()
+  target_include_directories(${NAME} PRIVATE
+      ${CMAKE_SOURCE_DIR}/misc/test_runner/include)
   target_compile_definitions(${NAME} PRIVATE
     IA2_ENABLE=1
     PKEY=${ARG_PKEY}
@@ -153,10 +148,11 @@ function(create_compile_commands NAME TYPE)
   # Copy target properties from the real target. We might need to add more properties.
   target_link_libraries(${COMPILE_COMMAND_TARGET} PRIVATE $<TARGET_PROPERTY:${NAME},LINK_LIBRARIES>)
   target_include_directories(${COMPILE_COMMAND_TARGET} PRIVATE ${INCLUDE_DIRECTORIES})
-  if (LIBIA2_AARCH64)
-    target_include_directories(${COMPILE_COMMAND_TARGET} PRIVATE
-        ${CMAKE_SOURCE_DIR}/misc/spoofed_criterion/include)
-  endif()
+  # The test runner is a static library that just defines main so we don't link it into
+  # the libraries defined by tests and instead just add the include flags for its
+  # assertions
+  target_include_directories(${COMPILE_COMMAND_TARGET} PRIVATE
+      ${CMAKE_SOURCE_DIR}/misc/test_runner/include)
   set(CMAKE_EXPORT_COMPILE_COMMANDS OFF)
 endfunction()
 

misc/CMakeLists.txt

@@ -1,3 +1 @@
-if (LIBIA2_AARCH64)
-    add_subdirectory(spoofed_criterion)
-endif()
+add_subdirectory(test_runner)

misc/test_runner/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_library(ia2_test_runner STATIC test_runner.c)
+target_include_directories(ia2_test_runner
+    PRIVATE ${CMAKE_SOURCE_DIR}/runtime/libia2/include)

misc/test_runner/include/ia2_test_runner.h

@@ -0,0 +1,66 @@
+#pragma once
+#include <ia2.h>
+#include <stdio.h>
+#include <stdbool.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <assert.h>
+
+/*
+ * Tests should include this header without defining the following macro to avoid rewriting function
+ * pointers that shouldn't be rewritten.
+ */
+#if !defined(IA2_TEST_RUNNER_SOURCE)
+typedef void *ia2_test_fn;
+#else
+typedef void (*ia2_test_fn)(void);
+#endif
+
+struct fake_criterion_test {
+  ia2_test_fn test;
+  ia2_test_fn init;
+  int exit_code;
+};
+
+/*
+ * Placing IA2_{BEGIN,END}_NO_WRAP between the function declaration stops the rewriter from creating a
+ * direct call gate for the test function or indirect call gates for function pointer expressions
+ * that reference it like the RHS when initializing struct fake_criterion_test's test field. The
+ * last line of this macro is the start of the test function's definition and should be followed by { }
+ */
+#define Test(suite, name, ...)                                                                                                 \
+  IA2_BEGIN_NO_WRAP                                                                                                            \
+  void fake_criterion_##suite##_##name(void);                                                                                  \
+  IA2_END_NO_WRAP                                                                                                              \
+  __attribute__((__section__("fake_criterion_tests"))) struct fake_criterion_test fake_criterion_##suite##_##name##_##test = { \
+      .test = fake_criterion_##suite##_##name,                                                                                 \
+      ##__VA_ARGS__};                                                                                                          \
+  void fake_criterion_##suite##_##name(void)
+
+#define cr_log_info(f, ...) printf(f "\n", ##__VA_ARGS__)
+#define cr_log_error(f, ...) fprintf(stderr, f "\n", ##__VA_ARGS__)
+
+#define cr_assert assert
+#define cr_assert_eq(a, b) cr_assert((a) == (b))
+#define cr_assert_lt(a, b) cr_assert((a) < (b))
+#define cr_fatal(s)          \
+  do {                       \
+    fprintf(stderr, s "\n"); \
+    exit(1);                 \
+  } while (0)
+
+/*
+ * Configure the signal handler to expect an mpk violation when `expr` is evaluated. If `expr`
+ * doesn't trigger a fault, the process exits with a non-zero exit status.
+ */
+#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

@@ -0,0 +1,123 @@
+/* Define this macro to make the pointers in struct fake_criterion_test function pointers */
+#define IA2_TEST_RUNNER_SOURCE
+#include "include/ia2_test_runner.h"
+
+#include <stdbool.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+extern struct fake_criterion_test __start_fake_criterion_tests;
+extern struct fake_criterion_test __stop_fake_criterion_tests;
+
+/* This is shared data to allow checking for violations from different compartments */
+bool expect_fault IA2_SHARED_DATA = false;
+
+/* Create a stack for the signal handler to use since the tests may trigger it from any compartment */
+char sighandler_stack[4 * 1024] IA2_SHARED_DATA __attribute__((aligned(16))) = {0};
+
+/* Pointer to the start of the signal handler stack */
+char *sighandler_sp 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__(
+      // This signal handler is defined in the main binary, but it doesn'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 callgate 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 is used for manual sanity-checks to ensure check whether a segfault occurred and
+ * if it was expected or not.
+ */
+void print_mpk_message(int sig) {
+  if (sig == SIGSEGV) {
+    /* segfault happened at the expected place so exit with status code zero */
+    const char *ok_msg = "CHECK_VIOLATION: seg faulted as expected\n";
+    /* segfault happened at an unexpected place so exit with non-zero status */
+    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 directly to stdout since printf is not async-signal-safe */
+    write(1, msg, strlen(msg));
+    if (!expect_fault) {
+      _exit(-1);
+    }
+  }
+  _exit(0);
+}
+
+int main() {
+  struct sigaction act = {
+      .sa_handler = handle_segfault,
+  };
+  /*
+   * Installs a signal handler that will be inherited by the child processes created for each
+   * invocation of the Test macro
+   */
+  sigaction(SIGSEGV, &act, NULL);
+  struct fake_criterion_test *test_info = &__start_fake_criterion_tests;
+  for (; test_info < &__stop_fake_criterion_tests; test_info++) {
+    if (!test_info->test) {
+      break;
+    }
+    pid_t pid = fork();
+    bool in_child = pid == 0;
+    if (in_child) {
+      /*
+       * This .c is not rewritten so these indirect callsites have no callgates and their callees must
+       * not be wrapped. That means the Test macro should not expose function pointer types to
+       * rewritten source files (i.e. the test sources).
+       */
+      if (test_info->init) {
+        (*test_info->init)();
+      }
+      (*test_info->test)();
+      return 0;
+    }
+    // otherwise, in parent
+    int stat;
+    pid_t waited_pid = waitpid(pid, &stat, 0);
+    if (waited_pid < 0) {
+      perror("waitpid");
+      return 2;
+    }
+    if WIFSIGNALED(stat) {
+      fprintf(stderr, "forked test child was terminated by signal %d\n", WTERMSIG(stat));
+      return 1;
+    }
+    int exit_status = WEXITSTATUS(stat);
+    if (exit_status != test_info->exit_code) {
+      fprintf(stderr, "forked test child exited with status %d, but %d was expected\n", exit_status, test_info->exit_code);
+      return 1;
+    }
+  }
+  return 0;
+}