Skip to content

Commit

Permalink
Simplify Secret Guard API
Browse files Browse the repository at this point in the history
  • Loading branch information
@Eyob94
Eyob94 committed Sep 8, 2024
1 parent 94c8bcc commit d99cfa8
Showing 1 changed file with 84 additions and 47 deletions.
131 changes: 84 additions & 47 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,14 @@ pub use zeroize;
/// Wrapper for the inner secret. Can be exposed by [`ExposeSecret`]
pub struct SecretBox<S: Zeroize> {
inner_secret: Box<S>,
protected: bool,
}

impl<S: Zeroize> Zeroize for SecretBox<S> {
fn zeroize(&mut self) {
let secret_ptr = self.inner_secret.as_ref() as *const S;
let len = mem::size_of_val(&*self.inner_secret);

let len = mem::size_of::<S>();
let secret_ptr = self.inner_secret.as_ref() as *const S;

unsafe {
if !munlock(secret_ptr as *mut u8, len) {
Expand Down Expand Up @@ -62,34 +63,37 @@ impl<S: Zeroize> From<Box<S>> for SecretBox<S> {
impl<S: Zeroize> SecretBox<S> {
/// Create a secret value using a pre-boxed value.
pub fn new(boxed_secret: Box<S>) -> Self {
let secret_ptr = Box::into_raw(boxed_secret);
let len = mem::size_of_val(&*boxed_secret);

let len = mem::size_of::<S>();
let secret_ptr = Box::into_raw(boxed_secret);

unsafe {
if !mlock(secret_ptr as *mut u8, len) {
eprintln!("Unable to mlock variable ")
panic!("Unable to mlock variable ")
}

if !mprotect(
NonNull::new(secret_ptr).expect("Unable to convert box to NonNull"),
PROT_NONE,
) {
eprintln!("Unable to protect secret")
panic!("Unable to protect secret")
}
}

let inner_secret = unsafe { Box::from_raw(secret_ptr) };

Self { inner_secret }
Self {
inner_secret,
protected: true,
}
}
}

impl<S: Zeroize + Default> SecretBox<S> {
/// Create a secret value using a function that can initialize the vale in-place.
pub fn new_with_mut(ctr: impl FnOnce(&mut S)) -> Self {
let mut secret = Self::default();
ctr(secret.expose_secret_mut());
ctr(secret.expose_secret().inner_secret_mut());
secret
}
}
Expand All @@ -105,9 +109,7 @@ impl<S: Zeroize + Clone> SecretBox<S> {
/// since this method's safety relies on empyric evidence and may be violated on some targets.
pub fn new_with_ctr(ctr: impl FnOnce() -> S) -> Self {
let mut data = ctr();
let secret = Self {
inner_secret: Box::new(data.clone()),
};
let secret = Self::new(Box::new(data.clone()));
data.zeroize();
secret
}
Expand All @@ -119,9 +121,7 @@ impl<S: Zeroize + Clone> SecretBox<S> {
/// since this method's safety relies on empyric evidence and may be violated on some targets.
pub fn try_new_with_ctr<E>(ctr: impl FnOnce() -> Result<S, E>) -> Result<Self, E> {
let mut data = ctr()?;
let secret = Self {
inner_secret: Box::new(data.clone()),
};
let secret = Self::new(Box::new(data.clone()));
data.zeroize();
Ok(secret)
}
Expand All @@ -131,6 +131,7 @@ impl<S: Zeroize + Default> Default for SecretBox<S> {
fn default() -> Self {
Self {
inner_secret: Box::<S>::default(),
protected: false,
}
}
}
Expand All @@ -148,55 +149,90 @@ where
fn clone(&self) -> Self {
SecretBox {
inner_secret: self.inner_secret.clone(),
protected: false,
}
}
}

impl<S: Zeroize> ExposeSecret<S> for SecretBox<S> {
fn expose_secret(&self) -> &S {
let secret_ptr = self.inner_secret.as_ref() as *const S;
fn expose_secret(&mut self) -> SecretGuard<'_, S> {
SecretGuard::new(self.inner_secret.as_mut(), self.protected)
}
}

unsafe {
if !mprotect(
NonNull::new(secret_ptr as *mut S).expect("Unable to convert ptr to NonNull"),
PROT_READ,
) {
eprintln!("Unable to unprotect variable")
/// Secret Guard that holds a mutable to reference to the secret
pub struct SecretGuard<'a, S>
where
S: Zeroize,
{
data: &'a mut S,
protected: bool,
}

impl<'a, S: Zeroize> SecretGuard<'a, S> {
/// Create a new SecretGuard instance.
pub fn new(data: &'a mut S, protected: bool) -> Self {
Self { data, protected }
}
/// Get a shared reference to the inner secret
pub fn inner_secret(&mut self) -> &S {
let secret_ptr = self.data as *const S;

if self.protected {
unsafe {
if !mprotect(
NonNull::new(secret_ptr as *mut S).expect("Unable to convert ptr to NonNull"),
PROT_READ,
) {
panic!("Unable to protect secret_guard")
}
}
}
self.inner_secret.as_ref()
self.protected = false;
self.data
}
}

impl<S: Zeroize> ExposeSecretMut<S> for SecretBox<S> {
fn expose_secret_mut(&mut self) -> &mut S {
let secret_ptr = self.inner_secret.as_ref() as *const S;
/// Get an exclusive reference to the inner secret
pub fn inner_secret_mut(&mut self) -> &mut S {
let secret_ptr = self.data as *const S;

unsafe {
if !mprotect(
NonNull::new(secret_ptr as *mut S).expect("Unable to convert ptr to NonNull"),
PROT_READ | PROT_WRITE,
) {
eprintln!("Unable to unprotect variable")
panic!("Unable to protect secret_guard")
}
}
self.protected = false;
self.data
}
}

impl<'a, S: Zeroize> Drop for SecretGuard<'a, S> {
fn drop(&mut self) {
let secret_ptr = self.data as *const S;

if !self.protected {
unsafe {
if !mprotect(
NonNull::new(secret_ptr as *mut S).expect("Unable to convert ptr to NonNull"),
PROT_NONE,
) {
panic!("Unable to protect memory")
}
}
}
self.inner_secret.as_mut()
}
}

/// Marker trait for secrets which are allowed to be cloned
pub trait CloneableSecret: Clone + Zeroize {}

/// Expose a reference to an inner secret
pub trait ExposeSecret<S> {
/// Expose secret: this is the only method providing access to a secret.
fn expose_secret(&self) -> &S;
}

/// Expose a mutable reference to an inner secret
pub trait ExposeSecretMut<S> {
/// Create a SecretGuard that holds a reference to the secret
pub trait ExposeSecret<S: Zeroize> {
/// Expose secret: this is the only method providing access to a secret.
fn expose_secret_mut(&mut self) -> &mut S;
fn expose_secret(&mut self) -> SecretGuard<'_, S>;
}

#[cfg(test)]
Expand Down Expand Up @@ -232,8 +268,8 @@ mod tests {
#[test]
fn test_secret_box_drop_zeroizes() {
let secret = Box::new(TestSecret::new(10));
let secret_box = SecretBox::new(secret);
assert!(secret_box.expose_secret().check_non_zero());
let mut secret_box = SecretBox::new(secret);
assert!(secret_box.expose_secret().inner_secret().check_non_zero());

drop(secret_box);

Expand All @@ -247,17 +283,18 @@ mod tests {
fn test_secret_box_expose_secret_mut() {
let secret = Box::new(TestSecret::new(10));
let mut secret_box = SecretBox::new(secret);
{
let mut exposed = secret_box.expose_secret();
exposed.inner_secret_mut().data[0] = 42;
}

let exposed = secret_box.expose_secret_mut();
exposed.data[0] = 42;

assert_eq!(secret_box.expose_secret().data[0], 42);
assert_eq!(secret_box.expose_secret().inner_secret().data[0], 42);
}

#[test]
fn test_secret_box_new_with_ctr() {
let secret_box = SecretBox::new_with_ctr(|| TestSecret::new(10));
assert!(secret_box.expose_secret().check_non_zero());
let mut secret_box = SecretBox::new_with_ctr(|| TestSecret::new(10));
assert!(secret_box.expose_secret().inner_secret().check_non_zero());
}

#[test]
Expand All @@ -266,7 +303,7 @@ mod tests {
SecretBox::try_new_with_ctr(|| Ok(TestSecret::new(10)));

match result {
Ok(secret_box) => assert!(secret_box.expose_secret().check_non_zero()),
Ok(mut secret_box) => assert!(secret_box.expose_secret().inner_secret().check_non_zero()),
Err(_) => panic!("Expected Ok variant"),
}
}
Expand Down

0 comments on commit d99cfa8

Please sign in to comment.