Fix KZG precompile I/O (#213)

* Fix parsing of calldata for KZG precompile

* Fix returndata

* Pacify mighty clippy

* Fix versioned_hash check

* Update test accordingly

evm_arithmetization/Cargo.toml

@@ -35,6 +35,7 @@ rand_chacha = "0.3.1"
 rlp = { workspace = true }
 rlp-derive = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
+sha2 = "0.10.6"
 static_assertions = "1.1.0"
 hashbrown = { version = "0.14.0" }
 tiny-keccak = "2.0.2"
@@ -50,7 +51,6 @@ jemallocator = "0.5.0"
 criterion = "0.5.1"
 hex = { workspace = true }
 ripemd = "0.1.3"
-sha2 = "0.10.6"
 
 [features]
 default = ["parallel"]

evm_arithmetization/src/cpu/kernel/asm/core/precompiles/kzg_peval.asm

@@ -66,10 +66,8 @@ global store_kzg_verification:
     PUSH @SEGMENT_RETURNDATA
     %build_address_no_offset
     // stack: addr, res_lo, res_hi, kexit_info
-    DUP1 %add_const(32)
-    %stack (addr_hi, addr_lo, res_lo, res_hi)
-        -> (addr_lo, res_lo, addr_hi, res_hi)
     MSTORE_32BYTES_32
+    // stack: addr', res_hi, kexit_info
     MSTORE_32BYTES_32
     // stack: kexit_info
 

evm_arithmetization/src/cpu/kernel/tests/bls381.rs

@@ -13,6 +13,7 @@ use crate::cpu::kernel::constants::context_metadata::ContextMetadata;
 use crate::cpu::kernel::interpreter::Interpreter;
 use crate::extension_tower::{Fp2, Stack, BLS381};
 use crate::memory::segments::Segment::{self, KernelGeneral};
+use crate::util::sha2;
 use crate::witness::errors::ProgramError;
 
 #[test]
@@ -116,16 +117,22 @@ const KZG_PRECOMPILE_TEST_SEQUENCES: [TestSequence; 10] = [
 fn test_kzg_peval_precompile() -> Result<()> {
     for (bytes, is_correct) in KZG_PRECOMPILE_TEST_SEQUENCES.iter() {
         let commitment_bytes = bytes.0;
-        let comm_hi = U256::from_big_endian(&commitment_bytes[0..16]);
-        let comm_lo = U256::from_big_endian(&commitment_bytes[16..48]);
-        let mut versioned_hash = keccak(commitment_bytes).0;
-        versioned_hash[0] = KZG_VERSIONED_HASH;
-        let versioned_hash = U256::from_big_endian(&versioned_hash);
+        let comm_hi = U256::from_big_endian(&commitment_bytes[0..32]);
+        let comm_lo = U256::from_big_endian(&commitment_bytes[32..48]);
+        let mut versioned_hash = sha2(commitment_bytes.to_vec());
+        const KZG_HASH_MASK: U256 = U256([
+            0xffffffffffffffff,
+            0xffffffffffffffff,
+            0xffffffffffffffff,
+            0x00ffffffffffffff,
+        ]);
+        versioned_hash &= KZG_HASH_MASK; // erase most significant byte
+        versioned_hash |= U256::from(KZG_VERSIONED_HASH) << 248; // append 1
         let z = U256::from_big_endian(&bytes.1);
         let y = U256::from_big_endian(&bytes.2);
         let proof_bytes = bytes.3;
-        let proof_hi = U256::from_big_endian(&proof_bytes[0..16]);
-        let proof_lo = U256::from_big_endian(&proof_bytes[16..48]);
+        let proof_hi = U256::from_big_endian(&proof_bytes[0..32]);
+        let proof_lo = U256::from_big_endian(&proof_bytes[32..48]);
 
         let mut stack = vec![
             versioned_hash,

evm_arithmetization/src/cpu/kernel/tests/hash.rs

@@ -4,11 +4,11 @@ use ethereum_types::U256;
 use plonky2::field::goldilocks_field::GoldilocksField as F;
 use rand::{thread_rng, Rng};
 use ripemd::{Digest, Ripemd160};
-use sha2::Sha256;
 
 use super::{run_interpreter_with_memory, InterpreterMemoryInitialization};
 use crate::cpu::kernel::interpreter::Interpreter;
 use crate::memory::segments::Segment::KernelGeneral;
+use crate::util::sha2;
 
 /// Standard RipeMD implementation.
 fn ripemd(input: Vec<u8>) -> U256 {
@@ -17,13 +17,6 @@ fn ripemd(input: Vec<u8>) -> U256 {
     U256::from(&hasher.finalize()[..])
 }
 
-/// Standard Sha2 implementation.
-fn sha2(input: Vec<u8>) -> U256 {
-    let mut hasher = Sha256::new();
-    hasher.update(input);
-    U256::from(&hasher.finalize()[..])
-}
-
 fn make_random_input() -> Vec<u8> {
     // Generate a random message, between 0 and 9999 bytes.
     let mut rng = thread_rng();

evm_arithmetization/src/generation/prover_input.rs

@@ -29,7 +29,7 @@ use crate::generation::prover_input::FieldOp::{Inverse, Sqrt};
 use crate::generation::state::GenerationState;
 use crate::memory::segments::Segment;
 use crate::memory::segments::Segment::BnPairing;
-use crate::util::{biguint_to_mem_vec, h2u, mem_vec_to_biguint, u256_to_u8, u256_to_usize};
+use crate::util::{biguint_to_mem_vec, h2u, mem_vec_to_biguint, sha2, u256_to_u8, u256_to_usize};
 use crate::witness::errors::ProverInputError::*;
 use crate::witness::errors::{ProgramError, ProverInputError};
 use crate::witness::memory::MemoryAddress;
@@ -429,18 +429,40 @@ impl<F: Field> GenerationState<F> {
             ));
         }
 
-        let mut comm_bytes = [0u8; 64];
+        let mut comm_bytes = [0u8; 48];
+        comm_lo.to_big_endian(&mut comm_bytes[16..48]); // only actually 16 bytes
+        if comm_bytes[16..32] != [0; 16] {
+            // Commitments must fit in 48 bytes.
+            return Err(ProgramError::ProverInputError(
+                ProverInputError::KzgEvalFailure(
+                    "Commitment does not fit in 48 bytes.".to_string(),
+                ),
+            ));
+        }
         comm_hi.to_big_endian(&mut comm_bytes[0..32]);
-        comm_lo.to_big_endian(&mut comm_bytes[32..64]); // only actually 16 bits
 
-        let mut proof_bytes = [0u8; 64];
+        let mut proof_bytes = [0u8; 48];
+        proof_lo.to_big_endian(&mut proof_bytes[16..48]); // only actually 16 bytes
+        if proof_bytes[16..32] != [0; 16] {
+            // Proofs must fit in 48 bytes.
+            return Err(ProgramError::ProverInputError(
+                ProverInputError::KzgEvalFailure("Proof does not fit in 48 bytes.".to_string()),
+            ));
+        }
         proof_hi.to_big_endian(&mut proof_bytes[0..32]);
-        proof_lo.to_big_endian(&mut proof_bytes[32..64]); // only actually 16 bits
 
-        let mut expected_versioned_hash = keccak(&comm_bytes[16..64]).0;
-        expected_versioned_hash[0] = KZG_VERSIONED_HASH;
+        let mut expected_versioned_hash = sha2(comm_bytes.to_vec());
 
-        if versioned_hash != U256::from_big_endian(&expected_versioned_hash) {
+        const KZG_HASH_MASK: U256 = U256([
+            0xffffffffffffffff,
+            0xffffffffffffffff,
+            0xffffffffffffffff,
+            0x00ffffffffffffff,
+        ]);
+        expected_versioned_hash &= KZG_HASH_MASK; // erase most significant byte
+        expected_versioned_hash |= U256::from(KZG_VERSIONED_HASH) << 248; // append 1
+
+        if versioned_hash != expected_versioned_hash {
             return Err(ProgramError::ProverInputError(
                 ProverInputError::KzgEvalFailure(
                     "Versioned hash does not match expected value.".to_string(),
@@ -475,21 +497,12 @@ impl<F: Field> GenerationState<F> {
     /// Verifies a KZG proof, i.e. that the commitment opens to y at z.
     fn verify_kzg_proof(
         &self,
-        comm_bytes: &[u8; 64],
+        comm_bytes: &[u8; 48],
         z: U256,
         y: U256,
-        proof_bytes: &[u8; 64],
+        proof_bytes: &[u8; 48],
     ) -> Result<U256, ProgramError> {
-        if comm_bytes[0..16] != [0; 16] || proof_bytes[0..16] != [0; 16] {
-            // Proofs and commitments must fit in 48 bytes to be deserializable.
-            return Err(ProgramError::ProverInputError(
-                ProverInputError::KzgEvalFailure(
-                    "Proof or commitment do not fit in 48 bytes.".to_string(),
-                ),
-            ));
-        }
-
-        let comm = if let Ok(c) = bls381::g1_from_bytes(comm_bytes[16..64].try_into().unwrap()) {
+        let comm = if let Ok(c) = bls381::g1_from_bytes(comm_bytes) {
             c
         } else {
             return Err(ProgramError::ProverInputError(
@@ -499,7 +512,7 @@ impl<F: Field> GenerationState<F> {
             ));
         };
 
-        let proof = if let Ok(p) = bls381::g1_from_bytes(proof_bytes[16..64].try_into().unwrap()) {
+        let proof = if let Ok(p) = bls381::g1_from_bytes(proof_bytes) {
             p
         } else {
             return Err(ProgramError::ProverInputError(

evm_arithmetization/src/util.rs

@@ -8,6 +8,7 @@ use plonky2::field::packed::PackedField;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
+use sha2::{Digest, Sha256};
 
 use crate::witness::errors::ProgramError;
 
@@ -255,3 +256,10 @@ pub(crate) fn get_h256<F: RichField>(slice: &[F]) -> H256 {
             .collect_vec(),
     )
 }
+
+/// Standard Sha2 implementation.
+pub(crate) fn sha2(input: Vec<u8>) -> U256 {
+    let mut hasher = Sha256::new();
+    hasher.update(input);
+    U256::from(&hasher.finalize()[..])
+}