wasmeld/components/kv-probe/src/lib.rs

mod bindings {
    wit_bindgen::generate!({
        path: "wit",
        world: "kv-probe-component",
        generate_all,
    });
}

use bindings::wasmeld::kv::store;

struct KvProbe;

impl bindings::Guest for KvProbe {
    fn init(_config: Vec<u8>) -> Result<(), bindings::ServiceError> {
        Ok(())
    }

    fn invoke(input: Vec<u8>) -> Result<Vec<u8>, bindings::ServiceError> {
        if let Some(key) = input.strip_prefix(b"get:") {
            let key = parse_key(key)?;
            return store::get(&key)
                .map(|value| value.unwrap_or_default())
                .map_err(kv_error);
        }

        if let Some(command) = input.strip_prefix(b"set:") {
            let separator = command
                .iter()
                .position(|byte| *byte == b':')
                .ok_or_else(|| {
                    bindings::ServiceError::CallFailed(
                        "set command must be set:<key>:<value>".to_owned(),
                    )
                })?;
            let key = parse_key(&command[..separator])?;
            let value = command[separator + 1..].to_vec();
            store::set(&key, &value).map_err(kv_error)?;
            return Ok(Vec::new());
        }

        if let Some(key) = input.strip_prefix(b"delete:") {
            let key = parse_key(key)?;
            store::delete(&key).map_err(kv_error)?;
            return Ok(Vec::new());
        }

        Err(bindings::ServiceError::CallFailed(
            "expected get:<key>, set:<key>:<value>, or delete:<key>".to_owned(),
        ))
    }
}

fn parse_key(bytes: &[u8]) -> Result<String, bindings::ServiceError> {
    core::str::from_utf8(bytes)
        .map(str::to_owned)
        .map_err(|_| bindings::ServiceError::CallFailed("key must be valid UTF-8".to_owned()))
}

fn kv_error(error: store::KvError) -> bindings::ServiceError {
    bindings::ServiceError::CallFailed(format!("KV operation failed: {error:?}"))
}

bindings::export!(KvProbe with_types_in bindings);