use crate::assert::assert_zeroize_on_drop;
use crate::collections::{
collection_zeroed, to_fast_zeroizable_dyn_mut, to_zeroization_probe_dyn_ref, zeroize_collection,
};
use crate::traits::{AssertZeroizeOnDrop, FastZeroizable, ZeroizationProbe};
use crate::zeroize_on_drop_sentinel::ZeroizeOnDropSentinel;
use crate::zeroizing_mut_guard::ZeroizingMutGuard;
struct Foo {
pub data: Vec<u8>,
__sentinel: ZeroizeOnDropSentinel,
}
impl Default for Foo {
fn default() -> Self {
Self {
data: vec![1, 2, 3, 4],
__sentinel: ZeroizeOnDropSentinel::default(),
}
}
}
impl FastZeroizable for Foo {
fn fast_zeroize(&mut self) {
self.data.fast_zeroize();
self.__sentinel.fast_zeroize();
}
}
impl Drop for Foo {
fn drop(&mut self) {
self.fast_zeroize();
}
}
impl ZeroizationProbe for Foo {
fn is_zeroized(&self) -> bool {
let fields: [&dyn ZeroizationProbe; 1] = [to_zeroization_probe_dyn_ref(&self.data)];
collection_zeroed(&mut fields.into_iter())
}
}
impl AssertZeroizeOnDrop for Foo {
fn clone_sentinel(&self) -> ZeroizeOnDropSentinel {
self.__sentinel.clone()
}
fn assert_zeroize_on_drop(self) {
assert_zeroize_on_drop(self);
}
}
struct FunctionalStruct<'a> {
pub bytes: Vec<u8>,
pub bytes_16: [u8; 16],
pub bytes_32: [u8; 32],
pub foo: ZeroizingMutGuard<'a, Foo>,
__sentinel: ZeroizeOnDropSentinel,
}
impl<'a> FunctionalStruct<'a> {
fn new(foo: &'a mut Foo) -> Self {
let mut bytes = Vec::new();
bytes.resize_with(128, || u8::MAX);
Self {
bytes,
bytes_16: [u8::MAX; 16],
bytes_32: [u8::MAX; 32],
foo: ZeroizingMutGuard::from(foo),
__sentinel: ZeroizeOnDropSentinel::default(),
}
}
}
impl<'a> FastZeroizable for FunctionalStruct<'a> {
fn fast_zeroize(&mut self) {
let fields: [&mut dyn FastZeroizable; 5] = [
to_fast_zeroizable_dyn_mut(&mut self.bytes_16),
to_fast_zeroizable_dyn_mut(&mut self.bytes_32),
to_fast_zeroizable_dyn_mut(&mut self.bytes),
to_fast_zeroizable_dyn_mut(&mut self.foo),
to_fast_zeroizable_dyn_mut(&mut self.__sentinel),
];
zeroize_collection(&mut fields.into_iter());
}
}
impl<'a> ZeroizationProbe for FunctionalStruct<'a> {
fn is_zeroized(&self) -> bool {
let fields: [&dyn ZeroizationProbe; 4] = [
to_zeroization_probe_dyn_ref(&self.bytes_16),
to_zeroization_probe_dyn_ref(&self.bytes_32),
to_zeroization_probe_dyn_ref(&self.bytes),
to_zeroization_probe_dyn_ref(&self.foo),
];
collection_zeroed(&mut fields.into_iter())
}
}
impl<'a> AssertZeroizeOnDrop for FunctionalStruct<'a> {
fn clone_sentinel(&self) -> ZeroizeOnDropSentinel {
self.__sentinel.clone()
}
fn assert_zeroize_on_drop(self) {
assert_zeroize_on_drop(self);
}
}
impl<'a> Drop for FunctionalStruct<'a> {
fn drop(&mut self) {
self.fast_zeroize();
}
}
#[test]
fn test_functionl_struct() {
let mut foo = Foo::default();
assert!(!foo.is_zeroized());
let mut fs = FunctionalStruct::new(&mut foo);
assert!(!fs.is_zeroized());
fs.fast_zeroize();
assert!(fs.is_zeroized());
fs.assert_zeroize_on_drop();
}