From ee8b27b185f12f4fd7a7fe4173ef85b8e0c366bc Mon Sep 17 00:00:00 2001
From: Mikael Voss <mvs@nyantec.com>
Date: Wed, 16 Mar 2022 19:49:23 +0100
Subject: [PATCH] WIP: Make memory locking strategy explicit

---
 src/alloc.rs | 36 +++++++++++++++++++++++++++---------
 src/lib.rs   |  1 +
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/src/alloc.rs b/src/alloc.rs
index 8628524..eae28a3 100644
--- a/src/alloc.rs
+++ b/src/alloc.rs
@@ -7,12 +7,28 @@ use std::alloc::{Allocator, AllocError, Layout, handle_alloc_error};
 use std::intrinsics::{likely, unlikely};
 use std::ptr::NonNull;
 
+/// Memory locking mode
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum LockingMode {
+	/// Do not attempt to lock memory
+	Neglect,
+
+	/// Attempt to lock memory, ignoring failure
+	Attempt,
+
+	/// Fail if memory cannot be locked
+	Enforce
+}
+
 /// Allocator for sensitive information
-pub struct Sensitive;
+pub struct Sensitive<const M: LockingMode = { LockingMode::Attempt }>;
 
-impl Sensitive {
+impl<const M: LockingMode> Sensitive<M> {
 	/// Number of guard pages
 	pub(crate) const GUARD_PAGES: usize = 1;
+
+	/// Memory locking mode
+	pub(crate) const LOCKING_MODE: LockingMode = M;
 }
 
 unsafe impl Allocator for Sensitive {
@@ -24,10 +40,10 @@ unsafe impl Allocator for Sensitive {
 
 		let alloc = GuardedAlloc::<{ Self::GUARD_PAGES }>::new(layout.size(), Protection::ReadWrite).map_err(|_| AllocError)?;
 
-		if likely(!alloc.inner().is_empty()) {
-			// Attempt to lock memory
-			#[allow(unused_must_use)] {
-				alloc.inner().lock();
+		if Self::LOCKING_MODE != LockingMode::Neglect && likely(!alloc.inner().is_empty()) {
+			// Lock memory
+			if alloc.inner().lock().is_err() && Self::LOCKING_MODE == LockingMode::Enforce {
+				return Err(AllocError)
 			}
 		}
 
@@ -53,9 +69,11 @@ unsafe impl Allocator for Sensitive {
 			// Zero memory before returning to OS
 			zero(ptr.as_ptr(), layout.size());
 
-			// Attempt to unlock memory
-			#[allow(unused_must_use)] {
-				alloc.inner().unlock();
+			if Self::LOCKING_MODE != LockingMode::Neglect {
+				// Unlock memory
+				if alloc.inner().unlock().is_err() && Self::LOCKING_MODE == LockingMode::Enforce {
+					handle_alloc_error(layout);
+				}
 			}
 		}
 	}
diff --git a/src/lib.rs b/src/lib.rs
index 4f2e163..701ff73 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,5 @@
 #![feature(
+	adt_const_params,
 	allocator_api,
 	core_intrinsics,
 	maybe_uninit_slice,