protect jit pages

3 files changed, 36 insertions, 9 deletions

diff --git a/src/common.h b/src/common.h
index b845e1c..6ee0df7 100644
--- a/src/common.h
+++ b/src/common.h
@@ -308,6 +308,6 @@ union interp_ret ejit_run(struct ejit_func *f, size_t argc,
 		     bool run,
 		     void ***labels_wb);
 
-bool ejit_compile(struct ejit_func *f, bool use_64);
+bool ejit_compile(struct ejit_func *f, bool use_64, bool im_scawed);
 
 #endif /* EJIT_COMMON_H */
diff --git a/src/compile/compile.c b/src/compile/compile.c
index e741a0f..490bc43 100644
--- a/src/compile/compile.c
+++ b/src/compile/compile.c
@@ -19,10 +19,10 @@ struct reloc_helper {
 #define VEC_NAME addrs
 #include "../vec.h"
 
-static void *alloc_arena(size_t size)
+static void *alloc_arena(size_t size, bool im_scawed)
 {
 	return mmap(NULL, size,
-	            PROT_EXEC | PROT_READ | PROT_WRITE,
+	            (!im_scawed ? PROT_EXEC : 0) | PROT_READ | PROT_WRITE,
 	            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
 }
 
@@ -2414,7 +2414,15 @@ static void assign_fprs(struct ejit_func *f)
 	}
 }
 
-bool ejit_compile(struct ejit_func *f, bool use_64)
+static size_t align_up(size_t a, size_t n)
+{
+	if (a % n == 0)
+		return a;
+
+	return a + (a % n);
+}
+
+bool ejit_compile(struct ejit_func *f, bool use_64, bool im_scawed)
 {
 	(void)use_64;
 #if __WORDSIZE == 32
@@ -2436,10 +2444,11 @@ bool ejit_compile(struct ejit_func *f, bool use_64)
 	assert(j);
 
 	void *arena = NULL;
-	size_t size = 4096;
+	size_t pagesize = sysconf(_SC_PAGE_SIZE);
+	size_t size = pagesize;
 
 	while (1) {
-		arena = alloc_arena(size);
+		arena = alloc_arena(size, im_scawed);
 		if (arena == (void *)(-1)) {
 			jit_destroy_state(j);
 			return false;
@@ -2450,10 +2459,16 @@ bool ejit_compile(struct ejit_func *f, bool use_64)
 			break;
 
 		free_arena(arena, size);
-		size = required_size + 4096;
+		size = align_up(required_size + pagesize, pagesize);
 	}
 
 	jit_destroy_state(j);
+
+	if (im_scawed && mprotect(arena, size, PROT_EXEC | PROT_READ)) {
+		free_arena(arena, size);
+		return false;
+	}
+
 	f->arena = arena;
 	f->size = size;
 	return true;
diff --git a/src/ejit.c b/src/ejit.c
index 660d476..265acea 100644
--- a/src/ejit.c
+++ b/src/ejit.c
@@ -363,17 +363,29 @@ struct ejit_func *ejit_create_func(enum ejit_type rtype, size_t argc,
 	return f;
 }
 
+void ejit_compile_func_unsafe(struct ejit_func *f)
+{
+	ejit_select_compile_func(f,
+			gpr_stats_len(&f->gpr),
+			fpr_stats_len(&f->fpr),
+			f->use_64,
+			true,
+			false);
+}
+
 void ejit_compile_func(struct ejit_func *f)
 {
+	/* might make sense to use flags instead of bools... */
 	ejit_select_compile_func(f,
 			gpr_stats_len(&f->gpr),
 			fpr_stats_len(&f->fpr),
 			f->use_64,
+			true,
 			true);
 }
 
 void ejit_select_compile_func(struct ejit_func *f, size_t gpr, size_t fpr,
-                              bool use_64, bool try_jit)
+                              bool use_64, bool try_jit, bool im_scawed)
 {
 	/* emit a final end instruction in case user didn't do a return */
 	emit_insn_o(f, END);
@@ -384,7 +396,7 @@ void ejit_select_compile_func(struct ejit_func *f, size_t gpr, size_t fpr,
 	assert(!f->use_64 || (use_64 == f->use_64));
 
 	/* try to jit compile if possible */
-	if (try_jit && ejit_compile(f, use_64))
+	if (try_jit && ejit_compile(f, use_64, im_scawed))
 		return;
 
 	/* otherwise, convert opcodes to address labels */