1
//! Define a wrapper for `Vec<u8>` that will act as Writer, but zeroize its
2
//! contents on drop or reallocation.
3

            
4
use crate::Writer;
5
use 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)]
20
pub 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`.
25
const DEFAULT_CAPACITY: usize = 384;
26

            
27
impl SecretBuf {
28
    /// Construct a new empty [`SecretBuf`]
29
3362
    pub fn new() -> Self {
30
3362
        Self::with_capacity(DEFAULT_CAPACITY)
31
3362
    }
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
5618
    pub fn with_capacity(capacity: usize) -> Self {
38
5618
        Self(Vec::with_capacity(capacity))
39
5618
    }
40

            
41
    /// Truncate this buffer to a given length.
42
1440
    pub fn truncate(&mut self, new_len: usize) {
43
1440
        self.0.truncate(new_len);
44
1440
    }
45

            
46
    /// Add all the bytes from `slice` to the end of this vector.
47
32368
    pub fn extend_from_slice(&mut self, slice: &[u8]) {
48
32368
        let new_len = self.0.len() + slice.len();
49
32368
        if new_len >= self.0.capacity() {
50
778
            // We will need to reallocate.  But in doing so we might reallocate,
51
778
            // which neglects to zero the previous contents.  So instead,
52
778
            // explicitly make a new vector and zeroize the old one.
53
778

            
54
778
            // Make sure we always at least double our capacity.
55
778
            let new_capacity = std::cmp::max(self.0.capacity() * 2, new_len);
56
778
            let mut new_vec = Vec::with_capacity(new_capacity);
57
778
            new_vec.extend_from_slice(&self.0[..]);
58
778

            
59
778
            let mut old_vec = std::mem::replace(&mut self.0, new_vec);
60
778
            old_vec.zeroize();
61
31590
        }
62
32368
        self.0.extend_from_slice(slice);
63
32368
        debug_assert_eq!(self.0.len(), new_len);
64
32368
    }
65
}
66

            
67
impl From<Vec<u8>> for SecretBuf {
68
2448
    fn from(v: Vec<u8>) -> Self {
69
2448
        Self(v)
70
2448
    }
71
}
72

            
73
impl Default for SecretBuf {
74
2
    fn default() -> Self {
75
2
        Self::new()
76
2
    }
77
}
78

            
79
impl AsMut<[u8]> for SecretBuf {
80
2256
    fn as_mut(&mut self) -> &mut [u8] {
81
2256
        &mut self.0[..]
82
2256
    }
83
}
84

            
85
// It's okay to implement `Deref` since all operations taking an _immutable_
86
// reference are still right here.
87
impl std::ops::Deref for SecretBuf {
88
    type Target = Vec<u8>;
89

            
90
16658
    fn deref(&self) -> &Self::Target {
91
16658
        &self.0
92
16658
    }
93
}
94

            
95
impl Writer for SecretBuf {
96
26224
    fn write_all(&mut self, b: &[u8]) {
97
26224
        self.extend_from_slice(b);
98
26224
    }
99
}
100

            
101
#[cfg(test)]
102
mod 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
}