fix(audit): refactor pad_variable_length for sha256 + sha512 (#342)

Author: ratankaliani

plonky2x/core/src/frontend/hash/sha/sha256/curta.rs

@@ -57,7 +57,7 @@ impl<L: PlonkParameters<D>, const D: usize> Hash<L, D, 64, false, 8> for SHA256
         input: &[ByteVariable],
         length: U32Variable,
     ) -> Vec<Self::IntVariable> {
-        let padded_bytes = builder.pad_message_sha256_variable(input, length);
+        let padded_bytes = builder.pad_sha256_variable_length(input, length);
 
         padded_bytes
             .chunks_exact(4)
@@ -148,6 +148,12 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
             0,
             "input length should be a multiple of 64"
         );
+
+        // Check that length <= input.len(). This is needed to ensure that users cannot prove the
+        // hash of a longer message than they supplied.
+        let supplied_input_length = self.constant::<U32Variable>(input.len() as u32);
+        self.lte(length, supplied_input_length);
+
         let last_chunk = self.compute_sha256_last_chunk(length);
         if self.sha256_accelerator.is_none() {
             self.sha256_accelerator = Some(SHA256Accelerator {

plonky2x/core/src/frontend/hash/sha/sha256/mod.rs

@@ -161,16 +161,13 @@ mod tests {
     use super::*;
     use crate::prelude::{ByteVariable, CircuitBuilder, DefaultParameters, U32Variable};
     use crate::utils::hash::sha256;
-    use crate::utils::setup_logger;
 
     type L = DefaultParameters;
     const D: usize = 2;
 
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha256_padding() {
-        setup_logger();
-
         let mut builder = CircuitBuilder::<L, D>::new();
 
         let max_len = 1024;
@@ -209,8 +206,6 @@ mod tests {
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha256_variable_padding() {
-        setup_logger();
-
         let mut builder = CircuitBuilder::<L, D>::new();
 
         let max_number_of_chunks = 5;
@@ -239,7 +234,7 @@ mod tests {
                 .map(|b| builder.constant::<ByteVariable>(*b))
                 .collect::<Vec<_>>();
 
-            let padding = builder.pad_message_sha256_variable(&message, length);
+            let padding = builder.pad_sha256_variable_length(&message, length);
 
             for (value, expected) in padding.iter().zip(expected_padding.iter()) {
                 builder.assert_is_equal(*value, *expected);

plonky2x/core/src/frontend/hash/sha/sha256/pad.rs

@@ -1,7 +1,14 @@
 use array_macro::array;
+use plonky2::util::ceil_div_usize;
 
 use crate::prelude::*;
 
+pub const SHA256_CHUNK_SIZE_BYTES: usize = 64;
+pub const SHA256_INPUT_LENGTH_BYTE_SIZE: usize = 8;
+
+pub const SHA256_CHUNK_SIZE_BITS: usize = SHA256_CHUNK_SIZE_BYTES * 8;
+pub const SHA256_INPUT_LENGTH_BIT_SIZE: usize = SHA256_INPUT_LENGTH_BYTE_SIZE * 8;
+
 impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
     /// Pad the given input according to the SHA-256 spec.
     /// The last chunk (each chunk is 64 bytes = 512 bits) gets padded.
@@ -13,7 +20,8 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
         bits.push(self.api._true());
 
         let l = bits.len() - 1;
-        let k = 512 - (l + 1 + 64) % 512;
+        let k = SHA256_CHUNK_SIZE_BITS
+            - (l + 1 + SHA256_INPUT_LENGTH_BIT_SIZE) % SHA256_CHUNK_SIZE_BITS;
         for _ in 0..k {
             bits.push(self.api._false());
         }
@@ -51,53 +59,66 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
     }
 
     /// Pad the given variable length input according to the SHA-256 spec.
-    ///
-    /// It is assumed that `input` has length MAX_NUM_CHUNKS * 64.
-    /// The true number of non-zero bytes in `input` is given by input_byte_length.
-    pub(crate) fn pad_message_sha256_variable(
+    /// input_byte_length gives the real length of the input in bytes.
+    pub(crate) fn pad_sha256_variable_length(
         &mut self,
         input: &[ByteVariable],
         input_byte_length: U32Variable,
     ) -> Vec<ByteVariable> {
-        let max_number_of_chunks = input.len() / 64;
-        assert_eq!(
-            max_number_of_chunks * 64,
-            input.len(),
-            "input length must be a multiple of 64 bytes"
-        );
+        let true_t = self._true();
+        let false_t = self._false();
+
         let last_chunk = self.compute_sha256_last_chunk(input_byte_length);
 
+        // Calculate the number of chunks needed to store the input. 9 bytes are added by the
+        // padding and LE length representation.
+        let max_num_chunks = ceil_div_usize(
+            input.len() + SHA256_INPUT_LENGTH_BYTE_SIZE + 1,
+            SHA256_CHUNK_SIZE_BYTES,
+        );
+
+        // Extend input to size max_num_chunks * 64 before padding.
+        let mut padded_input = input.to_vec();
+        padded_input.resize(max_num_chunks * SHA256_CHUNK_SIZE_BYTES, self.zero());
+
         // Compute the length bytes (big-endian representation of the length in bits).
         let zero_byte = self.constant::<ByteVariable>(0x00);
-        let mut length_bytes = vec![zero_byte; 4];
 
         let bits_per_byte = self.constant::<U32Variable>(8);
         let input_bit_length = self.mul(input_byte_length, bits_per_byte);
 
-        let mut length_bits = self.to_le_bits(input_bit_length);
+        // Get the length bits in LE order, padded to 64 bits.
+        let mut length_bits = self
+            .api
+            .split_le(input_bit_length.variable.0, SHA256_INPUT_LENGTH_BIT_SIZE);
+        // Convert length to BE bits
         length_bits.reverse();
 
-        // Prepend 4 zero bytes to length_bytes as abi.encodePacked(U32Variable) is 4 bytes.
-        length_bytes.extend_from_slice(
-            &length_bits
-                .chunks(8)
-                .map(|chunk| {
-                    let bits = array![x => chunk[x]; 8];
-                    ByteVariable(bits)
-                })
-                .collect::<Vec<_>>(),
-        );
+        let length_bytes = &length_bits
+            .chunks(8)
+            .map(|chunk| {
+                let bits = array![x => BoolVariable::from_targets(&[chunk[x].target]); 8];
+                ByteVariable(bits)
+            })
+            .collect::<Vec<_>>();
 
         let mut padded_bytes = Vec::new();
 
-        let mut message_byte_selector = self.constant::<BoolVariable>(true);
-        for i in 0..max_number_of_chunks {
-            let chunk_offset = 64 * i;
+        // Set to true if the last chunk has been reached. This is used to verify that
+        // input_byte_length is <= input.len().
+        let mut reached_last_chunk = false_t;
+
+        let mut message_byte_selector = true_t;
+        for i in 0..max_num_chunks {
+            let chunk_offset = SHA256_CHUNK_SIZE_BYTES * i;
             let curr_chunk = self.constant::<U32Variable>(i as u32);
 
+            // Check if this is the chunk where length should be added.
             let is_last_chunk = self.is_equal(curr_chunk, last_chunk);
+            reached_last_chunk = self.or(reached_last_chunk, is_last_chunk);
 
-            for j in 0..64 {
+            for j in 0..SHA256_CHUNK_SIZE_BYTES {
+                // First 64 - 8 bytes are either message | padding | nil bytes.
                 let idx = chunk_offset + j;
                 let idx_t = self.constant::<U32Variable>(idx as u32);
                 let is_last_msg_byte = self.is_equal(idx_t, input_byte_length);
@@ -112,21 +133,26 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
                 let padding_start_byte = self.constant::<ByteVariable>(0x80);
 
                 // If message_byte_selector is true, select the message byte.
-                let mut byte = self.select(message_byte_selector, input[idx], zero_byte);
+                let mut byte = self.select(message_byte_selector, padded_input[idx], zero_byte);
                 // If idx == length_bytes, select the padding start byte.
                 byte = self.select(is_last_msg_byte, padding_start_byte, byte);
-                if j >= 64 - 8 {
-                    // If in last chunk, select the length byte.
-                    byte = self.select(is_last_chunk, length_bytes[j % 8], byte);
+
+                if j >= SHA256_CHUNK_SIZE_BYTES - SHA256_INPUT_LENGTH_BYTE_SIZE {
+                    // If in last chunk, this is a length byte.
+                    byte = self.select(
+                        is_last_chunk,
+                        length_bytes[j % SHA256_INPUT_LENGTH_BYTE_SIZE],
+                        byte,
+                    );
                 }
 
                 padded_bytes.push(byte);
             }
         }
+        // These checks verify input_byte_length <= input.len().
+        self.is_equal(message_byte_selector, false_t);
+        self.is_equal(reached_last_chunk, true_t);
 
-        // self.watch_slice(&padded_bytes, "padded bytes");
-
-        assert_eq!(padded_bytes.len(), max_number_of_chunks * 64);
         padded_bytes
     }
 }

plonky2x/core/src/frontend/hash/sha/sha512/curta.rs

@@ -152,6 +152,11 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
         input: &[ByteVariable],
         length: U32Variable,
     ) -> BytesVariable<64> {
+        // Check that length <= input.len(). This is needed to ensure that users cannot prove the
+        // hash of a longer message than they supplied.
+        let supplied_input_length = self.constant::<U32Variable>(input.len() as u32);
+        self.lte(length, supplied_input_length);
+
         let last_chunk = self.compute_sha512_last_chunk(length);
 
         if self.sha512_accelerator.is_none() {
@@ -178,14 +183,14 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
 
 #[cfg(test)]
 mod tests {
+    use std::env;
+
     use rand::{thread_rng, Rng};
 
     use crate::prelude::*;
     use crate::utils::hash::sha512;
-    use crate::utils::setup_logger;
 
     fn test_sha512_fixed(msg: &[u8], expected_digest: [u8; 64]) {
-        setup_logger();
         let mut builder = DefaultBuilder::new();
         let message = msg
             .iter()
@@ -203,7 +208,6 @@ mod tests {
     }
 
     fn test_sha512_variable_length(message: &[u8], input_length: u32, expected_digest: [u8; 64]) {
-        setup_logger();
         let mut builder = DefaultBuilder::new();
 
         let input_length = builder.constant::<U32Variable>(input_length);
@@ -268,9 +272,14 @@ mod tests {
         test_sha512_variable_length(&msg, 0, expected_digest);
     }
 
+    // FAILED
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha512_curta_variable_large_message() {
+        env::set_var("RUST_LOG", "debug");
+        env_logger::try_init().unwrap_or_default();
+        dotenv::dotenv().ok();
+
         let mut msg : Vec<u8> = bytes!("35c323757c20640a294345c89c0bfcebe3d554fdb0c7b7a0bdb72222c531b1ecf7ec1c43f4de9d49556de87b86b26a98942cb078486fdb44de38b80864c3973153756363696e6374204c616273");
         let len = msg.len() as u32;
         msg.resize(256, 1);
@@ -279,9 +288,14 @@ mod tests {
         test_sha512_variable_length(&msg, len, expected_digest);
     }
 
+    // FAILED
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha512_curta_variable_short_message_same_slice() {
+        env::set_var("RUST_LOG", "debug");
+        env_logger::try_init().unwrap_or_default();
+        dotenv::dotenv().ok();
+
         let mut msg: Vec<u8> = b"plonky2".to_vec();
         let len = msg.len() as u32;
         msg.resize(128, 1);
@@ -290,9 +304,14 @@ mod tests {
         test_sha512_variable_length(&msg, len, expected_digest);
     }
 
+    // FAILED
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha512_curta_variable_short_message_different_slice() {
+        env::set_var("RUST_LOG", "debug");
+        env_logger::try_init().unwrap_or_default();
+        dotenv::dotenv().ok();
+
         let mut msg: Vec<u8> = b"plonky2".to_vec();
         let len = msg.len() as u32;
         msg.resize(512, 1);
@@ -304,7 +323,6 @@ mod tests {
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha512_fixed_length() {
-        setup_logger();
         let mut builder = DefaultBuilder::new();
 
         let max_len = 300;
@@ -331,7 +349,6 @@ mod tests {
     #[test]
     #[cfg_attr(feature = "ci", ignore)]
     fn test_sha512_variable_length_random() {
-        setup_logger();
         let mut builder = DefaultBuilder::new();
 
         let max_number_of_chunks = 2;
@@ -367,7 +384,6 @@ mod tests {
     fn test_sha512_variable_length_max_size() {
         // This test checks that sha512_variable_pad works as intended, especially when the max
         // input length is (length % 128 > 128 - 17).
-        setup_logger();
         let mut builder = DefaultBuilder::new();
 
         let max_number_of_chunks = 1;