Port C++ merkle tests to Rust

Cargo-guest.lock

@@ -51,9 +51,9 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "cpufeatures"
-version = "0.2.2"
+version = "0.2.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "59a6001667ab124aebae2a495118e11d30984c3a653e99d86d58971708cf5e4b"
+checksum = "dc948ebb96241bb40ab73effeb80d9f93afaad49359d159a5e61be51619fe813"
 dependencies = [
  "libc",
 ]
@@ -160,18 +160,18 @@ checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7"
 
 [[package]]
 name = "serde"
-version = "1.0.143"
+version = "1.0.144"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "53e8e5d5b70924f74ff5c6d64d9a5acd91422117c60f48c4e07855238a254553"
+checksum = "0f747710de3dcd43b88c9168773254e809d8ddbdf9653b84e2554ab219f17860"
 dependencies = [
  "serde_derive",
 ]
 
 [[package]]
 name = "serde_derive"
-version = "1.0.143"
+version = "1.0.144"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d3d8e8de557aee63c26b85b947f5e59b690d0454c753f3adeb5cd7835ab88391"
+checksum = "94ed3a816fb1d101812f83e789f888322c34e291f894f19590dc310963e87a00"
 dependencies = [
  "proc-macro2",
  "quote",

Cargo-host.lock

@@ -335,9 +335,9 @@ checksum = "5827cebf4670468b8772dd191856768aedcb1b0278a04f989f7766351917b9dc"
 
 [[package]]
 name = "cpufeatures"
-version = "0.2.2"
+version = "0.2.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "59a6001667ab124aebae2a495118e11d30984c3a653e99d86d58971708cf5e4b"
+checksum = "dc948ebb96241bb40ab73effeb80d9f93afaad49359d159a5e61be51619fe813"
 dependencies = [
  "libc",
 ]
@@ -997,9 +997,9 @@ checksum = "8371e4e5341c3a96db127eb2465ac681ced4c433e01dd0e938adbef26ba93ba5"
 
 [[package]]
 name = "link-cplusplus"
-version = "1.0.6"
+version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f8cae2cd7ba2f3f63938b9c724475dfb7b9861b545a90324476324ed21dbc8c8"
+checksum = "9272ab7b96c9046fbc5bc56c06c117cb639fe2d509df0c421cad82d2915cf369"
 dependencies = [
  "cc",
 ]
@@ -1273,9 +1273,9 @@ checksum = "1df8c4ec4b0627e53bdf214615ad287367e482558cf84b109250b37464dc03ae"
 
 [[package]]
 name = "plotters"
-version = "0.3.2"
+version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9428003b84df1496fb9d6eeee9c5f8145cb41ca375eb0dad204328888832811f"
+checksum = "716b4eeb6c4a1d3ecc956f75b43ec2e8e8ba80026413e70a3f41fd3313d3492b"
 dependencies = [
  "num-traits",
  "plotters-backend",
@@ -1292,9 +1292,9 @@ checksum = "193228616381fecdc1224c62e96946dfbc73ff4384fba576e052ff8c1bea8142"
 
 [[package]]
 name = "plotters-svg"
-version = "0.3.2"
+version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e0918736323d1baff32ee0eade54984f6f201ad7e97d5cfb5d6ab4a358529615"
+checksum = "f9a81d2759aae1dae668f783c308bc5c8ebd191ff4184aaa1b37f65a6ae5a56f"
 dependencies = [
  "plotters-backend",
 ]
@@ -1773,9 +1773,9 @@ dependencies = [
 
 [[package]]
 name = "security-framework"
-version = "2.6.1"
+version = "2.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2dc14f172faf8a0194a3aded622712b0de276821addc574fa54fc0a1167e10dc"
+checksum = "2bc1bb97804af6631813c55739f771071e0f2ed33ee20b68c86ec505d906356c"
 dependencies = [
  "bitflags",
  "core-foundation",
@@ -1805,9 +1805,9 @@ dependencies = [
 
 [[package]]
 name = "serde"
-version = "1.0.143"
+version = "1.0.144"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "53e8e5d5b70924f74ff5c6d64d9a5acd91422117c60f48c4e07855238a254553"
+checksum = "0f747710de3dcd43b88c9168773254e809d8ddbdf9653b84e2554ab219f17860"
 dependencies = [
  "serde_derive",
 ]
@@ -1824,9 +1824,9 @@ dependencies = [
 
 [[package]]
 name = "serde_derive"
-version = "1.0.143"
+version = "1.0.144"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d3d8e8de557aee63c26b85b947f5e59b690d0454c753f3adeb5cd7835ab88391"
+checksum = "94ed3a816fb1d101812f83e789f888322c34e291f894f19590dc310963e87a00"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -1835,9 +1835,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.83"
+version = "1.0.85"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38dd04e3c8279e75b31ef29dbdceebfe5ad89f4d0937213c53f7d49d01b3d5a7"
+checksum = "e55a28e3aaef9d5ce0506d0a14dbba8054ddc7e499ef522dd8b26859ec9d4a44"
 dependencies = [
  "itoa 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "ryu",

risc0/zkp/rust/src/prove/merkle.rs

@@ -125,3 +125,240 @@ impl<H: Hal> MerkleTreeProver<H> {
         out
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        core::sha_cpu,
+        hal::cpu::CpuHal,
+        verify::{merkle::MerkleTreeVerifier, read_iop::ReadIOP, VerificationError},
+    };
+    use rand::{Rng, RngCore};
+
+    use super::*;
+
+    fn init_prover<H: Hal>(
+        hal: &H,
+        rows: usize,
+        cols: usize,
+        queries: usize,
+    ) -> MerkleTreeProver<H> {
+        // Initialize a prover with leaves 0..size
+        let size: u32 = (rows * cols) as u32;
+        let mut data = Vec::<Fp>::new();
+        for val in 0..size {
+            data.push(Fp::from(val));
+        }
+        let matrix = hal.copy_fp_from(data.as_slice());
+
+        MerkleTreeProver::new(hal, &matrix, rows, cols, queries)
+    }
+
+    fn bad_row_access<H: Hal, S: Sha>(sha: &S, hal: &H, rows: usize, cols: usize, queries: usize) {
+        let prover = init_prover(hal, rows, cols, queries);
+        let mut iop = WriteIOP::new(sha);
+        prover.prove(&mut iop, rows);
+    }
+
+    fn possibly_bad_verify<H: Hal, S: Sha>(
+        sha: &S,
+        hal: &H,
+        rows: usize,
+        cols: usize,
+        queries: usize,
+        bad_query: usize,
+        manipulate_proof: bool,
+    ) {
+        let prover = init_prover(hal, rows, cols, queries);
+
+        let mut iop: WriteIOP<S> = WriteIOP::new(sha);
+        prover.commit(hal, &mut iop);
+        for _query in 0..queries {
+            let r_idx = (iop.rng.next_u32() as usize) % rows;
+            let col = prover.prove(&mut iop, r_idx);
+            for c_idx in 0..cols {
+                assert_eq!(col[c_idx], Fp::from((r_idx + c_idx * rows) as u32));
+            }
+        }
+        {
+            if manipulate_proof {
+                let mut rng = rand::thread_rng();
+                let manip_idx = rng.gen::<usize>() % iop.proof.len();
+                iop.proof[manip_idx] ^= 1;
+            }
+            let mut r_iop = ReadIOP::new(sha, &iop.proof);
+            let verifier = MerkleTreeVerifier::new(&mut r_iop, rows, cols, queries);
+            assert_eq!(verifier.root(), prover.root());
+            let mut err = false;
+            for query in 0..queries {
+                let r_idx = (r_iop.next_u32() as usize) % rows;
+                if query == bad_query {
+                    if rows == 1 {
+                        assert!(false, "Cannot test for bad query if there is only one row");
+                    }
+                    let r_idx = (r_idx + 1) % rows;
+                    let verification = verifier.verify(&mut r_iop, r_idx);
+                    match verification {
+                        Ok(_) => assert!(
+                            false,
+                            "Merkle tree wrongly passed verify when tested on the wrong row"
+                        ),
+                        Err(VerificationError::InvalidProof) => {}
+                        Err(_) => assert!(
+                            false,
+                            "Merkle tree failed validation for an unexpected reason"
+                        ),
+                    }
+                    err = true;
+                    break;
+                }
+                let col = verifier.verify(&mut r_iop, r_idx).unwrap();
+                for c_idx in 0..cols {
+                    assert_eq!(col[c_idx], Fp::from((r_idx + c_idx * rows) as u32));
+                }
+            }
+            if !err {
+                r_iop.verify_complete();
+            }
+        }
+    }
+
+    fn randomize_sizes() -> (usize, usize, usize) {
+        // Chooses random values of `rows`, `cols`, and `queries` such that:
+        // `rows` is a power of 2
+        // `cols` & `queries` have a wide distribution but tend to take small values
+        let mut rng = rand::thread_rng();
+        let rows = 1 << (rng.gen::<u8>() % 10);
+        let cols = (rng.gen::<usize>() % (1 << (rng.gen::<u8>() % 10))) + 1;
+        let queries = (rng.gen::<usize>() % (1 << (rng.gen::<u8>() % 10))) + 1;
+        (rows, cols, queries)
+    }
+
+    #[test]
+    #[should_panic(expected = "assertion failed: idx < self.params.row_size")]
+    fn merkle_cpu_1_1_1_bad_row_access() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        bad_row_access(&sha, &hal, 1, 1, 1);
+    }
+
+    #[test]
+    #[should_panic(expected = "assertion failed: idx < self.params.row_size")]
+    fn merkle_cpu_4_4_2_bad_row_access() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        bad_row_access(&sha, &hal, 4, 4, 2);
+    }
+
+    #[test]
+    #[should_panic(expected = "assertion failed: idx < self.params.row_size")]
+    fn merkle_cpu_randomized_bad_row_access() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        let (rows, cols, queries) = randomize_sizes();
+        bad_row_access(&sha, &hal, rows, cols, queries);
+    }
+
+    #[test]
+    fn merkle_cpu_1_1_1_verify() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        // Test a complete verification with no bad queries (by setting bad_query out of
+        // range)
+        possibly_bad_verify(&sha, &hal, 1, 1, 1, 4, false);
+    }
+
+    #[test]
+    fn merkle_cpu_4_4_2_verify() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        // Test a complete verification with no bad queries (by setting bad_query out of
+        // range)
+        possibly_bad_verify(&sha, &hal, 4, 4, 2, 4, false);
+    }
+
+    #[test]
+    fn merkle_cpu_randomized_verify() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        for _rep in 0..100 {
+            let (rows, cols, queries) = randomize_sizes();
+            // Test a complete verification with no bad queries (by setting bad_query out of
+            // range)
+            possibly_bad_verify(&sha, &hal, rows, cols, queries, queries + 1, false);
+        }
+    }
+
+    #[test]
+    fn merkle_cpu_2_1_1_bad_query() {
+        // n.b. since we test bad queries by incrementing the row, we can't test for a
+        // bad query with rows == 1
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        possibly_bad_verify(&sha, &hal, 2, 1, 1, 0, false);
+    }
+
+    #[test]
+    fn merkle_cpu_4_4_2_bad_query() {
+        let mut rng = rand::thread_rng();
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        let queries = 2;
+        // Test a complete verification with a bad query
+        let bad_query = rng.gen::<usize>() % queries;
+        possibly_bad_verify(&sha, &hal, 4, 4, queries, bad_query, false);
+    }
+
+    #[test]
+    fn merkle_cpu_randomized_bad_query() {
+        let mut rng = rand::thread_rng();
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        let (rows, cols, queries) = randomize_sizes();
+        // At least two rows are required to test querying an incorrect row
+        let rows = if rows == 1 { 2 } else { rows };
+        // Test a complete verification with a bad query
+        let bad_query = rng.gen::<usize>() % queries;
+        possibly_bad_verify(&sha, &hal, rows, cols, queries, bad_query, false);
+    }
+
+    #[test]
+    #[should_panic]
+    fn merkle_cpu_1_1_1_verify_manipulated() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        for _rep in 0..50 {
+            // Test a verification with a manipulated proof but no bad queries (by setting
+            // bad_query out of range) Do this multiple times as the
+            // manipulation location is random
+            possibly_bad_verify(&sha, &hal, 1, 1, 1, 2, true);
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn merkle_cpu_4_4_2_verify_manipulated() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        for _rep in 0..50 {
+            // Test a verification with a manipulated proof but no bad queries (by setting
+            // bad_query out of range) Do this multiple times as the
+            // manipulation location is random
+            possibly_bad_verify(&sha, &hal, 4, 4, 2, 4, true);
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn merkle_cpu_randomized_verify_manipulated() {
+        let sha = sha_cpu::Impl {};
+        let hal = CpuHal::new();
+        for _rep in 0..50 {
+            let (rows, cols, queries) = randomize_sizes();
+            // Test a verification with a manipulated proof but no bad queries (by setting
+            // bad_query out of range) Do this multiple times as the
+            // manipulation location is random
+            possibly_bad_verify(&sha, &hal, rows, cols, queries, queries + 1, true);
+        }
+    }
+}

risc0/zkp/rust/src/verify/mod.rs

@@ -14,7 +14,7 @@
 
 pub mod adapter;
 mod fri;
-mod merkle;
+pub(crate) mod merkle;
 pub mod read_iop;
 
 use alloc::vec;