secret_buffer.rs raw

use std::collections::TryReserveError;

use zeroize::Zeroizing;

/// The smallest capacity allocated when a `SecretBuffer` first grows.
const MIN_CAPACITY: usize = 256;

/// A growable byte buffer that never leaves copies of its contents in freed
/// memory.
///
/// A plain `Vec` abandons its old allocation, unwiped, whenever it grows. This
/// buffer grows by allocating a larger buffer, copying the contents across,
/// and zeroizing the old allocation before releasing it. The final allocation
/// is zeroized when the buffer is dropped.
#[derive(Default)]
pub(crate) struct SecretBuffer {
    bytes: Zeroizing<Vec<u8>>,
}

impl SecretBuffer {
    /// Creates an empty buffer without allocating.
    pub(crate) fn new() -> Self {
        Self::default()
    }

    /// Returns the buffer's contents.
    pub(crate) fn as_slice(&self) -> &[u8] {
        &self.bytes
    }

    /// Ensures there is room for `additional` more bytes without the
    /// underlying `Vec` reallocating on its own.
    fn reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
        let required = self.bytes.len().saturating_add(additional);

        if required <= self.bytes.capacity() {
            return Ok(());
        }

        let capacity = required
            .max(self.bytes.capacity().saturating_mul(2))
            .max(MIN_CAPACITY);

        let mut grown = Vec::new();
        grown.try_reserve_exact(capacity)?;
        grown.extend_from_slice(&self.bytes);

        // Replacing the `Zeroizing` wrapper drops the old one, which wipes the
        // old allocation before freeing it.
        self.bytes = Zeroizing::new(grown);

        Ok(())
    }
}

impl cbor4ii::core::enc::Write for SecretBuffer {
    type Error = TryReserveError;

    fn push(&mut self, input: &[u8]) -> Result<(), Self::Error> {
        self.reserve(input.len())?;
        self.bytes.extend_from_slice(input);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use cbor4ii::core::enc::Write as _;

    use super::*;

    #[test]
    fn growth_preserves_contents() {
        let mut buffer = SecretBuffer::new();
        let mut expected = Vec::new();

        for i in 0..1000u32 {
            let chunk = i.to_be_bytes();
            buffer.push(&chunk).unwrap();
            expected.extend_from_slice(&chunk);
        }

        assert_eq!(buffer.as_slice(), expected.as_slice());
    }

    #[test]
    fn extending_never_outgrows_the_reservation() {
        let mut buffer = SecretBuffer::new();

        for len in [1, 255, 256, 257, 1000, 5000] {
            buffer.reserve(len).unwrap();
            let reserved = buffer.bytes.as_ptr();
            buffer.push(&vec![0xa5; len]).unwrap();

            // The push must land in the allocation `reserve` prepared, not in
            // one the `Vec` obtained by reallocating behind our back.
            assert_eq!(buffer.bytes.as_ptr(), reserved);
        }
    }
}