wasmtime_wasi_keyvalue/
lib.rs

1//! # Wasmtime's [wasi-keyvalue] Implementation
2//!
3//! This crate provides a Wasmtime host implementation of the [wasi-keyvalue]
4//! API. With this crate, the runtime can run components that call APIs in
5//! [wasi-keyvalue] and provide components with access to key-value storages.
6//!
7//! Currently supported storage backends:
8//! * In-Memory (empty identifier)
9//!
10//! # Examples
11//!
12//! The usage of this crate is very similar to other WASI API implementations
13//! such as [wasi:cli] and [wasi:http].
14//!
15//! A common scenario is accessing KV store in a [wasi:cli] component.
16//! A standalone example of doing all this looks like:
17//!
18//! ```
19//! use wasmtime::{
20//!     component::{Linker, ResourceTable},
21//!     Engine, Result, Store,
22//! };
23//! use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView};
24//! use wasmtime_wasi_keyvalue::{WasiKeyValue, WasiKeyValueCtx, WasiKeyValueCtxBuilder};
25//!
26//! #[tokio::main]
27//! async fn main() -> Result<()> {
28//!     let engine = Engine::default();
29//!
30//!     let mut store = Store::new(&engine, Ctx {
31//!         table: ResourceTable::new(),
32//!         wasi_ctx: WasiCtx::builder().build(),
33//!         wasi_keyvalue_ctx: WasiKeyValueCtxBuilder::new().build(),
34//!     });
35//!
36//!     let mut linker = Linker::<Ctx>::new(&engine);
37//!     wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
38//!     // add `wasi-keyvalue` world's interfaces to the linker
39//!     wasmtime_wasi_keyvalue::add_to_linker(&mut linker, |h: &mut Ctx| {
40//!         WasiKeyValue::new(&h.wasi_keyvalue_ctx, &mut h.table)
41//!     })?;
42//!
43//!     // ... use `linker` to instantiate within `store` ...
44//!
45//!     Ok(())
46//! }
47//!
48//! struct Ctx {
49//!     table: ResourceTable,
50//!     wasi_ctx: WasiCtx,
51//!     wasi_keyvalue_ctx: WasiKeyValueCtx,
52//! }
53//!
54//! impl WasiView for Ctx {
55//!     fn ctx(&mut self) -> WasiCtxView<'_> {
56//!         WasiCtxView { ctx: &mut self.wasi_ctx, table: &mut self.table }
57//!     }
58//! }
59//! ```
60//!
61//! [wasi-keyvalue]: https://github.com/WebAssembly/wasi-keyvalue
62//! [wasi:cli]: https://docs.rs/wasmtime-wasi/latest
63//! [wasi:http]: https://docs.rs/wasmtime-wasi-http/latest
64
65#![deny(missing_docs)]
66
67mod generated {
68    wasmtime::component::bindgen!({
69        path: "wit",
70        world: "wasi:keyvalue/imports",
71        imports: { default: trappable },
72        with: {
73            "wasi:keyvalue/store.bucket": crate::Bucket,
74        },
75        trappable_error_type: {
76            "wasi:keyvalue/store.error" => crate::Error,
77        },
78    });
79}
80
81use self::generated::wasi::keyvalue;
82use std::collections::HashMap;
83use wasmtime::Result;
84use wasmtime::component::{HasData, Resource, ResourceTable, ResourceTableError};
85
86#[doc(hidden)]
87pub enum Error {
88    NoSuchStore,
89    AccessDenied,
90    Other(String),
91}
92
93impl From<ResourceTableError> for Error {
94    fn from(err: ResourceTableError) -> Self {
95        Self::Other(err.to_string())
96    }
97}
98
99#[doc(hidden)]
100pub struct Bucket {
101    in_memory_data: HashMap<String, Vec<u8>>,
102}
103
104/// Builder-style structure used to create a [`WasiKeyValueCtx`].
105#[derive(Default)]
106pub struct WasiKeyValueCtxBuilder {
107    in_memory_data: HashMap<String, Vec<u8>>,
108}
109
110impl WasiKeyValueCtxBuilder {
111    /// Creates a builder for a new context with default parameters set.
112    pub fn new() -> Self {
113        Default::default()
114    }
115
116    /// Preset data for the In-Memory provider.
117    pub fn in_memory_data<I, K, V>(mut self, data: I) -> Self
118    where
119        I: IntoIterator<Item = (K, V)>,
120        K: Into<String>,
121        V: Into<Vec<u8>>,
122    {
123        self.in_memory_data = data
124            .into_iter()
125            .map(|(k, v)| (k.into(), v.into()))
126            .collect();
127        self
128    }
129
130    /// Uses the configured context so far to construct the final [`WasiKeyValueCtx`].
131    pub fn build(self) -> WasiKeyValueCtx {
132        WasiKeyValueCtx {
133            in_memory_data: self.in_memory_data,
134        }
135    }
136}
137
138/// Capture the state necessary for use in the `wasi-keyvalue` API implementation.
139pub struct WasiKeyValueCtx {
140    in_memory_data: HashMap<String, Vec<u8>>,
141}
142
143impl WasiKeyValueCtx {
144    /// Convenience function for calling [`WasiKeyValueCtxBuilder::new`].
145    pub fn builder() -> WasiKeyValueCtxBuilder {
146        WasiKeyValueCtxBuilder::new()
147    }
148}
149
150/// A wrapper capturing the needed internal `wasi-keyvalue` state.
151pub struct WasiKeyValue<'a> {
152    ctx: &'a WasiKeyValueCtx,
153    table: &'a mut ResourceTable,
154}
155
156impl<'a> WasiKeyValue<'a> {
157    /// Create a new view into the `wasi-keyvalue` state.
158    pub fn new(ctx: &'a WasiKeyValueCtx, table: &'a mut ResourceTable) -> Self {
159        Self { ctx, table }
160    }
161}
162
163impl keyvalue::store::Host for WasiKeyValue<'_> {
164    fn open(&mut self, identifier: String) -> Result<Resource<Bucket>, Error> {
165        match identifier.as_str() {
166            "" => Ok(self.table.push(Bucket {
167                in_memory_data: self.ctx.in_memory_data.clone(),
168            })?),
169            _ => Err(Error::NoSuchStore),
170        }
171    }
172
173    fn convert_error(&mut self, err: Error) -> Result<keyvalue::store::Error> {
174        match err {
175            Error::NoSuchStore => Ok(keyvalue::store::Error::NoSuchStore),
176            Error::AccessDenied => Ok(keyvalue::store::Error::AccessDenied),
177            Error::Other(e) => Ok(keyvalue::store::Error::Other(e)),
178        }
179    }
180}
181
182impl keyvalue::store::HostBucket for WasiKeyValue<'_> {
183    fn get(&mut self, bucket: Resource<Bucket>, key: String) -> Result<Option<Vec<u8>>, Error> {
184        let bucket = self.table.get_mut(&bucket)?;
185        Ok(bucket.in_memory_data.get(&key).cloned())
186    }
187
188    fn set(&mut self, bucket: Resource<Bucket>, key: String, value: Vec<u8>) -> Result<(), Error> {
189        let bucket = self.table.get_mut(&bucket)?;
190        bucket.in_memory_data.insert(key, value);
191        Ok(())
192    }
193
194    fn delete(&mut self, bucket: Resource<Bucket>, key: String) -> Result<(), Error> {
195        let bucket = self.table.get_mut(&bucket)?;
196        bucket.in_memory_data.remove(&key);
197        Ok(())
198    }
199
200    fn exists(&mut self, bucket: Resource<Bucket>, key: String) -> Result<bool, Error> {
201        let bucket = self.table.get_mut(&bucket)?;
202        Ok(bucket.in_memory_data.contains_key(&key))
203    }
204
205    fn list_keys(
206        &mut self,
207        bucket: Resource<Bucket>,
208        cursor: Option<u64>,
209    ) -> Result<keyvalue::store::KeyResponse, Error> {
210        let bucket = self.table.get_mut(&bucket)?;
211        let keys: Vec<String> = bucket.in_memory_data.keys().cloned().collect();
212        // `cursor` is guest-controlled, so clamp it to the number of keys before
213        // slicing. `try_from` also guards against a `u64` that does not fit in a
214        // `usize` on 32-bit hosts, which the old `as usize` cast would truncate.
215        let cursor = usize::try_from(cursor.unwrap_or(0))
216            .unwrap_or(usize::MAX)
217            .min(keys.len());
218        let keys_slice = &keys[cursor..];
219        Ok(keyvalue::store::KeyResponse {
220            keys: keys_slice.to_vec(),
221            cursor: None,
222        })
223    }
224
225    fn drop(&mut self, bucket: Resource<Bucket>) -> Result<()> {
226        self.table.delete(bucket)?;
227        Ok(())
228    }
229}
230
231impl keyvalue::atomics::Host for WasiKeyValue<'_> {
232    fn increment(
233        &mut self,
234        bucket: Resource<Bucket>,
235        key: String,
236        delta: u64,
237    ) -> Result<u64, Error> {
238        let bucket = self.table.get_mut(&bucket)?;
239        let value = bucket
240            .in_memory_data
241            .entry(key.clone())
242            .or_insert("0".to_string().into_bytes());
243        let current_value = String::from_utf8(value.clone())
244            .map_err(|e| Error::Other(e.to_string()))?
245            .parse::<u64>()
246            .map_err(|e| Error::Other(e.to_string()))?;
247        let new_value = current_value
248            .checked_add(delta)
249            .ok_or_else(|| Error::Other("overflow incrementing value".to_string()))?;
250        *value = new_value.to_string().into_bytes();
251        Ok(new_value)
252    }
253}
254
255impl keyvalue::batch::Host for WasiKeyValue<'_> {
256    fn get_many(
257        &mut self,
258        bucket: Resource<Bucket>,
259        keys: Vec<String>,
260    ) -> Result<Vec<Option<(String, Vec<u8>)>>, Error> {
261        let bucket = self.table.get_mut(&bucket)?;
262        Ok(keys
263            .into_iter()
264            .map(|key| {
265                bucket
266                    .in_memory_data
267                    .get(&key)
268                    .map(|value| (key.clone(), value.clone()))
269            })
270            .collect())
271    }
272
273    fn set_many(
274        &mut self,
275        bucket: Resource<Bucket>,
276        key_values: Vec<(String, Vec<u8>)>,
277    ) -> Result<(), Error> {
278        let bucket = self.table.get_mut(&bucket)?;
279        for (key, value) in key_values {
280            bucket.in_memory_data.insert(key, value);
281        }
282        Ok(())
283    }
284
285    fn delete_many(&mut self, bucket: Resource<Bucket>, keys: Vec<String>) -> Result<(), Error> {
286        let bucket = self.table.get_mut(&bucket)?;
287        for key in keys {
288            bucket.in_memory_data.remove(&key);
289        }
290        Ok(())
291    }
292}
293
294/// Add all the `wasi-keyvalue` world's interfaces to a [`wasmtime::component::Linker`].
295pub fn add_to_linker<T: Send + 'static>(
296    l: &mut wasmtime::component::Linker<T>,
297    f: fn(&mut T) -> WasiKeyValue<'_>,
298) -> Result<()> {
299    keyvalue::store::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
300    keyvalue::atomics::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
301    keyvalue::batch::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
302    Ok(())
303}
304
305struct HasWasiKeyValue;
306
307impl HasData for HasWasiKeyValue {
308    type Data<'a> = WasiKeyValue<'a>;
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314    use keyvalue::atomics::Host as _;
315    use keyvalue::store::{Host as _, HostBucket as _};
316    use wasmtime::component::ResourceTable;
317
318    fn ctx_with(data: &[(&str, &str)]) -> WasiKeyValueCtx {
319        WasiKeyValueCtxBuilder::new()
320            .in_memory_data(data.iter().map(|(k, v)| (k.to_string(), v.to_string())))
321            .build()
322    }
323
324    #[test]
325    fn list_keys_cursor_past_end_returns_empty() {
326        let ctx = ctx_with(&[("k", "v")]);
327        let mut table = ResourceTable::new();
328        let mut kv = WasiKeyValue::new(&ctx, &mut table);
329        let bucket = match kv.open(String::new()) {
330            Ok(bucket) => bucket,
331            Err(_) => panic!("open failed"),
332        };
333        // A guest-supplied cursor past the last key must not panic.
334        match kv.list_keys(bucket, Some(1_000)) {
335            Ok(resp) => assert!(resp.keys.is_empty()),
336            Err(_) => panic!("list_keys returned an error"),
337        }
338    }
339
340    #[test]
341    fn increment_overflow_is_reported_as_error() {
342        let ctx = ctx_with(&[("c", "18446744073709551615")]); // u64::MAX
343        let mut table = ResourceTable::new();
344        let mut kv = WasiKeyValue::new(&ctx, &mut table);
345        let bucket = match kv.open(String::new()) {
346            Ok(bucket) => bucket,
347            Err(_) => panic!("open failed"),
348        };
349        assert!(kv.increment(bucket, "c".to_string(), 1).is_err());
350    }
351}
wasmtime_wasi_keyvalue/lib.rs

wasmtime_wasi_keyvalue/
lib.rs
