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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
use std::path::PathBuf;
use std::process::Command;
use anyhow::{Context, bail, ensure};
use atproto_identity::key::{KeyType, generate_key};
use wasmi::core::Trap;
use wasmi::{Caller, Engine, Linker, Module, Store};
pub fn seed() -> [u8; 32] {
generate_key(KeyType::K256Private)
.expect("k256 key generation")
.bytes()
.try_into()
.expect("32-byte k256 scalar")
}
pub fn stage_bytes(name: &str) -> anyhow::Result<Vec<u8>> {
let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("../../stages")
.join(name);
let manifest = dir.join("Cargo.toml");
ensure!(manifest.is_file(), "no stage {name}");
let status = Command::new(std::env::var_os("CARGO").unwrap_or_else(|| "cargo".into()))
.args([
"build",
"--quiet",
"--release",
"--target",
"wasm32-unknown-unknown",
"--manifest-path",
])
.arg(&manifest)
.env_remove("CARGO_TARGET_DIR")
.env_remove("RUSTFLAGS")
.env_remove("CARGO_ENCODED_RUSTFLAGS")
.status()
.context("running cargo")?;
ensure!(status.success(), "building stage {name}: {status}");
let wasm = dir.join(format!("target/wasm32-unknown-unknown/release/{name}.wasm"));
std::fs::read(&wasm).with_context(|| format!("reading {}", wasm.display()))
}
pub fn run_stage(wasm: &[u8], input: &[u8]) -> anyhow::Result<Vec<u8>> {
let engine = Engine::default();
let module = Module::new(&engine, wasm).map_err(anyhow::Error::msg)?;
let mut store = Store::new(&engine, ());
let mut linker = Linker::<()>::new(&engine);
linker
.func_wrap("env", "keygen", keygen)
.map_err(anyhow::Error::msg)?;
let instance = linker
.instantiate(&mut store, &module)
.and_then(|pre| pre.start(&mut store))
.map_err(anyhow::Error::msg)?;
let memory = instance
.get_memory(&store, "memory")
.context("no memory export")?;
let input_ptr = instance
.get_typed_func::<(), i32>(&store, "input_ptr")
.map_err(anyhow::Error::msg)?
.call(&mut store, ())
.map_err(anyhow::Error::msg)?;
let output_ptr = instance
.get_typed_func::<(), i32>(&store, "output_ptr")
.map_err(anyhow::Error::msg)?
.call(&mut store, ())
.map_err(anyhow::Error::msg)?;
memory
.write(&mut store, input_ptr as u32 as usize, input)
.map_err(anyhow::Error::msg)?;
let len = instance
.get_typed_func::<i32, i32>(&store, "run")
.map_err(anyhow::Error::msg)?
.call(&mut store, i32::try_from(input.len())?)
.map_err(anyhow::Error::msg)?;
let Ok(len) = usize::try_from(len) else {
bail!("negative output length {len}");
};
let mut output = vec![0; len];
memory
.read(&store, output_ptr as u32 as usize, &mut output)
.map_err(anyhow::Error::msg)?;
Ok(output)
}
fn keygen(mut caller: Caller<'_, ()>, seed_ptr: i32, out_ptr: i32) -> Result<i32, Trap> {
let memory = caller
.get_export("memory")
.and_then(|export| export.into_memory())
.ok_or_else(|| Trap::new("no memory export"))?;
let mut seed = [0u8; 32];
memory
.read(&caller, seed_ptr as u32 as usize, &mut seed)
.map_err(|e| Trap::new(e.to_string()))?;
let key = public_key(&seed).map_err(Trap::new)?;
memory
.write(&mut caller, out_ptr as u32 as usize, &key)
.map_err(|e| Trap::new(e.to_string()))?;
Ok(33)
}
pub(crate) fn public_key(seed: &[u8; 32]) -> Result<[u8; 33], String> {
let secret = k256::SecretKey::from_slice(seed).map_err(|e| e.to_string())?;
let point =
k256::elliptic_curve::sec1::ToEncodedPoint::to_encoded_point(&secret.public_key(), true);
point
.as_bytes()
.try_into()
.map_err(|_| "bad point".to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use atproto_identity::key::{KeyData, to_public};
#[test]
fn seed_is_k256_scalar() {
let a = seed();
assert_ne!(a, seed());
let key = to_public(&KeyData::new(KeyType::K256Private, a.to_vec())).unwrap();
assert_eq!(key.bytes(), public_key(&a).unwrap());
}
#[test]
fn invalid_seed_fails() {
assert!(public_key(&[0; 32]).is_err());
assert!(public_key(&[0xff; 32]).is_err());
}
#[test]
fn rejects_garbage_wasm() {
assert!(run_stage(b"not wasm", b"").is_err());
}
}
