lssa-zkvm-testing/shake256-33bytes-demo/src/main.rs

use hex::ToHex;
use risc0_zkvm::{default_prover, ExecutorEnv};

use shake256_33bytes_demo::{EncInput, enc_xor_shake256, nssa_kdf};      // now works via re-exports
use shake256_33bytes_demo::ser_bytes33::Bytes33;                        // for constructing wrapper
use methods::GUEST_ELF;                          // generated guest image

fn main() -> anyhow::Result<()> {
    
    let ss_bytes: [u8; 32] = [
    0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,
    0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F,
    0x10,0x11,0x12,0x13,0x14,0x15,0x16,0x17,
    0x18,0x19,0x1A,0x1B,0x1C,0x1D,0x1E,0x1F,
];

    let commitment: [u8; 32] = [
    0x80,0x81,0x82,0x83,0x84,0x85,0x86,0x87,
    0x88,0x89,0x8A,0x8B,0x8C,0x8D,0x8E,0x8F,
    0x90,0x91,0x92,0x93,0x94,0x95,0x96,0x97,
    0x98,0x99,0x9A,0x9B,0x9C,0x9D,0x9E,0x9F,
];

    let epk_raw: [u8; 33] = [
    0x02,
    0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,
    0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F,0x10,
    0x11,0x12,0x13,0x14,0x15,0x16,0x17,0x18,
    0x19,0x1A,0x1B,0x1C,0x1D,0x1E,0x1F,0x20,
];
    let ipk_raw: [u8; 33] = [
    0x03,
    0x21,0x22,0x23,0x24,0x25,0x26,0x27,0x28,
    0x29,0x2A,0x2B,0x2C,0x2D,0x2E,0x2F,0x30,
    0x31,0x32,0x33,0x34,0x35,0x36,0x37,0x38,
    0x39,0x3A,0x3B,0x3C,0x3D,0x3E,0x3F,0x40,
];

    let input = EncInput {
        ss_bytes,
        epk_bytes: Bytes33::from(epk_raw), 
        ipk_bytes: Bytes33::from(ipk_raw), 
        commitment,
        out_index: 7,
    };

    let env = ExecutorEnv::builder().write(&input)?.build()?;
    let prove_info = default_prover().prove(env, GUEST_ELF)?;
    let receipt = prove_info; 
    let guest_ct = receipt.receipt.journal.bytes.clone();


    // Locally recomputes host ct (optional sanity check)
    let info = {
        let mut v = Vec::with_capacity(66 + 32);
        v.extend_from_slice(input.epk_bytes.as_ref());
        v.extend_from_slice(input.ipk_bytes.as_ref());
        v.extend_from_slice(&input.commitment);
        v
    };
    let k_enc = nssa_kdf(
        input.ss_bytes,
        input.epk_bytes.as_ref(),
        input.ipk_bytes.as_ref(),
        &input.commitment,
        input.out_index,
    );

    let plaintext = b"hello";
    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>());

    Ok(())
}