MIsc comments (#747)

* fix comments

* fix comment

Co-authored-by: sinu.eth <[email protected]>

* describe all args

* change decrypted plaintext -> plaintext

* remove redundant comments

---------

Co-authored-by: sinu.eth <[email protected]>

Author: themighty1

crates/common/src/zk_aes.rs

@@ -78,11 +78,18 @@ impl ZkAesCtr {
 
     /// Proves the encryption of `len` bytes.
     ///
+    /// Here we only assign certain values in the VM but the actual proving
+    /// happens later when the plaintext is assigned and the VM is executed.
+    ///
     /// # Arguments
     ///
     /// * `vm` - Virtual machine.
     /// * `explicit_nonce` - Explicit nonce.
     /// * `len` - Length of the plaintext in bytes.
+    ///
+    /// # Returns
+    ///
+    /// A VM reference to the plaintext and the ciphertext.
     pub fn encrypt(
         &mut self,
         vm: &mut dyn Vm<Binary>,
@@ -132,6 +139,8 @@ impl ZkAesCtr {
         // Assign zeroes to the padding.
         if padding_len > 0 {
             let padding = input.split_off(input.len() - padding_len);
+            // To simplify the impl, we don't mark the padding as public, that's why only
+            // the prover assigns it.
             if let Role::Prover = self.role {
                 vm.assign(padding, vec![0; padding_len])
                     .map_err(ZkAesCtrError::vm)?;

crates/components/cipher/src/lib.rs

@@ -99,6 +99,7 @@ pub struct CtrBlock<N, C, O> {
 /// Can be used to XOR with the cipher input to operate the cipher in counter
 /// mode.
 pub struct Keystream<N, C, O> {
+    /// Sequential keystream blocks. Outputs are stored in contiguous memory.
     blocks: VecDeque<CtrBlock<N, C, O>>,
 }
 

crates/components/key-exchange/src/exchange.rs

@@ -297,7 +297,7 @@ where
         } = self.state.take()
         else {
             return Err(KeyExchangeError::state(
-                "can not compute shares before performing setup",
+                "cannot compute shares before performing setup",
             ));
         };
 

crates/mpc-tls/src/leader.rs

@@ -148,7 +148,7 @@ impl MpcTlsLeader {
 
         let client_random = Random::new().expect("rng is available");
 
-        // Allocate
+        // Allocate.
         let pms = ke.alloc(&mut (*vm_lock))?;
         let PrfOutput { keys, cf_vd, sf_vd } = prf.alloc(&mut (*vm_lock), pms)?;
         record_layer.set_keys(

crates/mpc-tls/src/leader/actor.rs

@@ -1728,7 +1728,7 @@ impl Wrap<BackendMsgServerClosed> for MpcTlsLeaderMsg {
     }
 }
 
-/// Message to start deferring the decryption
+/// Message to start deferring the decryption.
 #[allow(missing_docs)]
 #[derive(Debug)]
 pub struct DeferDecryption;

crates/mpc-tls/src/lib.rs

@@ -84,9 +84,9 @@ pub struct LeaderOutput {
     pub server_cert_details: ServerCertDetails,
     /// Key exchange details.
     pub server_kx_details: ServerKxDetails,
-    /// Client random
+    /// Client random.
     pub client_random: Random,
-    /// Server random
+    /// Server random.
     pub server_random: Random,
     /// TLS transcript.
     pub transcript: TlsTranscript,

crates/mpc-tls/src/record_layer.rs

@@ -178,7 +178,7 @@ impl RecordLayer {
             .try_lock_owned()
             .map_err(|_| MpcTlsError::other("decrypt lock is held"))?;
 
-        // Computes GHASH keys in parallel.
+        // Preprocesses GHASH keys in parallel.
         ctx.try_join(
             |ctx| async move { encrypt.preprocess(ctx).await }.scope_boxed(),
             |ctx| async move { decrypt.preprocess(ctx).await }.scope_boxed(),
@@ -486,7 +486,7 @@ impl RecordLayer {
 
         if !self.encrypt_buffer.is_empty() {
             return Err(MpcTlsError::state(
-                "record layer can not commit with pending encrypt operations",
+                "record layer cannot commit with pending encrypt operations",
             ));
         }
 

crates/mpc-tls/src/record_layer/aead/aes_gcm.rs

@@ -111,7 +111,7 @@ impl MpcAesGcm {
 
         // Allocate encryption/decryption.
 
-        // Round up the length to the nearest multiple of the block count.
+        // Round up the length to the nearest multiple of the block size.
         let len = 16 * len.div_ceil(16);
 
         let input = vm.alloc_vec::<U8>(len)?;
@@ -141,7 +141,7 @@ impl MpcAesGcm {
 
     pub(crate) async fn preprocess(&mut self, ctx: &mut Context) -> Result<(), AeadError> {
         let State::Setup { ghash, .. } = &mut self.state else {
-            return Err(AeadError::state("must be in setup state to allocate"));
+            return Err(AeadError::state("must be in setup state to preprocess"));
         };
 
         ghash.preprocess(ctx).await?;
@@ -216,7 +216,7 @@ impl MpcAesGcm {
         let explicit_nonce: [u8; 8] = explicit_nonce.try_into().map_err(|nonce: Vec<_>| {
             AeadError::cipher(format!(
                 "explicit nonce length: expected {}, got {}",
-                16,
+                8,
                 nonce.len()
             ))
         })?;
@@ -246,6 +246,8 @@ impl MpcAesGcm {
         // Assign zeroes to the padding.
         if padding_len > 0 {
             let padding = input.split_off(input.len() - padding_len);
+            // To simplify the impl, we don't mark the padding as public, that's why only
+            // the prover assigns it.
             if let Role::Leader = self.role {
                 vm.assign(padding, vec![0; padding_len])?;
             }
@@ -283,7 +285,7 @@ impl MpcAesGcm {
         let explicit_nonce: [u8; 8] = explicit_nonce.try_into().map_err(|nonce: Vec<_>| {
             AeadError::cipher(format!(
                 "explicit nonce length: expected {}, got {}",
-                16,
+                8,
                 nonce.len()
             ))
         })?;
@@ -375,7 +377,9 @@ impl MpcAesGcm {
     /// # Arguments
     ///
     /// * `vm` - Virtual machine.
-    /// * `inputs` - Data to verify the tags for.
+    /// * `data` - Tag data associated with `tags`.
+    /// * `ciphertexts` - Ciphertexts to verify the tags for.
+    /// * `tags` - Tags to verify.
     pub(crate) fn verify_tags(
         &mut self,
         vm: &mut dyn Vm<Binary>,

crates/mpc-tls/src/record_layer/aead/ghash.rs

@@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::record_layer::aead::AeadError;
 
-/// Maximum key share power.
+/// Maximum exponent used in GHASH.
 const MAX_POWER: usize = 1026;
 
 #[async_trait]
@@ -26,13 +26,13 @@ pub(crate) trait Ghash {
     /// Preprocesses GHASH.
     async fn preprocess(&mut self, ctx: &mut Context) -> Result<(), GhashError>;
 
-    /// Sets the key for the hash function.
+    /// Sets the additive key share for the hash function.
     fn set_key(&mut self, key: Vec<u8>) -> Result<(), GhashError>;
 
     /// Sets up GHASH, computing the key shares.
     async fn setup(&mut self, ctx: &mut Context) -> Result<(), GhashError>;
 
-    /// Computes the GHASH tag.
+    /// Computes the GHASH tag share.
     fn compute(&self, input: &[u8]) -> Result<Vec<u8>, GhashError>;
 }
 
@@ -80,12 +80,13 @@ where
 {
     fn alloc(&mut self) -> Result<(), GhashError> {
         if !self.alloc {
-            // We need only half the number of `MAX_POWER` M2As because of the free
-            // squaring trick and we need one extra A2M conversion in the beginning.
+            // Odd powers are computed using M2A, even powers are computed
+            // locally. We need one extra A2M conversion in the beginning.
             // Both M2A and A2M, each require a single OLE.
             AdditiveToMultiplicative::<Gf2_128>::alloc(&mut self.converter, 1)
                 .map_err(GhashError::conversion)?;
 
+            // -1 because the odd power H^1 is already known at this point.
             MultiplicativeToAdditive::<Gf2_128>::alloc(&mut self.converter, (MAX_POWER / 2) - 1)
                 .map_err(GhashError::conversion)?;
 
@@ -128,7 +129,7 @@ where
 
     async fn setup(&mut self, ctx: &mut Context) -> Result<(), GhashError> {
         let State::SetKey { key: add_key } = self.state.take() else {
-            return Err(GhashError::state("can not setup before key is set"));
+            return Err(GhashError::state("cannot setup before key is set"));
         };
 
         let mut mult_key = self
@@ -156,9 +157,9 @@ where
                 *acc = power_n * mult_key;
                 Some(power_n)
             })
-            // Start from H^3
+            // Start from H^3.
             .skip(2)
-            // Skip even powers
+            // Skip even powers.
             .step_by(2)
             .collect();
 
@@ -249,7 +250,7 @@ fn compute_shares(key: Gf2_128, odd_powers: &[Gf2_128]) -> Vec<Gf2_128> {
             let base = shares[i / 2 - 1];
             shares.push(base * base);
         } else {
-            // Odd power
+            // Odd power.
             shares.push(odd_powers[odd_idx]);
             odd_idx += 1;
         }
@@ -423,7 +424,7 @@ mod tests {
         let sender_key: u128 = rng.random();
         let receiver_key: u128 = h ^ sender_key;
 
-        // Message length is not a multiple of the block length
+        // Message length is not a multiple of the block length.
         let message: Vec<u8> = (0..14).map(|_| rng.random()).collect();
 
         let (mut sender, mut receiver) = create_pair();

crates/mpc-tls/src/record_layer/decrypt.rs

@@ -197,7 +197,9 @@ impl DecryptOp {
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub(crate) enum DecryptMode {
+    /// The plaintext is private.
     Private,
+    /// The plaintext is public.
     Public,
 }