Switch to compartment 0 when calling exit(3)

Destructors are wrapped to run in their respective compartments, but
transition back to compartment 0 after running. This means that the
caller stack must be writable in compartment 0. This change wraps
exit(3) and transitions to compartment 0 and the shared stack before
running exit() from libc.

Fixes #285

libia2/CMakeLists.txt

@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.12)
 project(libia2)
 
-add_library(libia2 ia2.c threads.c main.c)
+add_library(libia2 ia2.c threads.c main.c exit.c)
 target_compile_options(libia2 PRIVATE "-fPIC")
 
 if(LIBIA2_DEBUG)

libia2/exit.c

@@ -0,0 +1,33 @@
+#include "ia2.h"
+#include <dlfcn.h>
+
+__attribute__((used))
+static void call_libc_exit(int status) {
+  void (*exit_ptr)(int) = dlsym(RTLD_NEXT, "exit");
+  if (!exit_ptr) {
+    printf("Could not find exit(3) in the next DSO\n");
+    _exit(status);
+  }
+  exit_ptr(status);
+}
+
+__attribute__((naked)) void exit(int status) {
+  __asm__(
+      /* clang-format off */
+      "pushq %%rbp\n"
+      "movq %%rsp, %%rbp\n"
+      // Load the stack pointer for the shared compartment's stack.
+      "mov ia2_stackptr_0@GOTTPOFF(%%rip), %%r11\n"
+      "mov %%fs:(%%r11), %%rsp\n"
+      // Switch pkey to the appropriate compartment.
+      "xor %%ecx,%%ecx\n"
+      "mov %%ecx,%%edx\n"
+      "mov_pkru_eax 0\n"
+      "wrpkru\n"
+      // Align the stack before continuing
+      "subq $8, %%rsp\n"
+      // Call the real exit function.
+      "call call_libc_exit\n"
+      /* clang-format on */
+      ::);
+}

rewriter/tests/CMakeLists.txt

@@ -29,6 +29,7 @@ add_custom_target(check COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure)
 add_dependencies(check check-ia2)
 set_target_properties(check PROPERTIES FOLDER "tests")
 
+add_subdirectory(destructors)
 # add_subdirectory(ffmpeg)
 add_subdirectory(header_includes)
 add_subdirectory(heap_two_keys)

rewriter/tests/destructors/CMakeLists.txt

@@ -0,0 +1,18 @@
+# Build the plugin lib
+define_shared_lib(
+    SRCS plugin.c
+    INCLUDE_DIR include/main
+    NEEDS_LD_WRAP
+    PKEY 2
+)
+
+# Build the test
+define_test(
+    SRCS main.c
+    INCLUDE_DIR include/plugin
+    NEEDS_LD_WRAP
+    PKEY 1
+    CRITERION_TEST
+)
+
+define_ia2_wrapper()

rewriter/tests/destructors/include/main/exported_fn.h

@@ -0,0 +1,12 @@
+#pragma once
+#include <stdint.h>
+#include <stdbool.h>
+
+void print_message(void);
+
+// This is exported to avoid an implicit decl error when the plugin tries to
+// access it, but it's explicitly not shared to test that an MPK violation
+// occurs.
+extern uint32_t secret;
+
+extern bool debug_mode;

rewriter/tests/destructors/include/plugin/plugin.h

@@ -0,0 +1,13 @@
+// No need to execute the program here since the header exported by the main
+// binary does that.
+#pragma once
+#include <stdint.h>
+
+void start_plugin(void);
+
+void exit_from_plugin(void);
+
+// This is exported to avoid an implicit decl error when the main binary tries
+// to access it, but it's explicitly not shared to test that an MPK violation
+// occurs.
+extern uint32_t plugin_secret;

rewriter/tests/destructors/main.c

@@ -0,0 +1,56 @@
+#include <criterion/criterion.h>
+#include <criterion/logging.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#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);
+#define IA2_COMPARTMENT 1
+#include <ia2_compartment_init.inc>
+
+uint32_t secret = 0x09431233;
+
+static bool steal_plugin_secret = false;
+// Running in debug mode prints the addresses of the secrets defined in each
+// compartment. This is off by default to simplify the diff of stdout against
+// the expected output.
+bool debug_mode IA2_SHARED_DATA = false;
+
+bool clean_exit IA2_SHARED_DATA = false;
+
+void print_message(void) {
+    cr_log_info("this is defined in the main binary");
+    if (debug_mode) {
+        cr_log_info("the main secret is at %p", &secret);
+    }
+    cr_assert(secret == 0x09431233);
+    if (steal_plugin_secret) {
+        cr_assert(CHECK_VIOLATION(plugin_secret) == 0x78341244);
+    }
+}
+
+Test(two_keys, main) {
+    start_plugin();
+}
+
+Test(two_keys, plugin) {
+    steal_plugin_secret = true;
+    start_plugin();
+}
+
+Test(two_keys, clean_exit) {
+    clean_exit = true;
+    start_plugin();
+    exit(0);
+}
+
+Test(two_keys, plugin_clean_exit) {
+    clean_exit = true;
+    start_plugin();
+    exit_from_plugin();
+}

rewriter/tests/destructors/plugin.c

@@ -0,0 +1,29 @@
+#include <criterion/criterion.h>
+#include <criterion/logging.h>
+#include <stdio.h>
+#include <ia2.h>
+#include "exported_fn.h"
+#include "test_fault_handler.h"
+
+#define IA2_COMPARTMENT 2
+#include <ia2_compartment_init.inc>
+
+uint32_t plugin_secret = 0x78341244;
+
+extern bool clean_exit;
+
+void start_plugin(void) {
+    cr_log_info("this is defined in the plugin");
+    if (debug_mode) {
+        cr_log_info("the plugin secret is at %p", &plugin_secret);
+    }
+    cr_assert(plugin_secret == 0x78341244);
+    print_message();
+    if (!clean_exit) {
+        cr_assert(CHECK_VIOLATION(secret) == 0x09431233);
+    }
+}
+
+void exit_from_plugin(void) {
+    exit(0);
+}

rewriter/tests/lit.cfg.py

@@ -18,7 +18,8 @@
 # excludes: A list of directories to exclude from the testsuite. The 'Inputs'
 # subdirectories contain auxiliary inputs for various tests in their parent
 # directories.
-config.excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt', 'libusb', 'ffmpeg']
+config.excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt',
+                   'libusb', 'ffmpeg', 'destructors']
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)