shake256 demo using R0

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)'] }

shake256-demo/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    risc0_build::embed_methods();
+}

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"]

shake256-demo/methods/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    risc0_build::embed_methods();
+}

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 }

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<u8>,
+        pub k_enc: [u8; 32],
+        pub info: Vec<u8>,
+    }
+
+    pub fn enc_xor_shake256(k_enc: &[u8; 32], info: &[u8], pt: &[u8]) -> Vec<u8> {
+        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);
+}

shake256-demo/methods/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod methods {
+    include!(concat!(env!("OUT_DIR"), "/methods.rs"));
+}

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<u8>,
+    pub ss_bytes: [u8; 32],
+    pub epk_bytes: Vec<u8>,
+    pub ipk_bytes: Vec<u8>,
+    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<u8>) {
+    // 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::<Sha256>::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<u8> {
+    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);
+}

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<u8>,
+    pub k_enc: [u8; 32],  // 32-byte encryption key
+    pub info: Vec<u8>,    // 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<u8> {
+    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<u8>) {
+    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::<Sha256>::new(Some(&salt), &ss_bytes);
+
+    let mut k_enc = [0u8; 32];
+    hk.expand(&info, &mut k_enc).expect("HKDF expand");
+
+    (k_enc, info)
+}

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<u8> = 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::<String>());
+    println!("host  ct: {}", host_ct.encode_hex::<String>());
+
+    assert_eq!(guest_ct, host_ct);
+    Ok(())
+}