Skip to main content

secure_types/
vec.rs

1#[cfg(not(feature = "use_os"))]
2use alloc::{Layout, alloc, dealloc};
3
4#[cfg(feature = "use_os")]
5use std::vec::Vec;
6
7use super::{Error, SecureArray};
8use core::{
9   marker::PhantomData,
10   mem,
11   ops::{Bound, RangeBounds},
12   ptr::{self, NonNull},
13};
14use zeroize::{DefaultIsZeroes, Zeroize};
15
16#[cfg(feature = "use_os")]
17use super::{alloc, free};
18#[cfg(feature = "use_os")]
19use memsec::Prot;
20
21pub type SecureBytes = SecureVec<u8>;
22
23/// Unlocks the vector's memory on construction and re-locks it on drop —
24/// including when the drop happens because the fn closure panicked.
25struct UnlockGuard<'a, T: Zeroize> {
26   vec: &'a SecureVec<T>,
27}
28
29impl<'a, T: Zeroize> UnlockGuard<'a, T> {
30   fn new(vec: &'a SecureVec<T>) -> Self {
31      let ok = vec.unlock_memory();
32      debug_assert!(ok, "UnlockGuard::new: unlock_memory failed");
33
34      UnlockGuard { vec }
35   }
36}
37
38impl<'a, T: Zeroize> Drop for UnlockGuard<'a, T> {
39   fn drop(&mut self) {
40      let ok = self.vec.lock_memory();
41      debug_assert!(ok, "UnlockGuard::drop: lock_memory failed");
42   }
43}
44
45/// A securely allocated, growable vector, just like `std::vec::Vec`.
46///
47/// ## Security Model
48///
49/// When compiled with the `use_os` feature (the default), it provides several layers of protection:
50/// - **Zeroization on Drop**: The memory is zeroized when the vector is dropped.
51/// - **Memory Locking**: The underlying memory pages are locked using `mlock` & `madvise` for (Unix) or
52///   `VirtualLock` & `VirtualProtect` for (Windows) to prevent the OS from memory-dump/swap to disk or other processes accessing the memory.
53///
54/// In a `no_std` environment, it falls back to providing only the **zeroization-on-drop** guarantee.
55///
56/// ## Program Termination
57///
58/// Direct indexing (e.g., `vec[0]`) on a locked vector will cause the operating system
59/// to terminate the process with an access violation error.
60///
61/// Always use the provided scope methods (`unlock_slice`, `unlock_slice_mut`) for safe access.
62///
63/// # Notes
64///
65/// If you return a new allocated `Vec` from one of the unlock methods you are responsible for zeroizing the memory.
66///
67/// # Example
68///
69/// Using `SecureBytes` (a type alias for `SecureVec<u8>`) to handle a secret key.
70///
71/// ```
72/// use secure_types::{SecureBytes, Zeroize};
73///
74/// // Create a new, empty secure vector.
75/// let mut secret_key = SecureBytes::new().unwrap();
76///
77/// // Push some sensitive data into it.
78/// secret_key.push(0xAB);
79/// secret_key.push(0xCD);
80/// secret_key.push(0xEF);
81///
82/// // The memory is locked here.
83///
84/// // Use a scope to safely access the contents as a slice.
85/// secret_key.unlock_slice(|unlocked_slice| {
86///     assert_eq!(unlocked_slice, &[0xAB, 0xCD, 0xEF]);
87/// });
88///
89/// // Not recommended but if you allocate a new Vec make sure to zeroize it
90/// let mut exposed = secret_key.unlock_slice(|unlocked_slice| {
91///     Vec::from(unlocked_slice)
92/// });
93///
94/// // Do what you need to to do with the new vector
95/// // When you are done with it, zeroize it
96/// exposed.zeroize();
97///
98/// // The memory is automatically locked again when the scope ends.
99///
100/// // When `secret_key` is dropped, its memory is securely zeroized.
101/// ```
102pub struct SecureVec<T>
103where
104   T: Zeroize,
105{
106   ptr: NonNull<T>,
107   pub(crate) len: usize,
108   pub(crate) capacity: usize,
109   _marker: PhantomData<T>,
110}
111
112unsafe impl<T: Zeroize + Send> Send for SecureVec<T> {}
113unsafe impl<T: Zeroize + Send + Sync> Sync for SecureVec<T> {}
114
115impl<T: Zeroize> SecureVec<T> {
116   /// Create a new `SecureVec` with a capacity of 1
117   pub fn new() -> Result<Self, Error> {
118      let capacity = 1;
119      let size = capacity * mem::size_of::<T>();
120      let ptr = unsafe { alloc::<T>(size)? };
121
122      let secure = SecureVec {
123         ptr,
124         len: 0,
125         capacity,
126         _marker: PhantomData,
127      };
128
129      let locked = secure.lock_memory();
130
131      #[cfg(feature = "use_os")]
132      if !locked {
133         return Err(Error::LockFailed);
134      }
135
136      Ok(secure)
137   }
138
139   /// Create a new `SecureVec` with the given capacity
140   pub fn new_with_capacity(mut capacity: usize) -> Result<Self, Error> {
141      if capacity == 0 {
142         capacity = 1;
143      }
144
145      capacity
146         .checked_mul(size_of::<T>())
147         .ok_or(Error::AllocationFailed)?;
148
149      let size = capacity * mem::size_of::<T>();
150      let ptr = unsafe { alloc::<T>(size)? };
151
152      let secure = SecureVec {
153         ptr,
154         len: 0,
155         capacity,
156         _marker: PhantomData,
157      };
158
159      let locked = secure.lock_memory();
160
161      #[cfg(feature = "use_os")]
162      if !locked {
163         return Err(Error::LockFailed);
164      }
165
166      Ok(secure)
167   }
168
169   #[cfg(feature = "use_os")]
170   /// Create a new `SecureVec` from a `Vec`
171   ///
172   /// The `Vec` is zeroized afterwards
173   pub fn from_vec(mut vec: Vec<T>) -> Result<Self, Error> {
174      if vec.capacity() == 0 {
175         vec.reserve(1);
176      }
177
178      let capacity = vec.capacity();
179      let len = vec.len();
180
181      let size = match capacity.checked_mul(size_of::<T>()) {
182         Some(s) => s,
183         None => {
184            vec.zeroize();
185            return Err(Error::AllocationFailed);
186         }
187      };
188
189      let ptr = match unsafe { alloc::<T>(size) } {
190         Ok(ptr) => ptr,
191         Err(_) => {
192            vec.zeroize();
193            return Err(Error::AllocationFailed);
194         }
195      };
196
197      // Copy data from the old pointer to the new one
198      unsafe {
199         core::ptr::copy_nonoverlapping(vec.as_ptr(), ptr.as_ptr() as *mut T, len);
200      }
201
202      vec.zeroize();
203
204      let secure = SecureVec {
205         ptr,
206         len,
207         capacity,
208         _marker: PhantomData,
209      };
210
211      let locked = secure.lock_memory();
212
213      if !locked {
214         return Err(Error::LockFailed);
215      }
216
217      Ok(secure)
218   }
219
220   /// Create a new `SecureVec` from a mutable slice.
221   ///
222   /// The slice is zeroized afterwards
223   pub fn from_slice_mut(slice: &mut [T]) -> Result<Self, Error>
224   where
225      T: Clone + DefaultIsZeroes,
226   {
227      let mut secure_vec = match SecureVec::new_with_capacity(slice.len()) {
228         Ok(secure_vec) => secure_vec,
229         Err(e) => {
230            slice.zeroize();
231            return Err(e);
232         }
233      };
234
235      secure_vec.init_from_clone(slice);
236      slice.zeroize();
237
238      Ok(secure_vec)
239   }
240
241   /// Create a new `SecureVec` from a slice.
242   ///
243   /// The slice is not zeroized, you are responsible for zeroizing it
244   pub fn from_slice(slice: &[T]) -> Result<Self, Error>
245   where
246      T: Clone,
247   {
248      let mut secure_vec = SecureVec::new_with_capacity(slice.len())?;
249      secure_vec.init_from_clone(slice);
250      Ok(secure_vec)
251   }
252
253   pub fn len(&self) -> usize {
254      self.len
255   }
256
257   pub fn is_empty(&self) -> bool {
258      self.len() == 0
259   }
260
261   pub(crate) fn as_mut_ptr(&mut self) -> *mut u8 {
262      self.ptr.as_ptr() as *mut u8
263   }
264
265   pub(crate) fn lock_memory(&self) -> bool {
266      #[cfg(feature = "use_os")]
267      {
268         #[cfg(windows)]
269         {
270            super::mprotect(self.ptr, Prot::NoAccess)
271         }
272         #[cfg(unix)]
273         {
274            super::mprotect(self.ptr, Prot::NoAccess)
275         }
276      }
277      #[cfg(not(feature = "use_os"))]
278      {
279         true // No-op: always "succeeds"
280      }
281   }
282
283   pub(crate) fn unlock_memory(&self) -> bool {
284      #[cfg(feature = "use_os")]
285      {
286         #[cfg(windows)]
287         {
288            super::mprotect(self.ptr, Prot::ReadWrite)
289         }
290         #[cfg(unix)]
291         {
292            super::mprotect(self.ptr, Prot::ReadWrite)
293         }
294      }
295
296      #[cfg(not(feature = "use_os"))]
297      {
298         true // No-op: always "succeeds"
299      }
300   }
301
302   /// Immutable access to the `SecureVec`
303   pub fn unlock<F, R>(&self, f: F) -> R
304   where
305      F: FnOnce(&SecureVec<T>) -> R,
306   {
307      let _guard = UnlockGuard::new(self);
308      let result = f(self);
309      result
310   }
311
312   /// Immutable access to the `SecureVec` as `&[T]`
313   pub fn unlock_slice<F, R>(&self, f: F) -> R
314   where
315      F: FnOnce(&[T]) -> R,
316   {
317      let _guard = UnlockGuard::new(self);
318      let slice = unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.len) };
319      f(slice)
320   }
321
322   /// Mutable access to the `SecureVec` as `&mut [T]`
323   pub fn unlock_slice_mut<F, R>(&mut self, f: F) -> R
324   where
325      F: FnOnce(&mut [T]) -> R,
326   {
327      unsafe {
328         let _guard = UnlockGuard::new(self);
329         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
330         let result = f(slice);
331         result
332      }
333   }
334
335   /// Immutable access to the `SecureVec` as `Iter<T>`
336   pub fn unlock_iter<F, R>(&self, f: F) -> R
337   where
338      F: FnOnce(core::slice::Iter<T>) -> R,
339   {
340      unsafe {
341         let _guard = UnlockGuard::new(self);
342         let slice = core::slice::from_raw_parts(self.ptr.as_ptr(), self.len);
343         let iter = slice.iter();
344         let result = f(iter);
345         result
346      }
347   }
348
349   /// Mutable access to the `SecureVec` as `IterMut<T>`
350   pub fn unlock_iter_mut<F, R>(&mut self, f: F) -> R
351   where
352      F: FnOnce(core::slice::IterMut<T>) -> R,
353   {
354      unsafe {
355         let _guard = UnlockGuard::new(self);
356         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
357         let iter = slice.iter_mut();
358         let result = f(iter);
359         result
360      }
361   }
362
363   /// Erase the underlying data and clears the vector
364   ///
365   /// The memory is locked again and the capacity is preserved for reuse
366   pub fn erase(&mut self) {
367      unsafe {
368         let ok = self.unlock_memory();
369         debug_assert!(ok, "SecureVec::erase: unlock_memory failed");
370
371         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.capacity);
372         for elem in slice.iter_mut() {
373            elem.zeroize();
374         }
375
376         self.clear();
377
378         let ok = self.lock_memory();
379         debug_assert!(ok, "SecureVec::erase: lock_memory failed");
380      }
381   }
382
383   /// Clear the vector
384   ///
385   /// This just sets the vector's len to zero it does not erase the underlying data
386   pub fn clear(&mut self) {
387      self.len = 0;
388   }
389
390   pub fn push(&mut self, value: T) {
391      self.reserve(1);
392
393      let ok = self.unlock_memory();
394      debug_assert!(ok, "SecureVec::push: unlock_memory failed");
395
396      unsafe {
397         // Write the new value at the end of the vector.
398         core::ptr::write(self.ptr.as_ptr().add(self.len), value);
399
400         self.len += 1;
401      }
402
403      let ok = self.lock_memory();
404      debug_assert!(ok, "SecureVec::push: lock_memory failed");
405   }
406
407   /// Ensures that the vector has enough capacity for at least `additional` more elements.
408   ///
409   /// If more capacity is needed, it will reallocate. This may cause the buffer location to change.
410   ///
411   /// # Panics
412   ///
413   /// Panics if the new capacity overflows `usize` or if the allocation fails.
414   pub fn reserve(&mut self, additional: usize) {
415      if self.len() + additional <= self.capacity {
416         return;
417      }
418
419      // Use an amortized growth strategy to avoid reallocating on every push
420      let required_capacity = self.len() + additional;
421      let new_capacity = (self.capacity.max(1) * 2).max(required_capacity);
422
423      let new_size = new_capacity * mem::size_of::<T>();
424
425      // Safe to panic here because the memory is locked
426      let new_ptr = unsafe {
427         alloc::<T>(new_size).unwrap_or_else(|_| {
428            panic!(
429               "secure-types: failed to allocate {} bytes of locked memory \
430          (possibly RLIMIT_MEMLOCK exhausted); SecureVec left unchanged",
431               new_size
432            )
433         })
434      };
435
436      // Copy data to new pointer
437      unsafe {
438         let ok = self.unlock_memory();
439         debug_assert!(ok, "SecureVec::reserve: unlock_memory failed");
440
441         core::ptr::copy_nonoverlapping(
442            self.ptr.as_ptr(),
443            new_ptr.as_ptr() as *mut T,
444            self.len(),
445         );
446
447         // Erase and free the old memory
448         if self.capacity > 0 {
449            let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.capacity);
450            for elem in slice.iter_mut() {
451               elem.zeroize();
452            }
453         }
454
455         #[cfg(feature = "use_os")]
456         free(self.ptr);
457
458         #[cfg(not(feature = "use_os"))]
459         {
460            let old_size = self.capacity * mem::size_of::<T>();
461            let old_layout = Layout::from_size_align_unchecked(old_size, mem::align_of::<T>());
462            dealloc(self.ptr.as_ptr() as *mut u8, old_layout);
463         }
464      }
465
466      // Update pointer and capacity, then re-lock the new memory region
467      self.ptr = new_ptr;
468      self.capacity = new_capacity;
469      let ok = self.lock_memory();
470      debug_assert!(ok, "SecureVec::reserve: lock_memory failed");
471   }
472
473   /// Creates a draining iterator that removes the specified range from the vector
474   /// and yields the removed items.
475   ///
476   /// Note: The vector is unlocked during the lifetime of the `Drain` iterator.
477   /// The memory is relocked when the `Drain` iterator is dropped.
478   ///
479   /// # Panics
480   /// Panics if the starting point is greater than the end point or if the end point
481   /// is greater than the length of the vector.
482   pub fn drain<R>(&mut self, range: R) -> Drain<'_, T>
483   where
484      R: RangeBounds<usize>,
485   {
486      let original_len = self.len;
487
488      let (drain_start_idx, drain_end_idx) = resolve_range_indices(range, original_len);
489
490      let tail_len = original_len - drain_end_idx;
491
492      self.len = drain_start_idx;
493
494      let ok = self.unlock_memory();
495      debug_assert!(ok, "SecureVec::drain: unlock_memory failed");
496
497      Drain {
498         vec_ref: self,
499         drain_start_index: drain_start_idx,
500         current_drain_iter_index: drain_start_idx,
501         drain_end_index: drain_end_idx,
502         original_vec_len: original_len,
503         tail_len,
504         _marker: PhantomData,
505      }
506   }
507
508   /// Initializes a freshly-allocated (uninitialized) buffer by cloning `src`
509   /// into it. Uses `ptr::write` so the uninitialized destination slots are
510   /// never read, never dropped, and no `&mut [T]` is ever formed over them.
511   ///
512   /// `len` is set only after every write succeeds, so a panic from
513   /// `T::clone` leaves the vector at its previous length (0 for a fresh one).
514   pub(crate) fn init_from_clone(&mut self, src: &[T])
515   where
516      T: Clone,
517   {
518      debug_assert!(src.len() <= self.capacity);
519
520      let ok = self.unlock_memory();
521      debug_assert!(
522         ok,
523         "SecureVec::init_from_clone: unlock_memory failed"
524      );
525
526      unsafe {
527         let dst = self.ptr.as_ptr();
528         for (i, item) in src.iter().enumerate() {
529            core::ptr::write(dst.add(i), item.clone());
530         }
531      }
532
533      self.len = src.len();
534      let ok = self.lock_memory();
535      debug_assert!(
536         ok,
537         "SecureVec::init_from_clone: lock_memory failed"
538      );
539   }
540}
541
542impl<T: Clone + Zeroize> Clone for SecureVec<T> {
543   fn clone(&self) -> Self {
544      let mut new_vec = SecureVec::new_with_capacity(self.capacity).unwrap();
545      self.unlock_slice(|src_slice| {
546         new_vec.init_from_clone(src_slice);
547      });
548      new_vec
549   }
550}
551
552impl<const LENGTH: usize> From<SecureArray<u8, LENGTH>> for SecureVec<u8> {
553   fn from(array: SecureArray<u8, LENGTH>) -> Self {
554      let mut new_vec = SecureVec::new_with_capacity(LENGTH)
555         .expect("Failed to allocate SecureVec during conversion");
556      array.unlock(|array_slice| {
557         new_vec.init_from_clone(array_slice);
558      });
559      new_vec
560   }
561}
562
563impl<T: Zeroize> Drop for SecureVec<T> {
564   fn drop(&mut self) {
565      self.erase();
566      let ok = self.unlock_memory();
567      debug_assert!(ok, "SecureVec::drop: unlock_memory failed");
568
569      #[cfg(feature = "use_os")]
570      free(self.ptr);
571
572      #[cfg(not(feature = "use_os"))]
573      {
574         let layout =
575            Layout::from_size_align_unchecked(self.allocated_byte_size(), mem::align_of::<T>());
576         dealloc(self.ptr.as_ptr() as *mut u8, layout);
577      }
578   }
579}
580
581impl<T: Zeroize> core::ops::Index<usize> for SecureVec<T> {
582   type Output = T;
583
584   fn index(&self, index: usize) -> &Self::Output {
585      assert!(index < self.len, "Index out of bounds");
586      unsafe {
587         let ptr = self.ptr.as_ptr().add(index);
588         &*ptr
589      }
590   }
591}
592
593#[cfg(feature = "serde")]
594impl serde::Serialize for SecureVec<u8> {
595   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
596   where
597      S: serde::Serializer,
598   {
599      self.unlock_slice(|slice| serializer.collect_seq(slice.iter()))
600   }
601}
602
603#[cfg(feature = "serde")]
604impl<'de> serde::Deserialize<'de> for SecureVec<u8> {
605   fn deserialize<D>(deserializer: D) -> Result<SecureVec<u8>, D::Error>
606   where
607      D: serde::Deserializer<'de>,
608   {
609      struct SecureVecVisitor;
610      impl<'de> serde::de::Visitor<'de> for SecureVecVisitor {
611         type Value = SecureVec<u8>;
612         fn expecting(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
613            write!(formatter, "a sequence of bytes")
614         }
615         fn visit_seq<A>(
616            self,
617            mut seq: A,
618         ) -> Result<<Self as serde::de::Visitor<'de>>::Value, A::Error>
619         where
620            A: serde::de::SeqAccess<'de>,
621         {
622            let mut vec = SecureVec::new().map_err(serde::de::Error::custom)?;
623            while let Some(byte) = seq.next_element::<u8>()? {
624               vec.push(byte);
625            }
626            Ok(vec)
627         }
628      }
629      deserializer.deserialize_seq(SecureVecVisitor)
630   }
631}
632
633/// A draining iterator for `SecureVec<T>`.
634///
635/// This struct is created by the `drain` method on `SecureVec`.
636///
637/// # Safety
638/// The returned `Drain` iterator must not be forgotten (via `mem::forget`).
639/// Forgetting the iterator sets the len of `SecureVec` to 0 and the memory will remain unlocked
640pub struct Drain<'a, T: Zeroize + 'a> {
641   vec_ref: &'a mut SecureVec<T>,
642   drain_start_index: usize,
643   current_drain_iter_index: usize,
644   drain_end_index: usize,
645
646   original_vec_len: usize, // Original length of vec_ref before drain
647   tail_len: usize,         // Number of elements after the drain range in the original vec
648
649   _marker: PhantomData<&'a T>,
650}
651
652impl<'a, T: Zeroize> Iterator for Drain<'a, T> {
653   type Item = T;
654
655   fn next(&mut self) -> Option<T> {
656      if self.current_drain_iter_index < self.drain_end_index {
657         // SecureVec is already unlocked by the `drain` method.
658         unsafe {
659            let item_ptr = self.vec_ref.ptr.as_ptr().add(self.current_drain_iter_index);
660            let item = ptr::read(item_ptr);
661            self.current_drain_iter_index += 1;
662            Some(item)
663         }
664      } else {
665         None
666      }
667   }
668
669   fn size_hint(&self) -> (usize, Option<usize>) {
670      let remaining = self.drain_end_index - self.current_drain_iter_index;
671      (remaining, Some(remaining))
672   }
673}
674
675impl<'a, T: Zeroize> ExactSizeIterator for Drain<'a, T> {}
676
677impl<'a, T: Zeroize> Drop for Drain<'a, T> {
678   fn drop(&mut self) {
679      unsafe {
680         // The vec_ref's memory is currently unlocked.
681         if mem::needs_drop::<T>() {
682            let mut current_ptr = self.vec_ref.ptr.as_ptr().add(self.current_drain_iter_index);
683            let end_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_end_index);
684            while current_ptr < end_ptr {
685               ptr::drop_in_place(current_ptr);
686               current_ptr = current_ptr.add(1);
687            }
688         }
689
690         let hole_dst_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_start_index);
691         let tail_src_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_end_index);
692
693         if self.tail_len > 0 {
694            ptr::copy(tail_src_ptr, hole_dst_ptr, self.tail_len);
695         }
696
697         // The new length of the vector.
698         let new_len = self.drain_start_index + self.tail_len;
699
700         // Process the memory region that is no longer part of the active vector's content.
701         // This region is from `vec_ref.ptr + new_len` up to `vec_ref.ptr + original_vec_len`.
702         // It contains:
703         //    a) Original data of the latter part of the drained slice (if not overwritten by tail).
704         //       These were dropped in step 1 if T:Drop.
705         //    b) Original data of the tail items (which have now been copied).
706         //       These need to be dropped if T:Drop, as ptr::copy doesn't drop the source.
707         // After any necessary drops, this entire region must be zeroized.
708
709         let mut current_cleanup_ptr = self.vec_ref.ptr.as_ptr().add(new_len);
710         let end_cleanup_ptr = self.vec_ref.ptr.as_ptr().add(self.original_vec_len);
711
712         // Determine the start of the original tail's memory region
713         let original_tail_start_ptr_val = tail_src_ptr as usize;
714
715         while current_cleanup_ptr < end_cleanup_ptr {
716            if mem::needs_drop::<T>() {
717               let current_ptr_val = current_cleanup_ptr as usize;
718               let original_tail_end_ptr_val =
719                  original_tail_start_ptr_val + self.tail_len * mem::size_of::<T>();
720
721               if current_ptr_val >= original_tail_start_ptr_val
722                  && current_ptr_val < original_tail_end_ptr_val
723               {
724                  // This element was part of the original tail. ptr::copy moved its value.
725                  // The original instance here needs to be dropped.
726                  ptr::drop_in_place(current_cleanup_ptr);
727               }
728               // Else, it was part of the drained range (not covered by tail move).
729               // If it needed dropping, it was handled in step 1.
730            }
731
732            // Zeroize the memory of this element.
733            (*current_cleanup_ptr).zeroize();
734            current_cleanup_ptr = current_cleanup_ptr.add(1);
735         }
736
737         // Update the SecureVec's length.
738         self.vec_ref.len = new_len;
739
740         // Relock the SecureVec's memory.
741         let ok = self.vec_ref.lock_memory();
742         debug_assert!(ok, "Drain::drop: lock_memory failed");
743      }
744   }
745}
746
747// Helper function to resolve RangeBounds to (start, end) indices
748fn resolve_range_indices<R: RangeBounds<usize>>(range: R, len: usize) -> (usize, usize) {
749   let start_bound = range.start_bound();
750   let end_bound = range.end_bound();
751
752   let start = match start_bound {
753      Bound::Included(&s) => s,
754      Bound::Excluded(&s) => s
755         .checked_add(1)
756         .unwrap_or_else(|| panic!("attempted to start drain at Excluded(usize::MAX)")),
757      Bound::Unbounded => 0,
758   };
759
760   let end = match end_bound {
761      Bound::Included(&e) => e
762         .checked_add(1)
763         .unwrap_or_else(|| panic!("attempted to end drain at Included(usize::MAX)")),
764      Bound::Excluded(&e) => e,
765      Bound::Unbounded => len,
766   };
767
768   if start > end {
769      panic!(
770         "drain range start ({}) must be less than or equal to end ({})",
771         start, end
772      );
773   }
774   if end > len {
775      panic!(
776         "drain range end ({}) out of bounds for slice of length {}",
777         end, len
778      );
779   }
780
781   (start, end)
782}
783
784#[cfg(all(test, feature = "use_os"))]
785mod tests {
786   use super::*;
787   use std::process::{Command, Stdio};
788   use std::sync::{Arc, Mutex};
789
790   #[test]
791   fn test_creation() {
792      let vec: Vec<u8> = vec![1, 2, 3];
793      let secure_vec = SecureVec::from_vec(vec).unwrap();
794
795      secure_vec.unlock_slice(|slice| {
796         assert_eq!(slice, &[1, 2, 3]);
797      });
798
799      let exposed_slice = &mut [1, 2, 3];
800      let secure_slice = SecureVec::from_slice_mut(exposed_slice).unwrap();
801      assert_eq!(exposed_slice, &[0u8; 3]);
802
803      secure_slice.unlock_slice(|slice| {
804         assert_eq!(slice, &[1, 2, 3]);
805      });
806
807      let exposed_slice = [1, 2, 3];
808      let secure_slice = SecureVec::from_slice(&exposed_slice).unwrap();
809
810      secure_slice.unlock_slice(|slice| {
811         assert_eq!(slice, exposed_slice);
812      });
813   }
814
815   #[test]
816   fn test_from_secure_array() {
817      let exposed: &mut [u8; 3] = &mut [1, 2, 3];
818      let array: SecureArray<u8, 3> = SecureArray::from_slice_mut(exposed).unwrap();
819      let vec: SecureVec<u8> = array.into();
820      assert_eq!(vec.len(), 3);
821      vec.unlock_slice(|slice| {
822         assert_eq!(slice, &[1, 2, 3]);
823      });
824   }
825
826   #[test]
827   fn lock_unlock_works() {
828      let secure: SecureVec<u8> = SecureVec::new().unwrap();
829
830      let unlocked = secure.unlock_memory();
831      assert!(unlocked);
832
833      let locked = secure.lock_memory();
834      assert!(locked);
835   }
836
837   #[test]
838   fn test_thread_safety() {
839      let vec: Vec<u8> = vec![];
840      let secure = SecureVec::from_vec(vec).unwrap();
841      let secure = Arc::new(Mutex::new(secure));
842
843      let mut handles = Vec::new();
844      for i in 0..5u8 {
845         let secure_clone = secure.clone();
846         let handle = std::thread::spawn(move || {
847            let mut secure = secure_clone.lock().unwrap();
848            secure.push(i);
849         });
850         handles.push(handle);
851      }
852
853      for handle in handles {
854         handle.join().unwrap();
855      }
856
857      let mut sec = secure.lock().unwrap();
858      sec.unlock_slice_mut(|slice| {
859         slice.sort();
860         assert_eq!(slice.len(), 5);
861         assert_eq!(slice, &[0, 1, 2, 3, 4]);
862      });
863   }
864
865   #[test]
866   fn test_clone() {
867      let vec: Vec<u8> = vec![1, 2, 3];
868      let secure1 = SecureVec::from_vec(vec).unwrap();
869      let secure2 = secure1.clone();
870
871      secure1.unlock_slice(|slice| {
872         secure2.unlock_slice(|slice2| {
873            assert_eq!(slice, slice2);
874         });
875      });
876   }
877
878   #[test]
879   fn test_do_not_call_forget_on_drain() {
880      let vec: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
881      let mut secure = SecureVec::from_vec(vec).unwrap();
882      let drain = secure.drain(..3);
883      core::mem::forget(drain);
884      // we can still use secure vec but its state is unreachable
885      secure.unlock_slice(|secure| {
886         assert_eq!(secure.len(), 0);
887      });
888   }
889
890   #[test]
891   fn test_drain() {
892      let vec: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
893      let mut secure = SecureVec::from_vec(vec).unwrap();
894      let mut drain = secure.drain(..3);
895      assert_eq!(drain.next(), Some(1));
896      assert_eq!(drain.next(), Some(2));
897      assert_eq!(drain.next(), Some(3));
898      assert_eq!(drain.next(), None);
899      drop(drain);
900      secure.unlock_slice(|secure| {
901         assert_eq!(secure.len(), 7);
902         assert_eq!(secure, &[4, 5, 6, 7, 8, 9, 10]);
903      });
904   }
905
906   #[cfg(feature = "serde")]
907   #[test]
908   fn test_secure_vec_serde() {
909      let vec: Vec<u8> = vec![1, 2, 3];
910      let secure = SecureVec::from_vec(vec).unwrap();
911      let json = serde_json::to_vec(&secure).expect("Serialization failed");
912      let deserialized: SecureVec<u8> =
913         serde_json::from_slice(&json).expect("Deserialization failed");
914      deserialized.unlock_slice(|slice| {
915         assert_eq!(slice, &[1, 2, 3]);
916      });
917   }
918
919   #[test]
920   fn test_erase() {
921      let mut secure = SecureVec::new_with_capacity(10).unwrap();
922      for i in 0..9 {
923         secure.push(i);
924      }
925
926      secure.erase();
927
928      secure.unlock(|secure| {
929         assert_eq!(secure.len, 0);
930         assert_eq!(secure.capacity, 10);
931      });
932
933      secure.unlock_iter(|iter| {
934         for elem in iter {
935            assert_eq!(elem, &0);
936         }
937      });
938   }
939
940   #[test]
941   fn test_push() {
942      let vec: Vec<u8> = Vec::new();
943      let mut secure = SecureVec::from_vec(vec).unwrap();
944      for i in 0..10 {
945         secure.push(i);
946      }
947
948      assert_eq!(secure.len(), 10);
949
950      secure.unlock_slice(|slice| {
951         assert_eq!(slice, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
952      });
953   }
954
955   #[test]
956   fn test_reserve() {
957      let mut secure: SecureVec<u8> = SecureVec::new().unwrap();
958      secure.reserve(10);
959      assert_eq!(secure.capacity, 10);
960   }
961
962   #[test]
963   fn test_reserve_doubling() {
964      let mut secure: SecureVec<u8> = SecureVec::new().unwrap();
965      secure.reserve(10);
966
967      for i in 0..9 {
968         secure.push(i);
969      }
970
971      secure.push(9);
972      assert_eq!(secure.capacity, 10);
973      assert_eq!(secure.len(), 10);
974
975      secure.push(10);
976      assert_eq!(secure.capacity, 20);
977      assert_eq!(secure.len(), 11);
978   }
979
980   #[test]
981   fn test_index() {
982      let vec: Vec<u8> = vec![1, 2, 3];
983      let secure = SecureVec::from_vec(vec).unwrap();
984      secure.unlock(|secure| {
985         assert_eq!(secure[0], 1);
986         assert_eq!(secure[1], 2);
987         assert_eq!(secure[2], 3);
988      });
989   }
990
991   #[test]
992   fn test_unlock_slice() {
993      let vec: Vec<u8> = vec![1, 2, 3];
994      let secure = SecureVec::from_vec(vec).unwrap();
995      secure.unlock_slice(|slice| {
996         assert_eq!(slice, &[1, 2, 3]);
997      });
998   }
999
1000   #[test]
1001   fn test_unlock_slice_mut() {
1002      let vec: Vec<u8> = vec![1, 2, 3];
1003      let mut secure = SecureVec::from_vec(vec).unwrap();
1004
1005      secure.unlock_slice_mut(|slice| {
1006         slice[0] = 4;
1007         assert_eq!(slice, &mut [4, 2, 3]);
1008      });
1009   }
1010
1011   #[test]
1012   fn test_unlock_iter() {
1013      let vec: Vec<u8> = vec![1, 2, 3];
1014      let secure = SecureVec::from_vec(vec).unwrap();
1015      let sum: u8 = secure.unlock_iter(|iter| iter.map(|&x| x).sum());
1016
1017      assert_eq!(sum, 6);
1018
1019      let secure: SecureVec<u8> = SecureVec::new_with_capacity(3).unwrap();
1020      let sum: u8 = secure.unlock_iter(|iter| iter.map(|&x| x).sum());
1021
1022      assert_eq!(sum, 0);
1023   }
1024
1025   #[test]
1026   fn test_unlock_iter_mut() {
1027      let vec: Vec<u8> = vec![1, 2, 3];
1028      let mut secure = SecureVec::from_vec(vec).unwrap();
1029      secure.unlock_iter_mut(|iter| {
1030         for elem in iter {
1031            *elem += 1;
1032         }
1033      });
1034
1035      secure.unlock_slice(|slice| {
1036         assert_eq!(slice, &[2, 3, 4]);
1037      });
1038   }
1039
1040   #[test]
1041   fn test_index_should_fail_when_locked() {
1042      let arg = "CRASH_TEST_SECUREVEC_LOCKED";
1043
1044      if std::env::args().any(|a| a == arg) {
1045         let vec: Vec<u8> = vec![1, 2, 3];
1046         let secure = SecureVec::from_vec(vec).unwrap();
1047         let _value = core::hint::black_box(secure[0]);
1048
1049         std::process::exit(1);
1050      }
1051
1052      let child = Command::new(std::env::current_exe().unwrap())
1053         .arg("vec::tests::test_index_should_fail_when_locked")
1054         .arg(arg)
1055         .arg("--nocapture")
1056         .stdout(Stdio::piped())
1057         .stderr(Stdio::piped())
1058         .spawn()
1059         .expect("Failed to spawn child process");
1060
1061      let output = child.wait_with_output().expect("Failed to wait on child");
1062      let status = output.status;
1063
1064      assert!(
1065         !status.success(),
1066         "Process exited successfully with code {:?}, but it should have crashed.",
1067         status.code()
1068      );
1069
1070      #[cfg(unix)]
1071      {
1072         use std::os::unix::process::ExitStatusExt;
1073         let signal = status
1074            .signal()
1075            .expect("Process was not terminated by a signal on Unix.");
1076         assert!(
1077            signal == libc::SIGSEGV || signal == libc::SIGBUS,
1078            "Process terminated with unexpected signal: {}",
1079            signal
1080         );
1081         println!(
1082            "Test passed: Process correctly terminated with signal {}.",
1083            signal
1084         );
1085      }
1086
1087      #[cfg(windows)]
1088      {
1089         const STATUS_ACCESS_VIOLATION: i32 = 0xC0000005_u32 as i32;
1090         assert_eq!(
1091            status.code(),
1092            Some(STATUS_ACCESS_VIOLATION),
1093            "Process exited with unexpected code: {:x?}. Expected STATUS_ACCESS_VIOLATION.",
1094            status.code()
1095         );
1096         eprintln!("Test passed: Process correctly terminated with STATUS_ACCESS_VIOLATION.");
1097      }
1098   }
1099}