heliosdb_proxy/plugins/
host_imports.rs

1//! Wasmtime-side host imports exposed to WASM plugins.
2//!
3//! Plugins import these from the `env` module:
4//!
5//! ```wat
6//! (import "env" "kv_get"    (func (param i32 i32 i32 i32) (result i32)))
7//! (import "env" "kv_set"    (func (param i32 i32 i32 i32) (result i32)))
8//! (import "env" "kv_delete" (func (param i32 i32)         (result i32)))
9//! ```
10//!
11//! The KV namespace is per-plugin: each plugin sees only its own
12//! key-value store, keyed off `LoadedPlugin.metadata.name`. State
13//! survives across calls because the `KvBackend` is owned by the
14//! runtime, not the per-call `Store`.
15//!
16//! Return-value conventions (i32):
17//!
18//! - `kv_get`: bytes written, or `-1` for missing key, or `-2` if the
19//!   caller's output buffer is too small (caller can retry with a
20//!   larger buffer; the value is left intact).
21//! - `kv_set`: `0` on success, `-1` on internal error.
22//! - `kv_delete`: `0` (idempotent — no error if the key was absent).
23//!
24//! The implementation is in-process and in-memory. A future slice
25//! can swap the backend for a persistent store (sled, redb, …)
26//! without changing the import surface.
27
28use std::collections::HashMap;
29use std::sync::Arc;
30
31use parking_lot::RwLock;
32use wasmtime::{Caller, Linker, Memory};
33
34use super::runtime::PluginError;
35
36/// In-memory KV backend, namespaced by plugin name. The outer map
37/// is keyed by plugin name; the inner map by user-supplied key.
38#[derive(Clone, Default)]
39pub struct KvBackend {
40    inner: Arc<RwLock<HashMap<String, HashMap<Vec<u8>, Vec<u8>>>>>,
41}
42
43impl KvBackend {
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    /// Read a value. None if missing.
49    pub fn get(&self, plugin: &str, key: &[u8]) -> Option<Vec<u8>> {
50        let g = self.inner.read();
51        g.get(plugin).and_then(|m| m.get(key).cloned())
52    }
53
54    /// Insert / overwrite.
55    pub fn set(&self, plugin: &str, key: Vec<u8>, value: Vec<u8>) {
56        let mut g = self.inner.write();
57        g.entry(plugin.to_string())
58            .or_default()
59            .insert(key, value);
60    }
61
62    /// Delete; idempotent.
63    pub fn delete(&self, plugin: &str, key: &[u8]) {
64        let mut g = self.inner.write();
65        if let Some(m) = g.get_mut(plugin) {
66            m.remove(key);
67        }
68    }
69
70    /// Returns the number of keys in the plugin's namespace.
71    /// Useful for tests and the future admin endpoint.
72    pub fn len(&self, plugin: &str) -> usize {
73        self.inner
74            .read()
75            .get(plugin)
76            .map(|m| m.len())
77            .unwrap_or(0)
78    }
79}
80
81/// Per-call store data: the plugin name (so host imports route to
82/// the right KV namespace) and a clone of the shared KV backend.
83/// Carrying the Arc<KvBackend> by value here is cheap (one atomic
84/// inc) and lets the import functions call `caller.data()` to
85/// retrieve it.
86pub struct StoreCtx {
87    pub plugin_name: String,
88    pub kv: KvBackend,
89}
90
91/// Register all host imports under the `env` module against the
92/// supplied linker. Idempotent — calling twice replaces prior bindings.
93pub fn register_kv_imports(linker: &mut Linker<StoreCtx>) -> Result<(), PluginError> {
94    linker
95        .func_wrap(
96            "env",
97            "kv_get",
98            |mut caller: Caller<'_, StoreCtx>, key_ptr: i32, key_len: i32, val_out_ptr: i32, val_max_len: i32| -> i32 {
99                let memory = match get_memory(&mut caller) {
100                    Some(m) => m,
101                    None => return -1,
102                };
103                let key = match read_bytes(&memory, &caller, key_ptr, key_len) {
104                    Some(b) => b,
105                    None => return -1,
106                };
107                let plugin_name = caller.data().plugin_name.clone();
108                let kv = caller.data().kv.clone();
109                let value = match kv.get(&plugin_name, &key) {
110                    Some(v) => v,
111                    None => return -1,
112                };
113                if (value.len() as i32) > val_max_len {
114                    return -2;
115                }
116                if write_bytes(&memory, &mut caller, val_out_ptr, &value).is_err() {
117                    return -1;
118                }
119                value.len() as i32
120            },
121        )
122        .map_err(|e| PluginError::RuntimeError(format!("link kv_get: {}", e)))?;
123
124    linker
125        .func_wrap(
126            "env",
127            "kv_set",
128            |mut caller: Caller<'_, StoreCtx>, key_ptr: i32, key_len: i32, val_ptr: i32, val_len: i32| -> i32 {
129                let memory = match get_memory(&mut caller) {
130                    Some(m) => m,
131                    None => return -1,
132                };
133                let key = match read_bytes(&memory, &caller, key_ptr, key_len) {
134                    Some(b) => b,
135                    None => return -1,
136                };
137                let val = match read_bytes(&memory, &caller, val_ptr, val_len) {
138                    Some(b) => b,
139                    None => return -1,
140                };
141                let plugin_name = caller.data().plugin_name.clone();
142                let kv = caller.data().kv.clone();
143                kv.set(&plugin_name, key, val);
144                0
145            },
146        )
147        .map_err(|e| PluginError::RuntimeError(format!("link kv_set: {}", e)))?;
148
149    linker
150        .func_wrap(
151            "env",
152            "kv_delete",
153            |mut caller: Caller<'_, StoreCtx>, key_ptr: i32, key_len: i32| -> i32 {
154                let memory = match get_memory(&mut caller) {
155                    Some(m) => m,
156                    None => return -1,
157                };
158                let key = match read_bytes(&memory, &caller, key_ptr, key_len) {
159                    Some(b) => b,
160                    None => return -1,
161                };
162                let plugin_name = caller.data().plugin_name.clone();
163                let kv = caller.data().kv.clone();
164                kv.delete(&plugin_name, &key);
165                0
166            },
167        )
168        .map_err(|e| PluginError::RuntimeError(format!("link kv_delete: {}", e)))?;
169
170    Ok(())
171}
172
173/// Register the `env.sha256_hex` host import. Plugins call:
174///
175/// ```text
176/// env.sha256_hex(in_ptr: i32, in_len: i32, out_ptr: i32) -> i32
177/// ```
178///
179/// where `out_ptr` must point to at least 64 bytes inside plugin
180/// memory (the lower-case hex SHA-256 digest is exactly 64 ASCII
181/// chars). Returns 64 on success, -1 on memory error.
182///
183/// The host computes the digest over the plugin-supplied byte range
184/// using the production `sha2` crate; plugins no longer need to
185/// embed their own (placeholder) hash and stay small.
186pub fn register_crypto_imports(linker: &mut Linker<StoreCtx>) -> Result<(), PluginError> {
187    use sha2::{Digest, Sha256};
188
189    linker
190        .func_wrap(
191            "env",
192            "sha256_hex",
193            |mut caller: Caller<'_, StoreCtx>, in_ptr: i32, in_len: i32, out_ptr: i32| -> i32 {
194                let memory = match get_memory(&mut caller) {
195                    Some(m) => m,
196                    None => return -1,
197                };
198                let input = match read_bytes(&memory, &caller, in_ptr, in_len) {
199                    Some(b) => b,
200                    None => return -1,
201                };
202                let digest = Sha256::digest(&input);
203                // Hex-encode into a fixed 64-byte stack buffer so we
204                // don't allocate per call.
205                let mut hex = [0u8; 64];
206                const HEX: &[u8; 16] = b"0123456789abcdef";
207                for (i, b) in digest.iter().enumerate() {
208                    hex[i * 2]     = HEX[(b >> 4) as usize];
209                    hex[i * 2 + 1] = HEX[(b & 0x0f) as usize];
210                }
211                if write_bytes(&memory, &mut caller, out_ptr, &hex).is_err() {
212                    return -1;
213                }
214                64
215            },
216        )
217        .map_err(|e| PluginError::RuntimeError(format!("link sha256_hex: {}", e)))?;
218    Ok(())
219}
220
221fn get_memory(caller: &mut Caller<'_, StoreCtx>) -> Option<Memory> {
222    caller.get_export("memory").and_then(|e| e.into_memory())
223}
224
225fn read_bytes(memory: &Memory, caller: &Caller<'_, StoreCtx>, ptr: i32, len: i32) -> Option<Vec<u8>> {
226    if len < 0 {
227        return None;
228    }
229    let start = ptr as usize;
230    let end = start.checked_add(len as usize)?;
231    let data = memory.data(caller);
232    data.get(start..end).map(|s| s.to_vec())
233}
234
235fn write_bytes(
236    memory: &Memory,
237    caller: &mut Caller<'_, StoreCtx>,
238    ptr: i32,
239    bytes: &[u8],
240) -> Result<(), ()> {
241    let start = ptr as usize;
242    let end = start.checked_add(bytes.len()).ok_or(())?;
243    let data = memory.data_mut(caller);
244    let slot = data.get_mut(start..end).ok_or(())?;
245    slot.copy_from_slice(bytes);
246    Ok(())
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[test]
254    fn kv_namespaced_per_plugin() {
255        let kv = KvBackend::new();
256        kv.set("plugin-a", b"k".to_vec(), b"v1".to_vec());
257        kv.set("plugin-b", b"k".to_vec(), b"v2".to_vec());
258        assert_eq!(kv.get("plugin-a", b"k"), Some(b"v1".to_vec()));
259        assert_eq!(kv.get("plugin-b", b"k"), Some(b"v2".to_vec()));
260        assert_eq!(kv.get("plugin-c", b"k"), None);
261    }
262
263    #[test]
264    fn kv_overwrite_is_idempotent() {
265        let kv = KvBackend::new();
266        kv.set("p", b"k".to_vec(), b"v1".to_vec());
267        kv.set("p", b"k".to_vec(), b"v2".to_vec());
268        assert_eq!(kv.get("p", b"k"), Some(b"v2".to_vec()));
269        assert_eq!(kv.len("p"), 1);
270    }
271
272    #[test]
273    fn kv_delete_idempotent_on_missing() {
274        let kv = KvBackend::new();
275        kv.delete("p", b"never-set");
276        kv.set("p", b"k".to_vec(), b"v".to_vec());
277        kv.delete("p", b"k");
278        assert_eq!(kv.get("p", b"k"), None);
279    }
280}
heliosdb_proxy/plugins/host_imports.rs

heliosdb_proxy/plugins/
host_imports.rs
