shake256 with only 32 bytes

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"

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 }

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<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};
+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<u8>,
+    pub ipk_bytes: Vec<u8>,
+    pub commitment: [u8; 32],
+    pub out_index: u32,
+    pub pt: Vec<u8>,
 }
+
+
+#[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],
+    ipk: &[u8],
+    commitment: &[u8; 32],
+    out_index: u32,
+) -> ([u8; 32], Vec<u8>) {
+    // 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::<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_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);
+}

shake256-demo/methods/src/lib.rs

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

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<u8>,
+    pub ipk_bytes: Vec<u8>,
+    pub commitment: [u8; 32],
+    pub out_index: u32,
+    pub pt: Vec<u8>,
+}
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct KdfInputs {
+    pub epk_bytes: Vec<u8>,
+    pub ipk_bytes: Vec<u8>,
+    pub commitment: [u8; 32],
+    pub out_index: u32,
+}
+
+#[derive(Debug, serde::Serialize, serde::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)
+    pub ss_bytes: [u8; 32],     // ECDH x-coordinate (sender view)
+    pub epk_bytes: Vec<u8>,
+    pub ipk_bytes: Vec<u8>,
+    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<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"))]
+/// 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<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();
+    let mut hasher = Sha256::new();
+    Sha2Digest::update(&mut hasher, b"NSSA/v0.1/KDF-SHA256"); // avoids E0034
+    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::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::<Sha256>::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<u8> {
+    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()
+}

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<u8>,
+    pub ipk_bytes: Vec<u8>,
+    pub commitment: [u8; 32],
+    pub out_index: u32,
+    pub pt: Vec<u8>,
+}
 
 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<u8> = 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::<String>());
+
+    // 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::<String>());
 
     assert_eq!(guest_ct, host_ct);
+    println!("OK: guest and host ciphertexts match.");
     Ok(())
-}
+}