use extism_pdk::*;
use serde::{Deserialize, Serialize};
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
mod base64_bytes {
use serde::{self, Deserialize, Deserializer, Serializer};
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
pub fn serialize<S>(bytes: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&BASE64.encode(bytes))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let s = String::deserialize(deserializer)?;
BASE64.decode(&s).map_err(serde::de::Error::custom)
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreSetRequest {
key: String,
#[serde(with = "base64_bytes")]
value: Vec<u8>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreSetResponse {
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreSetWithTTLRequest {
key: String,
#[serde(with = "base64_bytes")]
value: Vec<u8>,
ttl_seconds: i64,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreSetWithTTLResponse {
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreGetRequest {
key: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreGetResponse {
#[serde(default)]
#[serde(with = "base64_bytes")]
value: Vec<u8>,
#[serde(default)]
exists: bool,
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreGetManyRequest {
keys: Vec<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreGetManyResponse {
#[serde(default)]
values: std::collections::HashMap<String, Vec<u8>>,
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreHasRequest {
key: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreHasResponse {
#[serde(default)]
exists: bool,
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreListRequest {
prefix: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreListResponse {
#[serde(default)]
keys: Vec<String>,
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreDeleteRequest {
key: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreDeleteResponse {
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreDeleteByPrefixRequest {
prefix: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreDeleteByPrefixResponse {
#[serde(default)]
deleted_count: i64,
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct KVStoreGetStorageUsedResponse {
#[serde(default)]
bytes: i64,
#[serde(default)]
error: Option<String>,
}
#[host_fn]
extern "ExtismHost" {
fn kvstore_set(input: Json<KVStoreSetRequest>) -> Json<KVStoreSetResponse>;
fn kvstore_setwithttl(input: Json<KVStoreSetWithTTLRequest>) -> Json<KVStoreSetWithTTLResponse>;
fn kvstore_get(input: Json<KVStoreGetRequest>) -> Json<KVStoreGetResponse>;
fn kvstore_getmany(input: Json<KVStoreGetManyRequest>) -> Json<KVStoreGetManyResponse>;
fn kvstore_has(input: Json<KVStoreHasRequest>) -> Json<KVStoreHasResponse>;
fn kvstore_list(input: Json<KVStoreListRequest>) -> Json<KVStoreListResponse>;
fn kvstore_delete(input: Json<KVStoreDeleteRequest>) -> Json<KVStoreDeleteResponse>;
fn kvstore_deletebyprefix(input: Json<KVStoreDeleteByPrefixRequest>) -> Json<KVStoreDeleteByPrefixResponse>;
fn kvstore_getstorageused(input: Json<serde_json::Value>) -> Json<KVStoreGetStorageUsedResponse>;
}
pub fn set(key: &str, value: Vec<u8>) -> Result<(), Error> {
let response = unsafe {
kvstore_set(Json(KVStoreSetRequest {
key: key.to_owned(),
value: value,
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(())
}
pub fn set_with_ttl(key: &str, value: Vec<u8>, ttl_seconds: i64) -> Result<(), Error> {
let response = unsafe {
kvstore_setwithttl(Json(KVStoreSetWithTTLRequest {
key: key.to_owned(),
value: value,
ttl_seconds: ttl_seconds,
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(())
}
pub fn get(key: &str) -> Result<Option<Vec<u8>>, Error> {
let response = unsafe {
kvstore_get(Json(KVStoreGetRequest {
key: key.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
if response.0.exists {
Ok(Some(response.0.value))
} else {
Ok(None)
}
}
pub fn get_many(keys: Vec<String>) -> Result<std::collections::HashMap<String, Vec<u8>>, Error> {
let response = unsafe {
kvstore_getmany(Json(KVStoreGetManyRequest {
keys: keys,
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(response.0.values)
}
pub fn has(key: &str) -> Result<bool, Error> {
let response = unsafe {
kvstore_has(Json(KVStoreHasRequest {
key: key.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(response.0.exists)
}
pub fn list(prefix: &str) -> Result<Vec<String>, Error> {
let response = unsafe {
kvstore_list(Json(KVStoreListRequest {
prefix: prefix.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(response.0.keys)
}
pub fn delete(key: &str) -> Result<(), Error> {
let response = unsafe {
kvstore_delete(Json(KVStoreDeleteRequest {
key: key.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(())
}
pub fn delete_by_prefix(prefix: &str) -> Result<i64, Error> {
let response = unsafe {
kvstore_deletebyprefix(Json(KVStoreDeleteByPrefixRequest {
prefix: prefix.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(response.0.deleted_count)
}
pub fn get_storage_used() -> Result<i64, Error> {
let response = unsafe {
kvstore_getstorageused(Json(serde_json::json!({})))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(response.0.bytes)
}