Adding Shake256

shake256-demo/Cargo.toml

@@ -12,10 +12,15 @@ 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"

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,15 @@
+[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/lib.rs

@@ -0,0 +1,2 @@
+#![no_std]
+#![no_main]

shake256-demo/methods/guest/src/main.rs

@@ -0,0 +1,95 @@
+#![no_std]
+#![no_main]
+
+extern crate alloc;
+use alloc::vec; 
+use alloc::vec::Vec;
+use risc0_zkvm::guest::env;
+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);
+
+
+#[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/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

@@ -1,37 +1,92 @@
 #[cfg(not(rust_analyzer))]
-mod generated {
+pub mod methods {
     include!(concat!(env!("OUT_DIR"), "/methods.rs"));
 }
-#[cfg(not(rust_analyzer))]
-pub use generated::{GUEST_ELF, GUEST_ID};
 
+#[cfg(not(rust_analyzer))]
 
 #[cfg(rust_analyzer)]
 pub const GUEST_ELF: &[u8] = &[];
 #[cfg(rust_analyzer)]
 pub const GUEST_ID: [u32; 8] = [0; 8];
 
-use anyhow::Result;
-use risc0_zkvm::{default_prover, ExecutorEnv, Receipt};
 
 
-pub fn prove_encrypt(
-    key: [u8; 32],
-    nonce: [u8; 12],
-    plaintext: &[u8],
-) -> Result<(Receipt, Vec<u8>)> {
-    let env = ExecutorEnv::builder()
-        .write(&key)?
-        .write(&nonce)?
-        .write(&plaintext.to_vec())?
-        .build()?;
+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;
 
-    let prover = default_prover();
-    let prove_info = prover.prove(env, GUEST_ELF)?;
-    let receipt = prove_info.receipt;
-    receipt.verify(GUEST_ID)?;
 
-    let ciphertext: Vec<u8> = receipt.journal.bytes.clone(); // if this errors, use .to_vec()
-    
-    Ok((receipt, ciphertext))
+#[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 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,
+}
+
+/// KDF per NSSA spec: HKDF-SHA256 with fixed salt and domain-separated info.
+pub 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();
+    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::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)
+        .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,44 +1,70 @@
 #[cfg(not(rust_analyzer))]
-include!(concat!(env!("OUT_DIR"), "/methods.rs"));
 
 #[cfg(rust_analyzer)]
-mod methods {
-    pub const GUEST_ELF: &[u8] = &[];
-    pub const GUEST_ID: [u32; 8] = [0; 8];
-}
+//mod methods {
+   // pub const GUEST_ELF: &[u8] = &[];
+   // pub const GUEST_ID: [u32; 8] = [0; 8];
+//}
 #[cfg(rust_analyzer)]
 use methods::*;
 
 
 use anyhow::Result; 
-use hex::encode;
 use risc0_zkvm::{default_prover, ExecutorEnv};
 
+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 key       = [0x42u8; 32];
-    let nonce     = [0x24u8; 12];
-    let plaintext = b"Hello, RISC Zero ChaCha20 demo!";
+    // 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;
 
-    let env = ExecutorEnv::builder()
-        .write(&key)?
-        .write(&nonce)?
-        .write(&plaintext.to_vec())?
-        .build()?;
-  
-   
-    let prover  = default_prover();
-    let prove_info = prover.prove(env, GUEST_ELF)?;
-    let receipt = prove_info.receipt; 
-   
-    // (Optionally) verify the proof
-    receipt.verify(GUEST_ID)?;
+    let input = EncInput {
+        plaintext: plaintext.clone(),
+        ss_bytes,
+        epk_bytes: epk.clone(),
+        ipk_bytes: ipk.clone(),
+        commitment,
+        out_index,
+    };
 
-    //  Extract and print the ciphertext
-    let ct: &[u8] = &receipt.journal.bytes;
-    println!("Ciphertext: {}", encode(ct));
+    // 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(())
 }
 
-