Skip to main content

secure_types/
vec.rs

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