tor_bytes/
secretbuf.rs

1//! Define a wrapper for `Vec<u8>` that will act as Writer, but zeroize its
2//! contents on drop or reallocation.
3
4use crate::Writer;
5use zeroize::{Zeroize, ZeroizeOnDrop};
6
7/// A [`Writer`] used for accumulating secret data, which gets cleared on drop.
8///
9/// Unlike `Zeroizing<Vec<u8>>`, this type makes sure that we always zeroize the
10/// contents of the buffer, even if the buffer has to be reallocated in order to
11/// grow.
12///
13/// We use this for cases when we're building the input to a key derivation
14/// function (KDF), and want to ensure that we don't expose the values we feed
15/// to it.
16///
17/// This struct is expected to have additional overhead beyond `Vec<u8>` only
18/// when it has to grow its capacity.
19#[derive(Zeroize, ZeroizeOnDrop, Debug, Clone, Eq, PartialEq)]
20pub struct SecretBuf(Vec<u8>);
21
22/// The default size of our buffer.
23///
24/// This is based on the size of a typical secret input in `tor-proto`.
25const DEFAULT_CAPACITY: usize = 384;
26
27impl SecretBuf {
28    /// Construct a new empty [`SecretBuf`]
29    pub fn new() -> Self {
30        Self::with_capacity(DEFAULT_CAPACITY)
31    }
32
33    /// Construct a new empty [`SecretBuf`] with a specified capacity.
34    ///
35    /// This buffer will not have to be reallocated until it uses `capacity`
36    /// bytes.
37    pub fn with_capacity(capacity: usize) -> Self {
38        Self(Vec::with_capacity(capacity))
39    }
40
41    /// Truncate this buffer to a given length.
42    pub fn truncate(&mut self, new_len: usize) {
43        self.0.truncate(new_len);
44    }
45
46    /// Add all the bytes from `slice` to the end of this vector.
47    pub fn extend_from_slice(&mut self, slice: &[u8]) {
48        let new_len = self.0.len() + slice.len();
49        if new_len >= self.0.capacity() {
50            // We will need to reallocate.  But in doing so we might reallocate,
51            // which neglects to zero the previous contents.  So instead,
52            // explicitly make a new vector and zeroize the old one.
53
54            // Make sure we always at least double our capacity.
55            let new_capacity = std::cmp::max(self.0.capacity() * 2, new_len);
56            let mut new_vec = Vec::with_capacity(new_capacity);
57            new_vec.extend_from_slice(&self.0[..]);
58
59            let mut old_vec = std::mem::replace(&mut self.0, new_vec);
60            old_vec.zeroize();
61        }
62        self.0.extend_from_slice(slice);
63        debug_assert_eq!(self.0.len(), new_len);
64    }
65}
66
67impl From<Vec<u8>> for SecretBuf {
68    fn from(v: Vec<u8>) -> Self {
69        Self(v)
70    }
71}
72
73impl Default for SecretBuf {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl AsMut<[u8]> for SecretBuf {
80    fn as_mut(&mut self) -> &mut [u8] {
81        &mut self.0[..]
82    }
83}
84
85// It's okay to implement `Deref` since all operations taking an _immutable_
86// reference are still right here.
87impl std::ops::Deref for SecretBuf {
88    type Target = Vec<u8>;
89
90    fn deref(&self) -> &Self::Target {
91        &self.0
92    }
93}
94
95impl Writer for SecretBuf {
96    fn write_all(&mut self, b: &[u8]) {
97        self.extend_from_slice(b);
98    }
99}
100
101#[cfg(test)]
102mod test {
103    // @@ begin test lint list maintained by maint/add_warning @@
104    #![allow(clippy::bool_assert_comparison)]
105    #![allow(clippy::clone_on_copy)]
106    #![allow(clippy::dbg_macro)]
107    #![allow(clippy::mixed_attributes_style)]
108    #![allow(clippy::print_stderr)]
109    #![allow(clippy::print_stdout)]
110    #![allow(clippy::single_char_pattern)]
111    #![allow(clippy::unwrap_used)]
112    #![allow(clippy::unchecked_duration_subtraction)]
113    #![allow(clippy::useless_vec)]
114    #![allow(clippy::needless_pass_by_value)]
115    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
116    use super::*;
117
118    #[test]
119    fn simple_case() -> crate::EncodeResult<()> {
120        // Sadly, there is no way in safe rust to test that the zeroization
121        // actually happened.  All we can test is that the data is correct.
122
123        let mut buf1 = SecretBuf::default();
124        let mut buf2 = Vec::new();
125        let xyz = b"Nine hundred pounds of sifted flax";
126
127        // This is enough to be sure that we'll reallocate.
128        for _ in 0..200 {
129            buf1.write(xyz)?;
130            buf2.write(xyz)?;
131        }
132        assert_eq!(&buf1[..], &buf2[..]);
133
134        Ok(())
135    }
136}