From 3df0c00d02e84966dabca0253a978cbfedd6d2fd 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 |  3 ++
 rewriter/tests/mmap_loop/include/mmap_loop.h  |  8 ++++
 rewriter/tests/mmap_loop/main.c               | 44 +++++++++++++++++++
 rewriter/tests/mmap_loop/mmap_loop.c          | 28 ++++++++++++
 6 files changed, 99 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 9da50c72d5..51c38bd002 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..a8185c24f4
--- /dev/null
+++ b/rewriter/tests/mmap_loop/Output/mmap_loop.out
@@ -0,0 +1,3 @@
+program started
+mmap buf
+res=0
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..08d6956f80
--- /dev/null
+++ b/rewriter/tests/mmap_loop/main.c
@@ -0,0 +1,44 @@
+/*
+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() {
+  puts("program started");
+
+  char *buf = malloc(4096);
+  buf[0] = 'r';
+  assert(buf[0] == 'r');
+
+  char *lib_buf_malloc = lib_malloc_buf(4096);
+
+  puts("mmap buf");
+  char *lib_buf_mmap = lib_mmap_buf(4096);
+
+  free(buf);
+
+  /* syscall tracing should forbid us to unmap compartment 2's buffer */
+  int res = munmap(lib_buf_mmap, 4096);
+  if (res < 0) {
+    perror("munmap");
+  }
+  printf("res=%d\n", res);
+
+  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..d6222d7cc3
--- /dev/null
+++ b/rewriter/tests/mmap_loop/mmap_loop.c
@@ -0,0 +1,28 @@
+/*
+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>
+
+// LINKARGS: --wrap=lib_malloc_buf
+char *lib_malloc_buf(size_t size) {
+  char *buf = malloc(size);
+  memset(buf, 0, size);
+  return buf;
+}
+
+// LINKARGS: --wrap=lib_mmap_buf
+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;
+}