use std::ops::Deref;
pub struct Secret {
buf: Vec<u8>,
}
impl Secret {
pub fn from_vec(buf: Vec<u8>) -> Self {
Self { buf }
}
pub fn with_capacity(n: usize) -> Self {
Self {
buf: Vec::with_capacity(n),
}
}
pub fn as_mut_vec(&mut self) -> &mut Vec<u8> {
&mut self.buf
}
pub fn len(&self) -> usize {
self.buf.len()
}
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
pub fn trim_trailing_newlines(&mut self) {
while matches!(self.buf.last(), Some(b'\n') | Some(b'\r')) {
self.buf.pop();
}
}
}
impl Deref for Secret {
type Target = [u8];
fn deref(&self) -> &[u8] {
&self.buf
}
}
unsafe fn volatile_zero(ptr: *mut u8, len: usize) {
for i in 0..len {
std::ptr::write_volatile(ptr.add(i), 0u8);
}
std::sync::atomic::compiler_fence(std::sync::atomic::Ordering::SeqCst);
}
impl Drop for Secret {
fn drop(&mut self) {
let cap = self.buf.capacity();
unsafe {
self.buf.set_len(0);
volatile_zero(self.buf.as_mut_ptr(), cap);
}
}
}
pub struct SecretString(pub String);
impl Deref for SecretString {
type Target = str;
fn deref(&self) -> &str {
&self.0
}
}
impl Drop for SecretString {
fn drop(&mut self) {
let cap = self.0.capacity();
unsafe {
let v = self.0.as_mut_vec();
v.set_len(0);
volatile_zero(v.as_mut_ptr(), cap);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn trims_only_trailing_newlines() {
let mut s = Secret::from_vec(b"a\nb\n\r\n".to_vec());
s.trim_trailing_newlines();
assert_eq!(&*s, b"a\nb");
}
#[test]
fn trimming_empty_is_safe() {
let mut s = Secret::from_vec(b"\n\n\n".to_vec());
s.trim_trailing_newlines();
assert!(s.is_empty());
}
#[test]
fn deref_exposes_bytes() {
let s = Secret::from_vec(vec![1, 2, 3]);
assert_eq!(s.len(), 3);
assert_eq!(&*s, &[1, 2, 3]);
}
#[test]
fn volatile_zero_actually_zeroes() {
let mut buf: Vec<u8> = vec![0xAA; 256];
unsafe { volatile_zero(buf.as_mut_ptr(), buf.len()) };
assert!(buf.iter().all(|&b| b == 0));
}
#[test]
fn volatile_zero_of_zero_length_is_a_no_op() {
let mut buf: Vec<u8> = vec![0xAA; 4];
unsafe { volatile_zero(buf.as_mut_ptr(), 1) };
assert_eq!(buf, vec![0x00, 0xAA, 0xAA, 0xAA]);
}
}