use crate::{
sodium::{rng, zeroize},
util::{cpy, mutating},
};
use lazy_static::lazy_static;
use libsodium_sys as libsodium;
use std::{
collections::HashMap,
convert::TryInto,
fmt,
ops::{Deref, DerefMut},
os::raw::c_void,
ptr::null_mut,
sync::Mutex,
};
lazy_static! {
static ref SECRET_CACHE: Mutex<SecretMemoryPool> = Mutex::new(SecretMemoryPool::new());
}
#[derive(Debug)] pub struct SecretMemoryPool {
pool: HashMap<usize, Vec<*mut c_void>>,
}
impl SecretMemoryPool {
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
let pool = HashMap::new();
Self { pool }
}
pub fn release<const N: usize>(&mut self, mut s: Secret<N>) {
unsafe {
self.release_by_ref(&mut s);
}
std::mem::forget(s);
}
unsafe fn release_by_ref<const N: usize>(&mut self, s: &mut Secret<N>) {
s.zeroize();
let Secret { ptr: secret } = s;
self.pool.entry(N).or_default().push(*secret);
}
pub fn take<const N: usize>(&mut self) -> Secret<N> {
let entry = self.pool.entry(N).or_default();
let secret = entry.pop().unwrap_or_else(|| {
let ptr = unsafe { libsodium::sodium_malloc(N) };
assert!(
!ptr.is_null(),
"libsodium::sodium_mallloc() returned a null ptr"
);
ptr
});
let mut s = Secret { ptr: secret };
s.zeroize();
s
}
}
impl Drop for SecretMemoryPool {
fn drop(&mut self) {
for ptr in self.pool.drain().flat_map(|(_, x)| x.into_iter()) {
unsafe {
libsodium::sodium_free(ptr);
}
}
}
}
unsafe impl Send for SecretMemoryPool {}
pub struct Secret<const N: usize> {
ptr: *mut c_void,
}
impl<const N: usize> Clone for Secret<N> {
fn clone(&self) -> Self {
let mut new = Self::zero();
new.secret_mut().clone_from_slice(self.secret());
new
}
}
impl<const N: usize> Drop for Secret<N> {
fn drop(&mut self) {
self.zeroize();
unsafe { SECRET_CACHE.lock().unwrap().release_by_ref(self) };
self.ptr = null_mut();
}
}
impl<const N: usize> Secret<N> {
pub fn from_slice(slice: &[u8]) -> Self {
let mut new_self = Self::zero();
new_self.secret_mut().copy_from_slice(slice);
new_self
}
pub fn zero() -> Self {
let s = SECRET_CACHE.lock().unwrap().take();
assert_eq!(s.secret(), &[0u8; N]);
s
}
pub fn random() -> Self {
mutating(Self::zero(), |r| r.randomize())
}
pub fn zeroize(&mut self) {
zeroize(self.secret_mut());
}
pub fn randomize(&mut self) {
rng(self.secret_mut());
}
pub fn secret(&self) -> &[u8; N] {
let ptr = self.ptr as *const u8;
let slice = unsafe { std::slice::from_raw_parts(ptr, N) };
slice.try_into().unwrap()
}
pub fn secret_mut(&mut self) -> &mut [u8; N] {
let ptr = self.ptr as *mut u8;
let slice = unsafe { std::slice::from_raw_parts_mut(ptr, N) };
slice.try_into().unwrap()
}
}
impl<const N: usize> fmt::Debug for Secret<N> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.write_str("<SECRET>")
}
}
#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Public<const N: usize> {
pub value: [u8; N],
}
impl<const N: usize> Public<N> {
pub fn from_slice(value: &[u8]) -> Self {
mutating(Self::zero(), |r| cpy(value, &mut r.value))
}
pub fn new(value: [u8; N]) -> Self {
Self { value }
}
pub fn zero() -> Self {
Self { value: [0u8; N] }
}
pub fn random() -> Self {
mutating(Self::zero(), |r| r.randomize())
}
pub fn randomize(&mut self) {
rng(&mut self.value);
}
}
pub fn debug_crypto_array(v: &[u8], fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.write_str("[{}]=")?;
if v.len() > 64 {
for byte in &v[..32] {
std::fmt::LowerHex::fmt(byte, fmt)?;
}
fmt.write_str("…")?;
for byte in &v[v.len() - 32..] {
std::fmt::LowerHex::fmt(byte, fmt)?;
}
} else {
for byte in v {
std::fmt::LowerHex::fmt(byte, fmt)?;
}
}
Ok(())
}
impl<const N: usize> fmt::Debug for Public<N> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
debug_crypto_array(&self.value, fmt)
}
}
impl<const N: usize> Deref for Public<N> {
type Target = [u8; N];
fn deref(&self) -> &[u8; N] {
&self.value
}
}
impl<const N: usize> DerefMut for Public<N> {
fn deref_mut(&mut self) -> &mut [u8; N] {
&mut self.value
}
}
#[cfg(test)]
mod test {
use super::*;
const SODIUM_MAGIC_BYTE: u8 = 0xdb;
fn init() {
unsafe { libsodium_sys::sodium_init() };
}
#[test]
fn sodium_malloc() {
init();
const N: usize = 8;
let ptr = unsafe { libsodium_sys::sodium_malloc(N) };
let mem = unsafe { std::slice::from_raw_parts(ptr as *mut u8, N) };
assert_eq!(mem, &[SODIUM_MAGIC_BYTE; N])
}
#[test]
fn sodium_free() {
init();
const N: usize = 8;
let ptr = unsafe { libsodium_sys::sodium_malloc(N) };
unsafe { libsodium_sys::sodium_free(ptr) }
}
#[test]
fn secret_memory_pool_take() {
init();
const N: usize = 0x100;
let mut pool = SecretMemoryPool::new();
let secret: Secret<N> = pool.take();
assert_eq!(secret.secret(), &[0; N]);
}
#[test]
fn secret_memory_pool_drop() {
init();
const N: usize = 0x100;
let mut pool = SecretMemoryPool::new();
let secret: Secret<N> = pool.take();
std::mem::drop(pool);
assert_eq!(secret.secret(), &[0; N]);
}
#[test]
fn secret_memory_pool_release() {
init();
const N: usize = 1;
let mut pool = SecretMemoryPool::new();
let mut secret: Secret<N> = pool.take();
let old_secret_ptr = secret.ptr;
secret.secret_mut()[0] = 0x13;
pool.release(secret);
let new_secret: Secret<N> = pool.take();
assert_eq!(old_secret_ptr, new_secret.ptr);
assert_eq!(new_secret.secret(), &[0; N]);
}
}