From 537a7be44f023ebf3db75855e8ba32d2b022ff74 Mon Sep 17 00:00:00 2001
From: Frances Wingerter <fw@immunant.com>
Date: Tue, 25 Jul 2023 12:28:44 -0400
Subject: [PATCH] tests: add mmap_loop test

---
 rewriter/tests/CMakeLists.txt                 |  1 +
 rewriter/tests/mmap_loop/CMakeLists.txt       | 15 ++++++
 rewriter/tests/mmap_loop/Output/mmap_loop.out |  1 +
 rewriter/tests/mmap_loop/include/mmap_loop.h  |  8 +++
 rewriter/tests/mmap_loop/main.c               | 53 +++++++++++++++++++
 rewriter/tests/mmap_loop/mmap_loop.c          | 25 +++++++++
 6 files changed, 103 insertions(+)
 create mode 100644 rewriter/tests/mmap_loop/CMakeLists.txt
 create mode 100644 rewriter/tests/mmap_loop/Output/mmap_loop.out
 create mode 100644 rewriter/tests/mmap_loop/include/mmap_loop.h
 create mode 100644 rewriter/tests/mmap_loop/main.c
 create mode 100644 rewriter/tests/mmap_loop/mmap_loop.c

diff --git a/rewriter/tests/CMakeLists.txt b/rewriter/tests/CMakeLists.txt
index 3067195c09..bf01f8b568 100644
--- a/rewriter/tests/CMakeLists.txt
+++ b/rewriter/tests/CMakeLists.txt
@@ -34,6 +34,7 @@ add_subdirectory(heap_two_keys)
 # add_subdirectory(libusb)
 add_subdirectory(macro_attr)
 add_subdirectory(minimal)
+add_subdirectory(mmap_loop)
 # TODO: support C++ namespaces
 #add_subdirectory(namespaces)
 add_subdirectory(recursion)
diff --git a/rewriter/tests/mmap_loop/CMakeLists.txt b/rewriter/tests/mmap_loop/CMakeLists.txt
new file mode 100644
index 0000000000..d45097fec5
--- /dev/null
+++ b/rewriter/tests/mmap_loop/CMakeLists.txt
@@ -0,0 +1,15 @@
+set(STDOUT_REF ${CMAKE_CURRENT_SOURCE_DIR}/Output/mmap_loop.out)
+configure_file(${STDOUT_REF} ${CMAKE_CURRENT_BINARY_DIR}/mmap_loop.out)
+
+define_shared_lib(
+    SRCS mmap_loop.c
+    PKEY 2
+)
+
+define_test(
+    SRCS main.c
+    PKEY 1
+    NEEDS_LD_WRAP
+)
+
+define_ia2_wrapper()
diff --git a/rewriter/tests/mmap_loop/Output/mmap_loop.out b/rewriter/tests/mmap_loop/Output/mmap_loop.out
new file mode 100644
index 0000000000..9766475a41
--- /dev/null
+++ b/rewriter/tests/mmap_loop/Output/mmap_loop.out
@@ -0,0 +1 @@
+ok
diff --git a/rewriter/tests/mmap_loop/include/mmap_loop.h b/rewriter/tests/mmap_loop/include/mmap_loop.h
new file mode 100644
index 0000000000..975775b765
--- /dev/null
+++ b/rewriter/tests/mmap_loop/include/mmap_loop.h
@@ -0,0 +1,8 @@
+#pragma once
+#include <stdlib.h>
+
+// LINKARGS: --wrap=lib_mmap_buf
+char *lib_mmap_buf(size_t size);
+
+// LINKARGS: --wrap=lib_malloc_buf
+char *lib_malloc_buf(size_t size);
diff --git a/rewriter/tests/mmap_loop/main.c b/rewriter/tests/mmap_loop/main.c
new file mode 100644
index 0000000000..b26d0d615a
--- /dev/null
+++ b/rewriter/tests/mmap_loop/main.c
@@ -0,0 +1,53 @@
+/*
+RUN: sh -c 'if [ ! -s "mmap_loop_call_gates_0.ld" ]; then echo "No link args as expected"; exit 0; fi; echo "Unexpected link args"; exit 1;'
+RUN: %binary_dir/tests/mmap_loop/mmap_loop_main_wrapped | diff %S/Output/mmap_loop.out -
+*/
+#include "mmap_loop.h"
+#include <assert.h>
+#include <ia2.h>
+#include <math.h>
+#include <stdio.h>
+
+#define IA2_DEFINE_TEST_HANDLER
+#include "test_fault_handler.h"
+
+/*
+    This program tests that mmap and heap allocations are handled properly.
+*/
+
+INIT_RUNTIME(2);
+#define IA2_COMPARTMENT 1
+#include <ia2_compartment_init.inc>
+
+int main() {
+  // wait to be traced
+  sleep(2);
+
+  char *buf = malloc(4096);
+  buf[0] = 'r';
+  assert(buf[0] == 'r');
+
+  char *lib_buf = lib_malloc_buf(4096);
+  /*printf("%p was r?\n", &lib_buf[0]);
+  bool x = CHECK_VIOLATION(lib_buf[0] = 'r');
+  printf("%s\n", x ? "true" : "false");*/
+
+  {
+    sleep(1);
+    lib_buf = lib_mmap_buf(4096);
+    printf("mmapped %p was r?\n", &lib_buf[0]);
+    // bool x = CHECK_VIOLATION(lib_buf[0] = 'r');
+
+    int res = munmap(buf, 4096);
+    if (res < 0) {
+      perror("munmap");
+    }
+    printf("res=%d\n", res);
+  }
+
+  printf("lib buf byte: %d", CHECK_VIOLATION(lib_buf[0]));
+
+  // CHECK_VIOLATION(library_access_int_ptr(x)));
+
+  return 0;
+}
diff --git a/rewriter/tests/mmap_loop/mmap_loop.c b/rewriter/tests/mmap_loop/mmap_loop.c
new file mode 100644
index 0000000000..45b7f4a7ce
--- /dev/null
+++ b/rewriter/tests/mmap_loop/mmap_loop.c
@@ -0,0 +1,25 @@
+/*
+RUN: cat mmap_loop_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
+*/
+#include "mmap_loop.h"
+#include "ia2.h"
+#include <stdio.h>
+#include <string.h>
+#include <sys/mman.h>
+
+#define IA2_COMPARTMENT 2
+#include <ia2_compartment_init.inc>
+
+char *lib_mmap_buf(size_t size) {
+  // mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
+  char *buf = mmap(NULL, size, PROT_NONE, MAP_ANON | MAP_PRIVATE, -1, 0);
+  pkey_mprotect(buf, size, PROT_READ | PROT_WRITE, 2);
+  memset(buf, 0, size);
+  return buf;
+}
+
+char *lib_malloc_buf(size_t size) {
+  char *buf = malloc(size);
+  memset(buf, 0, size);
+  return buf;
+}