Apply suggestions from code review

Co-authored-by: RJ Rybarczyk <[email protected]>

protocols/v2/noise-sv2/src/cipher_state.rs

@@ -19,33 +19,35 @@ where
     /// Retrieves a mutable reference to the 32-byte encryption key (`k`). 
     fn get_k(&mut self) -> &mut Option<[u8; 32]>;
 
-    /// Sets the 256-bit key for the cipher.
+    /// Sets the 32-byte encryption key to the optionally provided value (`k`).
     ///
-    /// This method is used to initialize or update the key within the cipher state.
+    /// Allows the encryption key to be explicitly set, typically after it has been derived or
+    /// initialized during the handshake process. If `None`, the encryption key is unset.
     fn set_k(&mut self, k: Option<[u8; 32]>);
 
-    /// Retrieves the current nonce value used by the cipher.
+    /// Retrieves the current nonce (`n`) used for encryption.
     ///
-    /// The nonce is incremented with each encryption/decryption operation to ensure
-    /// the uniqueness required by AEAD ciphers.
+    /// The nonce is a counter that is incremented with each encryption/decryption operations to
+    /// ensure that each encryption operation with the same key produces a unique ciphertext.
     fn get_n(&self) -> u64;
 
-    /// Sets the nonce value for the cipher.
+    /// Sets the nonce (`n`) to the provided value.
     ///
-    /// This method allows for manual control over the nonce, which can be useful
-    /// in specific scenarios where the nonce needs to be reset or initialized.
+    /// Allows the nonce to be explicitly set, typically after it has been initialized, incremented
+    /// during the encryption process, or reset.
     fn set_n(&mut self, n: u64);
 
-    /// Retrieves a mutable reference to the underlying cipher instance.
+    /// Retrieves a mutable reference to the optional cipher instance.
     ///
-    /// This method is used to access the cipher for encryption and decryption operations.
+    /// Provides access to the underlying AEAD cipher instance used for encryption and decryption
+    /// operations.
     fn get_cipher(&mut self) -> &mut Option<Cipher_>;
 
-    /// Converts the current nonce value to a 12-byte array
+    /// Converts the current 64-bit nonce value (`n`) to a 12-byte array.
     ///
-    /// This method prepares the nonce in a format suitable for use with AEAD ciphers.
-    /// The nonce is typically derived from a 64-bit counter, with the remaining bytes
-    /// set to zero.
+    /// Converts the 64-bit nonce value  to a 12-byte array suitable for use with AEAD ciphers,
+    /// which typically expect a 96-bit (12-byte) nonce. The result is a correctly formatted nonce
+    /// for use in encryption and decryption operations.
     fn nonce_to_bytes(&self) -> [u8; 12] {
         let mut res = [0u8; 12];
         let n = self.get_n();
@@ -55,10 +57,11 @@ where
         res
     }
 
-    /// Encrypts the given data in place with associated data (AD).
+    /// Encrypts the provided `data` in place using the cipher and AAD (`ad`).
     ///
-    /// This method uses the current cipher to encrypt the data, incrementing the nonce
-    /// with each operation. If encryption fails, the nonce is reverted to its previous value.
+    /// Performs authenticated encryption on the provided `data` buffer, modifying it in place to
+    /// contain the ciphertext. The encryption is performed using the current nonce and the AAD.
+    /// The nonce is incremented after each successful encryption.
     fn encrypt_with_ad<T: Buffer>(
         &mut self,
         ad: &[u8],
@@ -112,6 +115,12 @@ where
 /// `GenericCipher` abstracts over different AEAD cipher implementations, allowing
 /// for flexible use of either `ChaCha20Poly1305` or `Aes256Gcm`. It provides methods
 /// for encryption, decryption, and key erasure.
+///
+/// A generic cipher that can be either `ChaCha20Poly1305` or `Aes256Gcm`.
+///
+/// The `GenericCipher` enum allows for flexibility in choosing between different AEAD ciphers. It
+/// provides methods for encrypting, decrypting, and securely erasing the encryption key, ensuring
+/// that sensitive cryptographic material is handled correctly.
 #[allow(clippy::large_enum_variant)]
 pub enum GenericCipher {
     ChaCha20Poly1305(Cipher<ChaCha20Poly1305>),
@@ -120,40 +129,45 @@ pub enum GenericCipher {
 }
 
 impl Drop for GenericCipher {
-    /// Custom drop implementation to securely erase the key.
+    /// Securely erases the encryption key when the `GenericCipher` is dropped.
     ///
-    /// When a `GenericCipher` instance is dropped, the key is securely erased
-    /// from memory to prevent unauthorized access.
+    /// Ensures that the encryption key is securely erased from memory when the `GenericCipher`
+    /// instance is dropped, preventing any potential leakage of sensitive cryptographic material.
     fn drop(&mut self) {
         self.erase_k();
     }
 }
 
 impl GenericCipher {
-    /// Encrypts the given data in place.
+    /// Encrypts the data (`msg`) in place using the underlying cipher.
     ///
-    /// This method uses the current cipher to encrypt the data without any associated data.
+    /// Performs authenticated encryption on the provided data buffer, modifying it in place to
+    /// contain the ciphertext. The encryption is performed using the current nonce and an empty
+    /// additional associated data (AAD) buffer.
     pub fn encrypt<T: Buffer>(&mut self, msg: &mut T) -> Result<(), aes_gcm::Error> {
         match self {
             GenericCipher::ChaCha20Poly1305(c) => c.encrypt_with_ad(&[], msg),
             GenericCipher::Aes256Gcm(c) => c.encrypt_with_ad(&[], msg),
         }
     }
 
-    /// Decrypts the given data in place.
+    /// Decrypts the data (`msg`) in place using the underlying cipher.
     ///
-    /// This method uses the current cipher to decrypt the data without any associated data.
+    /// Performs authenticated decryption on the provided data buffer, modifying it in place to
+    /// contain the plaintext. The decryption is performed using the current nonce and an empty
+    /// additional associated data (AAD) buffer.
     pub fn decrypt<T: Buffer>(&mut self, msg: &mut T) -> Result<(), aes_gcm::Error> {
         match self {
             GenericCipher::ChaCha20Poly1305(c) => c.decrypt_with_ad(&[], msg),
             GenericCipher::Aes256Gcm(c) => c.decrypt_with_ad(&[], msg),
         }
     }
 
-    /// Securely erases the key from the memory.
+    /// Securely erases the encryption key (`k`) from memory.
     ///
-    /// This method overwrites the key with zeros in memory and sets it to `None`,
-    /// ensuring that the key cannot be recovered.
+    /// Overwrites the encryption key stored within the `GenericCipher` with zeros and sets it to
+    /// `None`, ensuring that the key cannot be recovered after the `GenericCipher` is dropped or
+    /// no longer needed.
     pub fn erase_k(&mut self) {
         match self {
             GenericCipher::ChaCha20Poly1305(c) => {
@@ -176,10 +190,13 @@ impl GenericCipher {
     }
 }
 
-/// Implementation of the `CipherState` trait for `GenericCipher` when using `Aes256Gcm`.
+/// Represents the state of an AEAD cipher when using `Aes256Gcm`, including the encryption key and
+/// nonce.
 ///
-/// This implementation provides the methods needed to manage and operate on an
-/// `Aes256Gcm` cipher within a `GenericCipher` instance.
+/// The `Cipher` struct manages the cryptographic state required to perform AEAD encryption and
+/// decryption operations. It stores the optional 32-byte encryption key (`k`), the nonce (`n`),
+/// and the optional cipher instance itself. The `CipherState` trait is implemented for this struct
+/// to provide a consistent interface for managing cipher state across different AEAD ciphers.
 impl CipherState<Aes256Gcm> for GenericCipher {
     fn get_k(&mut self) -> &mut Option<[u8; 32]> {
         match self {
@@ -217,12 +234,13 @@ impl CipherState<Aes256Gcm> for GenericCipher {
     }
 }
 
-/// Represents a cryptographic cipher state, including the key, nonce, and the cipher instance.
+/// Represents the state of an AEAD cipher, including the optional 32-byte encryption key (`k`),
+/// nonce (`n`), and optional cipher instance (`cipher`).
 ///
-/// This struct is designed to manage the state of an authenticated encryption algorithm,
-/// including its key and nonce. The key is stored as an optional 32-byte array, and the
-/// nonce `n` is a 64-bit unsigned integer. The cipher itself is stored as an optional
-/// instance of a type that implements the `AeadCipher` trait.
+/// Manages the cryptographic state required to perform AEAD encryption and decryption operations.
+/// It stores the optional encryption key, the nonce, and the optional cipher instance itself. The
+/// `CipherState` trait is implemented to provide a consistent interface for managing cipher state
+/// across different AEAD ciphers.
 pub struct Cipher<C: AeadCipher> {
     k: Option<[u8; 32]>,
     n: u64,
@@ -273,33 +291,3 @@ impl<C: AeadCipher> CipherState<C> for Cipher<C> {
         self.k = k;
     }
 }
-
-#[cfg(test)]
-mod test {
-    use aes_gcm::{aead::generic_array::GenericArray, KeyInit};
-    use chacha20poly1305::ChaCha20Poly1305;
-    use rand::{rngs::OsRng, RngCore};
-
-    use super::{Cipher, CipherState};
-
-    fn generate_random_key() -> [u8; 32] {
-        let mut key = [0u8; 32];
-        OsRng.fill_bytes(&mut key);
-        key
-    }
-
-    #[test]
-    fn test_cipher_nonce_handling() {
-        let key = generate_random_key();
-        let cipher_instance = ChaCha20Poly1305::new(GenericArray::from_slice(&key));
-        let mut cipher = Cipher::from_key_and_cipher(key, cipher_instance);
-
-        assert_eq!(cipher.get_n(), 0);
-
-        let mut data = b"nonce test".to_vec();
-        cipher
-            .encrypt_with_ad(&[], &mut data)
-            .expect("Encryption failed");
-        assert_eq!(cipher.get_n(), 1);
-    }
-}