From edaaf79d36b5e9de06881591a42b799c5de0f59a Mon Sep 17 00:00:00 2001
From: root <root@dw-6.eecs.uci.edu>
Date: Wed, 3 Jul 2019 13:06:02 -0700
Subject: [PATCH] Fork handler mitigations

---
 classlist.h |  2 ++
 common.h    |  3 ---
 config.h    |  3 +++
 model.cc    | 12 +++++++++++-
 mymemory.cc |  7 -------
 snapshot.cc |  5 +++++
 6 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/classlist.h b/classlist.h
index 641a1481..24c635fd 100644
--- a/classlist.h
+++ b/classlist.h
@@ -19,4 +19,6 @@ class Fuzzer;
 struct model_snapshot_members;
 struct bug_message;
 typedef SnapList<ModelAction *> action_list_t;
+
+extern volatile int forklock;
 #endif
diff --git a/common.h b/common.h
index 0f0ad64f..27f4c8ef 100644
--- a/common.h
+++ b/common.h
@@ -11,10 +11,8 @@
 #include "printf.h"
 
 extern int model_out;
-extern int switch_alloc;
 
 #define model_print(fmt, ...) do { \
-		switch_alloc = 1;              \
 		char mprintbuf[256];                                                \
 		int printbuflen=snprintf_(mprintbuf, 256, fmt, ## __VA_ARGS__);     \
 		int lenleft = printbuflen < 256 ? printbuflen : 256;                   \
@@ -24,7 +22,6 @@ extern int switch_alloc;
 			lenleft-=byteswritten;                                            \
 			totalwritten+=byteswritten;                                       \
 		}                                                                   \
-		switch_alloc = 0;                                                   \
 } while (0)
 
 #ifdef CONFIG_DEBUG
diff --git a/config.h b/config.h
index 43c04427..ed0b6bac 100644
--- a/config.h
+++ b/config.h
@@ -54,4 +54,7 @@
 /** Enable debugging assertions (via ASSERT()) */
 #define CONFIG_ASSERT
 
+/** Enable mitigations against fork handlers that call into locks...  */
+#define FORK_HANDLER_HACK
+
 #endif
diff --git a/model.cc b/model.cc
index 1be15edb..8454d0fb 100644
--- a/model.cc
+++ b/model.cc
@@ -319,7 +319,17 @@ void ModelChecker::switch_from_master(Thread *thread)
  */
 uint64_t ModelChecker::switch_to_master(ModelAction *act)
 {
-	DBG();
+  if (forklock) {
+    static bool fork_message_printed = false;
+
+    if (!fork_message_printed) {
+      model_print("Fork handler trying to call into model checker...\n");
+      fork_message_printed = true;
+    }
+    delete act;
+    return 0;
+  }
+  DBG();
 	Thread *old = thread_current();
 	scheduler->set_current_thread(NULL);
 	ASSERT(!old->get_pending());
diff --git a/mymemory.cc b/mymemory.cc
index a85c48c3..66a4fb97 100644
--- a/mymemory.cc
+++ b/mymemory.cc
@@ -17,7 +17,6 @@
 size_t allocatedReqs[REQUESTS_BEFORE_ALLOC] = { 0 };
 int nextRequest = 0;
 int howManyFreed = 0;
-int switch_alloc = 0;
 #if !USE_MPROTECT_SNAPSHOT
 static mspace sStaticSpace = NULL;
 #endif
@@ -181,9 +180,6 @@ static void * user_malloc(size_t size)
 void *malloc(size_t size)
 {
 	if (user_snapshot_space) {
-		if (switch_alloc) {
-			return model_malloc(size);
-		}
 		/* Only perform user allocations from user context */
 		ASSERT(!model || thread_current());
 		return user_malloc(size);
@@ -195,9 +191,6 @@ void *malloc(size_t size)
 void free(void * ptr)
 {
 	if (!DontFree(ptr)) {
-		if (switch_alloc) {
-			return model_free(ptr);
-		}
 		mspace_free(user_snapshot_space, ptr);
 	}
 }
diff --git a/snapshot.cc b/snapshot.cc
index dc42614e..6725a7ae 100644
--- a/snapshot.cc
+++ b/snapshot.cc
@@ -375,6 +375,8 @@ static void fork_snapshot_init(unsigned int numbackingpages,
 	model_snapshot_space = create_mspace(numheappages * PAGESIZE, 1);
 }
 
+volatile int forklock = 0;
+
 static void fork_loop() {
 	/* switch back here when takesnapshot is called */
 	snapshotid = fork_snap->currSnapShotID;
@@ -386,7 +388,10 @@ static void fork_loop() {
 	while (true) {
 		pid_t forkedID;
 		fork_snap->currSnapShotID = snapshotid + 1;
+
+		forklock = 1;
 		forkedID = fork();
+		forklock = 0;
 
 		if (0 == forkedID) {
 			setcontext(&fork_snap->shared_ctxt);
-- 
2.34.1