Apply suggestions from code review

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

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

@@ -5,28 +5,66 @@ use secp256k1::{
     rand, Keypair, Secp256k1,
 };
 
-/// A trait that represents the operations needed during a handshake process.
+/// Represents the operations needed during a Noise protocol handshake.
 ///
-/// This trait requires the implementation of various cryptographic operations,
-/// state management, and key exchange mechanisms needed to perform a secure
-/// handshake between parties.
+/// The `HandshakeOp` trait defines the necessary functions for managing the state and
+/// cryptographic operations required during the Noise protocol handshake. It provides methods for
+/// key generation, hash mixing, encryption, decryption, and key derivation, ensuring that the
+/// handshake process is secure and consistent.
+///
+/// ## Associated Types
+///
+/// * `Cipher: AeadCipher`: Represents the cipher used during the handshake, typically an
+///   authenticated encryption with associated data (AEAD) cipher like `ChaCha20Poly1305`.
+///
+/// ## Example
+///
+/// This trait is implemented for both the `Initiator` and `Responder` structs, allowing them
+/// to perform their respective roles in the Noise protocol handshake.
+///
+/// ```rs
+/// impl HandshakeOp<ChaCha20Poly1305> for Initiator {
+///     // Implementation for the Initiator...
+/// }
+///
+/// impl HandshakeOp<ChaCha20Poly1305> for Responder {
+///     // Implementation for the Responder...
+/// }
+/// ```
 pub trait HandshakeOp<Cipher: AeadCipher>: CipherState<Cipher> {
-    /// Returns the name of the handshake operation as a `String`.
+    /// Returns the name of the entity implementing the handshake operation.
+    ///
+    /// Provides a string that identifies the entity (e.g., "Initiator" or "Responder") that is
+    /// performing the handshake. It is primarily used for debugging or logging purposes.
     #[allow(dead_code)]
     fn name(&self) -> String;
 
-    /// Provides a mutable reference to the `h` value, which represents the hash
-    /// output in the handshake process.
+    /// Retrieves a mutable reference to the handshake hash (`h`).
+    ///
+    /// The handshake hash accumulates the state of the handshake, incorporating all exchanged
+    /// messages to ensure integrity and prevent tampering. This method provides access to the
+    /// current state of the handshake hash, allowing it to be updated as the handshake progresses.
     fn get_h(&mut self) -> &mut [u8; 32];
 
-    /// Provides a mutable reference to the `ck` value, which represents the
-    /// chaining key in the handshake process.
+    /// Retrieves a mutable reference to the chaining key (`ck`).
+    ///
+    /// The chaining key is used during the key derivation process to generate new keys throughout
+    /// the handshake. This method provides access to the current chaining key, which is updated
+    /// as the handshake progresses and new keys are derived.
     fn get_ck(&mut self) -> &mut [u8; 32];
 
-    /// Sets the `h` value to the provided data, updating the hash output.
+    /// Sets the handshake hash (`h`) to the provided value.
+    ///
+    /// This method allows the handshake hash to be explicitly set, typically after it has been
+    /// initialized or updated during the handshake process. The handshake hash ensures the
+    /// integrity of the handshake by incorporating all exchanged messages.
     fn set_h(&mut self, data: [u8; 32]);
 
-    /// Sets the `ck` value to the provided data, updating the chaining key.
+    /// Sets the chaining key (`ck`) to the provided value.
+    ///
+    /// This method allows the chaining key to be explicitly set, typically after it has been
+    /// initialized or updated during the handshake process. The chaining key is crucial for
+    /// deriving new keys as the handshake progresses.
     fn set_ck(&mut self, data: [u8; 32]);
 
     /// Mixes the data into the handshake hash (`h`).
@@ -54,14 +92,17 @@ pub trait HandshakeOp<Cipher: AeadCipher>: CipherState<Cipher> {
         }
     }
 
-    /// Computes an HMAC-SHA256 hash of the provided data using the given key.
+    /// Computes an HMAC-SHA256 (Hash-based Message Authentication Code) hash of the provided data
+    /// using the given key.
+    ///
+    /// This method implements the HMAC-SHA256 hashing algorithm, which combines a key and data to
+    /// produce a 32-byte hash. It is used during the handshake to securely derive new keys from
+    /// existing material, ensuring that the resulting keys are cryptographically strong.
     ///
-    /// The `hmac_hash` method implements the HMAC (Hash-based Message Authentication Code)
-    /// algorithm using SHA-256 as the underlying hash function. HMAC is used for key
-    /// derivation and message integrity, ensuring that the data has not been tampered
-    /// with and is authentic. This method uses a two-step process: first, the key is
-    /// XORed with an inner padding `ipad` and hashed with the data; then, the result is
-    /// XORed with an outer padding `opad` and hashed again to produce the final HMAC.
+    /// This method uses a two-step process:
+    /// 1. The key is XORed with an inner padding (`ipad`) and hashed with the data.
+    /// 2. The result is XORed with the outer padding (`opad`) and hashed again to produce the
+    ///    final HMAC.
     fn hmac_hash(key: &[u8; 32], data: &[u8]) -> [u8; 32] {
         #[allow(clippy::identity_op)]
         let mut ipad = [(0 ^ 0x36); 64];
@@ -86,14 +127,21 @@ pub trait HandshakeOp<Cipher: AeadCipher>: CipherState<Cipher> {
         Sha256Hash::hash(&to_hash).to_byte_array()
     }
 
-    /// Derive two keys using the HKDF (HMAC-based Key Derivation Function) algorithm.
+    /// Derives two new keys using the HKDF (HMAC-based Key Derivation Function) process.
+    ///
+    /// Performs the HKDF key derivation process, which uses an initial chaining key and input key
+    /// material to produce two new 32-byte keys. This process is used throughout the handshake to
+    /// generate fresh keys for encryption and authentication, ensuring that each step of the
+    /// handshake is securely linked.
     ///
-    /// The `hkdf_2` method derives two cryptographic keys from a given input key material
-    /// and a chaining key. HKDF is a widely used key derivation function that extracts
-    /// and expands keys in a secure manner. This method uses `hmac_hash` as the underlying
-    /// HMAC function to first derive a temporary key,  and then two output keys by
-    /// iteratively applying HMAC to the temporary key with specific byte sequences
-    /// (0x01, 0x02).
+    /// This method performs the following steps:
+    /// 1. Performs a HMAC hash on the chaining key and input key material to derive a temporary
+    ///    key.
+    /// 2. Performs a HMAC hash on the temporary key and specific byte sequence (`0x01`) to derive
+    ///    the first output.
+    /// 3. Performs a HMAC hash on the temporary key and the concatenation of the first output and
+    ///    a specific byte sequence (`0x02`).
+    /// 4. Returns both outputs.
     fn hkdf_2(chaining_key: &[u8; 32], input_key_material: &[u8]) -> ([u8; 32], [u8; 32]) {
         let temp_key = Self::hmac_hash(chaining_key, input_key_material);
         let out_1 = Self::hmac_hash(&temp_key, &[0x1]);

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

@@ -12,9 +12,11 @@ pub struct SignatureNoiseMessage {
 }
 
 impl From<[u8; 74]> for SignatureNoiseMessage {
-    /// Converts a 74-byte array into a `SignatureNoiseMessage` struct.
-    /// The first 10 bytes are interpreted as the version, valid_from, and not_valid_after fields,
-    /// while the remaining 64 bytes are used as the signature.
+    /// Converts a 74-byte array into a `SignatureNoiseMessage`.
+    ///
+    /// Allows a raw 74-byte array to be converted into a `SignatureNoiseMessage`, extracting the
+    /// version, validity periods, and signature from the provided data. Panics if the byte array
+    /// cannot be correctly converted into the struct fields.
     fn from(value: [u8; 74]) -> Self {
         let version = u16::from_le_bytes(value[0..2].try_into().unwrap());
         let valid_from = u32::from_le_bytes(value[2..6].try_into().unwrap());
@@ -62,9 +64,11 @@ impl SignatureNoiseMessage {
         }
     }
 
-    /// Signs a 74-byte message with a provided keypair and the static public key of the server.
-    /// The message is hashed using SHA-256 and the signature is generated using Schnorr
-    /// signature scheme. The resulting 64-byte signature is then appended to the original message.
+    /// Signs a `SignatureNoiseMessage` using the provided keypair (`kp`).
+    ///
+    /// Creates a Schnorr signature for the message, combining the version, validity period, and
+    /// the static public key of the server (`static_pk`). The resulting signature is then written
+    /// into the provided message buffer (`msg`).
     pub fn sign(msg: &mut [u8; 74], static_pk: &XOnlyPublicKey, kp: &Keypair) {
         let secp = Secp256k1::signing_only();
         let m = [&msg[0..10], &static_pk.serialize()].concat();
@@ -75,9 +79,12 @@ impl SignatureNoiseMessage {
         }
     }
 
-    /// Splits the `SignatureNoiseMessage` into two parts: the first 10 bytes, which include the
-    /// version, valid_from, and not_valid_after fields, and the 64-byte signature. This function
-    /// is primarily used internally during the verification process.
+    /// Splits the `SignatureNoiseMessage` into its component parts: the message hash and the
+    /// signature.
+    ///
+    /// Separates the message into the first 10 bytes (containing the version and validity period)
+    /// and the 64-byte Schnorr signature, returning them in a tuple. Used internally during the
+    /// verification process.
     fn split(self) -> ([u8; 10], [u8; 64]) {
         let mut m = [0; 10];
         m[0] = self.version.to_le_bytes()[0];