diff --git a/shake256-demo/Cargo.toml b/shake256-demo/Cargo.toml new file mode 100644 index 0000000..1713755 --- /dev/null +++ b/shake256-demo/Cargo.toml @@ -0,0 +1,36 @@ +[package] +name = "shake256-demo" +version = "0.1.0" +edition = "2021" +build = "build.rs" + +[[bin]] +name = "shake256-demo" +path = "src/main.rs" + +[dependencies] +risc0-zkvm = { version = "2.3.1", features = ["std", "prove"] } +anyhow = "1.0" +hex = "0.4" +serde = { version = "1.0", features = ["derive"] } +sha2 = "0.10" +hkdf = "0.12" +sha3 = "0.10" +serde-big-array = "0.5" + +[dev-dependencies] +rand = "0.8" +cipher = { version = "0.4", features = ["std"] } +serde = { version = "1", default-features = false, features = ["derive", "alloc"] } + + + + +[build-dependencies] +risc0-build = "2.3.1" + +[package.metadata.risc0] +methods = ["methods/guest"] + +[lints.rust] +unexpected_cfgs = { level = "allow", check-cfg = ['cfg(rust_analyzer)'] } \ No newline at end of file diff --git a/shake256-demo/build.rs b/shake256-demo/build.rs new file mode 100644 index 0000000..08a8a4e --- /dev/null +++ b/shake256-demo/build.rs @@ -0,0 +1,3 @@ +fn main() { + risc0_build::embed_methods(); +} diff --git a/shake256-demo/methods/Cargo.toml b/shake256-demo/methods/Cargo.toml new file mode 100644 index 0000000..9ab59da --- /dev/null +++ b/shake256-demo/methods/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "methods" +version = "0.1.0" +edition = "2021" + +[build-dependencies] +risc0-build = { version = "2.3.1" } + +[package.metadata.risc0] +methods = ["methods/guest"] diff --git a/shake256-demo/methods/build.rs b/shake256-demo/methods/build.rs new file mode 100644 index 0000000..08a8a4e --- /dev/null +++ b/shake256-demo/methods/build.rs @@ -0,0 +1,3 @@ +fn main() { + risc0_build::embed_methods(); +} diff --git a/shake256-demo/methods/guest/Cargo.toml b/shake256-demo/methods/guest/Cargo.toml new file mode 100644 index 0000000..18791c8 --- /dev/null +++ b/shake256-demo/methods/guest/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "guest" +version = "0.1.0" +edition = "2021" + +[workspace] + +[dependencies] +risc0-zkvm = { version = "2.3.1", default-features = false} +serde = { version = "1", default-features = false, features = ["derive"] } +serde-big-array = { version = "0.5", default-features = false } +sha2 = { version = "0.10", default-features = false } +hkdf = { version = "0.12", default-features = false } +sha3 = { version = "0.10", default-features = false } \ No newline at end of file diff --git a/shake256-demo/methods/guest/src/main.rs b/shake256-demo/methods/guest/src/main.rs new file mode 100644 index 0000000..34e3e15 --- /dev/null +++ b/shake256-demo/methods/guest/src/main.rs @@ -0,0 +1,45 @@ +#![no_main] +#![no_std] + +extern crate alloc; + +use risc0_zkvm::guest::env; + +mod shared { + extern crate alloc; + use alloc::vec::Vec; + use serde::{Deserialize, Serialize}; + use sha3::{ + digest::{ExtendableOutput, Update, XofReader}, + Shake256, + }; + + #[derive(Debug, Serialize, Deserialize, Clone)] + pub struct EncInput { + pub plaintext: Vec, + pub k_enc: [u8; 32], + pub info: Vec, + } + + pub fn enc_xor_shake256(k_enc: &[u8; 32], info: &[u8], pt: &[u8]) -> Vec { + let mut h = Shake256::default(); + h.update(k_enc); + h.update(info); + let mut xof = h.finalize_xof(); + + let mut ks = alloc::vec![0u8; pt.len()]; + xof.read(&mut ks); + + pt.iter().zip(ks).map(|(p, k)| p ^ k).collect() + } +} + +use shared::{enc_xor_shake256, EncInput}; + +risc0_zkvm::guest::entry!(main); + +pub fn main() { + let EncInput { plaintext, k_enc, info } = env::read(); + let ct = enc_xor_shake256(&k_enc, &info, &plaintext); + env::commit(&ct); +} diff --git a/shake256-demo/methods/src/lib.rs b/shake256-demo/methods/src/lib.rs new file mode 100644 index 0000000..833f494 --- /dev/null +++ b/shake256-demo/methods/src/lib.rs @@ -0,0 +1,3 @@ +pub mod methods { + include!(concat!(env!("OUT_DIR"), "/methods.rs")); +} \ No newline at end of file diff --git a/shake256-demo/methods/src/main.rs b/shake256-demo/methods/src/main.rs new file mode 100644 index 0000000..af9351d --- /dev/null +++ b/shake256-demo/methods/src/main.rs @@ -0,0 +1,83 @@ +#![no_std] +#![no_main] + +extern crate alloc; + +use alloc::vec::Vec; +use hkdf::Hkdf; +use risc0_zkvm::guest::env; +use risc0_zkvm::guest::entry; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use sha3::{digest::{ExtendableOutput, Update, XofReader}, Shake256}; + +entry!(main); + + +#[derive(Debug, Serialize, Deserialize)] +pub struct EncInput { + pub plaintext: Vec, + pub ss_bytes: [u8; 32], + pub epk_bytes: Vec, + pub ipk_bytes: Vec, + pub commitment: [u8; 32], + pub out_index: u32, +} + +fn nssa_kdf( + ss_bytes: &[u8; 32], + epk: &[u8; 32], + ipk: &[u8; 32], + commitment: &[u8; 32], + out_index: u32, +) -> ([u8; 32], Vec) { + // salt = SHA256("NSSA/v0.1/KDF-SHA256") + let mut hasher = Sha256::new(); + hasher.update(b"NSSA/v0.1/KDF-SHA256"); + let salt = hasher.finalize(); + + let hk = Hkdf::::new(Some(&salt), ss_bytes); + + // info = "NSSA/v0.1/enc" || Epk || Ipk || commitment || le(out_index) + let mut info = Vec::with_capacity(3 + 33 + 33 + 32 + 4 + 16); + info.extend_from_slice(b"NSSA/v0.1/enc"); + info.extend_from_slice(epk); + info.extend_from_slice(ipk); + info.extend_from_slice(commitment); + info.extend_from_slice(&out_index.to_le_bytes()); + + let mut k_enc = [0u8; 32]; + hk.expand(&info, &mut k_enc).unwrap(); + + (k_enc, info) +} + +fn enc_xor_shake256(k_enc: &[u8; 32], ad: &[u8], pt: &[u8]) -> Vec { + let mut shake = Shake256::default(); + shake.update(b"NSSA/v0.1/ENC/SHAKE256"); + shake.update(k_enc); + shake.update(ad); + let mut xof = shake.finalize_xof(); + + let mut ks = vec![0u8; pt.len()]; + xof.read(&mut ks); + + pt.iter().zip(ks).map(|(p, k)| p ^ k).collect() +} + +pub fn main() { + let input: EncInput = env::read(); + + let (k_enc, info) = nssa_kdf( + &input.ss_bytes, + &input.epk_compressed, + &input.ipk_compressed, + &input.commitment, + input.out_index, + ); + + let ct = enc_xor_shake256(&k_enc, &info, &input.plaintext); + + // Commit ciphertext to the journal + env::commit_slice(&ct); +} \ No newline at end of file diff --git a/shake256-demo/src/lib.rs b/shake256-demo/src/lib.rs new file mode 100644 index 0000000..4389ce2 --- /dev/null +++ b/shake256-demo/src/lib.rs @@ -0,0 +1,68 @@ +#![no_std] + +extern crate alloc; + +use alloc::vec; +use alloc::vec::Vec; +use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}}; +use serde::{Serialize, Deserialize}; + +/// Re-export the generated method symbols. +pub mod methods { + include!(concat!(env!("OUT_DIR"), "/methods.rs")); +} + +/// Inputs the host sends to the guest. +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct EncInput { + pub plaintext: Vec, + pub k_enc: [u8; 32], // 32-byte encryption key + pub info: Vec, // domain-separated info bytes (same on host & guest) +} + +/// XOR-encrypt using SHAKE256 keystream derived from k_enc || info. +pub fn enc_xor_shake256(k_enc: &[u8; 32], info: &[u8], pt: &[u8]) -> Vec { + let mut hasher = Shake256::default(); + hasher.update(k_enc); + hasher.update(info); + let mut xof = hasher.finalize_xof(); + let mut ks = vec![0u8; pt.len()]; + xof.read(&mut ks); + pt.iter().zip(ks).map(|(p, k)| p ^ k).collect() +} + +/// NSSA-style KDF: HKDF-SHA256(ss_bytes) with domain-separated salt+info. +/// Host-only (but compiles under no_std). +#[cfg(not(target_arch = "riscv32"))] +pub fn nssa_kdf( + ss_bytes: [u8; 32], + epk_bytes: &[u8], + ipk_bytes: &[u8], + commitment: [u8; 32], + out_index: u32, +) -> ([u8; 32], Vec) { + use hkdf::Hkdf; + use sha2::{Sha256, Digest as _}; + + // salt = SHA256("NSSA/v0.1/KDF-SHA256") + let mut h = Sha256::new(); + // Disambiguate to avoid the "multiple `update`" error: + sha2::Digest::update(&mut h, b"NSSA/v0.1/KDF-SHA256"); + let salt = h.finalize(); + + // info = "NSSA/v0.1/enc" || Epk || Ipk || commitment || le(out_index) + let mut info = Vec::new(); + info.extend_from_slice(b"NSSA/v0.1/enc"); + info.extend_from_slice(epk_bytes); + info.extend_from_slice(ipk_bytes); + info.extend_from_slice(&commitment); + info.extend_from_slice(&out_index.to_le_bytes()); + + // HKDF-Extract/Expand + let hk = Hkdf::::new(Some(&salt), &ss_bytes); + + let mut k_enc = [0u8; 32]; + hk.expand(&info, &mut k_enc).expect("HKDF expand"); + + (k_enc, info) +} diff --git a/shake256-demo/src/main.rs b/shake256-demo/src/main.rs new file mode 100644 index 0000000..4403eb1 --- /dev/null +++ b/shake256-demo/src/main.rs @@ -0,0 +1,39 @@ +use anyhow::Result; +use hex::ToHex; +use risc0_zkvm::{default_prover, ExecutorEnv}; +use shake256_demo::{EncInput, enc_xor_shake256, nssa_kdf, methods}; + +fn main() -> Result<()> { + // Example inputs. + let ss_bytes = [7u8; 32]; + let epk_bytes = vec![1u8; 33]; + let ipk_bytes = vec![2u8; 33]; + let commitment = [3u8; 32]; + let out_index = 0u32; + let plaintext = b"hello NSSA!".to_vec(); + + // Derive (k_enc, info) once on the host + let (k_enc, info) = nssa_kdf(ss_bytes, &epk_bytes, &ipk_bytes, commitment, out_index); + + // Build guest input + let input = EncInput { plaintext: plaintext.clone(), k_enc, info: info.clone() }; + + let env = ExecutorEnv::builder() + .write(&input)? + .build()?; + + // Prove + let receipt = default_prover().prove(env, methods::GUEST_ELF)?.receipt; + + // Get guest ciphertext + let guest_ct: Vec = receipt.journal.decode()?; + + // Compute host ciphertext the same way + let host_ct = enc_xor_shake256(&k_enc, &info, &plaintext); + + println!("guest ct: {}", guest_ct.encode_hex::()); + println!("host ct: {}", host_ct.encode_hex::()); + + assert_eq!(guest_ct, host_ct); + Ok(()) +}