Skip to main content

lib_q_aead/security/
stack_buffer.rs

1//! Stack Buffer Utilities
2//!
3//! This module provides fixed-size stack-allocated buffers for cryptographic operations,
4//! following libQ's zero dynamic allocation memory model.
5
6use core::mem::MaybeUninit;
7use core::ops::{
8    Deref,
9    DerefMut,
10};
11
12use crate::security::memory::secure_zero_slice;
13
14/// Maximum size for stack-allocated buffers
15pub const MAX_STACK_BUFFER_SIZE: usize = 32768; // 32KB max stack usage
16
17/// Memory-efficient uninitialized stack buffer
18///
19/// This buffer uses MaybeUninit to avoid zeroing memory that will be immediately overwritten,
20/// providing better performance for large buffers.
21#[derive(Debug)]
22pub struct UninitStackBuffer<const N: usize> {
23    data: [MaybeUninit<u8>; N],
24    used: usize,
25}
26
27impl<const N: usize> Default for UninitStackBuffer<N> {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl<const N: usize> UninitStackBuffer<N> {
34    /// Create a new uninitialized stack buffer
35    pub fn new() -> Self {
36        Self {
37            data: unsafe { MaybeUninit::uninit().assume_init() },
38            used: 0,
39        }
40    }
41
42    /// Get the maximum capacity of the buffer
43    pub fn capacity(&self) -> usize {
44        N
45    }
46
47    /// Get the number of bytes currently used
48    pub fn len(&self) -> usize {
49        self.used
50    }
51
52    /// Check if the buffer is empty
53    pub fn is_empty(&self) -> bool {
54        self.used == 0
55    }
56
57    /// Get the used portion as a slice
58    pub fn as_slice(&self) -> &[u8] {
59        unsafe { core::slice::from_raw_parts(self.data.as_ptr() as *const u8, self.used) }
60    }
61
62    /// Get the used portion as a mutable slice
63    pub fn as_mut_slice(&mut self) -> &mut [u8] {
64        unsafe { core::slice::from_raw_parts_mut(self.data.as_mut_ptr() as *mut u8, self.used) }
65    }
66
67    /// Resize the buffer to the specified length
68    pub fn resize(&mut self, new_len: usize) -> Result<(), &'static str> {
69        if new_len > N {
70            return Err("New length exceeds buffer capacity");
71        }
72        self.used = new_len;
73        Ok(())
74    }
75
76    /// Clear the buffer
77    pub fn clear(&mut self) {
78        if self.used > 0 {
79            secure_zero_slice(self.as_mut_slice());
80            self.used = 0;
81        }
82    }
83
84    /// Append data to the buffer
85    pub fn append(&mut self, data: &[u8]) -> Result<(), &'static str> {
86        if self.used + data.len() > N {
87            return Err("Not enough space in buffer");
88        }
89
90        unsafe {
91            core::ptr::copy_nonoverlapping(
92                data.as_ptr(),
93                self.data.as_mut_ptr().add(self.used) as *mut u8,
94                data.len(),
95            );
96        }
97        self.used += data.len();
98        Ok(())
99    }
100}
101
102impl<const N: usize> Drop for UninitStackBuffer<N> {
103    fn drop(&mut self) {
104        self.clear();
105    }
106}
107
108/// Fixed-size stack-allocated buffer for cryptographic operations
109///
110/// This buffer provides a safe, stack-allocated alternative to dynamic allocations
111/// for cryptographic operations. It automatically zeroes sensitive data on drop.
112#[derive(Debug)]
113pub struct StackBuffer<const N: usize> {
114    data: [u8; N],
115    used: usize,
116}
117
118impl<const N: usize> StackBuffer<N> {
119    /// Create a new stack buffer
120    pub fn new() -> Self {
121        Self {
122            data: [0u8; N],
123            used: 0,
124        }
125    }
126
127    /// Create a new stack buffer with initial data
128    pub fn from_slice(data: &[u8]) -> Result<Self, &'static str> {
129        if data.len() > N {
130            return Err("Data too large for buffer");
131        }
132
133        let mut buffer = Self::new();
134        buffer.data[..data.len()].copy_from_slice(data);
135        buffer.used = data.len();
136        Ok(buffer)
137    }
138
139    /// Get the maximum capacity of the buffer
140    pub fn capacity(&self) -> usize {
141        N
142    }
143
144    /// Get the number of bytes currently used
145    pub fn len(&self) -> usize {
146        self.used
147    }
148
149    /// Check if the buffer is empty
150    pub fn is_empty(&self) -> bool {
151        self.used == 0
152    }
153
154    /// Get the remaining capacity
155    pub fn remaining_capacity(&self) -> usize {
156        N - self.used
157    }
158
159    /// Clear the buffer and reset usage
160    pub fn clear(&mut self) {
161        secure_zero_slice(&mut self.data[..self.used]);
162        self.used = 0;
163    }
164
165    /// Resize the buffer to a new length
166    pub fn resize(&mut self, new_len: usize) -> Result<(), &'static str> {
167        if new_len > N {
168            return Err("New length exceeds buffer capacity");
169        }
170
171        if new_len < self.used {
172            // Zero the unused portion
173            secure_zero_slice(&mut self.data[new_len..self.used]);
174        }
175
176        self.used = new_len;
177        Ok(())
178    }
179
180    /// Append data to the buffer
181    pub fn append(&mut self, data: &[u8]) -> Result<(), &'static str> {
182        if self.used + data.len() > N {
183            return Err("Not enough space in buffer");
184        }
185
186        self.data[self.used..self.used + data.len()].copy_from_slice(data);
187        self.used += data.len();
188        Ok(())
189    }
190
191    /// Get a slice of the used portion
192    pub fn as_slice(&self) -> &[u8] {
193        &self.data[..self.used]
194    }
195
196    /// Get a mutable slice of the used portion
197    pub fn as_mut_slice(&mut self) -> &mut [u8] {
198        &mut self.data[..self.used]
199    }
200
201    /// Get a slice of the entire buffer (including unused portion)
202    pub fn as_full_slice(&self) -> &[u8] {
203        &self.data
204    }
205
206    /// Get a mutable slice of the entire buffer (including unused portion)
207    pub fn as_full_mut_slice(&mut self) -> &mut [u8] {
208        &mut self.data
209    }
210
211    /// Copy data from another buffer
212    pub fn copy_from(&mut self, other: &StackBuffer<N>) {
213        self.data.copy_from_slice(&other.data);
214        self.used = other.used;
215    }
216
217    /// Copy data from a slice
218    pub fn copy_from_slice(&mut self, data: &[u8]) -> Result<(), &'static str> {
219        if data.len() > N {
220            return Err("Data too large for buffer");
221        }
222
223        self.data[..data.len()].copy_from_slice(data);
224        self.used = data.len();
225        Ok(())
226    }
227
228    /// Copy data to a slice
229    pub fn copy_to_slice(&self, dest: &mut [u8]) -> Result<(), &'static str> {
230        if dest.len() < self.used {
231            return Err("Destination slice too small");
232        }
233
234        dest[..self.used].copy_from_slice(&self.data[..self.used]);
235        Ok(())
236    }
237}
238
239impl<const N: usize> Default for StackBuffer<N> {
240    fn default() -> Self {
241        Self::new()
242    }
243}
244
245impl<const N: usize> Clone for StackBuffer<N> {
246    fn clone(&self) -> Self {
247        let mut new_buffer = Self::new();
248        new_buffer.copy_from(self);
249        new_buffer
250    }
251}
252
253impl<const N: usize> Deref for StackBuffer<N> {
254    type Target = [u8];
255
256    fn deref(&self) -> &Self::Target {
257        self.as_slice()
258    }
259}
260
261impl<const N: usize> DerefMut for StackBuffer<N> {
262    fn deref_mut(&mut self) -> &mut Self::Target {
263        self.as_mut_slice()
264    }
265}
266
267impl<const N: usize> Drop for StackBuffer<N> {
268    fn drop(&mut self) {
269        // Securely zero the entire buffer on drop
270        secure_zero_slice(&mut self.data);
271    }
272}
273
274/// Predefined buffer sizes for common cryptographic operations
275pub const KEY_BUFFER_SIZE: usize = 64; // 64 bytes for keys
276pub const NONCE_BUFFER_SIZE: usize = 32; // 32 bytes for nonces
277pub const TAG_BUFFER_SIZE: usize = 64; // 64 bytes for authentication tags
278pub const HASH_BUFFER_SIZE: usize = 64; // 64 bytes for hash outputs
279pub const IV_BUFFER_SIZE: usize = 32; // 32 bytes for initialization vectors
280pub const CIPHERTEXT_BUFFER_SIZE: usize = 4096; // 4KB for ciphertext operations
281pub const PLAINTEXT_BUFFER_SIZE: usize = 4096; // 4KB for plaintext operations
282
283/// Type aliases for common buffer sizes
284pub type KeyBuffer = StackBuffer<KEY_BUFFER_SIZE>;
285pub type NonceBuffer = StackBuffer<NONCE_BUFFER_SIZE>;
286pub type TagBuffer = StackBuffer<TAG_BUFFER_SIZE>;
287pub type HashBuffer = StackBuffer<HASH_BUFFER_SIZE>;
288pub type IvBuffer = StackBuffer<IV_BUFFER_SIZE>;
289pub type CiphertextBuffer = StackBuffer<CIPHERTEXT_BUFFER_SIZE>;
290pub type PlaintextBuffer = StackBuffer<PLAINTEXT_BUFFER_SIZE>;
291
292/// Type aliases for memory-efficient uninitialized buffers
293pub type UninitKeyBuffer = UninitStackBuffer<KEY_BUFFER_SIZE>;
294pub type UninitNonceBuffer = UninitStackBuffer<NONCE_BUFFER_SIZE>;
295pub type UninitTagBuffer = UninitStackBuffer<TAG_BUFFER_SIZE>;
296pub type UninitHashBuffer = UninitStackBuffer<HASH_BUFFER_SIZE>;
297pub type UninitIvBuffer = UninitStackBuffer<IV_BUFFER_SIZE>;
298pub type UninitCiphertextBuffer = UninitStackBuffer<CIPHERTEXT_BUFFER_SIZE>;
299pub type UninitPlaintextBuffer = UninitStackBuffer<PLAINTEXT_BUFFER_SIZE>;
300
301/// Utility functions for working with stack buffers
302pub mod utils {
303    use super::*;
304
305    /// Create a key buffer from a slice
306    pub fn create_key_buffer(data: &[u8]) -> Result<KeyBuffer, &'static str> {
307        KeyBuffer::from_slice(data)
308    }
309
310    /// Create a nonce buffer from a slice
311    pub fn create_nonce_buffer(data: &[u8]) -> Result<NonceBuffer, &'static str> {
312        NonceBuffer::from_slice(data)
313    }
314
315    /// Create a tag buffer from a slice
316    pub fn create_tag_buffer(data: &[u8]) -> Result<TagBuffer, &'static str> {
317        TagBuffer::from_slice(data)
318    }
319
320    /// Create a hash buffer from a slice
321    pub fn create_hash_buffer(data: &[u8]) -> Result<HashBuffer, &'static str> {
322        HashBuffer::from_slice(data)
323    }
324
325    /// Create an IV buffer from a slice
326    pub fn create_iv_buffer(data: &[u8]) -> Result<IvBuffer, &'static str> {
327        IvBuffer::from_slice(data)
328    }
329
330    /// Create a ciphertext buffer from a slice
331    pub fn create_ciphertext_buffer(data: &[u8]) -> Result<CiphertextBuffer, &'static str> {
332        CiphertextBuffer::from_slice(data)
333    }
334
335    /// Create a plaintext buffer from a slice
336    pub fn create_plaintext_buffer(data: &[u8]) -> Result<PlaintextBuffer, &'static str> {
337        PlaintextBuffer::from_slice(data)
338    }
339
340    /// Copy data between buffers of different sizes
341    pub fn copy_between_buffers<const SRC_SIZE: usize, const DST_SIZE: usize>(
342        src: &StackBuffer<SRC_SIZE>,
343        dst: &mut StackBuffer<DST_SIZE>,
344    ) -> Result<(), &'static str> {
345        if src.len() > DST_SIZE {
346            return Err("Source buffer too large for destination");
347        }
348
349        dst.copy_from_slice(src.as_slice())
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356
357    #[test]
358    fn test_stack_buffer_creation() {
359        let buffer = StackBuffer::<32>::new();
360        assert_eq!(buffer.capacity(), 32);
361        assert_eq!(buffer.len(), 0);
362        assert!(buffer.is_empty());
363    }
364
365    #[test]
366    fn test_stack_buffer_from_slice() {
367        let data = b"hello world";
368        let buffer = StackBuffer::<32>::from_slice(data).unwrap();
369        assert_eq!(buffer.len(), data.len());
370        assert_eq!(buffer.as_slice(), data);
371    }
372
373    #[test]
374    fn test_stack_buffer_append() {
375        let mut buffer = StackBuffer::<32>::new();
376        buffer.append(b"hello").unwrap();
377        buffer.append(b" world").unwrap();
378        assert_eq!(buffer.as_slice(), b"hello world");
379    }
380
381    #[test]
382    fn test_stack_buffer_resize() {
383        let mut buffer = StackBuffer::<32>::new();
384        buffer.append(b"hello").unwrap();
385        buffer.resize(3).unwrap();
386        assert_eq!(buffer.as_slice(), b"hel");
387    }
388
389    #[test]
390    fn test_stack_buffer_clear() {
391        let mut buffer = StackBuffer::<32>::new();
392        buffer.append(b"hello").unwrap();
393        buffer.clear();
394        assert!(buffer.is_empty());
395    }
396
397    #[test]
398    fn test_stack_buffer_copy() {
399        let mut buffer1 = StackBuffer::<32>::new();
400        buffer1.append(b"hello").unwrap();
401
402        let mut buffer2 = StackBuffer::<32>::new();
403        buffer2.copy_from(&buffer1);
404        assert_eq!(buffer2.as_slice(), buffer1.as_slice());
405    }
406
407    #[test]
408    fn test_predefined_buffers() {
409        let key_buf = KeyBuffer::new();
410        assert_eq!(key_buf.capacity(), KEY_BUFFER_SIZE);
411
412        let nonce_buf = NonceBuffer::new();
413        assert_eq!(nonce_buf.capacity(), NONCE_BUFFER_SIZE);
414
415        let tag_buf = TagBuffer::new();
416        assert_eq!(tag_buf.capacity(), TAG_BUFFER_SIZE);
417    }
418
419    #[test]
420    fn test_utils() {
421        let data = b"test data";
422        let key_buf = utils::create_key_buffer(data).unwrap();
423        assert_eq!(key_buf.as_slice(), data);
424
425        let nonce_buf = utils::create_nonce_buffer(data).unwrap();
426        assert_eq!(nonce_buf.as_slice(), data);
427    }
428
429    #[test]
430    fn test_copy_between_buffers() {
431        let mut src = StackBuffer::<16>::new();
432        src.append(b"hello").unwrap();
433
434        let mut dst = StackBuffer::<32>::new();
435        utils::copy_between_buffers(&src, &mut dst).unwrap();
436        assert_eq!(dst.as_slice(), src.as_slice());
437    }
438}