secure-types 0.2.63

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

/// 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.
///
/// # 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 range from the string
   ///
   /// # 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| {
         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();

      // Perform the insertion in-place
      self.vec.unlock_memory();
      unsafe {
         let ptr = self.vec.as_mut_ptr();

         // 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,
         );

         self.vec.len += insert_len;
      }

      self.vec.lock_memory();

      chars_to_insert_count
   }

   /// Deletes the text in the given character range
   ///
   /// # 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| {
         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 0;
         }

         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 i in new_len..old_total_len {
            current_bytes[i].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.
   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.
   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,
   {
      let res = self.unlock_str(|str| serializer.serialize_str(str));
      res
   }
}

#[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))
         }
      }
      deserializer.deserialize_string(SecureStringVisitor)
   }
}

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
}

#[cfg(all(test, feature = "use_os"))]
mod tests {
   use super::*;

   #[test]
   fn test_creation() {
      let hello_world = "Hello, world!";
      let secure = SecureString::from(hello_world);

      secure.unlock_str(|str| {
         assert_eq!(str, hello_world);
      });
   }

   #[test]
   fn test_from_string() {
      let hello_world = String::from("Hello, world!");
      let string = SecureString::from(hello_world);

      string.unlock_str(|str| {
         assert_eq!(str, "Hello, world!");
      });
   }

   #[test]
   fn test_try_from_secure_vec() {
      let hello_world = "Hello, world!".to_string();
      let vec: SecureVec<u8> = SecureVec::from_slice(hello_world.as_bytes()).unwrap();

      let string = SecureString::from(hello_world);
      let string2 = SecureString::try_from(vec).unwrap();

      string.unlock_str(|str| {
         string2.unlock_str(|str2| {
            assert_eq!(str, str2);
         });
      });

      string.unlock_str_unchecked(|str| {
         string2.unlock_str_unchecked(|str2| {
            assert_eq!(str, str2);
         });
      });
   }

   #[test]
   fn test_clone() {
      let hello_world = "Hello, world!".to_string();
      let secure1 = SecureString::from(hello_world.clone());
      let secure2 = secure1.clone();

      secure2.unlock_str(|str| {
         assert_eq!(str, hello_world);
      });

      secure2.unlock_str_unchecked(|str| {
         assert_eq!(str, hello_world);
      });
   }

   #[test]
   fn test_insert_text_at_char_idx() {
      let hello_world = "My name is ";
      let mut secure = SecureString::from(hello_world);
      secure.insert_text_at_char_idx(12, "Mike");
      secure.unlock_str(|str| {
         assert_eq!(str, "My name is Mike");
      });

      secure.unlock_str_unchecked(|str| {
         assert_eq!(str, "My name is Mike");
      });
   }

   #[test]
   fn test_delete_text_char_range() {
      let hello_world = "My name is Mike";
      let mut secure = SecureString::from(hello_world);
      secure.delete_text_char_range(10..17);
      secure.unlock_str(|str| {
         assert_eq!(str, "My name is");
      });

      secure.unlock_str_unchecked(|str| {
         assert_eq!(str, "My name is");
      });
   }

   #[test]
   fn test_drain() {
      let hello_world = "Hello, world!";
      let mut secure = SecureString::from(hello_world);
      secure.drain(0..7);
      secure.unlock_str(|str| {
         assert_eq!(str, "world!");
      });

      secure.unlock_str_unchecked(|str| {
         assert_eq!(str, "world!");
      });
   }

   #[cfg(feature = "serde")]
   #[test]
   fn test_serde() {
      let hello_world = "Hello, world!";
      let secure = SecureString::from(hello_world);

      let json_string = serde_json::to_string(&secure).expect("Serialization failed");
      let json_bytes = serde_json::to_vec(&secure).expect("Serialization failed");

      let deserialized_string: SecureString =
         serde_json::from_str(&json_string).expect("Deserialization failed");

      let deserialized_bytes: SecureString =
         serde_json::from_slice(&json_bytes).expect("Deserialization failed");

      deserialized_string.unlock_str(|str| {
         assert_eq!(str, hello_world);
      });

      deserialized_string.unlock_str_unchecked(|str| {
         assert_eq!(str, hello_world);
      });

      deserialized_bytes.unlock_str(|str| {
         assert_eq!(str, hello_world);
      });

      deserialized_bytes.unlock_str_unchecked(|str| {
         assert_eq!(str, hello_world);
      });
   }

   #[test]
   fn test_unlock_str() {
      let hello_word = "Hello, world!";
      let string = SecureString::from(hello_word);
      let _exposed_string = string.unlock_str(|str| {
         assert_eq!(str, hello_word);
         String::from(str)
      });

      let _exposed_string = string.unlock_str_unchecked(|str| {
         assert_eq!(str, hello_word);
         String::from(str)
      });
   }

   #[test]
   fn test_push_str() {
      let hello_world = "Hello, world!";

      let mut string = SecureString::new().unwrap();
      string.push_str(hello_world);
      string.unlock_str(|str| {
         assert_eq!(str, hello_world);
      });

      string.unlock_str_unchecked(|str| {
         assert_eq!(str, hello_world);
      });
   }

   #[test]
   fn test_unlock_mut() {
      let hello_world = "Hello, world!";
      let mut string = SecureString::from("Hello, ");
      string.secure_mut(|string| {
         string.push_str("world!");
      });

      string.unlock_str(|str| {
         assert_eq!(str, hello_world);
      });

      string.unlock_str_unchecked(|str| {
         assert_eq!(str, hello_world);
      });
   }
}