From 532e02d03673d300a548947e70f5b14c44033707 Mon Sep 17 00:00:00 2001
From: Simon McFarlane <sim@immunant.com>
Date: Fri, 6 Oct 2023 17:03:57 -0700
Subject: [PATCH] heap_two_keys

---
 rewriter/tests/heap_two_keys/CMakeLists.txt |  1 +
 rewriter/tests/heap_two_keys/main.c         | 65 +++++----------------
 2 files changed, 17 insertions(+), 49 deletions(-)

diff --git a/rewriter/tests/heap_two_keys/CMakeLists.txt b/rewriter/tests/heap_two_keys/CMakeLists.txt
index 0fa58179c8..cd48fcd629 100644
--- a/rewriter/tests/heap_two_keys/CMakeLists.txt
+++ b/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()
diff --git a/rewriter/tests/heap_two_keys/main.c b/rewriter/tests/heap_two_keys/main.c
index e06c2ad554..2b14ab787e 100644
--- a/rewriter/tests/heap_two_keys/main.c
+++ b/rewriter/tests/heap_two_keys/main.c
@@ -2,10 +2,10 @@
 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;'
 TODO: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 0 | diff %S/Output/plugin.out -
 // TODO(src_rewriter_wip): had to change the output here, why?
-RUN: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 1 | diff %S/Output/main.out -
-TODO: %binary_dir/tests/heap_two_keys/heap_two_keys_main_wrapped 2 | 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>
@@ -21,81 +21,48 @@ INIT_RUNTIME(2);
 
 // Test that the program can exit without error
 // TODO(#112): it cannot.
-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_log_error("Failed to allocate memory on the heap");
+        exit(EXIT_FAILURE);
     }
     *x = 0x09431233;
     read_uint32_t_expect_fault(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_log_error("Failed to allocate memory on the heap");
+        exit(EXIT_FAILURE);
     }
     write_uint8_t_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_log_error("Failed to allocate memory on the heap");
+        exit(EXIT_FAILURE);
     }
     *x = 0xffed;
     assert(read_uint16_t(x) == 0xffed);
     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)
\ No newline at end of file