1use serde::Serialize;
19use serde::de::DeserializeOwned;
20use std::alloc::{GlobalAlloc, Layout, System};
21use std::sync::Once;
22
23pub mod action;
24pub mod base64;
25pub mod diagnose;
26pub mod export;
27pub mod http;
28pub mod images;
29pub mod ops;
30pub mod server;
31#[cfg(test)]
32mod ops_tests;
33mod word_classes;
34#[cfg(test)]
35mod text_tests;
36
37pub const API_VERSION: u32 = 1;
39
40pub const CACHE_MAX_KEY: usize = 1024;
41pub const CACHE_MAX_VALUE: usize = 1024 * 1024;
42
43#[link(wasm_import_module = "extism:host/env")]
46unsafe extern "C" {
47 fn input_length() -> u64;
48 fn input_load_u8(index: u64) -> u8;
49 fn input_load_u64(index: u64) -> u64;
50 fn alloc(length: u64) -> u64;
51 fn free(offset: u64);
52 fn length(offset: u64) -> u64;
53 fn load_u8(offset: u64) -> u8;
54 fn load_u64(offset: u64) -> u64;
55 fn store_u8(offset: u64, value: u8);
56 fn store_u64(offset: u64, value: u64);
57 fn output_set(offset: u64, length: u64);
58 fn error_set(offset: u64);
59 fn log_debug(offset: u64);
60 fn log_info(offset: u64);
61 fn log_warn(offset: u64);
62 fn log_error(offset: u64);
63}
64
65#[link(wasm_import_module = "extism:host/user")]
66unsafe extern "C" {
67 fn hammer_dispatch(request: u64) -> u64;
68 fn hammer_cache_get(key: u64) -> u64;
69 fn hammer_cache_set(key: u64, value: u64);
70 fn hammer_progress(report: u64);
71 fn hammer_resource(name: u64) -> u64;
72}
73
74fn to_kernel(bytes: &[u8]) -> u64 {
75 unsafe {
76 let offset = alloc(bytes.len() as u64);
77 let mut chunks = bytes.chunks_exact(8);
78 let mut at = offset;
79 for chunk in &mut chunks {
80 store_u64(at, u64::from_le_bytes(chunk.try_into().unwrap()));
81 at += 8;
82 }
83 for &byte in chunks.remainder() {
84 store_u8(at, byte);
85 at += 1;
86 }
87 offset
88 }
89}
90
91fn from_kernel(offset: u64) -> Vec<u8> {
92 unsafe {
93 let size = length(offset) as usize;
94 let mut bytes = Vec::with_capacity(size);
95 let mut at = 0;
96 while at + 8 <= size {
97 bytes.extend_from_slice(&load_u64(offset + at as u64).to_le_bytes());
98 at += 8;
99 }
100 while at < size {
101 bytes.push(load_u8(offset + at as u64));
102 at += 1;
103 }
104 bytes
105 }
106}
107
108pub type Error = Box<dyn std::error::Error>;
110
111pub type Result<T> = std::result::Result<T, Error>;
112
113pub fn run<O: AsRef<[u8]>>(body: impl FnOnce() -> Result<O>) -> i32 {
115 static PANIC_HOOK: Once = Once::new();
116 PANIC_HOOK.call_once(|| std::panic::set_hook(Box::new(|info| set_error(&info.to_string()))));
118 match body() {
119 Ok(output) => {
120 let output = output.as_ref();
121 unsafe { output_set(to_kernel(output), output.len() as u64) };
122 0
123 }
124 Err(error) => {
125 set_error(&error.to_string());
126 1
127 }
128 }
129}
130
131fn set_error(message: &str) {
132 unsafe { error_set(to_kernel(message.as_bytes())) };
133}
134
135pub fn input() -> Vec<u8> {
136 unsafe {
137 let size = input_length() as usize;
138 let mut bytes = Vec::with_capacity(size);
139 let mut at = 0;
140 while at + 8 <= size {
141 bytes.extend_from_slice(&input_load_u64(at as u64).to_le_bytes());
142 at += 8;
143 }
144 while at < size {
145 bytes.push(input_load_u8(at as u64));
146 at += 1;
147 }
148 bytes
149 }
150}
151
152pub fn input_json<T: DeserializeOwned>() -> Result<T> {
153 Ok(serde_json::from_slice(&input())?)
154}
155
156#[derive(serde::Deserialize)]
157struct Reply {
158 #[serde(default)]
159 output: serde_json::Value,
160 error: Option<ReplyError>,
161}
162
163#[derive(serde::Deserialize)]
164struct ReplyError {
165 kind: String,
166 #[serde(default)]
167 message: String,
168}
169
170pub fn call<T: DeserializeOwned>(operation: &str, input: &impl Serialize) -> Result<T> {
174 #[derive(Serialize)]
175 struct Request<'a, I> {
176 operation: &'a str,
177 input: &'a I,
178 }
179 let request = to_kernel(&serde_json::to_vec(&Request { operation, input })?);
180 let reply = unsafe { hammer_dispatch(request) };
181 unsafe { free(request) };
182 let bytes = from_kernel(reply);
183 unsafe { free(reply) };
184 let reply: Reply = serde_json::from_slice(&bytes)?;
185 match reply.error {
186 Some(error) => Err(format!("{}: {}", error.kind, error.message).into()),
187 None => Ok(serde_json::from_value(reply.output)?),
189 }
190}
191
192pub fn cache_get(key: &str) -> Option<Vec<u8>> {
199 if key.len() > CACHE_MAX_KEY {
200 return None;
201 }
202 let key = to_kernel(key.as_bytes());
203 let value = unsafe { hammer_cache_get(key) };
204 unsafe { free(key) };
205 if value == 0 {
206 return None;
207 }
208 let bytes = from_kernel(value);
209 unsafe { free(value) };
210 Some(bytes)
211}
212
213pub fn cache_set(key: &str, value: &[u8]) -> bool {
216 if key.len() > CACHE_MAX_KEY || value.len() > CACHE_MAX_VALUE {
217 return false;
218 }
219 let key = to_kernel(key.as_bytes());
220 let value = if value.is_empty() { 0 } else { to_kernel(value) };
221 unsafe {
222 hammer_cache_set(key, value);
223 free(key);
224 if value != 0 {
225 free(value);
226 }
227 }
228 true
229}
230
231pub fn resource(name: &str) -> Option<Vec<u8>> {
235 let name = to_kernel(name.as_bytes());
236 let value = unsafe { hammer_resource(name) };
237 unsafe { free(name) };
238 if value == 0 {
239 return None;
240 }
241 let bytes = from_kernel(value);
242 unsafe { free(value) };
243 Some(bytes)
244}
245
246pub fn progress(fraction: Option<f64>, message: Option<&str>) {
250 #[derive(Serialize)]
251 struct Report<'a> {
252 #[serde(skip_serializing_if = "Option::is_none")]
253 fraction: Option<f64>,
254 #[serde(skip_serializing_if = "Option::is_none")]
255 message: Option<&'a str>,
256 }
257 let fraction = fraction.map(|it| it.clamp(0.0, 1.0));
258 let Ok(report) = serde_json::to_vec(&Report { fraction, message }) else { return };
259 let handle = to_kernel(&report);
260 unsafe {
261 hammer_progress(handle);
262 free(handle);
263 }
264}
265
266pub fn char_byte_offsets(text: &str) -> Vec<usize> {
269 text.char_indices().map(|(at, _)| at).chain(std::iter::once(text.len())).collect()
270}
271
272#[derive(Debug, Clone, Copy, PartialEq, Eq)]
274pub struct Word<'a> {
275 pub text: &'a str,
276 pub start: usize,
277 pub end: usize,
278}
279
280pub fn words(text: &str) -> impl Iterator<Item = Word<'_>> {
293 let class_at = move |at: usize| text[at..].chars().next().map_or(word_class::OTHER, word_class::of);
294 let after = move |at: usize| at + text[at..].chars().next().map_or(0, char::len_utf8);
295 let mut at = 0;
296 std::iter::from_fn(move || {
297 while at < text.len() {
298 let start = at;
299 let class = class_at(at);
300 at = after(at);
301 match class {
302 word_class::IDEOGRAPH => {
303 while class_at(at) == word_class::MARK {
304 at = after(at);
305 }
306 }
307 word_class::LETTER => loop {
308 let next = class_at(at);
309 let inside = next == word_class::LETTER
310 || next == word_class::MARK
311 || (next == word_class::JOINER && class_at(after(at)) == word_class::LETTER);
312 if !inside {
313 break;
314 }
315 at = after(at);
316 },
317 _ => continue,
318 }
319 return Some(Word { text: &text[start..at], start, end: at });
320 }
321 None
322 })
323}
324
325mod word_class {
326 pub const OTHER: u8 = 0;
327 pub const LETTER: u8 = 1;
328 pub const MARK: u8 = 2;
329 pub const IDEOGRAPH: u8 = 3;
330 pub const JOINER: u8 = 4;
331
332 pub fn of(c: char) -> u8 {
333 match c {
334 'a'..='z' | 'A'..='Z' => LETTER,
335 '\'' => JOINER,
336 '\0'..='\x7F' => OTHER,
337 _ => {
338 let code = c as u32;
339 crate::word_classes::WORD_CLASSES
340 .binary_search_by(|&(first, last, _)| {
341 if last < code {
342 std::cmp::Ordering::Less
343 } else if first > code {
344 std::cmp::Ordering::Greater
345 } else {
346 std::cmp::Ordering::Equal
347 }
348 })
349 .map_or(OTHER, |i| crate::word_classes::WORD_CLASSES[i].2)
350 }
351 }
352 }
353}
354
355pub struct ZeroingAllocator;
363
364unsafe impl GlobalAlloc for ZeroingAllocator {
365 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
366 unsafe { System.alloc(layout) }
367 }
368
369 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
370 unsafe {
371 ptr.write_bytes(0, layout.size());
372 System.dealloc(ptr, layout)
373 }
374 }
375}
376
377#[cfg(feature = "getrandom")]
378mod random {
379 #[link(wasm_import_module = "wasi_snapshot_preview1")]
380 unsafe extern "C" {
381 fn random_get(buf: *mut u8, len: usize) -> u16;
382 }
383
384 #[unsafe(no_mangle)]
387 unsafe extern "Rust" fn __getrandom_v03_custom(dest: *mut u8, len: usize) -> Result<(), getrandom::Error> {
388 match unsafe { random_get(dest, len) } {
389 0 => Ok(()),
390 _ => Err(getrandom::Error::UNSUPPORTED),
391 }
392 }
393}
394
395pub mod log {
396 fn send(to: unsafe extern "C" fn(u64), message: &str) {
397 let handle = super::to_kernel(message.as_bytes());
398 unsafe {
399 to(handle);
400 super::free(handle);
401 }
402 }
403
404 pub fn debug(message: &str) {
405 send(super::log_debug, message)
406 }
407
408 pub fn info(message: &str) {
409 send(super::log_info, message)
410 }
411
412 pub fn warn(message: &str) {
413 send(super::log_warn, message)
414 }
415
416 pub fn error(message: &str) {
417 send(super::log_error, message)
418 }
419}