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