Skip to main content

submilli_engine/stdlib/
crypto.rs

1//! `submilli:crypto` — SHA-2 digests, HMAC, OS randomness, constant-time
2//! compare. Pure Rust host functions registered directly under the package
3//! name; the `string | Uint8Array` union params are discriminated host-side.
4
5use crate::runtime::effects;
6use crate::runtime::host::{
7    abi_arg, abi_result, read_string_arg, read_uint8_array_arg, register_host_fn,
8    register_host_value_async, write_submilli_uint8array_struct,
9};
10use hmac::{Hmac, Mac};
11use sha2::{Digest, Sha256, Sha512};
12use wasmtime::{Caller, FuncType, HeapType, Linker, RefType, Val, ValType};
13
14use crate::runtime::StoreData;
15use crate::runtime::fuel;
16use crate::runtime::intrinsic_types::{build_intrinsic_types, intrinsic_types};
17use crate::runtime::prelude::collection::is_a;
18use crate::{PackageDeclaration, Param, Span, Type, ValueKind, ValueSymbol};
19
20pub const MODULE_NAME: &str = "submilli:crypto";
21
22const MAX_RANDOM_BYTES: i32 = 1024 * 1024; // 1 MiB
23
24pub fn package_declaration() -> PackageDeclaration {
25    let mut defs = PackageDeclaration::with_package(MODULE_NAME);
26
27    let string_or_bytes = Type::union(vec![Type::String, Type::Uint8Array]);
28
29    insert_fn(
30        &mut defs,
31        "sha256",
32        vec![Param::new("input", string_or_bytes.clone())],
33        Type::Uint8Array,
34        "/**\n * SHA-256 digest of `input`. Returns the 32-byte digest as a `Uint8Array`. Pick `.toHex()` or `.toBase64()` at the call site to get text.\n * @param input UTF-8 text (string) or raw bytes (Uint8Array).\n */",
35    );
36    insert_fn(
37        &mut defs,
38        "sha512",
39        vec![Param::new("input", string_or_bytes.clone())],
40        Type::Uint8Array,
41        "/**\n * SHA-512 digest of `input`. Returns the 64-byte digest as a `Uint8Array`. Pick `.toHex()` or `.toBase64()` at the call site to get text.\n * @param input UTF-8 text (string) or raw bytes (Uint8Array).\n */",
42    );
43    insert_fn(
44        &mut defs,
45        "hmacSha256",
46        vec![
47            Param::new("key", Type::Uint8Array),
48            Param::new("message", string_or_bytes),
49        ],
50        Type::Uint8Array,
51        "/**\n * HMAC-SHA-256 of `message` under `key`. Returns the 32-byte tag as a `Uint8Array`.\n * @param key Secret key bytes (any length).\n * @param message UTF-8 text (string) or raw bytes (Uint8Array).\n */",
52    );
53    insert_fn(
54        &mut defs,
55        "randomBytes",
56        vec![Param::new("length", Type::Number)],
57        Type::Uint8Array,
58        "/**\n * Cryptographically secure random bytes from the OS entropy source. Traps if `length` is negative or above the per-call cap (1 MiB).\n * @param length Number of bytes to draw.\n */",
59    );
60    insert_fn(
61        &mut defs,
62        "timingSafeEqual",
63        vec![
64            Param::new("a", Type::Uint8Array),
65            Param::new("b", Type::Uint8Array),
66        ],
67        Type::Boolean,
68        "/**\n * Constant-time byte-array equality, suitable for comparing HMAC tags. Returns `false` for length mismatch without scanning; length is not assumed secret.\n * @param a First byte sequence.\n * @param b Second byte sequence.\n */",
69    );
70
71    defs
72}
73
74fn insert_fn(defs: &mut PackageDeclaration, name: &str, params: Vec<Param>, ret: Type, doc: &str) {
75    defs.values.insert(
76        name.to_string(),
77        ValueSymbol {
78            name: name.to_string(),
79            mangled_name: crate::mangle::package_symbol(MODULE_NAME, name),
80            declaration_span: Span::at(crate::FileId::CRYPTO),
81            kind: ValueKind::Function {
82                generics: Vec::new(),
83                params,
84                ret,
85                type_predicate: None,
86                doc: crate::doc(crate::FileId::CRYPTO, doc),
87            },
88        },
89    );
90}
91
92/// Read a `string | Uint8Array` union param — codegen passes it as the erased
93/// `(ref $Object)` — into the bytes to feed the digest (strings as UTF-8).
94fn read_string_or_bytes(
95    caller: &mut Caller<'_, StoreData>,
96    val: &Val,
97    context: &str,
98) -> wasmtime::Result<Vec<u8>> {
99    let string_ty = intrinsic_types(&mut *caller)?.string.clone();
100    if is_a(caller, val, &string_ty)? {
101        Ok(read_string_arg(caller, val, context)?.into_bytes())
102    } else {
103        read_uint8_array_arg(caller, val, context)
104    }
105}
106
107pub fn install(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
108    let engine = linker.engine().clone();
109    let intr = build_intrinsic_types(&engine)?;
110    let object = ValType::Ref(RefType::new(false, HeapType::ConcreteStruct(intr.object)));
111    let uint8 = ValType::Ref(RefType::new(
112        false,
113        HeapType::ConcreteStruct(intr.uint8_array),
114    ));
115
116    for (name, digest) in [
117        ("sha256", sha256_digest as fn(&[u8]) -> Vec<u8>),
118        ("sha512", sha512_digest),
119    ] {
120        let context = format!("crypto.{name}");
121        register_host_fn(
122            linker,
123            MODULE_NAME,
124            crate::mangle::package_symbol(MODULE_NAME, name),
125            FuncType::new(&engine, [object.clone()], [uint8.clone()]),
126            move |caller, params, results| {
127                let input = read_string_or_bytes(caller, abi_arg(params, 0)?, &context)?;
128                fuel::charge(&mut *caller, fuel::HASH, input.len() as u64)?;
129                let arr = write_submilli_uint8array_struct(caller, &digest(&input))?;
130                *abi_result(results, 0)? = Val::AnyRef(Some(arr.to_anyref()));
131                Ok(())
132            },
133        )?;
134    }
135
136    register_host_fn(
137        linker,
138        MODULE_NAME,
139        crate::mangle::package_symbol(MODULE_NAME, "hmacSha256"),
140        FuncType::new(&engine, [uint8.clone(), object], [uint8.clone()]),
141        |caller, params, results| {
142            let key =
143                read_uint8_array_arg(&mut *caller, abi_arg(params, 0)?, "crypto.hmacSha256 (key)")?;
144            let msg =
145                read_string_or_bytes(caller, abi_arg(params, 1)?, "crypto.hmacSha256 (message)")?;
146            fuel::charge(&mut *caller, fuel::HASH, (key.len() + msg.len()) as u64)?;
147            let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(&key)
148                .map_err(|e| crate::runtime::host::type_error(format!("crypto.hmacSha256: {e}")))?;
149            mac.update(&msg);
150            let arr = write_submilli_uint8array_struct(caller, &mac.finalize().into_bytes())?;
151            *abi_result(results, 0)? = Val::AnyRef(Some(arr.to_anyref()));
152            Ok(())
153        },
154    )?;
155
156    register_host_value_async(
157        linker,
158        MODULE_NAME,
159        crate::mangle::package_symbol(MODULE_NAME, "randomBytes"),
160        FuncType::new(&engine, [ValType::F64], [uint8.clone()]),
161        |caller, params| {
162            Box::pin(async move {
163                let Val::F64(bits) = *abi_arg(params, 0)? else {
164                    return Err(crate::runtime::host::type_error(
165                        "crypto.randomBytes: expected f64 length",
166                    ));
167                };
168                let len = f64::from_bits(bits) as i32;
169                if len < 0 {
170                    return Err(crate::runtime::host::range_error(format!(
171                        "crypto.randomBytes: length must be non-negative, got {len}"
172                    )));
173                }
174                if len > MAX_RANDOM_BYTES {
175                    return Err(crate::runtime::host::range_error(format!(
176                        "crypto.randomBytes: length {len} exceeds maximum {MAX_RANDOM_BYTES}"
177                    )));
178                }
179                let buf = effects::execute(caller, "crypto.randomBytes", len, |caller, _, len| {
180                    Box::pin(async move {
181                        fuel::charge(&mut *caller, fuel::SCAN, len as u64)?;
182                        let mut bytes = vec![0u8; len as usize];
183                        getrandom::getrandom(&mut bytes).map_err(|e| {
184                            wasmtime::Error::msg(format!("crypto.randomBytes: {e}"))
185                        })?;
186                        Ok(effects::OutcomeBytes(bytes))
187                    })
188                })
189                .await?;
190                let arr = write_submilli_uint8array_struct(caller, &buf.0)?;
191                Ok(Val::AnyRef(Some(arr.to_anyref())))
192            })
193        },
194    )?;
195
196    // Constant-time compare; length mismatch short-circuits (length is not secret).
197    register_host_fn(
198        linker,
199        MODULE_NAME,
200        crate::mangle::package_symbol(MODULE_NAME, "timingSafeEqual"),
201        FuncType::new(&engine, [uint8.clone(), uint8], [ValType::I32]),
202        |caller, params, results| {
203            let a = read_uint8_array_arg(
204                &mut *caller,
205                abi_arg(params, 0)?,
206                "crypto.timingSafeEqual (a)",
207            )?;
208            let b = read_uint8_array_arg(
209                &mut *caller,
210                abi_arg(params, 1)?,
211                "crypto.timingSafeEqual (b)",
212            )?;
213            fuel::charge(&mut *caller, fuel::SCAN, a.len().max(b.len()) as u64)?;
214            let equal = a.len() == b.len() && {
215                let mut diff: u8 = 0;
216                for (x, y) in a.iter().zip(b.iter()) {
217                    diff |= x ^ y;
218                }
219                diff == 0
220            };
221            *abi_result(results, 0)? = Val::I32(i32::from(equal));
222            Ok(())
223        },
224    )?;
225
226    Ok(())
227}
228
229fn sha256_digest(input: &[u8]) -> Vec<u8> {
230    let mut hasher = Sha256::new();
231    hasher.update(input);
232    hasher.finalize().to_vec()
233}
234
235fn sha512_digest(input: &[u8]) -> Vec<u8> {
236    let mut hasher = Sha512::new();
237    hasher.update(input);
238    hasher.finalize().to_vec()
239}