Abstract page allocation and locking in os.c

Simplifies ac.c just a little bit.

3 files changed, 64 insertions, 27 deletions

diff --git a/src/ac.c b/src/ac.c
index fc1d318..9d1430e 100644
--- a/src/ac.c
+++ b/src/ac.c
@@ -389,31 +389,18 @@ INIT {
 	struct hook_featsetup_ret_Key_Event h = hook_featsetup_Key_Event(Key_Event);
 	if_cold (h.err) return h.err;
 
-#ifdef _WIN32
-	keybox = VirtualAlloc(0, 4096, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+	keybox = os_pagealloc(4096, PAGE_READWRITE);
 	if_cold (!keybox) {
 		errmsg_errorsys("couldn't allocate memory for session state");
 		return FEAT_FAIL;
 	}
-	if_cold (!VirtualLock(keybox, 4096)) {
+	if_cold (!os_mlock(keybox, 4096)) {
 		errmsg_errorsys("couldn't secure session state");
 		goto e;
 	}
+#ifdef _WIN32
 	if_cold (!win32_init()) goto e;
 #else
-	keybox = mmap(0, 4096, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANON, -1, 0);
-	if_cold (keybox == MAP_FAILED) {
-		errmsg_errorstd("couldn't allocate memory for session state");
-		return FEAT_FAIL;
-	}
-	// linux-specific madvise stuff (there are some equivalents in OpenBSD and
-	// FreeBSD, if anyone's wondering, but we don't need to worry about those)
-	if_cold (madvise(keybox, 4096, MADV_DONTFORK) == -1 ||
-			madvise(keybox, 4096, MADV_DONTDUMP) == - 1 ||
-			mlock(keybox, 4096) == -1) {
-		errmsg_errorstd("couldn't secure session state");
-		goto e;
-	}
 	// TODO(linux): call other init things
 #endif
 
@@ -428,22 +415,17 @@ INIT {
 	hook_commit_Key_Event(h.hookpos, hookdest_Key_Event);
 	return FEAT_OK;
 
-#ifdef _WIN32
-e:	VirtualFree(keybox, 4096, MEM_RELEASE);
-#else
-e:	munmap(keybox, 4096);
-#endif
+e:	os_pagefree(keybox, 4096);
 	return FEAT_FAIL;
 }
 
 END {
 	// TODO(opt): *maybe* do the skip-on-quit stuff here. feels a bit scary...
 	ac_disable();
+	os_pagefree(keybox, 4096);
 #if defined(_WIN32)
-	VirtualFree(keybox, 4096, MEM_RELEASE);
 	win32_end();
 #else
-	munmap(keybox, 4096);
 	// TODO(linux): call other cleanup things
 #endif
 	unhook_Key_Event();
diff --git a/src/os.c b/src/os.c
index 2bc0e45..dce0d17 100644
--- a/src/os.c
+++ b/src/os.c
@@ -117,10 +117,20 @@ int os_dlfile(void *lib, ushort *buf, int sz) {
 	return n;
 }
 
+void *os_pagealloc(int sz, int prot) {
+	return VirtualAlloc(0, sz, MEM_RESERVE | MEM_COMMIT, prot);
+}
+void os_pagefree(void *addr, int sz) {
+	VirtualFree(addr, sz, MEM_RELEASE);
+}
+
 bool os_mprot(void *addr, int len, int mode) {
 	ulong old;
 	return !!VirtualProtect(addr, len, mode, &old);
 }
+bool os_mlock(void *addr, int sz) {
+	return VirtualLock(addr, 4096);
+}
 
 #else
 
@@ -155,13 +165,39 @@ void *os_dlsym(void *restrict lib, const char *restrict name) {
 	return dlsym(lib, name);
 }
 
+void *os_pagealloc(int sz, int prot) {
+	sz = sz + 4095 & ~4095; // round up size
+	void *ret = mmap(0, sz, prot, MAP_SHARED | MAP_ANON, -1, 0);
+	if (ret == MAP_FAILED) return 0;
+	return ret;
+}
+void os_pagefree(void *addr, int sz) {
+	// round down address and round up size
+	addr = (void *)((ulong)addr & ~4095);
+	len = len + 4095 & ~4095;
+	munmap(addr, len);
+}
+
 bool os_mprot(void *addr, int len, int mode) {
 	// round down address and round up size
-	addr = (void *)((ulong)addr & ~(4095));
-	len = len + 4095 & ~(4095);
+	addr = (void *)((ulong)addr & ~4095);
+	len = len + 4095 & ~4095;
 	return mprotect(addr, len, mode) != -1;
 }
 
+bool os_mlock(void *addr, int sz) {
+#ifdef __linux__
+	// Linux-specific madvise stuff (there are some equivalents in OpenBSD and
+	// FreeBSD, if anyone's wondering, but we don't use need them for our
+	// purposes).
+	if (madvise(addr, len, MADV_DONTFORK) == -1 ||
+			madvise(addr, len, MADV_DONTDUMP) == - 1) {
+		return false;
+	}
+#endif
+	return mlock(addr, len) != -1;
+}
+
 void os_randombytes(void *buf, int sz) { while (getentropy(buf, sz) == -1); }
 
 #endif
@@ -200,8 +236,8 @@ void *os_dlhandle(const char *name) {
 		int namelen = strlen(lm->l_name);
 		int sublen = strlen(name);
 		if (sublen >= namelen) continue;
-		if (lm->l_name[namelen - sublen - 1] == '/' && !memcmp(
-				lm->l_name + namelen - sublen, name, sublen)) {
+		if (lm->l_name[namelen - sublen - 1] == '/' &&
+				!memcmp(lm->l_name + namelen - sublen, name, sublen)) {
 			return lm;
 		}
 	}
diff --git a/src/os.h b/src/os.h
index 6466e45..9f8295f 100644
--- a/src/os.h
+++ b/src/os.h
@@ -228,6 +228,18 @@ int os_dlfile(void *lib, os_char *buf, int sz);
 #endif
 
 /*
+ * Allocates some memory from the operating system. Rounds the size up to the
+ * page size (4096). Returns an address on success, or null on failure.
+ */
+void *os_pagealloc(int sz, int prot);
+
+/*
+ * Frees a page range allocated using os_pagealloc(). Rounds the address and
+ * size to page boundaries (4096).
+ */
+void os_pagefree(void *addr, int sz);
+
+/*
  * Changes memory protection for the address range given by addr and len, using
  * one of the Win32-style PAGE_* flags specified above. Returns true on success
  * and false on failure.
@@ -235,6 +247,13 @@ int os_dlfile(void *lib, os_char *buf, int sz);
 bool os_mprot(void *addr, int len, int mode);
 
 /*
+ * Locks a page range in physical memory; prevents paging to disk. The address
+ * and size *must* be page-aligned, otherwise behaviour is undefined.
+ * Returns true on success, false on failure.
+ */
+bool os_mlock(void *addr, int sz);
+
+/*
  * Fills buf with up to sz cryptographically random bytes. sz has an OS-specific
  * upper limit - a safe value across all major operating systems is 256.
  */