update to do full encryption & hash plaintext inside zkVM

Author: nuke-web3

host/src/main.rs

@@ -5,7 +5,7 @@ use hex_literal::hex;
 // These constants represent the RISC-V ELF and the image ID generated by risc0-build.
 // The ELF is used for proving and the ID is used for verification.
 use methods::{GUEST_ENCRYPT_ELF, GUEST_ENCRYPT_ID};
-use risc0_zkvm::{default_prover, ExecutorEnv};
+use risc0_zkvm::{default_prover, sha::Digest, ExecutorEnv};
 
 fn main() {
     // Initialize tracing. In order to view logs, run `RUST_LOG=info cargo run`
@@ -22,19 +22,13 @@ fn main() {
     let plaintext = hex!("00010203 04050607 08090A0B 0C0D0E0F");
     println!("Plaintext bytes: {:?}", plaintext);
 
-    // Key and IV must be references to the `GenericArray` type.
-    // Here we use the `Into` trait to convert arrays into it.
-    let mut cipher = ChaCha20::new(&key.into(), &nonce.into());
-    // Generate raw keystream bytes
-    let mut keystream = [0u8; 16]; // ChaCha20 generates keystream in 64-byte blocks, but this can be any number of bits.
-    cipher.apply_keystream(&mut keystream); // keystream XOR with 0s = keystream
-    println!("Keystream bytes: {:?}", keystream);
-
     // zkVM
     let env = ExecutorEnv::builder()
-        .write(&plaintext)
+        .write(&key)
+        .unwrap()
+        .write(&nonce)
         .unwrap()
-        .write(&keystream)
+        .write(&plaintext)
         .unwrap()
         .build()
         .unwrap();
@@ -49,23 +43,26 @@ fn main() {
     // extract the receipt.
     let receipt = prove_info.receipt;
 
-    let mut output: [u8; 16] = receipt.journal.decode().unwrap();
-    println!("zkVM output bytes: {:?}", output);
+    let mut output: (Digest, [u8; 16]) = receipt.journal.decode().unwrap();
+    println!("zkVM -> plaintext hash:   {:?}", output.0);
+    println!("zkVM -> ciphertext bytes: {:?}", output.1);
 
     let ciphertext = hex!("e405626e 4f1236b3 670ee428 332ea20e");
     println!("expected output bytes: {:?}", ciphertext);
 
     // The receipt was verified at the end of proving, but the below code is an
     // example of how someone else could verify this receipt.
-    receipt
-        .verify(GUEST_ENCRYPT_ID)
-        .unwrap();
+    receipt.verify(GUEST_ENCRYPT_ID).unwrap();
     println!("Reciept from zkVM OK!");
-    
+
+    // Key and IV must be references to the `GenericArray` type.
+    // Here we use the `Into` trait to convert arrays into it.
+    let mut cipher = ChaCha20::new(&key.into(), &nonce.into());
+
     // ChaCha ciphers support seeking
     cipher.seek(0u32);
 
     // decrypt ciphertext by applying keystream again
-    cipher.apply_keystream(&mut output);
-    assert_eq!(output, plaintext);
+    cipher.apply_keystream(&mut output.1);
+    assert_eq!(output.1, plaintext);
 }

methods/guest/src/main.rs

@@ -1,29 +1,30 @@
-use risc0_zkvm::guest::env;
+use risc0_zkvm::{
+    guest::env,
+    sha::{Impl, Sha256},
+}
+;use chacha20::ChaCha20;
+use chacha20::cipher::{KeyIvInit, StreamCipher};
 
 fn main() {
-    let plaintext: [u8; 16] = env::read();
-    let keystream: [u8; 16] = env::read();
+    let key: [u8; 32] = env::read();
+    let nonce: [u8; 12] = env::read();
+    // Expects to be filled with plaintext to be encrypted
+    let mut buffer: [u8; 16] = env::read();
+
+    // Hash plaintext & commit
+    let digest = Impl::hash_bytes(&buffer);
+    env::commit(&digest);
 
     // TODO:
-    // - Hash plaintext & commit?
-    // - Hash keystream & commit?
-   
-    let mut buffer = plaintext.clone();
-    xor_arrays(&plaintext,&keystream, &mut buffer);
-    
-    // write public output to the journal
-    env::commit(&buffer);
-}
+    // - Hash key and/or nonce & commit?
+
+    // Key and IV must be references to the `GenericArray` type.
+    // Here we use the `Into` trait to convert arrays into it.
+    let mut cipher = ChaCha20::new(&key.into(), &nonce.into());
 
-/// UNSAFE: We demand input and output arrays here are the SAME length. 
-fn xor_arrays(arr1: &[u8], arr2: &[u8], output: &mut [u8]) {
-    // assert_eq!(arr1.len(), arr2.len(), "Arrays must have the same length");
-    // assert_eq!(arr1.len(), output.len(), "Output array must have the same length");
+    // Write ciphertext into buffer by applying keystream to plaintext
+    cipher.apply_keystream(&mut buffer);
 
-    for i in 0..arr1.len() {
-        // We assume that input arrays are the same length, skipping checks to save cycles
-        unsafe {
-            *output.get_unchecked_mut(i) = arr1.get_unchecked(i) ^ arr2.get_unchecked(i);
-        }
-    }
+    // write public output to the journal
+    env::commit(&buffer);
 }