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//!     Config, 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 mut config = Config::new();
29//!     config.async_support(true);
30//!     let engine = Engine::new(&config)?;
31//!
32//!     let mut store = Store::new(&engine, Ctx {
33//!         table: ResourceTable::new(),
34//!         wasi_ctx: WasiCtx::builder().build(),
35//!         wasi_keyvalue_ctx: WasiKeyValueCtxBuilder::new().build(),
36//!     });
37//!
38//!     let mut linker = Linker::<Ctx>::new(&engine);
39//!     wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
40//!     // add `wasi-keyvalue` world's interfaces to the linker
41//!     wasmtime_wasi_keyvalue::add_to_linker(&mut linker, |h: &mut Ctx| {
42//!         WasiKeyValue::new(&h.wasi_keyvalue_ctx, &mut h.table)
43//!     })?;
44//!
45//!     // ... use `linker` to instantiate within `store` ...
46//!
47//!     Ok(())
48//! }
49//!
50//! struct Ctx {
51//!     table: ResourceTable,
52//!     wasi_ctx: WasiCtx,
53//!     wasi_keyvalue_ctx: WasiKeyValueCtx,
54//! }
55//!
56//! impl WasiView for Ctx {
57//!     fn ctx(&mut self) -> WasiCtxView<'_> {
58//!         WasiCtxView { ctx: &mut self.wasi_ctx, table: &mut self.table }
59//!     }
60//! }
61//! ```
62//!
63//! [wasi-keyvalue]: https://github.com/WebAssembly/wasi-keyvalue
64//! [wasi:cli]: https://docs.rs/wasmtime-wasi/latest
65//! [wasi:http]: https://docs.rs/wasmtime-wasi-http/latest
66
67#![deny(missing_docs)]
68
69mod generated {
70    wasmtime::component::bindgen!({
71        path: "wit",
72        world: "wasi:keyvalue/imports",
73        imports: { default: trappable },
74        with: {
75            "wasi:keyvalue/store/bucket": crate::Bucket,
76        },
77        trappable_error_type: {
78            "wasi:keyvalue/store/error" => crate::Error,
79        },
80    });
81}
82
83use self::generated::wasi::keyvalue;
84use anyhow::Result;
85use std::collections::HashMap;
86use wasmtime::component::{HasData, Resource, ResourceTable, ResourceTableError};
87
88#[doc(hidden)]
89pub enum Error {
90    NoSuchStore,
91    AccessDenied,
92    Other(String),
93}
94
95impl From<ResourceTableError> for Error {
96    fn from(err: ResourceTableError) -> Self {
97        Self::Other(err.to_string())
98    }
99}
100
101#[doc(hidden)]
102pub struct Bucket {
103    in_memory_data: HashMap<String, Vec<u8>>,
104}
105
106/// Builder-style structure used to create a [`WasiKeyValueCtx`].
107#[derive(Default)]
108pub struct WasiKeyValueCtxBuilder {
109    in_memory_data: HashMap<String, Vec<u8>>,
110}
111
112impl WasiKeyValueCtxBuilder {
113    /// Creates a builder for a new context with default parameters set.
114    pub fn new() -> Self {
115        Default::default()
116    }
117
118    /// Preset data for the In-Memory provider.
119    pub fn in_memory_data<I, K, V>(mut self, data: I) -> Self
120    where
121        I: IntoIterator<Item = (K, V)>,
122        K: Into<String>,
123        V: Into<Vec<u8>>,
124    {
125        self.in_memory_data = data
126            .into_iter()
127            .map(|(k, v)| (k.into(), v.into()))
128            .collect();
129        self
130    }
131
132    /// Uses the configured context so far to construct the final [`WasiKeyValueCtx`].
133    pub fn build(self) -> WasiKeyValueCtx {
134        WasiKeyValueCtx {
135            in_memory_data: self.in_memory_data,
136        }
137    }
138}
139
140/// Capture the state necessary for use in the `wasi-keyvalue` API implementation.
141pub struct WasiKeyValueCtx {
142    in_memory_data: HashMap<String, Vec<u8>>,
143}
144
145impl WasiKeyValueCtx {
146    /// Convenience function for calling [`WasiKeyValueCtxBuilder::new`].
147    pub fn builder() -> WasiKeyValueCtxBuilder {
148        WasiKeyValueCtxBuilder::new()
149    }
150}
151
152/// A wrapper capturing the needed internal `wasi-keyvalue` state.
153pub struct WasiKeyValue<'a> {
154    ctx: &'a WasiKeyValueCtx,
155    table: &'a mut ResourceTable,
156}
157
158impl<'a> WasiKeyValue<'a> {
159    /// Create a new view into the `wasi-keyvalue` state.
160    pub fn new(ctx: &'a WasiKeyValueCtx, table: &'a mut ResourceTable) -> Self {
161        Self { ctx, table }
162    }
163}
164
165impl keyvalue::store::Host for WasiKeyValue<'_> {
166    fn open(&mut self, identifier: String) -> Result<Resource<Bucket>, Error> {
167        match identifier.as_str() {
168            "" => Ok(self.table.push(Bucket {
169                in_memory_data: self.ctx.in_memory_data.clone(),
170            })?),
171            _ => Err(Error::NoSuchStore),
172        }
173    }
174
175    fn convert_error(&mut self, err: Error) -> Result<keyvalue::store::Error> {
176        match err {
177            Error::NoSuchStore => Ok(keyvalue::store::Error::NoSuchStore),
178            Error::AccessDenied => Ok(keyvalue::store::Error::AccessDenied),
179            Error::Other(e) => Ok(keyvalue::store::Error::Other(e)),
180        }
181    }
182}
183
184impl keyvalue::store::HostBucket for WasiKeyValue<'_> {
185    fn get(&mut self, bucket: Resource<Bucket>, key: String) -> Result<Option<Vec<u8>>, Error> {
186        let bucket = self.table.get_mut(&bucket)?;
187        Ok(bucket.in_memory_data.get(&key).cloned())
188    }
189
190    fn set(&mut self, bucket: Resource<Bucket>, key: String, value: Vec<u8>) -> Result<(), Error> {
191        let bucket = self.table.get_mut(&bucket)?;
192        bucket.in_memory_data.insert(key, value);
193        Ok(())
194    }
195
196    fn delete(&mut self, bucket: Resource<Bucket>, key: String) -> Result<(), Error> {
197        let bucket = self.table.get_mut(&bucket)?;
198        bucket.in_memory_data.remove(&key);
199        Ok(())
200    }
201
202    fn exists(&mut self, bucket: Resource<Bucket>, key: String) -> Result<bool, Error> {
203        let bucket = self.table.get_mut(&bucket)?;
204        Ok(bucket.in_memory_data.contains_key(&key))
205    }
206
207    fn list_keys(
208        &mut self,
209        bucket: Resource<Bucket>,
210        cursor: Option<u64>,
211    ) -> Result<keyvalue::store::KeyResponse, Error> {
212        let bucket = self.table.get_mut(&bucket)?;
213        let keys: Vec<String> = bucket.in_memory_data.keys().cloned().collect();
214        let cursor = cursor.unwrap_or(0) as usize;
215        let keys_slice = &keys[cursor..];
216        Ok(keyvalue::store::KeyResponse {
217            keys: keys_slice.to_vec(),
218            cursor: None,
219        })
220    }
221
222    fn drop(&mut self, bucket: Resource<Bucket>) -> Result<()> {
223        self.table.delete(bucket)?;
224        Ok(())
225    }
226}
227
228impl keyvalue::atomics::Host for WasiKeyValue<'_> {
229    fn increment(
230        &mut self,
231        bucket: Resource<Bucket>,
232        key: String,
233        delta: u64,
234    ) -> Result<u64, Error> {
235        let bucket = self.table.get_mut(&bucket)?;
236        let value = bucket
237            .in_memory_data
238            .entry(key.clone())
239            .or_insert("0".to_string().into_bytes());
240        let current_value = String::from_utf8(value.clone())
241            .map_err(|e| Error::Other(e.to_string()))?
242            .parse::<u64>()
243            .map_err(|e| Error::Other(e.to_string()))?;
244        let new_value = current_value + delta;
245        *value = new_value.to_string().into_bytes();
246        Ok(new_value)
247    }
248}
249
250impl keyvalue::batch::Host for WasiKeyValue<'_> {
251    fn get_many(
252        &mut self,
253        bucket: Resource<Bucket>,
254        keys: Vec<String>,
255    ) -> Result<Vec<Option<(String, Vec<u8>)>>, Error> {
256        let bucket = self.table.get_mut(&bucket)?;
257        Ok(keys
258            .into_iter()
259            .map(|key| {
260                bucket
261                    .in_memory_data
262                    .get(&key)
263                    .map(|value| (key.clone(), value.clone()))
264            })
265            .collect())
266    }
267
268    fn set_many(
269        &mut self,
270        bucket: Resource<Bucket>,
271        key_values: Vec<(String, Vec<u8>)>,
272    ) -> Result<(), Error> {
273        let bucket = self.table.get_mut(&bucket)?;
274        for (key, value) in key_values {
275            bucket.in_memory_data.insert(key, value);
276        }
277        Ok(())
278    }
279
280    fn delete_many(&mut self, bucket: Resource<Bucket>, keys: Vec<String>) -> Result<(), Error> {
281        let bucket = self.table.get_mut(&bucket)?;
282        for key in keys {
283            bucket.in_memory_data.remove(&key);
284        }
285        Ok(())
286    }
287}
288
289/// Add all the `wasi-keyvalue` world's interfaces to a [`wasmtime::component::Linker`].
290pub fn add_to_linker<T: Send + 'static>(
291    l: &mut wasmtime::component::Linker<T>,
292    f: fn(&mut T) -> WasiKeyValue<'_>,
293) -> Result<()> {
294    keyvalue::store::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
295    keyvalue::atomics::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
296    keyvalue::batch::add_to_linker::<_, HasWasiKeyValue>(l, f)?;
297    Ok(())
298}
299
300struct HasWasiKeyValue;
301
302impl HasData for HasWasiKeyValue {
303    type Data<'a> = WasiKeyValue<'a>;
304}
wasmtime_wasi_keyvalue/lib.rs

wasmtime_wasi_keyvalue/
lib.rs
