Improve testing framework

cmake/define-test.cmake

@@ -138,9 +138,11 @@ endfunction()
 #       wrapped.
 # UNWRAPPED_LIBRARY_DIRS (optional) - extra library directories.
 # UNWRAPPED_LIBS (optional) - extra libraries that are not wrapped.
+# CRITERION_TEST - If present, link against criterion and add the test to the
+#                  cmake test infrastructure.
 function(define_test)
     # Parse options
-    set(options NEEDS_LD_WRAP NOT_IN_CHECK_IA2 NO_LIBS)
+    set(options NEEDS_LD_WRAP NOT_IN_CHECK_IA2 NO_LIBS CRITERION_TEST)
     set(oneValueArgs PKEY)
     set(multiValueArgs LIBS SRCS INCLUDE_DIR
         UNWRAPPED_INCLUDE_DIRS UNWRAPPED_LIBRARY_DIRS UNWRAPPED_LIBS)
@@ -160,6 +162,14 @@ function(define_test)
         set(DEFINE_TEST_LIBS ${TEST_NAME}_lib)
     endif()
 
+    if (DEFINE_TEST_CRITERION_TEST)
+        list(APPEND DEFINE_TEST_UNWRAPPED_LIBS criterion)
+        if (NOT DEFINE_TEST_NOT_IN_CHECK_IA2)
+            add_test(${TEST_NAME} ${WRAPPED_MAIN})
+            add_dependencies(check ${WRAPPED_MAIN})
+        endif()
+    endif()
+
     if(DEFINED DEFINE_TEST_INCLUDE_DIR)
         set(RELATIVE_INCLUDE_DIR ${DEFINE_TEST_INCLUDE_DIR})
     else()

rewriter/tests/CMakeLists.txt

@@ -1,3 +1,5 @@
+include(CTest)
+
 enable_testing()
 
 include("../../cmake/define-ia2-wrapper.cmake")
@@ -22,13 +24,12 @@ set_target_properties(check-ia2 PROPERTIES FOLDER "tests")
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
-# Set up 'check' alias
-add_custom_target(check)
+# Set up 'check' target
+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(ffmpeg)
-add_subdirectory(function_allowlist)
 add_subdirectory(header_includes)
 add_subdirectory(heap_two_keys)
 # add_subdirectory(libusb)

rewriter/tests/global_fn_ptr/CMakeLists.txt

@@ -7,6 +7,7 @@ define_shared_lib(
 define_test(
     SRCS main.c
     NEEDS_LD_WRAP
+    CRITERION_TEST
 )
 
 # Build the wrapper lib

rewriter/tests/global_fn_ptr/include/operations.h

@@ -1,14 +1,13 @@
 /*
 RUN: cat %t.c.args | FileCheck --check-prefix=LINKARGS %s
-RUN: %binary_dir/tests/global_fn_ptr/global_fn_ptr-main | diff %S/../Output/operations.out -
 */
 
 #pragma once
 #include <stdlib.h>
 #include <stdint.h>
 
 // LINKARGS: --wrap=call_operations
-void call_operations(void);
+uint32_t call_operation(size_t i);
 
 // CHECK: typedef struct IA2_fnptr__ZTSPFjjjE WordFn;
 typedef uint32_t (*WordFn)(uint32_t, uint32_t);

rewriter/tests/global_fn_ptr/main.c

@@ -1,9 +1,9 @@
 /*
 RUN: sh -c 'if [ ! -s "global_fn_ptr_call_gates_0.ld" ]; then echo "No link args as expected"; exit 0; fi; echo "Unexpected link args"; exit 1;'
-RUN: %binary_dir/tests/global_fn_ptr/global_fn_ptr_main_wrapped | diff %S/Output/operations.out -
 */
 #include "operations.h"
 #include <ia2.h>
+#include <criterion/criterion.h>
 
 uint32_t add(uint32_t x, uint32_t y) { return x + y; }
 uint16_t sub(uint16_t x, uint16_t y) { return x - y; }
@@ -39,6 +39,7 @@ Op operations[2] IA2_SHARED_DATA = {
     },
 };
 
-int main() {
-    call_operations();
+Test(global_fn_ptr, main) {
+    cr_assert(call_operation(0) == 43312);
+    cr_assert(call_operation(1) == 461513047);
 }

rewriter/tests/global_fn_ptr/operations.c

@@ -2,18 +2,17 @@
 RUN: cat global_fn_ptr_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
 */
 #include "operations.h"
+#include <criterion/criterion.h>
 #include <stdio.h>
 
 extern Op operations[2];
 
-// LINKARGS: --wrap=call_operations
-void call_operations(void) {
-    for (int i = 0; i < 2; i++) {
-        // TODO: Add a way to share strings between compartments
-        //printf("%s\n", operations[i].desc.data);
-        uint32_t x = 18923;
-        uint32_t y = 24389;
-        uint32_t res = operations[i].function(x, y);
-        printf("%d = f(%d, %d)\n", res, x, y);
-    }
+// LINKARGS: --wrap=call_operation
+uint32_t call_operation(size_t i) {
+    // TODO: Add a way to share strings between compartments
+    //printf("%s\n", operations[i].desc.data);
+    uint32_t x = 18923;
+    uint32_t y = 24389;
+    uint32_t res = operations[i].function(x, y);
+    return res;
 }

rewriter/tests/header_includes/CMakeLists.txt

@@ -7,6 +7,7 @@ define_shared_lib(
 define_test(
     SRCS main.c
     NEEDS_LD_WRAP
+    CRITERION_TEST
 )
 
 # Build the wrapper lib

rewriter/tests/header_includes/liboption.c

@@ -1,13 +1,13 @@
 /*
 RUN: cat header_includes_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
 */
-#include <stdio.h>
+#include <criterion/logging.h>
 #include "liboption.h"
 #include "types.h"
 
 // LINKARGS: --wrap=None
 Option None() {
-    printf("returning `None`\n");
+    cr_log_info("returning `None`");
     Option none = {
         .value = 0,
         .present = false,
@@ -17,7 +17,7 @@ Option None() {
 
 // LINKARGS: --wrap=Some
 Option Some(int x) {
-    printf("returning `Some(%d)`\n", x);
+    cr_log_info("returning `Some(%d)`", x);
     Option opt = {
         .value = x,
         .present = true,

rewriter/tests/header_includes/main.c

@@ -3,16 +3,17 @@ RUN: sh -c 'if [ ! -s "header_includes_call_gates_0.ld" ]; then echo "No link ar
 */
 #include "liboption.h"
 #include "types.h"
-#include <stdio.h>
+#include <criterion/criterion.h>
+#include <criterion/logging.h>
 #include <ia2.h>
 
 INIT_RUNTIME(1);
 #define IA2_COMPARTMENT 1
 #include <ia2_compartment_init.inc>
 
-int main() {
+Test(header_includes, main) {
     Option x = Some(3);
     Option none = None();
-    printf("`x` has value %d\n", unwrap_or(x, -1));
-    printf("`none` has no value %d\n", unwrap_or(none, -1));
+    cr_assert_eq(unwrap_or(x, -1), 3);
+    cr_assert_eq(unwrap_or(none, -1), -1);
 }

rewriter/tests/heap_two_keys/CMakeLists.txt

@@ -11,6 +11,7 @@ define_test(
     PKEY 1
     NEEDS_LD_WRAP
     INCLUDE_DIR include/plugin
+    CRITERION_TEST
 )
 
 define_ia2_wrapper()

rewriter/tests/heap_two_keys/main.c

@@ -1,11 +1,9 @@
 /*
 RUN: sh -c 'if [ ! -s "heap_two_keys_call_gates_0.ld" ]; then echo "No link args as expected"; exit 0; fi; echo "Unexpected link args"; exit 1;'
-RUN: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 0 | diff %S/Output/clean_exit.out -
-RUN: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 1 | diff %S/Output/fault.out -
-RUN: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 2 | diff %S/Output/fault.out -
-RUN: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 3 | diff %S/Output/clean_exit.out -
 */
-#include <stdio.h>
+#include <criterion/criterion.h>
+#include <criterion/logging.h>
+#include <criterion/new/assert.h>
 #include <unistd.h>
 #include <assert.h>
 #include <ia2.h>
@@ -20,81 +18,45 @@ INIT_RUNTIME(2);
 #include <ia2_compartment_init.inc>
 
 // Test that the program can exit without error
-int test_0() {
-    return 0;
+Test(heap_two_keys, 0, .init = trigger_compartment_init) {
 }
 
 // Test that the main binary's heap can't be read
-int test_1() {
+Test(heap_two_keys, 1, .init = trigger_compartment_init) {
     uint32_t *x = (uint32_t *)malloc(sizeof(uint32_t));
     if (!x) {
-        LOG("Failed to allocate memory on the heap");
-        return -1;
+        cr_fatal("Failed to allocate memory on the heap");
     }
     *x = 0x09431233;
     read_from_plugin_expect_fault((uint8_t*)x);
     free(x);
     // This test shouldn't return
-    return -1;
+    cr_fatal("Should have segfaulted but didn't");
 }
 
 // Test that the main binary's heap can't be written to
-int test_2() {
+Test(heap_two_keys, 2, .init = trigger_compartment_init) {
     // This zeroes out the allocated memory
     uint8_t *x = (uint8_t *)calloc(sizeof(uint8_t), 12);
     if (!x) {
-        LOG("Failed to allocate memory on the heap");
-        return -1;
+        cr_fatal("Failed to allocate memory on the heap");
     }
     write_from_plugin_expect_fault(x, 12);
     free(x);
     // This test shouldn't return
-    return -1;
+    cr_fatal("Should have segfaulted but didn't");
 }
 
 // Test that the main binary's shared data can be read
-int test_3() {
+Test(heap_two_keys, 3, .init = trigger_compartment_init) {
     uint16_t *x = (uint16_t *)shared_malloc(sizeof(uint16_t));
     if (!x) {
-        LOG("Failed to allocate memory on the heap");
-        return -1;
+        cr_fatal("Failed to allocate memory on the heap");
     }
     *x = 0xffed;
     assert(read_from_plugin((uint8_t*)x) == 0xed);
     shared_free(x);
-    return 0;
 }
 
 // TODO: Add tests for free, realloc, the plugin's heap and reserving more than
-// once from the gigacage (> 2MB in allocations)
-
-int main(int argc, char **argv) {
-    if (argc < 2) {
-        LOG("Run with an integer (0-3) as the first argument to select a test");
-        return -1;
-    }
-
-    // Call a no-op function to switch to this compartment's PKRU
-    trigger_compartment_init();
-
-    int test_num = *argv[1] - '0';
-
-    switch (test_num) {
-        case 0: {
-            return test_0();
-        }
-        case 1: {
-            return test_1();
-        }
-        case 2: {
-            return test_2();
-        }
-        case 3: {
-            return test_3();
-        }
-        default: {
-            LOG("Unknown test selected");
-            return -1;
-        }
-    }
-}
+// once from the gigacage (> 2MB in allocations)

rewriter/tests/macro_attr/CMakeLists.txt

@@ -7,6 +7,7 @@ define_shared_lib(
 define_test(
     SRCS main.c
     NEEDS_LD_WRAP
+    CRITERION_TEST
 )
 
 # Build the wrapper lib