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/// ## Security Note on Direct Access
60///
61/// We intentionally do **not** implement `Index` / `IndexMut`.
62/// Using `secure_vec[0]` is a compile error.
63///
64/// This is by design: direct indexing would allow bypassing the explicit
65/// unlock mechanism. Always use `unlock_slice()` / `unlock_slice_mut()` (or
66/// the `unlock*` family of methods) to access the contents.
67///
68/// # Notes
69///
70/// If you return a new allocated `Vec` from one of the unlock methods you are responsible for zeroizing the memory.
71///
72/// # Example
73///
74/// Using `SecureBytes` (a type alias for `SecureVec<u8>`) to handle a secret key.
75///
76/// ```
77/// use secure_types::{SecureBytes, Zeroize};
78///
79/// // Create a new, empty secure vector.
80/// let mut secret_key = SecureBytes::new().unwrap();
81///
82/// // Push some sensitive data into it.
83/// secret_key.push(0xAB);
84/// secret_key.push(0xCD);
85/// secret_key.push(0xEF);
86///
87/// // The memory is locked here.
88///
89/// // Use a scope to safely access the contents as a slice.
90/// secret_key.unlock_slice(|unlocked_slice| {
91///     assert_eq!(unlocked_slice, &[0xAB, 0xCD, 0xEF]);
92/// });
93///
94/// // Not recommended but if you allocate a new Vec make sure to zeroize it
95/// let mut exposed = secret_key.unlock_slice(|unlocked_slice| {
96///     Vec::from(unlocked_slice)
97/// });
98///
99/// // Do what you need to to do with the new vector
100/// // When you are done with it, zeroize it
101/// exposed.zeroize();
102///
103/// // The memory is automatically locked again when the scope ends.
104///
105/// // When `secret_key` is dropped, its memory is securely zeroized.
106/// ```
107pub struct SecureVec<T>
108where
109   T: Zeroize,
110{
111   ptr: NonNull<T>,
112   pub(crate) len: usize,
113   pub(crate) capacity: usize,
114   _marker: PhantomData<T>,
115}
116
117unsafe impl<T: Zeroize + Send> Send for SecureVec<T> {}
118unsafe impl<T: Zeroize + Send + Sync> Sync for SecureVec<T> {}
119
120impl<T: Zeroize> SecureVec<T> {
121   /// Create a new `SecureVec` with a capacity of 1
122   pub fn new() -> Result<Self, Error> {
123      let capacity = 1;
124      let size = capacity * mem::size_of::<T>();
125      let ptr = unsafe { alloc::<T>(size)? };
126
127      let secure = SecureVec {
128         ptr,
129         len: 0,
130         capacity,
131         _marker: PhantomData,
132      };
133
134      let _locked = secure.lock_memory();
135
136      #[cfg(feature = "use_os")]
137      if !_locked {
138         return Err(Error::LockFailed);
139      }
140
141      Ok(secure)
142   }
143
144   /// Create a new `SecureVec` with the given capacity
145   pub fn new_with_capacity(mut capacity: usize) -> Result<Self, Error> {
146      if capacity == 0 {
147         capacity = 1;
148      }
149
150      capacity
151         .checked_mul(size_of::<T>())
152         .ok_or(Error::AllocationFailed)?;
153
154      let size = capacity * mem::size_of::<T>();
155      let ptr = unsafe { alloc::<T>(size)? };
156
157      let secure = SecureVec {
158         ptr,
159         len: 0,
160         capacity,
161         _marker: PhantomData,
162      };
163
164      let _locked = secure.lock_memory();
165
166      #[cfg(feature = "use_os")]
167      if !_locked {
168         return Err(Error::LockFailed);
169      }
170
171      Ok(secure)
172   }
173
174   #[cfg(feature = "use_os")]
175   /// Create a new `SecureVec` from a `Vec`
176   ///
177   /// The `Vec` is zeroized afterwards
178   pub fn from_vec(mut vec: Vec<T>) -> Result<Self, Error> {
179      if vec.capacity() == 0 {
180         vec.reserve(1);
181      }
182
183      let capacity = vec.capacity();
184      let len = vec.len();
185
186      let size = match capacity.checked_mul(size_of::<T>()) {
187         Some(s) => s,
188         None => {
189            vec.zeroize();
190            return Err(Error::AllocationFailed);
191         }
192      };
193
194      let ptr = match unsafe { alloc::<T>(size) } {
195         Ok(ptr) => ptr,
196         Err(_) => {
197            vec.zeroize();
198            return Err(Error::AllocationFailed);
199         }
200      };
201
202      // Move data from the old vec into the secure allocation using ptr::read / ptr::write
203      // This correctly transfers ownership for non-Copy types (e.g. structs containing String).
204      // We then zero the *bytes* of the source buffer (after moving values out) to avoid
205      // leaving sensitive data, and prevent double-drop by clearing the vec length.
206      unsafe {
207         let src = vec.as_ptr();
208         let dst = ptr.as_ptr();
209         for i in 0..len {
210            let value = core::ptr::read(src.add(i));
211            core::ptr::write(dst.add(i), value);
212         }
213      }
214
215      // Prevent the Vec from dropping the now-moved-from elements (would be UB)
216      // and securely erase whatever representation bytes remain in its buffer.
217      //
218      // We use set_len(0) + zeroize on a &mut [u8] view of the allocation
219      // (instead of calling vec.zeroize()) because the Ts have been moved out
220      // via ptr::read. The normal Vec::zeroize impl would zeroize+drop the
221      // moved-from elements, which is UB (and often SIGABRT for a non-copy type).
222      let old_byte_size = capacity * mem::size_of::<T>();
223      unsafe {
224         vec.set_len(0);
225      }
226      if old_byte_size > 0 {
227         // SAFETY: after set_len(0) the allocation bytes are still valid,
228         // we own them exclusively, and no Ts will be dropped by the Vec.
229         let bytes =
230            unsafe { core::slice::from_raw_parts_mut(vec.as_mut_ptr() as *mut u8, old_byte_size) };
231         bytes.zeroize();
232      }
233
234      let secure = SecureVec {
235         ptr,
236         len,
237         capacity,
238         _marker: PhantomData,
239      };
240
241      let locked = secure.lock_memory();
242
243      if !locked {
244         return Err(Error::LockFailed);
245      }
246
247      Ok(secure)
248   }
249
250   /// Create a new `SecureVec` from a mutable slice.
251   ///
252   /// The slice is zeroized afterwards
253   pub fn from_slice_mut(slice: &mut [T]) -> Result<Self, Error>
254   where
255      T: Clone + DefaultIsZeroes,
256   {
257      let mut secure_vec = match SecureVec::new_with_capacity(slice.len()) {
258         Ok(secure_vec) => secure_vec,
259         Err(e) => {
260            slice.zeroize();
261            return Err(e);
262         }
263      };
264
265      secure_vec.init_from_clone(slice);
266      slice.zeroize();
267
268      Ok(secure_vec)
269   }
270
271   /// Create a new `SecureVec` from a slice.
272   ///
273   /// The slice is not zeroized, you are responsible for zeroizing it
274   pub fn from_slice(slice: &[T]) -> Result<Self, Error>
275   where
276      T: Clone,
277   {
278      let mut secure_vec = SecureVec::new_with_capacity(slice.len())?;
279      secure_vec.init_from_clone(slice);
280      Ok(secure_vec)
281   }
282
283   pub fn len(&self) -> usize {
284      self.len
285   }
286
287   pub fn is_empty(&self) -> bool {
288      self.len() == 0
289   }
290
291   /// Returns the pointer to the locked memory region
292   ///
293   /// # DANGER
294   ///
295   /// This is a low-level API, which should be used only for
296   /// testing purposes. If you need to access the locked memory
297   /// region, use one of the unlock methods.
298   #[cfg(feature = "expose-ptr")]
299   #[deprecated(
300      since = "0.3.0",
301      note = "This method is intended only for testing/crash reproduction. Use one of the unlock methods instead."
302   )]
303   pub fn ptr(&self) -> NonNull<T> {
304      self.ptr
305   }
306
307   /// Returns the total number of bytes currently allocated for this vector.
308   #[cfg(not(feature = "use_os"))]
309   pub(crate) fn allocated_byte_size(&self) -> usize {
310      self.capacity * mem::size_of::<T>()
311   }
312
313   pub(crate) fn as_mut_ptr(&mut self) -> *mut u8 {
314      self.ptr.as_ptr() as *mut u8
315   }
316
317   pub(crate) fn lock_memory(&self) -> bool {
318      #[cfg(feature = "use_os")]
319      {
320         #[cfg(windows)]
321         {
322            super::mprotect(self.ptr, Prot::NoAccess)
323         }
324         #[cfg(unix)]
325         {
326            super::mprotect(self.ptr, Prot::NoAccess)
327         }
328      }
329      #[cfg(not(feature = "use_os"))]
330      {
331         true // No-op: always "succeeds"
332      }
333   }
334
335   pub(crate) fn unlock_memory(&self) -> bool {
336      #[cfg(feature = "use_os")]
337      {
338         #[cfg(windows)]
339         {
340            super::mprotect(self.ptr, Prot::ReadWrite)
341         }
342         #[cfg(unix)]
343         {
344            super::mprotect(self.ptr, Prot::ReadWrite)
345         }
346      }
347
348      #[cfg(not(feature = "use_os"))]
349      {
350         true // No-op: always "succeeds"
351      }
352   }
353
354   /// Immutable access to the `SecureVec`
355   pub fn unlock<F, R>(&self, f: F) -> R
356   where
357      F: FnOnce(&SecureVec<T>) -> R,
358   {
359      let _guard = UnlockGuard::new(self);
360      let result = f(self);
361      result
362   }
363
364   /// Immutable access to the `SecureVec` as `&[T]`
365   pub fn unlock_slice<F, R>(&self, f: F) -> R
366   where
367      F: FnOnce(&[T]) -> R,
368   {
369      let _guard = UnlockGuard::new(self);
370      let slice = unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.len) };
371      f(slice)
372   }
373
374   /// Mutable access to the `SecureVec` as `&mut [T]`
375   pub fn unlock_slice_mut<F, R>(&mut self, f: F) -> R
376   where
377      F: FnOnce(&mut [T]) -> R,
378   {
379      unsafe {
380         let _guard = UnlockGuard::new(self);
381         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
382         let result = f(slice);
383         result
384      }
385   }
386
387   /// Immutable access to the `SecureVec` as `Iter<T>`
388   pub fn unlock_iter<F, R>(&self, f: F) -> R
389   where
390      F: FnOnce(core::slice::Iter<T>) -> R,
391   {
392      unsafe {
393         let _guard = UnlockGuard::new(self);
394         let slice = core::slice::from_raw_parts(self.ptr.as_ptr(), self.len);
395         let iter = slice.iter();
396         let result = f(iter);
397         result
398      }
399   }
400
401   /// Mutable access to the `SecureVec` as `IterMut<T>`
402   pub fn unlock_iter_mut<F, R>(&mut self, f: F) -> R
403   where
404      F: FnOnce(core::slice::IterMut<T>) -> R,
405   {
406      unsafe {
407         let _guard = UnlockGuard::new(self);
408         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
409         let iter = slice.iter_mut();
410         let result = f(iter);
411         result
412      }
413   }
414
415   /// Erase the underlying data and clears the vector
416   ///
417   /// The memory is locked again and the capacity is preserved for reuse
418   pub fn erase(&mut self) {
419      unsafe {
420         let ok = self.unlock_memory();
421         debug_assert!(ok, "SecureVec::erase: unlock_memory failed");
422
423         // Only zero the initialized elements. Zeroizing capacity would try to
424         // zeroize uninitialized memory as T, which for Drop types (eg. String)
425         // is UB and causes SIGSEGV/SIGABRT.
426         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
427         for elem in slice.iter_mut() {
428            elem.zeroize();
429         }
430
431         self.clear();
432
433         let ok = self.lock_memory();
434         debug_assert!(ok, "SecureVec::erase: lock_memory failed");
435      }
436   }
437
438   /// Clear the vector
439   ///
440   /// This just sets the vector's len to zero it does not erase the underlying data
441   pub fn clear(&mut self) {
442      self.len = 0;
443   }
444
445   pub fn push(&mut self, value: T) {
446      self.reserve(1);
447
448      let ok = self.unlock_memory();
449      debug_assert!(ok, "SecureVec::push: unlock_memory failed");
450
451      unsafe {
452         // Write the new value at the end of the vector.
453         core::ptr::write(self.ptr.as_ptr().add(self.len), value);
454
455         self.len += 1;
456      }
457
458      let ok = self.lock_memory();
459      debug_assert!(ok, "SecureVec::push: lock_memory failed");
460   }
461
462   /// Ensures that the vector has enough capacity for at least `additional` more elements.
463   ///
464   /// If more capacity is needed, it will reallocate. This may cause the buffer location to change.
465   ///
466   /// # Panics
467   ///
468   /// Panics if the new capacity overflows `usize` or if the allocation fails.
469   pub fn reserve(&mut self, additional: usize) {
470      if self.len() + additional <= self.capacity {
471         return;
472      }
473
474      // Use an amortized growth strategy to avoid reallocating on every push
475      let required_capacity = self.len() + additional;
476      let new_capacity = (self.capacity.max(1) * 2).max(required_capacity);
477
478      let new_size = new_capacity * mem::size_of::<T>();
479
480      // Safe to panic here because the memory is locked
481      let new_ptr = unsafe {
482         alloc::<T>(new_size).unwrap_or_else(|_| {
483            panic!(
484               "secure-types: failed to allocate {} bytes of locked memory \
485          (possibly RLIMIT_MEMLOCK exhausted); SecureVec left unchanged",
486               new_size
487            )
488         })
489      };
490
491      // Copy data to new pointer
492      unsafe {
493         let ok = self.unlock_memory();
494         debug_assert!(ok, "SecureVec::reserve: unlock_memory failed");
495
496         // Move (not copy) elements to new buffer to support non-Copy T correctly.
497         // Using read+write transfers ownership of e.g. Strings.
498         let len = self.len();
499         for i in 0..len {
500            let val = core::ptr::read(self.ptr.as_ptr().add(i));
501            core::ptr::write(new_ptr.as_ptr().add(i), val);
502         }
503
504         // Erase old buffer bytes (after move-out)
505         if self.capacity > 0 {
506            let old_bytes = self.capacity * mem::size_of::<T>();
507            let bytes = core::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut u8, old_bytes);
508            bytes.zeroize();
509         }
510
511         #[cfg(feature = "use_os")]
512         free(self.ptr);
513
514         #[cfg(not(feature = "use_os"))]
515         {
516            let old_size = self.capacity * mem::size_of::<T>();
517            let old_layout = Layout::from_size_align_unchecked(old_size, mem::align_of::<T>());
518            alloc::alloc::dealloc(self.ptr.as_ptr() as *mut u8, old_layout);
519         }
520      }
521
522      // Update pointer and capacity, then re-lock the new memory region
523      self.ptr = new_ptr;
524      self.capacity = new_capacity;
525      let ok = self.lock_memory();
526      debug_assert!(ok, "SecureVec::reserve: lock_memory failed");
527   }
528
529   /// Creates a draining iterator that removes the specified range from the vector
530   /// and yields the removed items.
531   ///
532   /// Note: The vector is unlocked during the lifetime of the `Drain` iterator.
533   /// The memory is relocked when the `Drain` iterator is dropped.
534   ///
535   /// # Panics
536   /// Panics if the starting point is greater than the end point or if the end point
537   /// is greater than the length of the vector.
538   pub fn drain<R>(&mut self, range: R) -> Drain<'_, T>
539   where
540      R: RangeBounds<usize>,
541   {
542      let original_len = self.len;
543
544      let (drain_start_idx, drain_end_idx) = resolve_range_indices(range, original_len);
545
546      let tail_len = original_len - drain_end_idx;
547
548      self.len = drain_start_idx;
549
550      let ok = self.unlock_memory();
551      debug_assert!(ok, "SecureVec::drain: unlock_memory failed");
552
553      Drain {
554         vec_ref: self,
555         drain_start_index: drain_start_idx,
556         current_drain_iter_index: drain_start_idx,
557         drain_end_index: drain_end_idx,
558         original_vec_len: original_len,
559         tail_len,
560         _marker: PhantomData,
561      }
562   }
563
564   /// Initializes a freshly-allocated (uninitialized) buffer by cloning `src`
565   /// into it. Uses `ptr::write` so the uninitialized destination slots are
566   /// never read, never dropped, and no `&mut [T]` is ever formed over them.
567   ///
568   /// `len` is set only after every write succeeds, so a panic from
569   /// `T::clone` leaves the vector at its previous length (0 for a fresh one).
570   pub(crate) fn init_from_clone(&mut self, src: &[T])
571   where
572      T: Clone,
573   {
574      debug_assert!(src.len() <= self.capacity);
575
576      let ok = self.unlock_memory();
577      debug_assert!(
578         ok,
579         "SecureVec::init_from_clone: unlock_memory failed"
580      );
581
582      unsafe {
583         let dst = self.ptr.as_ptr();
584         for (i, item) in src.iter().enumerate() {
585            core::ptr::write(dst.add(i), item.clone());
586         }
587      }
588
589      self.len = src.len();
590      let ok = self.lock_memory();
591      debug_assert!(
592         ok,
593         "SecureVec::init_from_clone: lock_memory failed"
594      );
595   }
596}
597
598impl<T: Clone + Zeroize> Clone for SecureVec<T> {
599   fn clone(&self) -> Self {
600      let mut new_vec = SecureVec::new_with_capacity(self.capacity).unwrap();
601      self.unlock_slice(|src_slice| {
602         new_vec.init_from_clone(src_slice);
603      });
604      new_vec
605   }
606}
607
608impl<const LENGTH: usize> From<SecureArray<u8, LENGTH>> for SecureVec<u8> {
609   fn from(array: SecureArray<u8, LENGTH>) -> Self {
610      let mut new_vec = SecureVec::new_with_capacity(LENGTH)
611         .expect("Failed to allocate SecureVec during conversion");
612      array.unlock(|array_slice| {
613         new_vec.init_from_clone(array_slice);
614      });
615      new_vec
616   }
617}
618
619impl<T: Zeroize> Drop for SecureVec<T> {
620   fn drop(&mut self) {
621      unsafe {
622         let ok = self.unlock_memory();
623         debug_assert!(ok, "SecureVec::erase: unlock_memory failed");
624
625         // Only zero the initialized elements. Zeroizing capacity would try to
626         // zeroize uninitialized memory as T, which for Drop types (eg. String)
627         // is UB and causes SIGSEGV/SIGABRT.
628         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
629         for elem in slice.iter_mut() {
630            elem.zeroize();
631         }
632      }
633
634      #[cfg(feature = "use_os")]
635      free(self.ptr);
636
637      #[cfg(not(feature = "use_os"))]
638      unsafe {
639         let layout =
640            Layout::from_size_align_unchecked(self.allocated_byte_size(), mem::align_of::<T>());
641         alloc::alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
642      }
643   }
644}
645
646// Note: We intentionally do **not** implement Index / IndexMut.
647// Direct indexing (`vec[0]`) would bypass the unlock mechanism and
648// access locked memory, causing a segfault. This is by design.
649// Always use unlock_slice() / unlock_slice_mut().
650
651#[cfg(feature = "serde")]
652impl serde::Serialize for SecureVec<u8> {
653   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
654   where
655      S: serde::Serializer,
656   {
657      self.unlock_slice(|slice| serializer.collect_seq(slice.iter()))
658   }
659}
660
661#[cfg(feature = "serde")]
662impl<'de> serde::Deserialize<'de> for SecureVec<u8> {
663   fn deserialize<D>(deserializer: D) -> Result<SecureVec<u8>, D::Error>
664   where
665      D: serde::Deserializer<'de>,
666   {
667      struct SecureVecVisitor;
668      impl<'de> serde::de::Visitor<'de> for SecureVecVisitor {
669         type Value = SecureVec<u8>;
670         fn expecting(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
671            write!(formatter, "a sequence of bytes")
672         }
673         fn visit_seq<A>(
674            self,
675            mut seq: A,
676         ) -> Result<<Self as serde::de::Visitor<'de>>::Value, A::Error>
677         where
678            A: serde::de::SeqAccess<'de>,
679         {
680            let mut vec = SecureVec::new().map_err(serde::de::Error::custom)?;
681            while let Some(byte) = seq.next_element::<u8>()? {
682               vec.push(byte);
683            }
684            Ok(vec)
685         }
686      }
687      deserializer.deserialize_seq(SecureVecVisitor)
688   }
689}
690
691/// A draining iterator for `SecureVec<T>`.
692///
693/// This struct is created by the `drain` method on `SecureVec`.
694///
695/// # Safety
696/// The returned `Drain` iterator must not be forgotten (via `mem::forget`).
697/// Forgetting the iterator sets the len of `SecureVec` to 0 and the memory will remain unlocked
698pub struct Drain<'a, T: Zeroize + 'a> {
699   vec_ref: &'a mut SecureVec<T>,
700   drain_start_index: usize,
701   current_drain_iter_index: usize,
702   drain_end_index: usize,
703
704   original_vec_len: usize, // Original length of vec_ref before drain
705   tail_len: usize,         // Number of elements after the drain range in the original vec
706
707   _marker: PhantomData<&'a T>,
708}
709
710impl<'a, T: Zeroize> Iterator for Drain<'a, T> {
711   type Item = T;
712
713   fn next(&mut self) -> Option<T> {
714      if self.current_drain_iter_index < self.drain_end_index {
715         // SecureVec is already unlocked by the `drain` method.
716         unsafe {
717            let item_ptr = self.vec_ref.ptr.as_ptr().add(self.current_drain_iter_index);
718            let item = ptr::read(item_ptr);
719            self.current_drain_iter_index += 1;
720            Some(item)
721         }
722      } else {
723         None
724      }
725   }
726
727   fn size_hint(&self) -> (usize, Option<usize>) {
728      let remaining = self.drain_end_index - self.current_drain_iter_index;
729      (remaining, Some(remaining))
730   }
731}
732
733impl<'a, T: Zeroize> ExactSizeIterator for Drain<'a, T> {}
734
735impl<'a, T: Zeroize> Drop for Drain<'a, T> {
736   fn drop(&mut self) {
737      unsafe {
738         // The vec_ref's memory is currently unlocked.
739         if mem::needs_drop::<T>() {
740            let mut current_ptr = self.vec_ref.ptr.as_ptr().add(self.current_drain_iter_index);
741            let end_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_end_index);
742            while current_ptr < end_ptr {
743               ptr::drop_in_place(current_ptr);
744               current_ptr = current_ptr.add(1);
745            }
746         }
747
748         let hole_dst_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_start_index);
749         let tail_src_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_end_index);
750
751         if self.tail_len > 0 {
752            ptr::copy(tail_src_ptr, hole_dst_ptr, self.tail_len);
753         }
754
755         // The new length of the vector.
756         let new_len = self.drain_start_index + self.tail_len;
757
758         // Process the memory region that is no longer part of the active vector's content.
759         // This region is from `vec_ref.ptr + new_len` up to `vec_ref.ptr + original_vec_len`.
760         // It contains:
761         //    a) Original data of the latter part of the drained slice (if not overwritten by tail).
762         //       These were dropped in step 1 if T:Drop.
763         //    b) Original data of the tail items (which have now been copied).
764         //       These need to be dropped if T:Drop, as ptr::copy doesn't drop the source.
765         // After any necessary drops, this entire region must be zeroized.
766
767         let mut current_cleanup_ptr = self.vec_ref.ptr.as_ptr().add(new_len);
768         let end_cleanup_ptr = self.vec_ref.ptr.as_ptr().add(self.original_vec_len);
769
770         // Determine the start of the original tail's memory region
771         let original_tail_start_ptr_val = tail_src_ptr as usize;
772
773         while current_cleanup_ptr < end_cleanup_ptr {
774            if mem::needs_drop::<T>() {
775               let current_ptr_val = current_cleanup_ptr as usize;
776               let original_tail_end_ptr_val =
777                  original_tail_start_ptr_val + self.tail_len * mem::size_of::<T>();
778
779               if current_ptr_val >= original_tail_start_ptr_val
780                  && current_ptr_val < original_tail_end_ptr_val
781               {
782                  // This element was part of the original tail. ptr::copy moved its value.
783                  // The original instance here needs to be dropped.
784                  ptr::drop_in_place(current_cleanup_ptr);
785               }
786               // Else, it was part of the drained range (not covered by tail move).
787               // If it needed dropping, it was handled in step 1.
788            }
789
790            // Zeroize the memory of this element.
791            (*current_cleanup_ptr).zeroize();
792            current_cleanup_ptr = current_cleanup_ptr.add(1);
793         }
794
795         // Update the SecureVec's length.
796         self.vec_ref.len = new_len;
797
798         // Relock the SecureVec's memory.
799         let ok = self.vec_ref.lock_memory();
800         debug_assert!(ok, "Drain::drop: lock_memory failed");
801      }
802   }
803}
804
805// Helper function to resolve RangeBounds to (start, end) indices
806fn resolve_range_indices<R: RangeBounds<usize>>(range: R, len: usize) -> (usize, usize) {
807   let start_bound = range.start_bound();
808   let end_bound = range.end_bound();
809
810   let start = match start_bound {
811      Bound::Included(&s) => s,
812      Bound::Excluded(&s) => s
813         .checked_add(1)
814         .unwrap_or_else(|| panic!("attempted to start drain at Excluded(usize::MAX)")),
815      Bound::Unbounded => 0,
816   };
817
818   let end = match end_bound {
819      Bound::Included(&e) => e
820         .checked_add(1)
821         .unwrap_or_else(|| panic!("attempted to end drain at Included(usize::MAX)")),
822      Bound::Excluded(&e) => e,
823      Bound::Unbounded => len,
824   };
825
826   if start > end {
827      panic!(
828         "drain range start ({}) must be less than or equal to end ({})",
829         start, end
830      );
831   }
832   if end > len {
833      panic!(
834         "drain range end ({}) out of bounds for slice of length {}",
835         end, len
836      );
837   }
838
839   (start, end)
840}
841
842#[cfg(all(test, feature = "use_os"))]
843mod tests {
844   use super::*;
845   use std::fmt::Debug;
846   use std::process::{Command, Stdio};
847   use std::sync::{Arc, Mutex};
848   use zeroize::Zeroize;
849
850   // Test helper types for variety (different sizes, alignments, complex data)
851
852   #[derive(Clone, Debug, PartialEq)]
853   struct SmallStruct {
854      a: u8,
855      b: u16,
856   }
857
858   impl Zeroize for SmallStruct {
859      fn zeroize(&mut self) {
860         self.a.zeroize();
861         self.b.zeroize();
862      }
863   }
864
865   #[derive(Clone, Debug, PartialEq)]
866   struct LargeStruct {
867      data: [u64; 4],
868      flag: bool,
869   }
870
871   impl Zeroize for LargeStruct {
872      fn zeroize(&mut self) {
873         self.data.zeroize();
874         self.flag.zeroize();
875      }
876   }
877
878   #[derive(Clone, Debug, PartialEq)]
879   #[repr(align(64))]
880   struct AlignedStruct {
881      value: u64,
882   }
883
884   impl Zeroize for AlignedStruct {
885      fn zeroize(&mut self) {
886         self.value.zeroize();
887      }
888   }
889
890   #[derive(Clone, Debug, PartialEq)]
891   struct Person {
892      name: String,
893      age: u32,
894      notes: String,
895   }
896
897   impl Zeroize for Person {
898      fn zeroize(&mut self) {
899         self.name.zeroize();
900         self.age.zeroize();
901         self.notes.zeroize();
902      }
903   }
904
905   impl Person {
906      fn new(name: impl Into<String>, age: u32, notes: impl Into<String>) -> Self {
907         Self {
908            name: name.into(),
909            age,
910            notes: notes.into(),
911         }
912      }
913   }
914
915   fn create_test_person(id: usize) -> Person {
916      Person::new(
917         format!("Person{}", id),
918         (id % 100) as u32,
919         format!("Some secret notes for person #{}", id),
920      )
921   }
922
923   fn test_vec_generic_basics<T: Zeroize + Clone + PartialEq + Debug>(initial: &[T]) {
924      if initial.is_empty() {
925         return;
926      }
927
928      // from_slice
929      let secure = SecureVec::from_slice(initial).unwrap();
930      assert_eq!(secure.len(), initial.len());
931      secure.unlock_slice(|slice| {
932         assert_eq!(slice, initial);
933      });
934
935      // new + push
936      let mut secure_push = SecureVec::new().unwrap();
937      for item in initial {
938         secure_push.push(item.clone());
939      }
940      secure_push.unlock_slice(|slice| {
941         assert_eq!(slice, initial);
942      });
943
944      // clone
945      let cloned = secure.clone();
946      secure.unlock_slice(|s| {
947         cloned.unlock_slice(|c| {
948            assert_eq!(s, c);
949         });
950      });
951
952      // reserve + push more
953      let mut res = SecureVec::new().unwrap();
954      res.reserve(initial.len() + 2);
955      for item in initial {
956         res.push(item.clone());
957      }
958      assert!(res.capacity >= initial.len() + 2 || res.capacity >= initial.len());
959      res.unlock_slice(|slice| {
960         assert_eq!(slice, initial);
961      });
962
963      // erase
964      res.erase();
965      res.unlock(|v| {
966         assert_eq!(v.len, 0);
967         // capacity should be preserved
968         assert!(v.capacity > 0);
969      });
970
971      // from_vec
972      let vec_data: Vec<T> = initial.to_vec();
973      let from_vec_secure = SecureVec::from_vec(vec_data).unwrap();
974      from_vec_secure.unlock_slice(|slice| {
975         assert_eq!(slice, initial);
976      });
977   }
978
979   #[test]
980   fn test_creation() {
981      let vec: Vec<u8> = vec![1, 2, 3];
982      let secure_vec = SecureVec::from_vec(vec).unwrap();
983
984      secure_vec.unlock_slice(|slice| {
985         assert_eq!(slice, &[1, 2, 3]);
986      });
987
988      let exposed_slice = &mut [1, 2, 3];
989      let secure_slice = SecureVec::from_slice_mut(exposed_slice).unwrap();
990      assert_eq!(exposed_slice, &[0u8; 3]);
991
992      secure_slice.unlock_slice(|slice| {
993         assert_eq!(slice, &[1, 2, 3]);
994      });
995
996      let exposed_slice = [1, 2, 3];
997      let secure_slice = SecureVec::from_slice(&exposed_slice).unwrap();
998
999      secure_slice.unlock_slice(|slice| {
1000         assert_eq!(slice, exposed_slice);
1001      });
1002   }
1003
1004   #[test]
1005   fn test_from_secure_array() {
1006      let exposed: &mut [u8; 3] = &mut [1, 2, 3];
1007      let array: SecureArray<u8, 3> = SecureArray::from_slice_mut(exposed).unwrap();
1008      let vec: SecureVec<u8> = array.into();
1009      assert_eq!(vec.len(), 3);
1010      vec.unlock_slice(|slice| {
1011         assert_eq!(slice, &[1, 2, 3]);
1012      });
1013   }
1014
1015   #[test]
1016   fn lock_unlock_works() {
1017      let secure: SecureVec<u8> = SecureVec::new().unwrap();
1018
1019      let unlocked = secure.unlock_memory();
1020      assert!(unlocked);
1021
1022      let locked = secure.lock_memory();
1023      assert!(locked);
1024   }
1025
1026   #[test]
1027   fn test_thread_safety() {
1028      let vec: Vec<u8> = vec![];
1029      let secure = SecureVec::from_vec(vec).unwrap();
1030      let secure = Arc::new(Mutex::new(secure));
1031
1032      let mut handles = Vec::new();
1033      for i in 0..5u8 {
1034         let secure_clone = secure.clone();
1035         let handle = std::thread::spawn(move || {
1036            let mut secure = secure_clone.lock().unwrap();
1037            secure.push(i);
1038         });
1039         handles.push(handle);
1040      }
1041
1042      for handle in handles {
1043         handle.join().unwrap();
1044      }
1045
1046      let mut sec = secure.lock().unwrap();
1047      sec.unlock_slice_mut(|slice| {
1048         slice.sort();
1049         assert_eq!(slice.len(), 5);
1050         assert_eq!(slice, &[0, 1, 2, 3, 4]);
1051      });
1052   }
1053
1054   #[test]
1055   fn test_clone() {
1056      let vec: Vec<u8> = vec![1, 2, 3];
1057      let secure1 = SecureVec::from_vec(vec).unwrap();
1058      let secure2 = secure1.clone();
1059
1060      secure1.unlock_slice(|slice| {
1061         secure2.unlock_slice(|slice2| {
1062            assert_eq!(slice, slice2);
1063         });
1064      });
1065   }
1066
1067   #[test]
1068   fn test_do_not_call_forget_on_drain() {
1069      let vec: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
1070      let mut secure = SecureVec::from_vec(vec).unwrap();
1071      let drain = secure.drain(..3);
1072      core::mem::forget(drain);
1073      // we can still use secure vec but its state is unreachable
1074      secure.unlock_slice(|secure| {
1075         assert_eq!(secure.len(), 0);
1076      });
1077   }
1078
1079   #[test]
1080   fn test_drain() {
1081      let vec: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
1082      let mut secure = SecureVec::from_vec(vec).unwrap();
1083      let mut drain = secure.drain(..3);
1084      assert_eq!(drain.next(), Some(1));
1085      assert_eq!(drain.next(), Some(2));
1086      assert_eq!(drain.next(), Some(3));
1087      assert_eq!(drain.next(), None);
1088      drop(drain);
1089      secure.unlock_slice(|secure| {
1090         assert_eq!(secure.len(), 7);
1091         assert_eq!(secure, &[4, 5, 6, 7, 8, 9, 10]);
1092      });
1093   }
1094
1095   #[cfg(feature = "serde")]
1096   #[test]
1097   fn test_secure_vec_serde() {
1098      let vec: Vec<u8> = vec![1, 2, 3];
1099      let secure = SecureVec::from_vec(vec).unwrap();
1100      let json = serde_json::to_vec(&secure).expect("Serialization failed");
1101      let deserialized: SecureVec<u8> =
1102         serde_json::from_slice(&json).expect("Deserialization failed");
1103      deserialized.unlock_slice(|slice| {
1104         assert_eq!(slice, &[1, 2, 3]);
1105      });
1106   }
1107
1108   #[test]
1109   fn test_erase() {
1110      let mut secure = SecureVec::new_with_capacity(10).unwrap();
1111      for i in 0..9 {
1112         secure.push(i);
1113      }
1114
1115      secure.erase();
1116
1117      secure.unlock(|secure| {
1118         assert_eq!(secure.len, 0);
1119         assert_eq!(secure.capacity, 10);
1120      });
1121
1122      secure.unlock_iter(|iter| {
1123         for elem in iter {
1124            assert_eq!(elem, &0);
1125         }
1126      });
1127   }
1128
1129   #[test]
1130   fn test_push() {
1131      let vec: Vec<u8> = Vec::new();
1132      let mut secure = SecureVec::from_vec(vec).unwrap();
1133      for i in 0..10 {
1134         secure.push(i);
1135      }
1136
1137      assert_eq!(secure.len(), 10);
1138
1139      secure.unlock_slice(|slice| {
1140         assert_eq!(slice, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
1141      });
1142   }
1143
1144   #[test]
1145   fn test_reserve() {
1146      let mut secure: SecureVec<u8> = SecureVec::new().unwrap();
1147      secure.reserve(10);
1148      assert_eq!(secure.capacity, 10);
1149   }
1150
1151   #[test]
1152   fn test_reserve_doubling() {
1153      let mut secure: SecureVec<u8> = SecureVec::new().unwrap();
1154      secure.reserve(10);
1155
1156      for i in 0..9 {
1157         secure.push(i);
1158      }
1159
1160      secure.push(9);
1161      assert_eq!(secure.capacity, 10);
1162      assert_eq!(secure.len(), 10);
1163
1164      secure.push(10);
1165      assert_eq!(secure.capacity, 20);
1166      assert_eq!(secure.len(), 11);
1167   }
1168
1169   #[test]
1170   fn test_unlock_gives_access() {
1171      let vec: Vec<u8> = vec![1, 2, 3];
1172      let secure = SecureVec::from_vec(vec).unwrap();
1173      secure.unlock_slice(|slice| {
1174         assert_eq!(slice[0], 1);
1175         assert_eq!(slice[1], 2);
1176         assert_eq!(slice[2], 3);
1177      });
1178   }
1179
1180   #[test]
1181   fn test_unlock_slice() {
1182      let vec: Vec<u8> = vec![1, 2, 3];
1183      let secure = SecureVec::from_vec(vec).unwrap();
1184      secure.unlock_slice(|slice| {
1185         assert_eq!(slice, &[1, 2, 3]);
1186      });
1187   }
1188
1189   #[test]
1190   fn test_unlock_slice_mut() {
1191      let vec: Vec<u8> = vec![1, 2, 3];
1192      let mut secure = SecureVec::from_vec(vec).unwrap();
1193
1194      secure.unlock_slice_mut(|slice| {
1195         slice[0] = 4;
1196         assert_eq!(slice, &mut [4, 2, 3]);
1197      });
1198   }
1199
1200   #[test]
1201   fn test_unlock_iter() {
1202      let vec: Vec<u8> = vec![1, 2, 3];
1203      let secure = SecureVec::from_vec(vec).unwrap();
1204      let sum: u8 = secure.unlock_iter(|iter| iter.map(|&x| x).sum());
1205
1206      assert_eq!(sum, 6);
1207
1208      let secure: SecureVec<u8> = SecureVec::new_with_capacity(3).unwrap();
1209      let sum: u8 = secure.unlock_iter(|iter| iter.map(|&x| x).sum());
1210
1211      assert_eq!(sum, 0);
1212   }
1213
1214   #[test]
1215   fn test_unlock_iter_mut() {
1216      let vec: Vec<u8> = vec![1, 2, 3];
1217      let mut secure = SecureVec::from_vec(vec).unwrap();
1218      secure.unlock_iter_mut(|iter| {
1219         for elem in iter {
1220            *elem += 1;
1221         }
1222      });
1223
1224      secure.unlock_slice(|slice| {
1225         assert_eq!(slice, &[2, 3, 4]);
1226      });
1227   }
1228
1229   #[test]
1230   fn test_index_should_fail_when_locked() {
1231      let arg = "CRASH_TEST_SECUREVEC_LOCKED";
1232
1233      if std::env::args().any(|a| a == arg) {
1234         let vec: Vec<u8> = vec![1, 2, 3];
1235         let secure = SecureVec::from_vec(vec).unwrap();
1236         // Deliberately dereference the locked pointer to test that
1237         // the security model (mlock + no normal access) works as expected.
1238         let _value = unsafe { core::hint::black_box(*secure.ptr.as_ptr()) };
1239
1240         std::process::exit(1);
1241      }
1242
1243      let child = Command::new(std::env::current_exe().unwrap())
1244         .arg("vec::tests::test_index_should_fail_when_locked")
1245         .arg(arg)
1246         .arg("--nocapture")
1247         .stdout(Stdio::piped())
1248         .stderr(Stdio::piped())
1249         .spawn()
1250         .expect("Failed to spawn child process");
1251
1252      let output = child.wait_with_output().expect("Failed to wait on child");
1253      let status = output.status;
1254
1255      assert!(
1256         !status.success(),
1257         "Process exited successfully with code {:?}, but it should have crashed.",
1258         status.code()
1259      );
1260
1261      #[cfg(unix)]
1262      {
1263         use std::os::unix::process::ExitStatusExt;
1264         let signal = status
1265            .signal()
1266            .expect("Process was not terminated by a signal on Unix.");
1267         assert!(
1268            signal == libc::SIGSEGV || signal == libc::SIGBUS,
1269            "Process terminated with unexpected signal: {}",
1270            signal
1271         );
1272         println!(
1273            "Test passed: Process correctly terminated with signal {}.",
1274            signal
1275         );
1276      }
1277
1278      #[cfg(windows)]
1279      {
1280         const STATUS_ACCESS_VIOLATION: i32 = 0xC0000005_u32 as i32;
1281         assert_eq!(
1282            status.code(),
1283            Some(STATUS_ACCESS_VIOLATION),
1284            "Process exited with unexpected code: {:x?}. Expected STATUS_ACCESS_VIOLATION.",
1285            status.code()
1286         );
1287         eprintln!("Test passed: Process correctly terminated with STATUS_ACCESS_VIOLATION.");
1288      }
1289   }
1290
1291   #[test]
1292   fn test_vec_u8_variety() {
1293      let data: Vec<u8> = vec![1, 2, 3, 4, 5];
1294      test_vec_generic_basics(&data);
1295   }
1296
1297   #[test]
1298   fn test_vec_u16() {
1299      let data: Vec<u16> = vec![1000, 2000, 3000];
1300      test_vec_generic_basics(&data);
1301   }
1302
1303   #[test]
1304   fn test_vec_u64() {
1305      let data: Vec<u64> = vec![0xDEADBEEF_CAFEBABE, 1, 2, 3, 4];
1306      test_vec_generic_basics(&data);
1307   }
1308
1309   #[test]
1310   fn test_vec_byte_array() {
1311      let data: Vec<[u8; 32]> = vec![[0xAB; 32], [0xCD; 32]];
1312      test_vec_generic_basics(&data);
1313   }
1314
1315   #[test]
1316   fn test_vec_small_struct() {
1317      let data = vec![
1318         SmallStruct { a: 10, b: 20 },
1319         SmallStruct { a: 30, b: 40 },
1320         SmallStruct { a: 50, b: 60 },
1321      ];
1322      test_vec_generic_basics(&data);
1323   }
1324
1325   #[test]
1326   fn test_vec_large_struct() {
1327      let data = vec![
1328         LargeStruct {
1329            data: [1, 2, 3, 4],
1330            flag: true,
1331         },
1332         LargeStruct {
1333            data: [10, 20, 30, 40],
1334            flag: false,
1335         },
1336      ];
1337      test_vec_generic_basics(&data);
1338   }
1339
1340   #[test]
1341   fn test_vec_person() {
1342      let data = vec![
1343         create_test_person(1),
1344         create_test_person(42),
1345         create_test_person(99),
1346      ];
1347      // test push which triggers reserve for > initial cap
1348      let mut pvec = SecureVec::new().unwrap();
1349      for p in &data {
1350         pvec.push(p.clone());
1351      }
1352      pvec.unlock_slice(|slice| {
1353         assert_eq!(slice.len(), 3);
1354         assert_eq!(slice[0].name, "Person1");
1355         assert_eq!(slice[2].notes, "Some secret notes for person #99");
1356      });
1357      println!("person push+realloc ok");
1358   }
1359
1360   #[test]
1361   fn test_vec_aligned_struct() {
1362      let data = vec![
1363         AlignedStruct {
1364            value: 0x1234_5678_9ABC_DEF0,
1365         },
1366         AlignedStruct { value: 42 },
1367      ];
1368      test_vec_generic_basics(&data);
1369   }
1370
1371   #[test]
1372   fn test_vec_mixed_sizes() {
1373      // Push many to force multiple reallocs with larger type
1374      let mut vec: SecureVec<u64> = SecureVec::new().unwrap();
1375      for i in 0..20u64 {
1376         vec.push(i * 1000);
1377      }
1378      vec.unlock_slice(|slice| {
1379         assert_eq!(slice.len(), 20);
1380         assert_eq!(slice[0], 0);
1381         assert_eq!(slice[19], 19000);
1382      });
1383   }
1384}