secure-types 0.3.0

// No_std: we only need `Layout` for computing allocation sizes.
// We call `alloc::alloc::dealloc` via fully-qualified path to avoid
// shadowing the crate-level `alloc::<T>()` helper.
#[cfg(not(feature = "use_os"))]
use alloc::alloc::Layout;

#[cfg(feature = "use_os")]
use std::vec::Vec;

use super::{Error, SecureArray, alloc};
use core::{
   marker::PhantomData,
   mem,
   ops::{Bound, RangeBounds},
   ptr::{self, NonNull},
};
use zeroize::{DefaultIsZeroes, Zeroize};

#[cfg(feature = "use_os")]
use super::free;
#[cfg(feature = "use_os")]
use memsec::Prot;

pub type SecureBytes = SecureVec<u8>;

/// Unlocks the vector's memory on construction and re-locks it on drop —
/// including when the drop happens because the fn closure panicked.
struct UnlockGuard<'a, T: Zeroize> {
   vec: &'a SecureVec<T>,
}

impl<'a, T: Zeroize> UnlockGuard<'a, T> {
   fn new(vec: &'a SecureVec<T>) -> Self {
      let ok = vec.unlock_memory();
      debug_assert!(ok, "UnlockGuard::new: unlock_memory failed");

      UnlockGuard { vec }
   }
}

impl<'a, T: Zeroize> Drop for UnlockGuard<'a, T> {
   fn drop(&mut self) {
      let ok = self.vec.lock_memory();
      debug_assert!(ok, "UnlockGuard::drop: lock_memory failed");
   }
}

/// A securely allocated, growable vector, just like `std::vec::Vec`.
///
/// ## Security Model
///
/// When compiled with the `use_os` feature (the default), it provides several layers of protection:
/// - **Zeroization on Drop**: The memory is zeroized when the vector is dropped.
/// - **Memory Locking**: The underlying memory pages are locked using `mlock` & `madvise` for (Unix) or
///   `VirtualLock` & `VirtualProtect` for (Windows) to prevent the OS from memory-dump/swap to disk or other processes accessing the memory.
///
/// In a `no_std` environment, it falls back to providing only the **zeroization-on-drop** guarantee.
///
/// ## Security Note on Direct Access
///
/// We intentionally do **not** implement `Index` / `IndexMut`.
/// Using `secure_vec[0]` is a compile error.
///
/// This is by design: direct indexing would allow bypassing the explicit
/// unlock mechanism. Always use `unlock_slice()` / `unlock_slice_mut()` (or
/// the `unlock*` family of methods) to access the contents.
///
/// # Notes
///
/// If you return a new allocated `Vec` from one of the unlock methods you are responsible for zeroizing the memory.
///
/// # Example
///
/// Using `SecureBytes` (a type alias for `SecureVec<u8>`) to handle a secret key.
///
/// ```
/// use secure_types::{SecureBytes, Zeroize};
///
/// // Create a new, empty secure vector.
/// let mut secret_key = SecureBytes::new().unwrap();
///
/// // Push some sensitive data into it.
/// secret_key.push(0xAB);
/// secret_key.push(0xCD);
/// secret_key.push(0xEF);
///
/// // The memory is locked here.
///
/// // Use a scope to safely access the contents as a slice.
/// secret_key.unlock_slice(|unlocked_slice| {
///     assert_eq!(unlocked_slice, &[0xAB, 0xCD, 0xEF]);
/// });
///
/// // Not recommended but if you allocate a new Vec make sure to zeroize it
/// let mut exposed = secret_key.unlock_slice(|unlocked_slice| {
///     Vec::from(unlocked_slice)
/// });
///
/// // Do what you need to to do with the new vector
/// // When you are done with it, zeroize it
/// exposed.zeroize();
///
/// // The memory is automatically locked again when the scope ends.
///
/// // When `secret_key` is dropped, its memory is securely zeroized.
/// ```
pub struct SecureVec<T>
where
   T: Zeroize,
{
   ptr: NonNull<T>,
   pub(crate) len: usize,
   pub(crate) capacity: usize,
   _marker: PhantomData<T>,
}

unsafe impl<T: Zeroize + Send> Send for SecureVec<T> {}
unsafe impl<T: Zeroize + Send + Sync> Sync for SecureVec<T> {}

impl<T: Zeroize> SecureVec<T> {
   /// Create a new `SecureVec` with a capacity of 1
   pub fn new() -> Result<Self, Error> {
      let capacity = 1;
      let size = capacity * mem::size_of::<T>();
      let ptr = unsafe { alloc::<T>(size)? };

      let secure = SecureVec {
         ptr,
         len: 0,
         capacity,
         _marker: PhantomData,
      };

      let _locked = secure.lock_memory();

      #[cfg(feature = "use_os")]
      if !_locked {
         return Err(Error::LockFailed);
      }

      Ok(secure)
   }

   /// Create a new `SecureVec` with the given capacity
   pub fn new_with_capacity(mut capacity: usize) -> Result<Self, Error> {
      if capacity == 0 {
         capacity = 1;
      }

      capacity
         .checked_mul(size_of::<T>())
         .ok_or(Error::AllocationFailed)?;

      let size = capacity * mem::size_of::<T>();
      let ptr = unsafe { alloc::<T>(size)? };

      let secure = SecureVec {
         ptr,
         len: 0,
         capacity,
         _marker: PhantomData,
      };

      let _locked = secure.lock_memory();

      #[cfg(feature = "use_os")]
      if !_locked {
         return Err(Error::LockFailed);
      }

      Ok(secure)
   }

   #[cfg(feature = "use_os")]
   /// Create a new `SecureVec` from a `Vec`
   ///
   /// The `Vec` is zeroized afterwards
   pub fn from_vec(mut vec: Vec<T>) -> Result<Self, Error> {
      if vec.capacity() == 0 {
         vec.reserve(1);
      }

      let capacity = vec.capacity();
      let len = vec.len();

      let size = match capacity.checked_mul(size_of::<T>()) {
         Some(s) => s,
         None => {
            vec.zeroize();
            return Err(Error::AllocationFailed);
         }
      };

      let ptr = match unsafe { alloc::<T>(size) } {
         Ok(ptr) => ptr,
         Err(_) => {
            vec.zeroize();
            return Err(Error::AllocationFailed);
         }
      };

      // Move data from the old vec into the secure allocation using ptr::read / ptr::write
      // This correctly transfers ownership for non-Copy types (e.g. structs containing String).
      // We then zero the *bytes* of the source buffer (after moving values out) to avoid
      // leaving sensitive data, and prevent double-drop by clearing the vec length.
      unsafe {
         let src = vec.as_ptr();
         let dst = ptr.as_ptr();
         for i in 0..len {
            let value = core::ptr::read(src.add(i));
            core::ptr::write(dst.add(i), value);
         }
      }

      // Prevent the Vec from dropping the now-moved-from elements (would be UB)
      // and securely erase whatever representation bytes remain in its buffer.
      //
      // We use set_len(0) + zeroize on a &mut [u8] view of the allocation
      // (instead of calling vec.zeroize()) because the Ts have been moved out
      // via ptr::read. The normal Vec::zeroize impl would zeroize+drop the
      // moved-from elements, which is UB (and often SIGABRT for a non-copy type).
      let old_byte_size = capacity * mem::size_of::<T>();
      unsafe {
         vec.set_len(0);
      }
      if old_byte_size > 0 {
         // SAFETY: after set_len(0) the allocation bytes are still valid,
         // we own them exclusively, and no Ts will be dropped by the Vec.
         let bytes =
            unsafe { core::slice::from_raw_parts_mut(vec.as_mut_ptr() as *mut u8, old_byte_size) };
         bytes.zeroize();
      }

      let secure = SecureVec {
         ptr,
         len,
         capacity,
         _marker: PhantomData,
      };

      let locked = secure.lock_memory();

      if !locked {
         return Err(Error::LockFailed);
      }

      Ok(secure)
   }

   /// Create a new `SecureVec` from a mutable slice.
   ///
   /// The slice is zeroized afterwards
   pub fn from_slice_mut(slice: &mut [T]) -> Result<Self, Error>
   where
      T: Clone + DefaultIsZeroes,
   {
      let mut secure_vec = match SecureVec::new_with_capacity(slice.len()) {
         Ok(secure_vec) => secure_vec,
         Err(e) => {
            slice.zeroize();
            return Err(e);
         }
      };

      secure_vec.init_from_clone(slice);
      slice.zeroize();

      Ok(secure_vec)
   }

   /// Create a new `SecureVec` from a slice.
   ///
   /// The slice is not zeroized, you are responsible for zeroizing it
   pub fn from_slice(slice: &[T]) -> Result<Self, Error>
   where
      T: Clone,
   {
      let mut secure_vec = SecureVec::new_with_capacity(slice.len())?;
      secure_vec.init_from_clone(slice);
      Ok(secure_vec)
   }

   pub fn len(&self) -> usize {
      self.len
   }

   pub fn is_empty(&self) -> bool {
      self.len() == 0
   }

   /// Returns the pointer to the locked memory region
   ///
   /// # DANGER
   ///
   /// This is a low-level API, which should be used only for
   /// testing purposes. If you need to access the locked memory
   /// region, use one of the unlock methods.
   #[cfg(feature = "expose-ptr")]
   #[deprecated(
      since = "0.3.0",
      note = "This method is intended only for testing/crash reproduction. Use one of the unlock methods instead."
   )]
   pub fn ptr(&self) -> NonNull<T> {
      self.ptr
   }

   /// Returns the total number of bytes currently allocated for this vector.
   #[cfg(not(feature = "use_os"))]
   pub(crate) fn allocated_byte_size(&self) -> usize {
      self.capacity * mem::size_of::<T>()
   }

   pub(crate) fn as_mut_ptr(&mut self) -> *mut u8 {
      self.ptr.as_ptr() as *mut u8
   }

   pub(crate) fn lock_memory(&self) -> bool {
      #[cfg(feature = "use_os")]
      {
         #[cfg(windows)]
         {
            super::mprotect(self.ptr, Prot::NoAccess)
         }
         #[cfg(unix)]
         {
            super::mprotect(self.ptr, Prot::NoAccess)
         }
      }
      #[cfg(not(feature = "use_os"))]
      {
         true // No-op: always "succeeds"
      }
   }

   pub(crate) fn unlock_memory(&self) -> bool {
      #[cfg(feature = "use_os")]
      {
         #[cfg(windows)]
         {
            super::mprotect(self.ptr, Prot::ReadWrite)
         }
         #[cfg(unix)]
         {
            super::mprotect(self.ptr, Prot::ReadWrite)
         }
      }

      #[cfg(not(feature = "use_os"))]
      {
         true // No-op: always "succeeds"
      }
   }

   /// Immutable access to the `SecureVec`
   pub fn unlock<F, R>(&self, f: F) -> R
   where
      F: FnOnce(&SecureVec<T>) -> R,
   {
      let _guard = UnlockGuard::new(self);
      let result = f(self);
      result
   }

   /// Immutable access to the `SecureVec` as `&[T]`
   pub fn unlock_slice<F, R>(&self, f: F) -> R
   where
      F: FnOnce(&[T]) -> R,
   {
      let _guard = UnlockGuard::new(self);
      let slice = unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.len) };
      f(slice)
   }

   /// Mutable access to the `SecureVec` as `&mut [T]`
   pub fn unlock_slice_mut<F, R>(&mut self, f: F) -> R
   where
      F: FnOnce(&mut [T]) -> R,
   {
      unsafe {
         let _guard = UnlockGuard::new(self);
         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
         let result = f(slice);
         result
      }
   }

   /// Immutable access to the `SecureVec` as `Iter<T>`
   pub fn unlock_iter<F, R>(&self, f: F) -> R
   where
      F: FnOnce(core::slice::Iter<T>) -> R,
   {
      unsafe {
         let _guard = UnlockGuard::new(self);
         let slice = core::slice::from_raw_parts(self.ptr.as_ptr(), self.len);
         let iter = slice.iter();
         let result = f(iter);
         result
      }
   }

   /// Mutable access to the `SecureVec` as `IterMut<T>`
   pub fn unlock_iter_mut<F, R>(&mut self, f: F) -> R
   where
      F: FnOnce(core::slice::IterMut<T>) -> R,
   {
      unsafe {
         let _guard = UnlockGuard::new(self);
         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
         let iter = slice.iter_mut();
         let result = f(iter);
         result
      }
   }

   /// Erase the underlying data and clears the vector
   ///
   /// The memory is locked again and the capacity is preserved for reuse
   pub fn erase(&mut self) {
      unsafe {
         let ok = self.unlock_memory();
         debug_assert!(ok, "SecureVec::erase: unlock_memory failed");

         // Only zero the initialized elements. Zeroizing capacity would try to
         // zeroize uninitialized memory as T, which for Drop types (eg. String)
         // is UB and causes SIGSEGV/SIGABRT.
         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
         for elem in slice.iter_mut() {
            elem.zeroize();
         }

         self.clear();

         let ok = self.lock_memory();
         debug_assert!(ok, "SecureVec::erase: lock_memory failed");
      }
   }

   /// Clear the vector
   ///
   /// This just sets the vector's len to zero it does not erase the underlying data
   pub fn clear(&mut self) {
      self.len = 0;
   }

   pub fn push(&mut self, value: T) {
      self.reserve(1);

      let ok = self.unlock_memory();
      debug_assert!(ok, "SecureVec::push: unlock_memory failed");

      unsafe {
         // Write the new value at the end of the vector.
         core::ptr::write(self.ptr.as_ptr().add(self.len), value);

         self.len += 1;
      }

      let ok = self.lock_memory();
      debug_assert!(ok, "SecureVec::push: lock_memory failed");
   }

   /// Ensures that the vector has enough capacity for at least `additional` more elements.
   ///
   /// If more capacity is needed, it will reallocate. This may cause the buffer location to change.
   ///
   /// # Panics
   ///
   /// Panics if the new capacity overflows `usize` or if the allocation fails.
   pub fn reserve(&mut self, additional: usize) {
      if self.len() + additional <= self.capacity {
         return;
      }

      // Use an amortized growth strategy to avoid reallocating on every push
      let required_capacity = self.len() + additional;
      let new_capacity = (self.capacity.max(1) * 2).max(required_capacity);

      let new_size = new_capacity * mem::size_of::<T>();

      // Safe to panic here because the memory is locked
      let new_ptr = unsafe {
         alloc::<T>(new_size).unwrap_or_else(|_| {
            panic!(
               "secure-types: failed to allocate {} bytes of locked memory \
          (possibly RLIMIT_MEMLOCK exhausted); SecureVec left unchanged",
               new_size
            )
         })
      };

      // Copy data to new pointer
      unsafe {
         let ok = self.unlock_memory();
         debug_assert!(ok, "SecureVec::reserve: unlock_memory failed");

         // Move (not copy) elements to new buffer to support non-Copy T correctly.
         // Using read+write transfers ownership of e.g. Strings.
         let len = self.len();
         for i in 0..len {
            let val = core::ptr::read(self.ptr.as_ptr().add(i));
            core::ptr::write(new_ptr.as_ptr().add(i), val);
         }

         // Erase old buffer bytes (after move-out)
         if self.capacity > 0 {
            let old_bytes = self.capacity * mem::size_of::<T>();
            let bytes = core::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut u8, old_bytes);
            bytes.zeroize();
         }

         #[cfg(feature = "use_os")]
         free(self.ptr);

         #[cfg(not(feature = "use_os"))]
         {
            let old_size = self.capacity * mem::size_of::<T>();
            let old_layout = Layout::from_size_align_unchecked(old_size, mem::align_of::<T>());
            alloc::alloc::dealloc(self.ptr.as_ptr() as *mut u8, old_layout);
         }
      }

      // Update pointer and capacity, then re-lock the new memory region
      self.ptr = new_ptr;
      self.capacity = new_capacity;
      let ok = self.lock_memory();
      debug_assert!(ok, "SecureVec::reserve: lock_memory failed");
   }

   /// Creates a draining iterator that removes the specified range from the vector
   /// and yields the removed items.
   ///
   /// Note: The vector is unlocked during the lifetime of the `Drain` iterator.
   /// The memory is relocked when the `Drain` iterator is dropped.
   ///
   /// # Panics
   /// Panics if the starting point is greater than the end point or if the end point
   /// is greater than the length of the vector.
   pub fn drain<R>(&mut self, range: R) -> Drain<'_, T>
   where
      R: RangeBounds<usize>,
   {
      let original_len = self.len;

      let (drain_start_idx, drain_end_idx) = resolve_range_indices(range, original_len);

      let tail_len = original_len - drain_end_idx;

      self.len = drain_start_idx;

      let ok = self.unlock_memory();
      debug_assert!(ok, "SecureVec::drain: unlock_memory failed");

      Drain {
         vec_ref: self,
         drain_start_index: drain_start_idx,
         current_drain_iter_index: drain_start_idx,
         drain_end_index: drain_end_idx,
         original_vec_len: original_len,
         tail_len,
         _marker: PhantomData,
      }
   }

   /// Initializes a freshly-allocated (uninitialized) buffer by cloning `src`
   /// into it. Uses `ptr::write` so the uninitialized destination slots are
   /// never read, never dropped, and no `&mut [T]` is ever formed over them.
   ///
   /// `len` is set only after every write succeeds, so a panic from
   /// `T::clone` leaves the vector at its previous length (0 for a fresh one).
   pub(crate) fn init_from_clone(&mut self, src: &[T])
   where
      T: Clone,
   {
      debug_assert!(src.len() <= self.capacity);

      let ok = self.unlock_memory();
      debug_assert!(
         ok,
         "SecureVec::init_from_clone: unlock_memory failed"
      );

      unsafe {
         let dst = self.ptr.as_ptr();
         for (i, item) in src.iter().enumerate() {
            core::ptr::write(dst.add(i), item.clone());
         }
      }

      self.len = src.len();
      let ok = self.lock_memory();
      debug_assert!(
         ok,
         "SecureVec::init_from_clone: lock_memory failed"
      );
   }
}

impl<T: Clone + Zeroize> Clone for SecureVec<T> {
   fn clone(&self) -> Self {
      let mut new_vec = SecureVec::new_with_capacity(self.capacity).unwrap();
      self.unlock_slice(|src_slice| {
         new_vec.init_from_clone(src_slice);
      });
      new_vec
   }
}

impl<const LENGTH: usize> From<SecureArray<u8, LENGTH>> for SecureVec<u8> {
   fn from(array: SecureArray<u8, LENGTH>) -> Self {
      let mut new_vec = SecureVec::new_with_capacity(LENGTH)
         .expect("Failed to allocate SecureVec during conversion");
      array.unlock(|array_slice| {
         new_vec.init_from_clone(array_slice);
      });
      new_vec
   }
}

impl<T: Zeroize> Drop for SecureVec<T> {
   fn drop(&mut self) {
      unsafe {
         let ok = self.unlock_memory();
         debug_assert!(ok, "SecureVec::erase: unlock_memory failed");

         // Only zero the initialized elements. Zeroizing capacity would try to
         // zeroize uninitialized memory as T, which for Drop types (eg. String)
         // is UB and causes SIGSEGV/SIGABRT.
         let slice = core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len);
         for elem in slice.iter_mut() {
            elem.zeroize();
         }
      }

      #[cfg(feature = "use_os")]
      free(self.ptr);

      #[cfg(not(feature = "use_os"))]
      unsafe {
         let layout =
            Layout::from_size_align_unchecked(self.allocated_byte_size(), mem::align_of::<T>());
         alloc::alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
      }
   }
}

// Note: We intentionally do **not** implement Index / IndexMut.
// Direct indexing (`vec[0]`) would bypass the unlock mechanism and
// access locked memory, causing a segfault. This is by design.
// Always use unlock_slice() / unlock_slice_mut().

#[cfg(feature = "serde")]
impl serde::Serialize for SecureVec<u8> {
   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
   where
      S: serde::Serializer,
   {
      self.unlock_slice(|slice| serializer.collect_seq(slice.iter()))
   }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for SecureVec<u8> {
   fn deserialize<D>(deserializer: D) -> Result<SecureVec<u8>, D::Error>
   where
      D: serde::Deserializer<'de>,
   {
      struct SecureVecVisitor;
      impl<'de> serde::de::Visitor<'de> for SecureVecVisitor {
         type Value = SecureVec<u8>;
         fn expecting(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
            write!(formatter, "a sequence of bytes")
         }
         fn visit_seq<A>(
            self,
            mut seq: A,
         ) -> Result<<Self as serde::de::Visitor<'de>>::Value, A::Error>
         where
            A: serde::de::SeqAccess<'de>,
         {
            let mut vec = SecureVec::new().map_err(serde::de::Error::custom)?;
            while let Some(byte) = seq.next_element::<u8>()? {
               vec.push(byte);
            }
            Ok(vec)
         }
      }
      deserializer.deserialize_seq(SecureVecVisitor)
   }
}

/// A draining iterator for `SecureVec<T>`.
///
/// This struct is created by the `drain` method on `SecureVec`.
///
/// # Safety
/// The returned `Drain` iterator must not be forgotten (via `mem::forget`).
/// Forgetting the iterator sets the len of `SecureVec` to 0 and the memory will remain unlocked
pub struct Drain<'a, T: Zeroize + 'a> {
   vec_ref: &'a mut SecureVec<T>,
   drain_start_index: usize,
   current_drain_iter_index: usize,
   drain_end_index: usize,

   original_vec_len: usize, // Original length of vec_ref before drain
   tail_len: usize,         // Number of elements after the drain range in the original vec

   _marker: PhantomData<&'a T>,
}

impl<'a, T: Zeroize> Iterator for Drain<'a, T> {
   type Item = T;

   fn next(&mut self) -> Option<T> {
      if self.current_drain_iter_index < self.drain_end_index {
         // SecureVec is already unlocked by the `drain` method.
         unsafe {
            let item_ptr = self.vec_ref.ptr.as_ptr().add(self.current_drain_iter_index);
            let item = ptr::read(item_ptr);
            self.current_drain_iter_index += 1;
            Some(item)
         }
      } else {
         None
      }
   }

   fn size_hint(&self) -> (usize, Option<usize>) {
      let remaining = self.drain_end_index - self.current_drain_iter_index;
      (remaining, Some(remaining))
   }
}

impl<'a, T: Zeroize> ExactSizeIterator for Drain<'a, T> {}

impl<'a, T: Zeroize> Drop for Drain<'a, T> {
   fn drop(&mut self) {
      unsafe {
         // The vec_ref's memory is currently unlocked.
         if mem::needs_drop::<T>() {
            let mut current_ptr = self.vec_ref.ptr.as_ptr().add(self.current_drain_iter_index);
            let end_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_end_index);
            while current_ptr < end_ptr {
               ptr::drop_in_place(current_ptr);
               current_ptr = current_ptr.add(1);
            }
         }

         let hole_dst_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_start_index);
         let tail_src_ptr = self.vec_ref.ptr.as_ptr().add(self.drain_end_index);

         if self.tail_len > 0 {
            ptr::copy(tail_src_ptr, hole_dst_ptr, self.tail_len);
         }

         // The new length of the vector.
         let new_len = self.drain_start_index + self.tail_len;

         // Process the memory region that is no longer part of the active vector's content.
         // This region is from `vec_ref.ptr + new_len` up to `vec_ref.ptr + original_vec_len`.
         // It contains:
         //    a) Original data of the latter part of the drained slice (if not overwritten by tail).
         //       These were dropped in step 1 if T:Drop.
         //    b) Original data of the tail items (which have now been copied).
         //       These need to be dropped if T:Drop, as ptr::copy doesn't drop the source.
         // After any necessary drops, this entire region must be zeroized.

         let mut current_cleanup_ptr = self.vec_ref.ptr.as_ptr().add(new_len);
         let end_cleanup_ptr = self.vec_ref.ptr.as_ptr().add(self.original_vec_len);

         // Determine the start of the original tail's memory region
         let original_tail_start_ptr_val = tail_src_ptr as usize;

         while current_cleanup_ptr < end_cleanup_ptr {
            if mem::needs_drop::<T>() {
               let current_ptr_val = current_cleanup_ptr as usize;
               let original_tail_end_ptr_val =
                  original_tail_start_ptr_val + self.tail_len * mem::size_of::<T>();

               if current_ptr_val >= original_tail_start_ptr_val
                  && current_ptr_val < original_tail_end_ptr_val
               {
                  // This element was part of the original tail. ptr::copy moved its value.
                  // The original instance here needs to be dropped.
                  ptr::drop_in_place(current_cleanup_ptr);
               }
               // Else, it was part of the drained range (not covered by tail move).
               // If it needed dropping, it was handled in step 1.
            }

            // Zeroize the memory of this element.
            (*current_cleanup_ptr).zeroize();
            current_cleanup_ptr = current_cleanup_ptr.add(1);
         }

         // Update the SecureVec's length.
         self.vec_ref.len = new_len;

         // Relock the SecureVec's memory.
         let ok = self.vec_ref.lock_memory();
         debug_assert!(ok, "Drain::drop: lock_memory failed");
      }
   }
}

// Helper function to resolve RangeBounds to (start, end) indices
fn resolve_range_indices<R: RangeBounds<usize>>(range: R, len: usize) -> (usize, usize) {
   let start_bound = range.start_bound();
   let end_bound = range.end_bound();

   let start = match start_bound {
      Bound::Included(&s) => s,
      Bound::Excluded(&s) => s
         .checked_add(1)
         .unwrap_or_else(|| panic!("attempted to start drain at Excluded(usize::MAX)")),
      Bound::Unbounded => 0,
   };

   let end = match end_bound {
      Bound::Included(&e) => e
         .checked_add(1)
         .unwrap_or_else(|| panic!("attempted to end drain at Included(usize::MAX)")),
      Bound::Excluded(&e) => e,
      Bound::Unbounded => len,
   };

   if start > end {
      panic!(
         "drain range start ({}) must be less than or equal to end ({})",
         start, end
      );
   }
   if end > len {
      panic!(
         "drain range end ({}) out of bounds for slice of length {}",
         end, len
      );
   }

   (start, end)
}

#[cfg(all(test, feature = "use_os"))]
mod tests {
   use super::*;
   use std::fmt::Debug;
   use std::process::{Command, Stdio};
   use std::sync::{Arc, Mutex};
   use zeroize::Zeroize;

   // Test helper types for variety (different sizes, alignments, complex data)

   #[derive(Clone, Debug, PartialEq)]
   struct SmallStruct {
      a: u8,
      b: u16,
   }

   impl Zeroize for SmallStruct {
      fn zeroize(&mut self) {
         self.a.zeroize();
         self.b.zeroize();
      }
   }

   #[derive(Clone, Debug, PartialEq)]
   struct LargeStruct {
      data: [u64; 4],
      flag: bool,
   }

   impl Zeroize for LargeStruct {
      fn zeroize(&mut self) {
         self.data.zeroize();
         self.flag.zeroize();
      }
   }

   #[derive(Clone, Debug, PartialEq)]
   #[repr(align(64))]
   struct AlignedStruct {
      value: u64,
   }

   impl Zeroize for AlignedStruct {
      fn zeroize(&mut self) {
         self.value.zeroize();
      }
   }

   #[derive(Clone, Debug, PartialEq)]
   struct Person {
      name: String,
      age: u32,
      notes: String,
   }

   impl Zeroize for Person {
      fn zeroize(&mut self) {
         self.name.zeroize();
         self.age.zeroize();
         self.notes.zeroize();
      }
   }

   impl Person {
      fn new(name: impl Into<String>, age: u32, notes: impl Into<String>) -> Self {
         Self {
            name: name.into(),
            age,
            notes: notes.into(),
         }
      }
   }

   fn create_test_person(id: usize) -> Person {
      Person::new(
         format!("Person{}", id),
         (id % 100) as u32,
         format!("Some secret notes for person #{}", id),
      )
   }

   fn test_vec_generic_basics<T: Zeroize + Clone + PartialEq + Debug>(initial: &[T]) {
      if initial.is_empty() {
         return;
      }

      // from_slice
      let secure = SecureVec::from_slice(initial).unwrap();
      assert_eq!(secure.len(), initial.len());
      secure.unlock_slice(|slice| {
         assert_eq!(slice, initial);
      });

      // new + push
      let mut secure_push = SecureVec::new().unwrap();
      for item in initial {
         secure_push.push(item.clone());
      }
      secure_push.unlock_slice(|slice| {
         assert_eq!(slice, initial);
      });

      // clone
      let cloned = secure.clone();
      secure.unlock_slice(|s| {
         cloned.unlock_slice(|c| {
            assert_eq!(s, c);
         });
      });

      // reserve + push more
      let mut res = SecureVec::new().unwrap();
      res.reserve(initial.len() + 2);
      for item in initial {
         res.push(item.clone());
      }
      assert!(res.capacity >= initial.len() + 2 || res.capacity >= initial.len());
      res.unlock_slice(|slice| {
         assert_eq!(slice, initial);
      });

      // erase
      res.erase();
      res.unlock(|v| {
         assert_eq!(v.len, 0);
         // capacity should be preserved
         assert!(v.capacity > 0);
      });

      // from_vec
      let vec_data: Vec<T> = initial.to_vec();
      let from_vec_secure = SecureVec::from_vec(vec_data).unwrap();
      from_vec_secure.unlock_slice(|slice| {
         assert_eq!(slice, initial);
      });
   }

   #[test]
   fn test_creation() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let secure_vec = SecureVec::from_vec(vec).unwrap();

      secure_vec.unlock_slice(|slice| {
         assert_eq!(slice, &[1, 2, 3]);
      });

      let exposed_slice = &mut [1, 2, 3];
      let secure_slice = SecureVec::from_slice_mut(exposed_slice).unwrap();
      assert_eq!(exposed_slice, &[0u8; 3]);

      secure_slice.unlock_slice(|slice| {
         assert_eq!(slice, &[1, 2, 3]);
      });

      let exposed_slice = [1, 2, 3];
      let secure_slice = SecureVec::from_slice(&exposed_slice).unwrap();

      secure_slice.unlock_slice(|slice| {
         assert_eq!(slice, exposed_slice);
      });
   }

   #[test]
   fn test_from_secure_array() {
      let exposed: &mut [u8; 3] = &mut [1, 2, 3];
      let array: SecureArray<u8, 3> = SecureArray::from_slice_mut(exposed).unwrap();
      let vec: SecureVec<u8> = array.into();
      assert_eq!(vec.len(), 3);
      vec.unlock_slice(|slice| {
         assert_eq!(slice, &[1, 2, 3]);
      });
   }

   #[test]
   fn lock_unlock_works() {
      let secure: SecureVec<u8> = SecureVec::new().unwrap();

      let unlocked = secure.unlock_memory();
      assert!(unlocked);

      let locked = secure.lock_memory();
      assert!(locked);
   }

   #[test]
   fn test_thread_safety() {
      let vec: Vec<u8> = vec![];
      let secure = SecureVec::from_vec(vec).unwrap();
      let secure = Arc::new(Mutex::new(secure));

      let mut handles = Vec::new();
      for i in 0..5u8 {
         let secure_clone = secure.clone();
         let handle = std::thread::spawn(move || {
            let mut secure = secure_clone.lock().unwrap();
            secure.push(i);
         });
         handles.push(handle);
      }

      for handle in handles {
         handle.join().unwrap();
      }

      let mut sec = secure.lock().unwrap();
      sec.unlock_slice_mut(|slice| {
         slice.sort();
         assert_eq!(slice.len(), 5);
         assert_eq!(slice, &[0, 1, 2, 3, 4]);
      });
   }

   #[test]
   fn test_clone() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let secure1 = SecureVec::from_vec(vec).unwrap();
      let secure2 = secure1.clone();

      secure1.unlock_slice(|slice| {
         secure2.unlock_slice(|slice2| {
            assert_eq!(slice, slice2);
         });
      });
   }

   #[test]
   fn test_do_not_call_forget_on_drain() {
      let vec: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
      let mut secure = SecureVec::from_vec(vec).unwrap();
      let drain = secure.drain(..3);
      core::mem::forget(drain);
      // we can still use secure vec but its state is unreachable
      secure.unlock_slice(|secure| {
         assert_eq!(secure.len(), 0);
      });
   }

   #[test]
   fn test_drain() {
      let vec: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
      let mut secure = SecureVec::from_vec(vec).unwrap();
      let mut drain = secure.drain(..3);
      assert_eq!(drain.next(), Some(1));
      assert_eq!(drain.next(), Some(2));
      assert_eq!(drain.next(), Some(3));
      assert_eq!(drain.next(), None);
      drop(drain);
      secure.unlock_slice(|secure| {
         assert_eq!(secure.len(), 7);
         assert_eq!(secure, &[4, 5, 6, 7, 8, 9, 10]);
      });
   }

   #[cfg(feature = "serde")]
   #[test]
   fn test_secure_vec_serde() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let secure = SecureVec::from_vec(vec).unwrap();
      let json = serde_json::to_vec(&secure).expect("Serialization failed");
      let deserialized: SecureVec<u8> =
         serde_json::from_slice(&json).expect("Deserialization failed");
      deserialized.unlock_slice(|slice| {
         assert_eq!(slice, &[1, 2, 3]);
      });
   }

   #[test]
   fn test_erase() {
      let mut secure = SecureVec::new_with_capacity(10).unwrap();
      for i in 0..9 {
         secure.push(i);
      }

      secure.erase();

      secure.unlock(|secure| {
         assert_eq!(secure.len, 0);
         assert_eq!(secure.capacity, 10);
      });

      secure.unlock_iter(|iter| {
         for elem in iter {
            assert_eq!(elem, &0);
         }
      });
   }

   #[test]
   fn test_push() {
      let vec: Vec<u8> = Vec::new();
      let mut secure = SecureVec::from_vec(vec).unwrap();
      for i in 0..10 {
         secure.push(i);
      }

      assert_eq!(secure.len(), 10);

      secure.unlock_slice(|slice| {
         assert_eq!(slice, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
      });
   }

   #[test]
   fn test_reserve() {
      let mut secure: SecureVec<u8> = SecureVec::new().unwrap();
      secure.reserve(10);
      assert_eq!(secure.capacity, 10);
   }

   #[test]
   fn test_reserve_doubling() {
      let mut secure: SecureVec<u8> = SecureVec::new().unwrap();
      secure.reserve(10);

      for i in 0..9 {
         secure.push(i);
      }

      secure.push(9);
      assert_eq!(secure.capacity, 10);
      assert_eq!(secure.len(), 10);

      secure.push(10);
      assert_eq!(secure.capacity, 20);
      assert_eq!(secure.len(), 11);
   }

   #[test]
   fn test_unlock_gives_access() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let secure = SecureVec::from_vec(vec).unwrap();
      secure.unlock_slice(|slice| {
         assert_eq!(slice[0], 1);
         assert_eq!(slice[1], 2);
         assert_eq!(slice[2], 3);
      });
   }

   #[test]
   fn test_unlock_slice() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let secure = SecureVec::from_vec(vec).unwrap();
      secure.unlock_slice(|slice| {
         assert_eq!(slice, &[1, 2, 3]);
      });
   }

   #[test]
   fn test_unlock_slice_mut() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let mut secure = SecureVec::from_vec(vec).unwrap();

      secure.unlock_slice_mut(|slice| {
         slice[0] = 4;
         assert_eq!(slice, &mut [4, 2, 3]);
      });
   }

   #[test]
   fn test_unlock_iter() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let secure = SecureVec::from_vec(vec).unwrap();
      let sum: u8 = secure.unlock_iter(|iter| iter.map(|&x| x).sum());

      assert_eq!(sum, 6);

      let secure: SecureVec<u8> = SecureVec::new_with_capacity(3).unwrap();
      let sum: u8 = secure.unlock_iter(|iter| iter.map(|&x| x).sum());

      assert_eq!(sum, 0);
   }

   #[test]
   fn test_unlock_iter_mut() {
      let vec: Vec<u8> = vec![1, 2, 3];
      let mut secure = SecureVec::from_vec(vec).unwrap();
      secure.unlock_iter_mut(|iter| {
         for elem in iter {
            *elem += 1;
         }
      });

      secure.unlock_slice(|slice| {
         assert_eq!(slice, &[2, 3, 4]);
      });
   }

   #[test]
   fn test_index_should_fail_when_locked() {
      let arg = "CRASH_TEST_SECUREVEC_LOCKED";

      if std::env::args().any(|a| a == arg) {
         let vec: Vec<u8> = vec![1, 2, 3];
         let secure = SecureVec::from_vec(vec).unwrap();
         // Deliberately dereference the locked pointer to test that
         // the security model (mlock + no normal access) works as expected.
         let _value = unsafe { core::hint::black_box(*secure.ptr.as_ptr()) };

         std::process::exit(1);
      }

      let child = Command::new(std::env::current_exe().unwrap())
         .arg("vec::tests::test_index_should_fail_when_locked")
         .arg(arg)
         .arg("--nocapture")
         .stdout(Stdio::piped())
         .stderr(Stdio::piped())
         .spawn()
         .expect("Failed to spawn child process");

      let output = child.wait_with_output().expect("Failed to wait on child");
      let status = output.status;

      assert!(
         !status.success(),
         "Process exited successfully with code {:?}, but it should have crashed.",
         status.code()
      );

      #[cfg(unix)]
      {
         use std::os::unix::process::ExitStatusExt;
         let signal = status
            .signal()
            .expect("Process was not terminated by a signal on Unix.");
         assert!(
            signal == libc::SIGSEGV || signal == libc::SIGBUS,
            "Process terminated with unexpected signal: {}",
            signal
         );
         println!(
            "Test passed: Process correctly terminated with signal {}.",
            signal
         );
      }

      #[cfg(windows)]
      {
         const STATUS_ACCESS_VIOLATION: i32 = 0xC0000005_u32 as i32;
         assert_eq!(
            status.code(),
            Some(STATUS_ACCESS_VIOLATION),
            "Process exited with unexpected code: {:x?}. Expected STATUS_ACCESS_VIOLATION.",
            status.code()
         );
         eprintln!("Test passed: Process correctly terminated with STATUS_ACCESS_VIOLATION.");
      }
   }

   #[test]
   fn test_vec_u8_variety() {
      let data: Vec<u8> = vec![1, 2, 3, 4, 5];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_u16() {
      let data: Vec<u16> = vec![1000, 2000, 3000];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_u64() {
      let data: Vec<u64> = vec![0xDEADBEEF_CAFEBABE, 1, 2, 3, 4];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_byte_array() {
      let data: Vec<[u8; 32]> = vec![[0xAB; 32], [0xCD; 32]];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_small_struct() {
      let data = vec![
         SmallStruct { a: 10, b: 20 },
         SmallStruct { a: 30, b: 40 },
         SmallStruct { a: 50, b: 60 },
      ];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_large_struct() {
      let data = vec![
         LargeStruct {
            data: [1, 2, 3, 4],
            flag: true,
         },
         LargeStruct {
            data: [10, 20, 30, 40],
            flag: false,
         },
      ];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_person() {
      let data = vec![
         create_test_person(1),
         create_test_person(42),
         create_test_person(99),
      ];
      // test push which triggers reserve for > initial cap
      let mut pvec = SecureVec::new().unwrap();
      for p in &data {
         pvec.push(p.clone());
      }
      pvec.unlock_slice(|slice| {
         assert_eq!(slice.len(), 3);
         assert_eq!(slice[0].name, "Person1");
         assert_eq!(slice[2].notes, "Some secret notes for person #99");
      });
      println!("person push+realloc ok");
   }

   #[test]
   fn test_vec_aligned_struct() {
      let data = vec![
         AlignedStruct {
            value: 0x1234_5678_9ABC_DEF0,
         },
         AlignedStruct { value: 42 },
      ];
      test_vec_generic_basics(&data);
   }

   #[test]
   fn test_vec_mixed_sizes() {
      // Push many to force multiple reallocs with larger type
      let mut vec: SecureVec<u64> = SecureVec::new().unwrap();
      for i in 0..20u64 {
         vec.push(i * 1000);
      }
      vec.unlock_slice(|slice| {
         assert_eq!(slice.len(), 20);
         assert_eq!(slice[0], 0);
         assert_eq!(slice[19], 19000);
      });
   }
}