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        let cursor = cursor.unwrap_or(0) as usize;
213        let keys_slice = &keys[cursor..];
214        Ok(keyvalue::store::KeyResponse {
215            keys: keys_slice.to_vec(),
216            cursor: None,
217        })
218    }
219
220    fn drop(&mut self, bucket: Resource<Bucket>) -> Result<()> {
221        self.table.delete(bucket)?;
222        Ok(())
223    }
224}
225
226impl keyvalue::atomics::Host for WasiKeyValue<'_> {
227    fn increment(
228        &mut self,
229        bucket: Resource<Bucket>,
230        key: String,
231        delta: u64,
232    ) -> Result<u64, Error> {
233        let bucket = self.table.get_mut(&bucket)?;
234        let value = bucket
235            .in_memory_data
236            .entry(key.clone())
237            .or_insert("0".to_string().into_bytes());
238        let current_value = String::from_utf8(value.clone())
239            .map_err(|e| Error::Other(e.to_string()))?
240            .parse::<u64>()
241            .map_err(|e| Error::Other(e.to_string()))?;
242        let new_value = current_value + delta;
243        *value = new_value.to_string().into_bytes();
244        Ok(new_value)
245    }
246}
247
248impl keyvalue::batch::Host for WasiKeyValue<'_> {
249    fn get_many(
250        &mut self,
251        bucket: Resource<Bucket>,
252        keys: Vec<String>,
253    ) -> Result<Vec<Option<(String, Vec<u8>)>>, Error> {
254        let bucket = self.table.get_mut(&bucket)?;
255        Ok(keys
256            .into_iter()
257            .map(|key| {
258                bucket
259                    .in_memory_data
260                    .get(&key)
261                    .map(|value| (key.clone(), value.clone()))
262            })
263            .collect())
264    }
265
266    fn set_many(
267        &mut self,
268        bucket: Resource<Bucket>,
269        key_values: Vec<(String, Vec<u8>)>,
270    ) -> Result<(), Error> {
271        let bucket = self.table.get_mut(&bucket)?;
272        for (key, value) in key_values {
273            bucket.in_memory_data.insert(key, value);
274        }
275        Ok(())
276    }
277
278    fn delete_many(&mut self, bucket: Resource<Bucket>, keys: Vec<String>) -> Result<(), Error> {
279        let bucket = self.table.get_mut(&bucket)?;
280        for key in keys {
281            bucket.in_memory_data.remove(&key);
282        }
283        Ok(())
284    }
285}
286
287/// Add all the `wasi-keyvalue` world's interfaces to a [`wasmtime::component::Linker`].
288pub fn add_to_linker<T: Send + 'static>(
289    l: &mut wasmtime::component::Linker<T>,
290    f: fn(&mut T) -> WasiKeyValue<'_>,
291) -> Result<()> {
292    keyvalue::store::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
293    keyvalue::atomics::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
294    keyvalue::batch::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
295    Ok(())
296}
297
298struct HasWasiKeyValue;
299
300impl HasData for HasWasiKeyValue {
301    type Data<'a> = WasiKeyValue<'a>;
302}
wasmtime_wasi_keyvalue/lib.rs

wasmtime_wasi_keyvalue/
lib.rs
