use std::sync::Arc;
use async_trait::async_trait;
use crate::{PluginError, PluginResult};
#[async_trait]
pub trait PluginStore: Send + Sync {
async fn get(&self, vault: &str, key: &str) -> PluginResult<Option<Vec<u8>>>;
async fn put(&self, vault: &str, key: &str, value: &[u8]) -> PluginResult<()>;
async fn delete(&self, vault: &str, key: &str) -> PluginResult<()>;
async fn list(&self, vault: &str, prefix: &str) -> PluginResult<Vec<String>>;
}
#[derive(Clone)]
pub struct PluginStorage {
store: Arc<dyn PluginStore>,
}
impl std::fmt::Debug for PluginStorage {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("PluginStorage")
.finish_non_exhaustive()
}
}
impl PluginStorage {
pub fn new(store: Arc<dyn PluginStore>) -> Self {
Self { store }
}
pub async fn get(&self, vault: &str, key: &str) -> PluginResult<Option<Vec<u8>>> {
self.store.get(vault, &validate_key(key)?).await
}
pub async fn put(&self, vault: &str, key: &str, value: &[u8]) -> PluginResult<()> {
self.store.put(vault, &validate_key(key)?, value).await
}
pub async fn delete(&self, vault: &str, key: &str) -> PluginResult<()> {
self.store.delete(vault, &validate_key(key)?).await
}
pub async fn list(&self, vault: &str, prefix: &str) -> PluginResult<Vec<String>> {
let prefix = if prefix.is_empty() {
String::new()
} else {
validate_key(prefix)?
};
self.store.list(vault, &prefix).await
}
}
pub(crate) fn validate_key(key: &str) -> PluginResult<String> {
if key.is_empty() {
return Err(PluginError::invalid_input("storage key must not be empty"));
}
if key.len() > MAX_KEY_LEN {
return Err(PluginError::invalid_input(format!(
"storage key is {} bytes; the maximum is {MAX_KEY_LEN}",
key.len()
)));
}
if key.starts_with('/') || key.starts_with('\\') {
return Err(PluginError::invalid_input(format!(
"storage key {key:?} must be relative"
)));
}
for segment in key.split('/') {
if segment.is_empty() || segment == "." || segment == ".." {
return Err(PluginError::invalid_input(format!(
"storage key {key:?} must not contain empty or relative segments"
)));
}
if segment
.chars()
.any(|character| character == '\\' || character == '\0' || character.is_control())
{
return Err(PluginError::invalid_input(format!(
"storage key {key:?} contains an unsupported character"
)));
}
}
Ok(key.to_string())
}
const MAX_KEY_LEN: usize = 200;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn keys_may_group_with_slashes() {
for valid in ["index", "embeddings/v2/shard-0.bin", "watermark.json"] {
assert_eq!(validate_key(valid).expect("valid key"), valid);
}
}
#[test]
fn keys_cannot_escape_the_namespace() {
for rejected in [
"",
"/absolute",
"..",
"../sibling/data",
"nested/../../escape",
"trailing/",
"back\\slash",
] {
let error = validate_key(rejected).expect_err("must reject");
assert_eq!(error.code, crate::PluginErrorCode::InvalidInput);
}
assert!(validate_key(&"k".repeat(MAX_KEY_LEN + 1)).is_err());
}
}