1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
use k256::SecretKey;
use k256::elliptic_curve::sec1::ToEncodedPoint;
use risc0_zkvm::guest::env;
use risc0_zkvm::sha::{Impl, Sha256};
use wasmi::{Caller, Engine, Extern, Linker, Memory, Module, Store};
use zkwid_zkvm_core::{Journal, ProveInput};
fn memory(caller: &Caller<'_, ()>) -> Memory {
match caller.get_export("memory") {
Some(Extern::Memory(m)) => m,
_ => panic!("no memory"),
}
}
fn keygen(mut caller: Caller<'_, ()>, seed_ptr: i32, out_ptr: i32) -> i32 {
let mem = memory(&caller);
let mut seed = [0u8; 32];
mem.read(&caller, seed_ptr as u32 as usize, &mut seed)
.expect("seed");
let key = SecretKey::from_slice(&seed).expect("secret key");
let point = key.public_key().to_encoded_point(true);
let bytes = point.as_bytes();
mem.write(&mut caller, out_ptr as u32 as usize, bytes)
.expect("out");
bytes.len() as i32
}
fn run_shim(wasm: &[u8], seed: &[u8; 32]) -> Vec<u8> {
let engine = Engine::default();
let module = Module::new(&engine, wasm).expect("module");
let mut store = Store::new(&engine, ());
let mut linker = <Linker<()>>::new(&engine);
linker.func_wrap("env", "keygen", keygen).expect("link");
let instance = linker
.instantiate(&mut store, &module)
.expect("instantiate")
.start(&mut store)
.expect("start");
let mem = instance.get_memory(&store, "memory").expect("memory");
let input_ptr = instance
.get_typed_func::<(), i32>(&store, "input_ptr")
.expect("input_ptr");
let output_ptr = instance
.get_typed_func::<(), i32>(&store, "output_ptr")
.expect("output_ptr");
let run = instance
.get_typed_func::<i32, i32>(&store, "run")
.expect("run");
let inp = input_ptr.call(&mut store, ()).expect("input_ptr call") as u32 as usize;
mem.write(&mut store, inp, seed).expect("write seed");
let len = run.call(&mut store, seed.len() as i32).expect("run call");
assert_eq!(len, 33, "output length");
let out = output_ptr.call(&mut store, ()).expect("output_ptr call") as u32 as usize;
let mut key = vec![0u8; 33];
mem.read(&store, out, &mut key).expect("read key");
key
}
fn main() {
let input: ProveInput = env::read();
assert!(!input.stages.is_empty(), "no stages");
let measurements = input
.stages
.iter()
.map(|s| Impl::hash_bytes(s).as_bytes().try_into().expect("digest"))
.collect();
let key = run_shim(&input.stages[0], &input.seed);
env::commit(&Journal {
measurements,
key,
nonce: input.nonce,
aud: input.aud,
});
}
