refactor(cipher): Improve security of the cipher module by encapsulating the decrypted string within the expose method. Implement serde's deserialization trait to make sure data is encrypted as soon as it is received

Author: Archasion

src/cipher.rs

@@ -14,9 +14,13 @@ static CIPHER: LazyLock<Aes256Gcm> = LazyLock::new(|| {
 });
 
 /// Represents encrypted data with the encrypted string and the nonce used for encryption.
-pub(crate) struct Encrypted {
+#[derive(Clone)]
+pub(crate) struct EncryptedString {
+    /// The cached decrypted data as a UTF-8 string.
+    /// This field is only set if exposed using the `expose` method.
+    decrypted: Option<String>,
     /// The encrypted data as a base64 encoded string.
-    pub(crate) data: String,
+    pub(crate) encrypted: String,
     /// The nonce used for encryption as a base64 encoded string.
     pub(crate) nonce: String,
 }
@@ -49,64 +53,106 @@ fn get_encryption_key() -> Key<Aes256Gcm> {
     *Key::<Aes256Gcm>::from_slice(key.as_bytes())
 }
 
-/// Encrypts the given data using AES-256-GCM.
-///
-/// This function encrypts the provided data using the AES-256-GCM algorithm and a randomly generated nonce.
-/// The encrypted data and nonce are returned as base64 encoded strings.
-///
-/// # Arguments
-///
-/// * `data` - A byte slice of the data to be encrypted.
-///
-/// # Returns
-///
-/// A [`Result`] containing an [`Encrypted`] struct on success, or a [`String`] error message on failure.
-pub(crate) fn encrypt(data: &[u8]) -> Result<Encrypted, String> {
-    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
-    let encrypted_bytes = CIPHER
-        .encrypt(&nonce, data)
-        .map_err(|e| format!("Failed to encrypt data: {}", e))?;
-
-    Ok(Encrypted {
-        data: base64::encode(&encrypted_bytes),
-        nonce: base64::encode(nonce.as_slice()),
-    })
+impl EncryptedString {
+    /// Creates a new [`EncryptedString`] struct with the provided data and nonce.
+    pub(crate) fn new<S>(data: S, nonce: S) -> Self
+    where
+        S: Into<String>,
+    {
+        EncryptedString {
+            encrypted: data.into(),
+            nonce: nonce.into(),
+            decrypted: None,
+        }
+    }
+
+    /// Encrypts the given data using AES-256-GCM.
+    ///
+    /// This function encrypts the provided data using the AES-256-GCM algorithm and a randomly generated nonce.
+    /// The encrypted data and nonce are returned as base64 encoded strings.
+    ///
+    /// # Arguments
+    ///
+    /// * `data` - A byte slice of the data to be encrypted.
+    ///
+    /// # Returns
+    ///
+    /// A [`Result`] containing an [`EncryptedString`] struct on success, or a [`String`] error message on failure.
+    pub(crate) fn encrypt(data: &[u8]) -> Result<Self, String> {
+        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
+        let encrypted_bytes = CIPHER
+            .encrypt(&nonce, data)
+            .map_err(|e| format!("Failed to encrypt data: {}", e))?;
+
+        Ok(EncryptedString::new(
+            base64::encode(&encrypted_bytes),
+            base64::encode(nonce.as_slice()),
+        ))
+    }
+
+    /// Decrypts the given encrypted data using AES-256-GCM.
+    ///
+    /// This function decrypts the provided [`EncryptedString`] using the AES-256-GCM algorithm and the nonce.
+    /// The decrypted data is returned as a UTF-8 string.
+    ///
+    /// # Arguments
+    ///
+    /// * `data` - A reference to the [`EncryptedString`] struct containing the encrypted data and nonce.
+    ///
+    /// # Returns
+    ///
+    /// A [`Result`] containing the decrypted data as a [`String`] on success, or a [`String`] error message on failure.
+    fn decrypt(&self) -> Result<String, String> {
+        // Decode the base64 encoded data and nonce
+        let encrypted_bytes = base64::decode(&self.encrypted)
+            .map_err(|e| format!("Failed to decode base64 data: {}", e))?;
+        let nonce = base64::decode(&self.nonce)
+            .map_err(|e| format!("Failed to decode base64 nonce: {}", e))?;
+        let nonce = Nonce::<Aes256Gcm>::from_slice(nonce.as_slice());
+
+        // Decrypt the data and convert it to a UTF-8 string
+        let decrypted_bytes = CIPHER
+            .decrypt(nonce, encrypted_bytes.as_slice())
+            .map_err(|e| format!("Failed to decrypt data: {}", e))?;
+        let decrypted_data = String::from_utf8(decrypted_bytes)
+            .map_err(|e| format!("Failed to convert decrypted data to UTF-8: {}", e))?;
+
+        Ok(decrypted_data)
+    }
+
+    /// Exposes the decrypted data to the provided function.
+    /// This method should only be used if the program intends to use the decrypted data **once**.
+    pub(crate) fn expose_once<T>(&self, f: impl FnOnce(&str) -> T) -> Result<T, String> {
+        let decrypted_data = self.decrypt()?;
+        Ok(f(&decrypted_data))
+    }
+
+    /// Exposes the decrypted data to the provided function.
+    /// This method should only be used if the program intends to use the decrypted data **multiple times**.
+    pub(crate) fn expose<T>(&mut self, f: impl FnOnce(&str) -> T) -> Result<T, String> {
+        if self.decrypted.is_none() {
+            let decrypted_data = self.decrypt()?;
+            self.decrypted = Some(decrypted_data);
+        }
+
+        Ok(f(self.decrypted.as_ref().unwrap()))
+    }
 }
 
-/// Decrypts the given encrypted data using AES-256-GCM.
-///
-/// This function decrypts the provided [`Encrypted`] using the AES-256-GCM algorithm and the nonce.
-/// The decrypted data is returned as a UTF-8 string.
-///
-/// # Arguments
-///
-/// * `data` - A reference to the [`Encrypted`] struct containing the encrypted data and nonce.
-///
-/// # Returns
-///
-/// A [`Result`] containing the decrypted data as a [`String`] on success, or a [`String`] error message on failure.
-pub(crate) fn decrypt(data: &Encrypted) -> Result<String, String> {
-    // Decode the base64 encoded data and nonce
-    let encrypted_bytes =
-        base64::decode(&data.data).map_err(|e| format!("Failed to decode base64 data: {}", e))?;
-    let nonce =
-        base64::decode(&data.nonce).map_err(|e| format!("Failed to decode base64 nonce: {}", e))?;
-    let nonce = Nonce::<Aes256Gcm>::from_slice(nonce.as_slice());
-
-    // Decrypt the data and convert it to a UTF-8 string
-    let decrypted_bytes = CIPHER
-        .decrypt(nonce, encrypted_bytes.as_slice())
-        .map_err(|e| format!("Failed to decrypt data: {}", e))?;
-    let decrypted_data = std::str::from_utf8(&decrypted_bytes)
-        .map_err(|e| format!("Failed to convert decrypted data to UTF-8: {}", e))?;
-
-    Ok(decrypted_data.to_string())
+impl<'de> rocket::serde::Deserialize<'de> for EncryptedString {
+    fn deserialize<D>(deserializer: D) -> Result<EncryptedString, D::Error>
+    where
+        D: rocket::serde::Deserializer<'de>,
+    {
+        let b = <&[u8]>::deserialize(deserializer)?;
+        Self::encrypt(b).map_err(rocket::serde::de::Error::custom)
+    }
 }
 
 #[cfg(test)]
 #[serial_test::file_serial(env)]
 mod tests {
-    use super::{decrypt, encrypt, get_encryption_key, Encrypted};
+    use super::{get_encryption_key, EncryptedString};
     use crate::constants;
 
     const DATA: &[u8] = b"Hello, world!";
@@ -158,10 +204,10 @@ mod tests {
             constants::env::ENCRYPTION_KEY,
             constants::test::ENCRYPTION_KEY,
         );
-        let encrypted_data = encrypt(DATA).unwrap();
-        let decrypted_data = decrypt(&encrypted_data).unwrap();
 
-        assert_eq!(decrypted_data.as_bytes(), DATA);
+        let data = EncryptedString::encrypt(DATA).unwrap().decrypt().unwrap();
+        assert_eq!(data.as_bytes(), DATA);
+
         std::env::remove_var(constants::env::ENCRYPTION_KEY);
     }
 
@@ -171,13 +217,10 @@ mod tests {
             constants::env::ENCRYPTION_KEY,
             constants::test::ENCRYPTION_KEY,
         );
-        let encrypted_data = Encrypted {
-            data: "valid+data".to_string(),
-            nonce: "valid+nonce".to_string(),
-        };
-        let result = decrypt(&encrypted_data);
 
+        let result = EncryptedString::new("valid+data", "valid+nonce").decrypt();
         assert!(result.is_err());
+
         std::env::remove_var(constants::env::ENCRYPTION_KEY);
     }
 
@@ -188,11 +231,40 @@ mod tests {
             constants::test::ENCRYPTION_KEY,
         );
 
-        let mut encrypted_data = encrypt(DATA).unwrap();
-        encrypted_data.nonce = "@".to_string();
-        let result = decrypt(&encrypted_data);
-
+        let result = EncryptedString::new("valid", "@").decrypt();
         assert!(result.is_err());
+
+        std::env::remove_var(constants::env::ENCRYPTION_KEY);
+    }
+
+    #[test]
+    fn expose_once() {
+        std::env::set_var(
+            constants::env::ENCRYPTION_KEY,
+            constants::test::ENCRYPTION_KEY,
+        );
+
+        let encrypted_data = EncryptedString::encrypt(DATA).unwrap();
+        let decrypted_data = encrypted_data.expose_once(|data| data.to_owned()).unwrap();
+        assert_eq!(decrypted_data.as_bytes(), DATA);
+
+        std::env::remove_var(constants::env::ENCRYPTION_KEY);
+    }
+
+    #[test]
+    fn expose() {
+        std::env::set_var(
+            constants::env::ENCRYPTION_KEY,
+            constants::test::ENCRYPTION_KEY,
+        );
+
+        let mut encrypted_data = EncryptedString::encrypt(DATA).unwrap();
+        assert!(encrypted_data.decrypted.is_none());
+
+        let decrypted_data = encrypted_data.expose(|data| data.to_owned()).unwrap();
+        assert_eq!(decrypted_data.as_bytes(), DATA);
+        assert_eq!(encrypted_data.decrypted.unwrap().as_bytes(), DATA);
+
         std::env::remove_var(constants::env::ENCRYPTION_KEY);
     }
 }