From 45655fada2804b7cb4533f3efda941778ce1b685 Mon Sep 17 00:00:00 2001 From: Moudy Date: Wed, 13 Aug 2025 09:04:36 +0200 Subject: [PATCH] shake256 with only 32 bytes --- shake256-demo/Cargo.toml | 4 - shake256-demo/methods/guest/Cargo.toml | 3 +- shake256-demo/methods/guest/src/main.rs | 124 +++++++++++++++++------- shake256-demo/methods/src/lib.rs | 3 - shake256-demo/src/lib.rs | 115 +++++++++++++--------- shake256-demo/src/main.rs | 81 ++++++++++------ 6 files changed, 211 insertions(+), 119 deletions(-) delete mode 100644 shake256-demo/methods/src/lib.rs diff --git a/shake256-demo/Cargo.toml b/shake256-demo/Cargo.toml index 1713755..e79678c 100644 --- a/shake256-demo/Cargo.toml +++ b/shake256-demo/Cargo.toml @@ -4,9 +4,6 @@ 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"] } @@ -16,7 +13,6 @@ 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" diff --git a/shake256-demo/methods/guest/Cargo.toml b/shake256-demo/methods/guest/Cargo.toml index 18791c8..7c0bee6 100644 --- a/shake256-demo/methods/guest/Cargo.toml +++ b/shake256-demo/methods/guest/Cargo.toml @@ -7,8 +7,7 @@ edition = "2021" [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 } +serde = { version = "1", default-features = false, features = ["derive", "alloc"] } 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 index 34e3e15..d5b1b90 100644 --- a/shake256-demo/methods/guest/src/main.rs +++ b/shake256-demo/methods/guest/src/main.rs @@ -1,45 +1,95 @@ -#![no_main] #![no_std] +#![no_main] extern crate alloc; - +use alloc::vec; +use alloc::vec::Vec; 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}; +use hkdf::Hkdf; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use sha3::{digest::{ExtendableOutput, Update, XofReader}, Shake256}; +//use serde_big_array::BigArray; 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); + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct KeystreamRequest { + pub kdf_salt: [u8; 32], + pub ss_bytes: [u8; 32], + pub epk_bytes: Vec, + pub ipk_bytes: Vec, + pub commitment: [u8; 32], + pub out_index: u32, + pub pt: Vec, } + + +#[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], + ipk: &[u8], + commitment: &[u8; 32], + out_index: u32, +) -> ([u8; 32], Vec) { + // salt = SHA256("NSSA/v0.1/KDF-SHA256") + let mut hasher = Sha256::new(); + sha2::Digest::update(&mut hasher, 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_bytes, + &input.ipk_bytes, + &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/methods/src/lib.rs b/shake256-demo/methods/src/lib.rs deleted file mode 100644 index 833f494..0000000 --- a/shake256-demo/methods/src/lib.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod methods { - include!(concat!(env!("OUT_DIR"), "/methods.rs")); -} \ No newline at end of file diff --git a/shake256-demo/src/lib.rs b/shake256-demo/src/lib.rs index 4389ce2..7e553c7 100644 --- a/shake256-demo/src/lib.rs +++ b/shake256-demo/src/lib.rs @@ -1,68 +1,93 @@ -#![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. +#[cfg(not(rust_analyzer))] pub mod methods { include!(concat!(env!("OUT_DIR"), "/methods.rs")); } -/// Inputs the host sends to the guest. +#[cfg(not(rust_analyzer))] + +#[cfg(rust_analyzer)] +pub const GUEST_ELF: &[u8] = &[]; +#[cfg(rust_analyzer)] +pub const GUEST_ID: [u32; 8] = [0; 8]; + + + +use hkdf::Hkdf; +use sha2::{Digest as Sha2Digest, Sha256}; +use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}}; +use serde::{Deserialize, Serialize}; +//use serde_big_array::BigArray; + + #[derive(Debug, Serialize, Deserialize, Clone)] +pub struct KeystreamRequest { + pub kdf_salt: [u8; 32], + pub ss_bytes: [u8; 32], + pub epk_bytes: Vec, + pub ipk_bytes: Vec, + pub commitment: [u8; 32], + pub out_index: u32, + pub pt: Vec, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KdfInputs { + pub epk_bytes: Vec, + pub ipk_bytes: Vec, + pub commitment: [u8; 32], + pub out_index: u32, +} + +#[derive(Debug, serde::Serialize, serde::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) + pub ss_bytes: [u8; 32], // ECDH x-coordinate (sender view) + pub epk_bytes: Vec, + pub ipk_bytes: Vec, + pub commitment: [u8; 32], + pub out_index: u32, } -/// 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"))] +/// KDF per NSSA spec: HKDF-SHA256 with fixed salt and domain-separated info. pub fn nssa_kdf( - ss_bytes: [u8; 32], - epk_bytes: &[u8], - ipk_bytes: &[u8], - commitment: [u8; 32], + ss_bytes: &[u8; 32], + epk: &[u8], + ipk: &[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(); + let mut hasher = Sha256::new(); + Sha2Digest::update(&mut hasher, b"NSSA/v0.1/KDF-SHA256"); // avoids E0034 + 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::new(); + 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_bytes); - info.extend_from_slice(ipk_bytes); - info.extend_from_slice(&commitment); + info.extend_from_slice(epk); + info.extend_from_slice(ipk); + 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"); + hk.expand(&info, &mut k_enc) + .expect("HKDF-Expand must produce 32 bytes"); (k_enc, info) } + +/// SHAKE256-XOF keystream XOR (symmetric). +pub 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"); // domain sep + 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() +} diff --git a/shake256-demo/src/main.rs b/shake256-demo/src/main.rs index 4403eb1..4287b18 100644 --- a/shake256-demo/src/main.rs +++ b/shake256-demo/src/main.rs @@ -1,39 +1,64 @@ -use anyhow::Result; -use hex::ToHex; +#[cfg(not(rust_analyzer))] + +#[cfg(rust_analyzer)] +use methods::*; + + +use anyhow::Result; use risc0_zkvm::{default_prover, ExecutorEnv}; -use shake256_demo::{EncInput, enc_xor_shake256, nssa_kdf, methods}; + +use shake256_demo::{enc_xor_shake256, nssa_kdf, EncInput}; +//ßuse methods::GUEST_ELF; + +use serde::{Serialize, Deserialize}; +//use serde_big_array::BigArray; +use hex::ToHex; // for encode_hex +use shake256_demo::methods::GUEST_ELF; // <— crate name uses underscore + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct KeystreamRequest { + pub kdf_salt: [u8; 32], + pub ss_bytes: [u8; 32], + pub epk_bytes: Vec, + pub ipk_bytes: Vec, + pub commitment: [u8; 32], + pub out_index: u32, + pub pt: Vec, +} 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(); + // Tiny demo inputs (in real NSSA you’d compute ss via ECDH). + let plaintext = b"hello NSSA (shake256)".to_vec(); + let ss_bytes = [1u8; 32]; + let epk = vec![2u8; 33]; + let ipk = vec![3u8; 33]; + let commitment = [4u8; 32]; + let out_index = 0u32; - // Derive (k_enc, info) once on the host - let (k_enc, info) = nssa_kdf(ss_bytes, &epk_bytes, &ipk_bytes, commitment, out_index); + let input = EncInput { + plaintext: plaintext.clone(), + ss_bytes, + epk_bytes: epk.clone(), + ipk_bytes: ipk.clone(), + 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); + // Prove inside zkVM + let env = ExecutorEnv::builder().write(&input)?.build()?; + let prove_info = default_prover().prove(env, GUEST_ELF)?; + println!("WARNING: proving in dev mode. Not production-safe."); + // Ciphertext from journal + let guest_ct = prove_info.receipt.journal.bytes.to_vec(); println!("guest ct: {}", guest_ct.encode_hex::()); + + // Host recompute (sanity) + let (k_enc, info) = nssa_kdf(&ss_bytes, &epk, &ipk, &commitment, out_index); + let host_ct = enc_xor_shake256(&k_enc, &info, &plaintext); println!("host ct: {}", host_ct.encode_hex::()); assert_eq!(guest_ct, host_ct); + println!("OK: guest and host ciphertexts match."); Ok(()) -} +} \ No newline at end of file