secure-types 0.5.16

use super::{
   Error,
   vec::{SecureVec, UnlockGuard},
};
use core::ops::Range;
use zeroize::Zeroize;

// `String` is only used by the serde visitor below; in a `no_std` build it has to come
// from `alloc` (with `use_os` the prelude provides it).
#[cfg(all(feature = "serde", not(feature = "use_os")))]
use alloc::string::String;

/// A securely allocated, growable UTF-8 string, just like `std::string::String`.
///
/// It is a wrapper around [SecureVec<u8>] and inherits all of its security guarantees.
///
/// Access to the string contents is provided through scoped methods like `unlock_str`,
/// which ensure the memory is only unlocked for the briefest possible time.
///
/// # Thread Safety
///
/// Same as [`SecureVec`]: `Send` but not `Sync`. Share as `Arc<Mutex<SecureString>>`.
///
/// # Notes
///
/// If you return a new allocated `String` from one of the unlock methods you are responsible for zeroizing the memory.
///
/// # Example
///
/// ```
/// use secure_types::{SecureString, Zeroize};
///
/// // Create a SecureString
/// let mut secret = SecureString::from("my_super_secret");
///
/// // The memory is locked here
///
/// // Safely append more data.
/// secret.push_str("_password");
///
/// // The memory is locked here.
///
/// // Use a scope to safely access the content as a &str.
/// secret.unlock_str(|exposed_str| {
///     assert_eq!(exposed_str, "my_super_secret_password");
/// });
///
/// // Not recommended but if you allocate a new String make sure to zeroize it
/// let mut exposed = secret.unlock_str(|exposed_str| {
///     String::from(exposed_str)
/// });
///
/// // Do what you need to to do with the new string
/// // When you are done with it, zeroize it
/// exposed.zeroize();
///
/// // When `secret` is dropped, its data zeroized.
/// ```
#[derive(Clone)]
pub struct SecureString {
   vec: SecureVec<u8>,
}

impl SecureString {
   pub fn new() -> Result<Self, Error> {
      let vec = SecureVec::new()?;
      Ok(SecureString { vec })
   }

   pub fn new_with_capacity(capacity: usize) -> Result<Self, Error> {
      let vec = SecureVec::new_with_capacity(capacity)?;
      Ok(SecureString { vec })
   }

   /// Creates a `SecureString` from a `SecureVec<u8>` without checking UTF-8.
   ///
   /// # Safety
   /// The caller must guarantee `vec` holds valid UTF-8. Violating this breaks
   /// the `SecureString` invariant and will make `unlock_str`/`char_len`/serde
   /// panic.
   pub unsafe fn from_utf8_unchecked(vec: SecureVec<u8>) -> SecureString {
      SecureString { vec }
   }

   pub fn erase(&mut self) {
      self.vec.erase();
   }

   /// Returns the length of the inner `SecureVec`
   ///
   /// If you want the character length use [`char_len`](Self::char_len)
   pub fn byte_len(&self) -> usize {
      self.vec.len()
   }

   pub fn is_empty(&self) -> bool {
      self.vec.is_empty()
   }

   /// Removes the specified byte range from the string.
   ///
   /// `range` is a **byte** range, not a character range: take it from `&str`
   /// byte offsets. For a character range use
   /// [`delete_text_char_range`](Self::delete_text_char_range).
   ///
   /// # Panics
   /// Panics if the range is not on UTF-8 char boundaries.
   pub fn drain(&mut self, range: Range<usize>) {
      self.unlock_str(|s| {
         assert!(
            s.is_char_boundary(range.start) && s.is_char_boundary(range.end),
            "SecureString::drain: range {:?} does not lie on UTF-8 char boundaries",
            range
         );
      });
      let _d = self.vec.drain(range);
   }

   /// Returns the number of chars in the string
   ///
   /// # Panics
   /// Panics if the string is not valid UTF-8.
   pub fn char_len(&self) -> usize {
      self.unlock_str(|s| s.chars().count())
   }

   /// Returns the number of chars in the string
   ///
   /// # Safety
   /// The caller must guarantee that the string is valid UTF-8.
   pub fn char_len_unchecked(&self) -> usize {
      self.unlock_str_unchecked(|s| s.chars().count())
   }

   /// Push a `&str` into the `SecureString`
   pub fn push_str(&mut self, string: &str) {
      let slice = string.as_bytes();
      for s in slice.iter() {
         self.vec.push(*s);
      }
   }

   /// Immutable access as `&str`
   ///
   /// It uses the `from_utf8` function to check the validity of the internal
   /// bytes. If the bytes are not valid UTF-8, the function panics.
   pub fn unlock_str<F, R>(&self, f: F) -> R
   where
      F: FnOnce(&str) -> R,
   {
      self.vec.unlock_slice(|slice| {
         let str = core::str::from_utf8(slice)
            .expect("SecureString invariant violated: internal bytes are not valid UTF-8");
         f(str)
      })
   }

   /// Immutable access as `&str`
   ///
   /// It uses the `from_utf8_unchecked` function to bypass the validity check.
   ///
   /// # Safety
   /// The caller must guarantee that the internal bytes are valid UTF-8.
   pub fn unlock_str_unchecked<F, R>(&self, f: F) -> R
   where
      F: FnOnce(&str) -> R,
   {
      self.vec.unlock_slice(|slice| {
         // SAFETY: this is the `unchecked` variant — the caller promised the
         // internal bytes are valid UTF-8.
         let str = unsafe { core::str::from_utf8_unchecked(slice) };
         f(str)
      })
   }

   /// Mutable access to the `SecureString`
   ///
   /// This method does not unlock the memory.
   pub fn secure_mut<F, R>(&mut self, f: F) -> R
   where
      F: FnOnce(&mut SecureString) -> R,
   {
      f(self)
   }

   /// Inserts text at the given character index
   ///
   /// # Returns
   ///
   /// The number of characters inserted
   ///
   /// # Example
   ///
   /// ```
   /// use secure_types::SecureString;
   ///
   /// let mut string = SecureString::from("GreekFeta");
   /// string.insert_text_at_char_idx(9, "Cheese");
   /// string.unlock_str(|str| {
   ///     assert_eq!(str, "GreekFetaCheese");
   /// });
   /// ```
   pub fn insert_text_at_char_idx(&mut self, char_idx: usize, text_to_insert: &str) -> usize {
      let chars_to_insert_count = text_to_insert.chars().count();
      if chars_to_insert_count == 0 {
         return 0;
      }

      let bytes_to_insert = text_to_insert.as_bytes();
      let insert_len = bytes_to_insert.len();

      // Get the byte index corresponding to the character index
      let byte_idx = self
         .vec
         .unlock_slice(|current_bytes| char_to_byte_idx(current_bytes, char_idx));

      self.vec.reserve(insert_len);

      let old_byte_len = self.vec.len();

      // Taken before the guard borrows the vector: raw pointers do not keep the
      // borrow alive, and the guard must stay alive while writing through it.
      let ptr = self.vec.as_mut_ptr();

      // Perform the insertion in-place. `UnlockGuard` re-locks the memory on
      // drop, including when the block unwinds.
      let new_byte_len = {
         let _guard = UnlockGuard::new(&self.vec);

         // SAFETY: the guard has unprotected `self.vec`'s live buffer for this
         // block. `reserve(insert_len)` above guaranteed room for
         // `old_byte_len + insert_len` bytes and `byte_idx <= old_byte_len`, so
         // both the shifted tail and the inserted span stay inside the
         // allocation. `ptr::copy` tolerates source/destination overlap; the
         // `copy_nonoverlapping` source is `text_to_insert`, a distinct `&str`.
         unsafe {
            // Shift the "tail" of the string (from the insertion point to the end)
            // to the right to make a gap for the new content.
            if byte_idx < old_byte_len {
               core::ptr::copy(
                  ptr.add(byte_idx),
                  ptr.add(byte_idx + insert_len),
                  old_byte_len - byte_idx,
               );
            }

            // Copy the new text into the newly created gap.
            core::ptr::copy_nonoverlapping(
               bytes_to_insert.as_ptr(),
               ptr.add(byte_idx),
               insert_len,
            );
         }

         old_byte_len + insert_len
      };

      self.vec.len = new_byte_len;

      chars_to_insert_count
   }

   /// Deletes the text in the given **character** range.
   ///
   /// `char_range` is a char range, not a byte range — unlike
   /// [`drain`](Self::drain), which takes byte offsets.
   ///
   /// # Example
   ///
   /// ```
   /// use secure_types::SecureString;
   ///
   /// let mut string = SecureString::from("GreekFetaCheese");
   /// string.delete_text_char_range(9..15);
   /// string.unlock_str(|str| {
   ///     assert_eq!(str, "GreekFeta");
   /// });
   /// ```
   pub fn delete_text_char_range(&mut self, char_range: core::ops::Range<usize>) {
      if char_range.start >= char_range.end {
         return;
      }

      let new_len = self.vec.unlock_slice_mut(|current_bytes| {
         // SAFETY: `SecureString` upholds the invariant that its bytes are valid
         // UTF-8 — every constructor but the `unsafe` `from_utf8_unchecked`
         // validates it. `current_bytes` is that buffer, unlocked by the guard.
         let current_text = unsafe { core::str::from_utf8_unchecked(current_bytes) };
         let byte_start = char_to_byte_idx(current_text.as_bytes(), char_range.start);
         let byte_end = char_to_byte_idx(current_text.as_bytes(), char_range.end);

         if byte_start >= byte_end || byte_end > current_bytes.len() {
            return current_bytes.len();
         }

         let remove_len = byte_end - byte_start;
         let old_total_len = current_bytes.len();

         // Shift elements left
         current_bytes.copy_within(byte_end..old_total_len, byte_start);

         let new_len = old_total_len - remove_len;
         // Zeroize the tail end that is now unused
         for byte in current_bytes[new_len..old_total_len].iter_mut() {
            byte.zeroize();
         }
         new_len
      });
      self.vec.len = new_len;
   }
}

#[cfg(feature = "use_os")]
impl From<String> for SecureString {
   /// Creates a new `SecureString` from a `String`.
   ///
   /// The `String` is zeroized afterwards.
   ///
   /// # Panics
   /// Panics if the secure allocation cannot be made or locked — `From` cannot
   /// return an error. Use [`SecureVec::from_vec`] with
   /// [`SecureString::try_from`] for a fallible path.
   fn from(s: String) -> SecureString {
      let vec = SecureVec::from_vec(s.into_bytes()).unwrap();
      SecureString { vec }
   }
}

impl From<&str> for SecureString {
   /// Creates a new `SecureString` from a `&str`.
   ///
   /// The `&str` is not zeroized, you are responsible for zeroizing it.
   ///
   /// # Panics
   /// Panics if the secure allocation cannot be made or locked — `From` cannot
   /// return an error.
   fn from(s: &str) -> SecureString {
      let bytes = s.as_bytes();
      // new_with_capacity bumps 0 -> 1 internally, so empty &str is fine.
      let mut new_vec = SecureVec::new_with_capacity(bytes.len()).unwrap();
      new_vec.init_from_clone(bytes);
      SecureString { vec: new_vec }
   }
}

impl TryFrom<SecureVec<u8>> for SecureString {
   type Error = Error;

   /// Creates a `SecureString` from a `SecureVec<u8>`, validating UTF-8.
   ///
   /// The `SecureVec` is consumed. On invalid UTF-8 it is dropped (and thus
   /// zeroized) and `Error::InvalidUtf8` is returned.
   fn try_from(vec: SecureVec<u8>) -> Result<Self, Self::Error> {
      let valid = vec.unlock_slice(|slice| core::str::from_utf8(slice).is_ok());
      if valid {
         Ok(SecureString { vec })
      } else {
         Err(Error::InvalidUtf8)
      }
   }
}

#[cfg(feature = "serde")]
impl serde::Serialize for SecureString {
   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
   where
      S: serde::Serializer,
   {
      self.unlock_str(|str| serializer.serialize_str(str))
   }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for SecureString {
   fn deserialize<D>(deserializer: D) -> Result<SecureString, D::Error>
   where
      D: serde::Deserializer<'de>,
   {
      struct SecureStringVisitor;
      impl<'de> serde::de::Visitor<'de> for SecureStringVisitor {
         type Value = SecureString;
         fn expecting(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
            write!(formatter, "an utf-8 encoded string")
         }
         fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
         where
            E: serde::de::Error,
         {
            Ok(SecureString::from(v))
         }

         /// Formats that build an owned `String` (unescaping, normalization) hand it
         /// over here. serde's default implementation would copy out of it and then
         /// drop it with the plaintext still inside, so wipe it after the copy.
         fn visit_string<E>(self, mut v: String) -> Result<Self::Value, E>
         where
            E: serde::de::Error,
         {
            let secure_string = SecureString::from(v.as_str());
            v.zeroize();
            Ok(secure_string)
         }
      }
      deserializer.deserialize_string(SecureStringVisitor)
   }
}

/// Maps a character index to a byte index.
///
/// An index at or past the end of the string resolves to the string's *length*, so callers
/// clamp rather than fail: `insert_text_at_char_idx` appends and `delete_text_char_range`
/// becomes a no-op. That has been the behaviour since the first release, and it matches what
/// the callers expect (Zeus's text field behaves the same way over egui's `TextEdit`), so it
/// is kept deliberately. Revisit if an out-of-range index should be rejected instead.
fn char_to_byte_idx(s_bytes: &[u8], char_idx: usize) -> usize {
   core::str::from_utf8(s_bytes)
      .ok()
      .and_then(|s| s.char_indices().nth(char_idx).map(|(idx, _)| idx))
      .unwrap_or(s_bytes.len()) // Fallback to end if char_idx is out of bounds
}