Skip to main content

dryoc/
types.rs

1use std::fmt;
2use std::ops::{Deref, DerefMut};
3
4use subtle::ConstantTimeEq;
5use zeroize::{Zeroize, ZeroizeOnDrop};
6
7use crate::rng::copy_randombytes;
8
9/// A stack-allocated fixed-length byte array for working with data, with
10/// optional [Serde](https://serde.rs) features.
11#[derive(Zeroize, ZeroizeOnDrop, Clone)]
12pub struct StackByteArray<const LENGTH: usize>([u8; LENGTH]);
13
14impl<const LENGTH: usize> fmt::Debug for StackByteArray<LENGTH> {
15    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
16        f.debug_struct("StackByteArray")
17            .field("len", &LENGTH)
18            .field("contents", &"[REDACTED]")
19            .finish()
20    }
21}
22
23impl<const LENGTH: usize> PartialEq for StackByteArray<LENGTH> {
24    fn eq(&self, other: &Self) -> bool {
25        self.0.ct_eq(&other.0).into()
26    }
27}
28
29impl<const LENGTH: usize> Eq for StackByteArray<LENGTH> {}
30
31/// Fixed-length byte array.
32pub trait ByteArray<const LENGTH: usize>: Bytes {
33    /// Returns a reference to the underlying fixed-length byte array.
34    fn as_array(&self) -> &[u8; LENGTH];
35}
36
37/// Arbitrary-length array of bytes.
38pub trait Bytes {
39    /// Returns a slice of the underlying bytes.
40    fn as_slice(&self) -> &[u8];
41    /// Shorthand to retrieve the underlying length of the byte array.
42    fn len(&self) -> usize;
43    /// Returns true if the array is empty.
44    fn is_empty(&self) -> bool;
45}
46
47/// Fixed-length mutable byte array.
48pub trait MutByteArray<const LENGTH: usize>: ByteArray<LENGTH> + MutBytes {
49    /// Returns a mutable reference to the underlying fixed-length byte array.
50    fn as_mut_array(&mut self) -> &mut [u8; LENGTH];
51}
52
53/// Fixed-length byte array that can be created and initialized.
54pub trait NewByteArray<const LENGTH: usize>: MutByteArray<LENGTH> + NewBytes {
55    /// Returns a new fixed-length byte array, initialized with zeroes.
56    fn new_byte_array() -> Self;
57    /// Returns a new fixed-length byte array, filled with random values.
58    #[allow(deprecated)]
59    fn generate() -> Self
60    where
61        Self: Sized,
62    {
63        Self::r#gen()
64    }
65    /// Returns a new fixed-length byte array, filled with random values.
66    ///
67    /// Prefer [`generate`](Self::generate). `gen` is retained for compatibility
68    /// with older Rust editions.
69    #[deprecated(note = "use generate() instead")]
70    fn r#gen() -> Self;
71}
72
73/// Arbitrary-length array of mutable bytes.
74pub trait MutBytes: Bytes {
75    /// Returns a mutable slice to the underlying bytes.
76    fn as_mut_slice(&mut self) -> &mut [u8];
77    /// Copies into the underlying slice from `other`. Panics if lengths do not
78    /// match.
79    fn copy_from_slice(&mut self, other: &[u8]);
80}
81
82/// Arbitrary-length byte array that can be created and initialized.
83pub trait NewBytes: MutBytes {
84    /// Returns an empty, unallocated, arbitrary-length byte array.
85    fn new_bytes() -> Self;
86}
87
88/// A byte array which can be resized.
89pub trait ResizableBytes {
90    /// Resizes `self` with `new_len` elements, populating new values with
91    /// `value`.
92    fn resize(&mut self, new_len: usize, value: u8);
93}
94
95impl<const LENGTH: usize> ByteArray<LENGTH> for StackByteArray<LENGTH> {
96    #[inline]
97    fn as_array(&self) -> &[u8; LENGTH] {
98        &self.0
99    }
100}
101
102impl<const LENGTH: usize> Bytes for StackByteArray<LENGTH> {
103    #[inline]
104    fn as_slice(&self) -> &[u8] {
105        &self.0
106    }
107
108    #[inline]
109    fn len(&self) -> usize {
110        self.0.len()
111    }
112
113    #[inline]
114    fn is_empty(&self) -> bool {
115        self.0.is_empty()
116    }
117}
118
119impl<const LENGTH: usize> NewBytes for StackByteArray<LENGTH> {
120    fn new_bytes() -> Self {
121        Self::default()
122    }
123}
124
125impl<const LENGTH: usize> NewByteArray<LENGTH> for StackByteArray<LENGTH> {
126    fn new_byte_array() -> Self {
127        Self::default()
128    }
129
130    /// Returns a new byte array filled with random data.
131    fn r#gen() -> Self {
132        let mut res = Self::default();
133        copy_randombytes(&mut res.0);
134        res
135    }
136}
137
138impl<const LENGTH: usize> MutByteArray<LENGTH> for StackByteArray<LENGTH> {
139    #[inline]
140    fn as_mut_array(&mut self) -> &mut [u8; LENGTH] {
141        &mut self.0
142    }
143}
144
145impl<const LENGTH: usize> MutBytes for StackByteArray<LENGTH> {
146    #[inline]
147    fn as_mut_slice(&mut self) -> &mut [u8] {
148        &mut self.0
149    }
150
151    fn copy_from_slice(&mut self, other: &[u8]) {
152        self.0.copy_from_slice(other)
153    }
154}
155
156impl<const LENGTH: usize> NewByteArray<LENGTH> for Vec<u8> {
157    fn new_byte_array() -> Self {
158        vec![0u8; LENGTH]
159    }
160
161    /// Returns a new byte array filled with random data.
162    fn r#gen() -> Self {
163        let mut res = vec![0u8; LENGTH];
164        copy_randombytes(&mut res);
165        res
166    }
167}
168
169impl<const LENGTH: usize> MutByteArray<LENGTH> for Vec<u8> {
170    #[inline]
171    fn as_mut_array(&mut self) -> &mut [u8; LENGTH] {
172        assert!(
173            self.len() >= LENGTH,
174            "invalid vec length {}, expecting at least {}",
175            self.len(),
176            LENGTH
177        );
178        let arr = self.as_mut_ptr() as *mut [u8; LENGTH];
179        // SAFETY: The assertion above guarantees the vector has at least
180        // `LENGTH` initialized bytes. `[u8; LENGTH]` has alignment 1, and the
181        // exclusive `&mut self` borrow prevents aliasing the returned prefix.
182        unsafe { &mut *arr }
183    }
184}
185
186impl<const LENGTH: usize> ByteArray<LENGTH> for Vec<u8> {
187    #[inline]
188    fn as_array(&self) -> &[u8; LENGTH] {
189        assert!(
190            self.len() >= LENGTH,
191            "invalid vec length {}, expecting at least {}",
192            self.len(),
193            LENGTH
194        );
195        let arr = self.as_ptr() as *const [u8; LENGTH];
196        // SAFETY: The assertion above guarantees the vector has at least
197        // `LENGTH` initialized bytes. `[u8; LENGTH]` has alignment 1, so the
198        // first `LENGTH` bytes can be viewed as a fixed-size byte array.
199        unsafe { &*arr }
200    }
201}
202
203impl<const LENGTH: usize> NewBytes for [u8; LENGTH] {
204    fn new_bytes() -> Self {
205        [0u8; LENGTH]
206    }
207}
208
209impl<const LENGTH: usize> NewByteArray<LENGTH> for [u8; LENGTH] {
210    fn new_byte_array() -> Self {
211        [0u8; LENGTH]
212    }
213
214    /// Returns a new byte array filled with random data.
215    fn r#gen() -> Self {
216        let mut res = Self::new_byte_array();
217        copy_randombytes(&mut res);
218        res
219    }
220}
221
222impl<const LENGTH: usize> MutByteArray<LENGTH> for [u8; LENGTH] {
223    #[inline]
224    fn as_mut_array(&mut self) -> &mut [u8; LENGTH] {
225        self
226    }
227}
228
229impl<const LENGTH: usize> MutBytes for [u8; LENGTH] {
230    #[inline]
231    fn as_mut_slice(&mut self) -> &mut [u8] {
232        self
233    }
234
235    fn copy_from_slice(&mut self, other: &[u8]) {
236        <[u8]>::copy_from_slice(self, other)
237    }
238}
239
240impl Bytes for Vec<u8> {
241    #[inline]
242    fn as_slice(&self) -> &[u8] {
243        self.as_slice()
244    }
245
246    #[inline]
247    fn len(&self) -> usize {
248        <[u8]>::len(self)
249    }
250
251    #[inline]
252    fn is_empty(&self) -> bool {
253        <[u8]>::is_empty(self)
254    }
255}
256
257impl NewBytes for Vec<u8> {
258    fn new_bytes() -> Self {
259        vec![]
260    }
261}
262
263impl MutBytes for Vec<u8> {
264    #[inline]
265    fn as_mut_slice(&mut self) -> &mut [u8] {
266        self.as_mut_slice()
267    }
268
269    fn copy_from_slice(&mut self, other: &[u8]) {
270        <[u8]>::copy_from_slice(self, other)
271    }
272}
273
274impl ResizableBytes for Vec<u8> {
275    fn resize(&mut self, new_len: usize, value: u8) {
276        self.resize(new_len, value);
277    }
278}
279
280impl Bytes for [u8] {
281    #[inline]
282    fn as_slice(&self) -> &[u8] {
283        self
284    }
285
286    #[inline]
287    fn len(&self) -> usize {
288        <[u8]>::len(self)
289    }
290
291    #[inline]
292    fn is_empty(&self) -> bool {
293        <[u8]>::is_empty(self)
294    }
295}
296
297impl Bytes for &[u8] {
298    #[inline]
299    fn as_slice(&self) -> &[u8] {
300        self
301    }
302
303    #[inline]
304    fn len(&self) -> usize {
305        <[u8]>::len(self)
306    }
307
308    #[inline]
309    fn is_empty(&self) -> bool {
310        <[u8]>::is_empty(self)
311    }
312}
313
314impl Bytes for &mut [u8] {
315    #[inline]
316    fn as_slice(&self) -> &[u8] {
317        self
318    }
319
320    #[inline]
321    fn len(&self) -> usize {
322        <[u8]>::len(self)
323    }
324
325    #[inline]
326    fn is_empty(&self) -> bool {
327        <[u8]>::is_empty(self)
328    }
329}
330
331impl<const LENGTH: usize> Bytes for [u8; LENGTH] {
332    #[inline]
333    fn as_slice(&self) -> &[u8] {
334        self
335    }
336
337    #[inline]
338    fn len(&self) -> usize {
339        <[u8]>::len(self)
340    }
341
342    #[inline]
343    fn is_empty(&self) -> bool {
344        <[u8]>::is_empty(self)
345    }
346}
347
348#[allow(suspicious_double_ref_op)]
349impl<const LENGTH: usize> Bytes for &[u8; LENGTH] {
350    #[inline]
351    fn as_slice(&self) -> &[u8] {
352        self.deref()
353    }
354
355    #[inline]
356    fn len(&self) -> usize {
357        <[u8]>::len(self.deref())
358    }
359
360    #[inline]
361    fn is_empty(&self) -> bool {
362        <[u8]>::is_empty(self.deref())
363    }
364}
365
366impl<const LENGTH: usize> ByteArray<LENGTH> for [u8; LENGTH] {
367    #[inline]
368    fn as_array(&self) -> &[u8; LENGTH] {
369        self
370    }
371}
372
373/// Provided for convenience. Panics if the input array size doesn't match
374/// `LENGTH`.
375impl<const LENGTH: usize> ByteArray<LENGTH> for &[u8] {
376    #[inline]
377    fn as_array(&self) -> &[u8; LENGTH] {
378        assert!(
379            self.len() >= LENGTH,
380            "invalid slice length {}, expecting at least {}",
381            self.len(),
382            LENGTH
383        );
384        let arr = self.as_ptr() as *const [u8; LENGTH];
385        // SAFETY: The assertion above guarantees the slice has at least
386        // `LENGTH` initialized bytes. `[u8; LENGTH]` has alignment 1, so the
387        // first `LENGTH` bytes can be viewed as a fixed-size byte array.
388        unsafe { &*arr }
389    }
390}
391
392impl<const LENGTH: usize> ByteArray<LENGTH> for [u8] {
393    #[inline]
394    fn as_array(&self) -> &[u8; LENGTH] {
395        assert!(
396            self.len() >= LENGTH,
397            "invalid slice length {}, expecting at least {}",
398            self.len(),
399            LENGTH
400        );
401        let arr = self.as_ptr() as *const [u8; LENGTH];
402        // SAFETY: The assertion above guarantees the slice has at least
403        // `LENGTH` initialized bytes. `[u8; LENGTH]` has alignment 1, so the
404        // first `LENGTH` bytes can be viewed as a fixed-size byte array.
405        unsafe { &*arr }
406    }
407}
408
409impl<const LENGTH: usize> MutByteArray<LENGTH> for [u8] {
410    fn as_mut_array(&mut self) -> &mut [u8; LENGTH] {
411        assert!(
412            self.len() >= LENGTH,
413            "invalid slice length {}, expecting at least {}",
414            self.len(),
415            LENGTH
416        );
417        let arr = self.as_mut_ptr() as *mut [u8; LENGTH];
418        // SAFETY: The assertion above guarantees the slice has at least
419        // `LENGTH` initialized bytes. `[u8; LENGTH]` has alignment 1, and
420        // `&mut self` provides exclusive access to the returned prefix.
421        unsafe { &mut *arr }
422    }
423}
424
425impl MutBytes for [u8] {
426    #[inline]
427    fn as_mut_slice(&mut self) -> &mut [u8] {
428        self
429    }
430
431    fn copy_from_slice(&mut self, other: &[u8]) {
432        self.copy_from_slice(other)
433    }
434}
435
436impl<const LENGTH: usize> StackByteArray<LENGTH> {
437    /// Returns a new fixed-length stack-allocated array
438    pub fn new() -> Self {
439        Self::default()
440    }
441}
442
443impl<const LENGTH: usize> std::convert::AsRef<[u8; LENGTH]> for StackByteArray<LENGTH> {
444    fn as_ref(&self) -> &[u8; LENGTH] {
445        let arr = self.0.as_ptr() as *const [u8; LENGTH];
446        // SAFETY: `StackByteArray<LENGTH>` stores exactly `[u8; LENGTH]` in
447        // `self.0`, so this cast preserves size, alignment, and initialization.
448        unsafe { &*arr }
449    }
450}
451
452impl<const LENGTH: usize> std::convert::AsMut<[u8; LENGTH]> for StackByteArray<LENGTH> {
453    fn as_mut(&mut self) -> &mut [u8; LENGTH] {
454        let arr = self.0.as_mut_ptr() as *mut [u8; LENGTH];
455        // SAFETY: `StackByteArray<LENGTH>` stores exactly `[u8; LENGTH]` in
456        // `self.0`, and `&mut self` provides exclusive access to it.
457        unsafe { &mut *arr }
458    }
459}
460
461impl<const LENGTH: usize> std::convert::AsRef<[u8]> for StackByteArray<LENGTH> {
462    fn as_ref(&self) -> &[u8] {
463        self.0.as_ref()
464    }
465}
466
467impl<const LENGTH: usize> std::convert::AsMut<[u8]> for StackByteArray<LENGTH> {
468    fn as_mut(&mut self) -> &mut [u8] {
469        self.0.as_mut()
470    }
471}
472
473impl<const LENGTH: usize> Deref for StackByteArray<LENGTH> {
474    type Target = [u8];
475
476    fn deref(&self) -> &Self::Target {
477        &self.0
478    }
479}
480
481impl<const LENGTH: usize> DerefMut for StackByteArray<LENGTH> {
482    fn deref_mut(&mut self) -> &mut Self::Target {
483        &mut self.0
484    }
485}
486
487impl<const LENGTH: usize> std::ops::Index<usize> for StackByteArray<LENGTH> {
488    type Output = u8;
489
490    #[inline]
491    fn index(&self, index: usize) -> &Self::Output {
492        &self.0[index]
493    }
494}
495impl<const LENGTH: usize> std::ops::IndexMut<usize> for StackByteArray<LENGTH> {
496    #[inline]
497    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
498        &mut self.0[index]
499    }
500}
501
502macro_rules! impl_index {
503    ($range:ty) => {
504        impl<const LENGTH: usize> std::ops::Index<$range> for StackByteArray<LENGTH> {
505            type Output = [u8];
506
507            #[inline]
508            fn index(&self, index: $range) -> &Self::Output {
509                &self.0[index]
510            }
511        }
512        impl<const LENGTH: usize> std::ops::IndexMut<$range> for StackByteArray<LENGTH> {
513            #[inline]
514            fn index_mut(&mut self, index: $range) -> &mut Self::Output {
515                &mut self.0[index]
516            }
517        }
518    };
519}
520
521impl_index!(std::ops::Range<usize>);
522impl_index!(std::ops::RangeFull);
523impl_index!(std::ops::RangeFrom<usize>);
524impl_index!(std::ops::RangeInclusive<usize>);
525impl_index!(std::ops::RangeTo<usize>);
526impl_index!(std::ops::RangeToInclusive<usize>);
527
528impl<const LENGTH: usize> Default for StackByteArray<LENGTH> {
529    fn default() -> Self {
530        Self([0u8; LENGTH])
531    }
532}
533
534impl<const LENGTH: usize> From<&[u8; LENGTH]> for StackByteArray<LENGTH> {
535    fn from(src: &[u8; LENGTH]) -> Self {
536        let mut arr = Self::default();
537        arr.0.copy_from_slice(src);
538        arr
539    }
540}
541
542impl<const LENGTH: usize> From<[u8; LENGTH]> for StackByteArray<LENGTH> {
543    fn from(src: [u8; LENGTH]) -> Self {
544        Self::from(&src)
545    }
546}
547
548impl<const LENGTH: usize> TryFrom<&[u8]> for StackByteArray<LENGTH> {
549    type Error = crate::error::Error;
550
551    fn try_from(src: &[u8]) -> Result<Self, Self::Error> {
552        if src.len() != LENGTH {
553            Err(length_error!(crate::ErrorContext::Slice, src.len(), exact LENGTH))
554        } else {
555            let mut arr = Self::default();
556            arr.0.copy_from_slice(src);
557            Ok(arr)
558        }
559    }
560}
561
562#[cfg(test)]
563mod tests {
564    use super::*;
565
566    #[test]
567    #[should_panic(expected = "invalid vec length 2, expecting at least 3")]
568    fn test_vec_as_array_out_of_bounds_panic() {
569        let vec = vec![1, 2];
570        let _ = <Vec<u8> as ByteArray<3>>::as_array(&vec)[2];
571    }
572
573    #[test]
574    fn test_vec_as_array_out_of_bounds_ok() {
575        let vec = vec![1, 2];
576        let _ = <Vec<u8> as ByteArray<2>>::as_array(&vec)[1];
577    }
578
579    #[test]
580    #[should_panic(expected = "invalid vec length 2, expecting at least 3")]
581    fn test_vec_as_mut_array_out_of_bounds_panic() {
582        let mut vec = vec![1, 2];
583        let _ = <Vec<u8> as MutByteArray<3>>::as_mut_array(&mut vec)[2];
584    }
585
586    #[test]
587    fn test_vec_as_mut_array_out_of_bounds_ok() {
588        let mut vec = vec![1, 2];
589        let _ = <Vec<u8> as MutByteArray<2>>::as_mut_array(&mut vec)[1];
590    }
591
592    #[test]
593    fn stack_byte_array_debug_redacts_contents() {
594        let bytes = StackByteArray::from([0xabu8; 4]);
595        let debug = format!("{bytes:?}");
596
597        assert!(debug.contains("[REDACTED]"));
598        assert!(!debug.contains("171"));
599    }
600}