Address comments

.github/workflows/main.yml

@@ -736,22 +736,31 @@ jobs:
   # Test the `wasmtime-wasi-keyvalue` crate. Split out from the main tests
   # because it needs additional database service.
   test_wasi_keyvalue:
-    strategy:
-      matrix:
-        os: [ "ubuntu-latest", "macOS-latest" ]
-    name: Test wasi-keyvalue on ${{ matrix.os }}
-    runs-on: ${{ matrix.os }}
+    name: Test wasi-keyvalue
+    runs-on: ubuntu-latest
     needs: determine
     if: needs.determine.outputs.run-wasi-keyvalue
+    # Setup redis server
+    services:
+      redis:
+        # Docker Hub image
+        image: redis
+        # Set health checks to wait until redis has started
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          # Maps port 6379 on service container to the host
+          - 6379:6379
     steps:
     - uses: actions/checkout@v4
       with:
         submodules: true
     - uses: ./.github/actions/install-rust
     # Install Rust targets
     - run: rustup target add wasm32-wasi
-    # Setup redis server
-    - uses: shogo82148/actions-setup-redis@v1
     - run: cargo test --all-features -p wasmtime-wasi-keyvalue
       env:
         RUST_BACKTRACE: 1

crates/wasi-keyvalue/src/lib.rs

@@ -87,6 +87,7 @@ mod generated {
 use self::generated::wasi::keyvalue;
 use anyhow::Result;
 use async_trait::async_trait;
+use std::fmt::Display;
 use url::Url;
 use wasmtime::component::{Resource, ResourceTable, ResourceTableError};
 
@@ -103,6 +104,10 @@ impl From<ResourceTableError> for Error {
     }
 }
 
+pub(crate) fn to_other_error(e: impl Display) -> Error {
+    Error::Other(e.to_string())
+}
+
 #[doc(hidden)]
 pub struct Bucket {
     inner: Box<dyn Host + Send>,
@@ -138,6 +143,7 @@ trait Host {
 /// Builder-style structure used to create a [`WasiKeyValueCtx`].
 #[derive(Default)]
 pub struct WasiKeyValueCtxBuilder {
+    allowed_hosts: Vec<String>,
     #[cfg(feature = "redis")]
     redis_connection_timeout: Option<std::time::Duration>,
     #[cfg(feature = "redis")]
@@ -150,6 +156,27 @@ impl WasiKeyValueCtxBuilder {
         Default::default()
     }
 
+    /// Appends a list of hosts to the allow-listed set each component gets
+    /// access to. It can be in the format `<hostname>[:port]` or a unix domain
+    /// socket path.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use wasmtime_wasi_keyvalue::WasiKeyValueCtxBuilder;
+    ///
+    /// # fn main() {
+    /// let ctx = WasiKeyValueCtxBuilder::new()
+    ///     .allow_hosts(&["localhost:1234", "/var/run/redis.sock"])
+    ///     .build();
+    /// # }
+    /// ```
+    pub fn allow_hosts(mut self, hosts: &[impl AsRef<str>]) -> Self {
+        self.allowed_hosts
+            .extend(hosts.iter().map(|a| a.as_ref().to_owned()));
+        self
+    }
+
     /// Sets the connection timeout parameter for the Redis provider.
     #[cfg(feature = "redis")]
     pub fn redis_connection_timeout(mut self, t: std::time::Duration) -> Self {
@@ -167,6 +194,7 @@ impl WasiKeyValueCtxBuilder {
     /// Uses the configured context so far to construct the final [`WasiKeyValueCtx`].
     pub fn build(self) -> WasiKeyValueCtx {
         WasiKeyValueCtx {
+            allowed_hosts: self.allowed_hosts,
             #[cfg(feature = "redis")]
             redis_connection_timeout: self.redis_connection_timeout,
             #[cfg(feature = "redis")]
@@ -177,6 +205,7 @@ impl WasiKeyValueCtxBuilder {
 
 /// Capture the state necessary for use in the `wasi-keyvalue` API implementation.
 pub struct WasiKeyValueCtx {
+    allowed_hosts: Vec<String>,
     #[cfg(feature = "redis")]
     redis_connection_timeout: Option<std::time::Duration>,
     #[cfg(feature = "redis")]
@@ -188,6 +217,18 @@ impl WasiKeyValueCtx {
     pub fn builder() -> WasiKeyValueCtxBuilder {
         WasiKeyValueCtxBuilder::new()
     }
+
+    fn allow_host(&self, u: &Url) -> bool {
+        let host = match u.host() {
+            Some(h) => match u.port() {
+                Some(port) => format!("{}:{}", h, port),
+                None => h.to_string(),
+            },
+            // unix domain socket path
+            None => u.path().to_string(),
+        };
+        self.allowed_hosts.contains(&host)
+    }
 }
 
 /// A wrapper capturing the needed internal `wasi-keyvalue` state.
@@ -212,15 +253,18 @@ impl keyvalue::store::Host for WasiKeyValue<'_> {
             })?);
         }
 
-        let u = Url::parse(&identifier).map_err(|e| Error::Other(e.to_string()))?;
+        let u = Url::parse(&identifier).map_err(to_other_error)?;
+        if !self.ctx.allow_host(&u) {
+            return Err(Error::Other(format!(
+                "the identifier {} is not in the allowed list",
+                identifier
+            )));
+        }
+
         match u.scheme() {
             "redis" | "redis+unix" => {
                 #[cfg(not(feature = "redis"))]
                 {
-                    // Silence the unused warning for `ctx` as it is only used in
-                    // the conditionally-compiled.
-                    let _ = self.ctx;
-
                     return Err(Error::Other(
                         "Cannot enable Redis support when the crate is not compiled with this feature."
                             .to_string(),
@@ -351,3 +395,22 @@ pub fn add_to_linker<T: Send>(
     keyvalue::batch::add_to_linker_get_host(l, f)?;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_allow_host() -> Result<()> {
+        let ctx = WasiKeyValueCtx::builder()
+            .allow_hosts(&["127.0.0.1:1234", "localhost", "/var/run/redis.sock"])
+            .build();
+        assert!(ctx.allow_host(&Url::parse("redis://127.0.0.1:1234/db")?));
+        assert!(ctx.allow_host(&Url::parse("redis://localhost")?));
+        assert!(!ctx.allow_host(&Url::parse("redis://192.168.0.1")?));
+        assert!(ctx.allow_host(&Url::parse("redis+unix:///var/run/redis.sock?db=db")?));
+        assert!(!ctx.allow_host(&Url::parse("redis+unix:///var/local/redis.sock?db=db")?));
+
+        Ok(())
+    }
+}

crates/wasi-keyvalue/src/provider/inmemory.rs

@@ -1,4 +1,4 @@
-use crate::{generated::wasi::keyvalue::store::KeyResponse, Error, Host};
+use crate::{generated::wasi::keyvalue::store::KeyResponse, to_other_error, Error, Host};
 use async_trait::async_trait;
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
@@ -45,11 +45,13 @@ impl Host for InMemory {
 
     async fn increment(&mut self, key: String, delta: u64) -> Result<u64, Error> {
         let mut store = self.store.lock().unwrap();
-        let value = store.entry(key.clone()).or_insert(vec![0]);
+        let value = store
+            .entry(key.clone())
+            .or_insert("0".to_string().into_bytes());
         let current_value = String::from_utf8(value.clone())
-            .unwrap_or_else(|_| "0".to_string())
+            .map_err(to_other_error)?
             .parse::<u64>()
-            .unwrap_or(0);
+            .map_err(to_other_error)?;
         let new_value = current_value + delta;
         *value = new_value.to_string().into_bytes();
         Ok(new_value)

crates/wasi-keyvalue/tests/main.rs

@@ -77,6 +77,7 @@ async fn keyvalue_redis() -> Result<()> {
                 .env("IDENTIFIER", "redis://127.0.0.1/")
                 .build(),
             wasi_keyvalue_ctx: WasiKeyValueCtx::builder()
+                .allow_hosts(&["127.0.0.1"])
                 .redis_connection_timeout(std::time::Duration::from_secs(5))
                 .redis_response_timeout(std::time::Duration::from_secs(5))
                 .build(),

crates/wasi-runtime-config/Cargo.toml

@@ -18,8 +18,3 @@ wasmtime = { workspace = true, features = ["runtime", "component-model"] }
 test-programs-artifacts = { workspace = true }
 wasmtime-wasi = { workspace = true }
 tokio = { workspace = true, features = ["macros"] }
-
-[features]
-# Do not use, this is only used for testing in CI,
-# see https://github.com/bytecodealliance/wasmtime/issues/8997
-test = ["wasmtime/wmemcheck"]