use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use serde_json::Value;
pub trait StateStore: Send + Sync {
fn get(&self, ns: &str, key: &str) -> Result<Option<Value>, String>;
fn set(&self, ns: &str, key: &str, value: Value) -> Result<(), String>;
fn delete(&self, ns: &str, key: &str) -> Result<bool, String>;
fn keys(&self, ns: &str) -> Result<Vec<String>, String>;
fn has(&self, ns: &str, key: &str) -> Result<bool, String>;
fn set_nx(&self, ns: &str, key: &str, value: Value) -> Result<bool, String>;
fn incr(&self, ns: &str, key: &str, delta: f64, default: f64) -> Result<f64, String>;
}
pub struct JsonFileStore {
root: PathBuf,
locks: Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>,
}
impl JsonFileStore {
pub fn new(root: PathBuf) -> Self {
Self {
root,
locks: Mutex::new(HashMap::new()),
}
}
fn ns_lock(&self, path: &Path) -> Result<Arc<Mutex<()>>, String> {
let mut map = self
.locks
.lock()
.map_err(|_| "state: locks map poisoned".to_string())?;
Ok(Arc::clone(
map.entry(path.to_path_buf())
.or_insert_with(|| Arc::new(Mutex::new(()))),
))
}
pub fn root(&self) -> &Path {
&self.root
}
fn ensure_root(&self) -> Result<&Path, String> {
if !self.root.exists() {
fs::create_dir_all(&self.root)
.map_err(|e| format!("Failed to create state dir: {e}"))?;
}
Ok(&self.root)
}
pub fn state_path(&self, ns: &str) -> Result<PathBuf, String> {
if ns.contains('/')
|| ns.contains('\\')
|| ns.contains("..")
|| ns.contains('\0')
|| ns.is_empty()
{
return Err(format!("Invalid namespace: '{ns}'"));
}
let dir = self.ensure_root()?;
Ok(dir.join(format!("{ns}.json")))
}
fn load(&self, ns: &str) -> Result<HashMap<String, Value>, String> {
let path = self.state_path(ns)?;
if !path.exists() {
return Ok(HashMap::new());
}
let content =
fs::read_to_string(&path).map_err(|e| format!("Failed to read state '{ns}': {e}"))?;
serde_json::from_str(&content).map_err(|e| format!("Failed to parse state '{ns}': {e}"))
}
fn save(&self, ns: &str, data: &HashMap<String, Value>) -> Result<(), String> {
let path = self.state_path(ns)?;
let tmp = path.with_extension("json.tmp");
let content = serde_json::to_string_pretty(data)
.map_err(|e| format!("Failed to serialize state: {e}"))?;
fs::write(&tmp, &content).map_err(|e| format!("Failed to write state tmp: {e}"))?;
fs::rename(&tmp, &path).map_err(|e| format!("Failed to rename state file: {e}"))?;
Ok(())
}
}
impl StateStore for JsonFileStore {
fn get(&self, ns: &str, key: &str) -> Result<Option<Value>, String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let state = self.load(ns)?;
Ok(state.get(key).cloned())
}
fn set(&self, ns: &str, key: &str, value: Value) -> Result<(), String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let mut state = self.load(ns)?;
state.insert(key.to_string(), value);
self.save(ns, &state)
}
fn delete(&self, ns: &str, key: &str) -> Result<bool, String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let mut state = self.load(ns)?;
let existed = state.remove(key).is_some();
if existed {
self.save(ns, &state)?;
}
Ok(existed)
}
fn keys(&self, ns: &str) -> Result<Vec<String>, String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let state = self.load(ns)?;
Ok(state.keys().cloned().collect())
}
fn has(&self, ns: &str, key: &str) -> Result<bool, String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let state = self.load(ns)?;
Ok(state.contains_key(key))
}
fn set_nx(&self, ns: &str, key: &str, value: Value) -> Result<bool, String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let mut state = self.load(ns)?;
if state.contains_key(key) {
return Ok(false);
}
state.insert(key.to_string(), value);
self.save(ns, &state)?;
Ok(true)
}
fn incr(&self, ns: &str, key: &str, delta: f64, default: f64) -> Result<f64, String> {
let path = self.state_path(ns)?;
let lock = self.ns_lock(&path)?;
let _guard = lock
.lock()
.map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
let mut state = self.load(ns)?;
let current = match state.get(key) {
Some(v) => v
.as_f64()
.ok_or_else(|| format!("incr: value at '{key}' is not a number"))?,
None => default,
};
let new_val = current + delta;
state.insert(key.to_string(), serde_json::json!(new_val));
self.save(ns, &state)?;
Ok(new_val)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn new_store() -> (JsonFileStore, TempDir) {
let tmp = tempfile::tempdir().unwrap();
let store = JsonFileStore::new(tmp.path().to_path_buf());
(store, tmp)
}
#[test]
fn roundtrip() {
let (store, _tmp) = new_store();
let ns = "rt";
store.set(ns, "count", serde_json::json!(42)).unwrap();
store
.set(ns, "name", serde_json::json!("algocline"))
.unwrap();
assert_eq!(store.get(ns, "count").unwrap(), Some(serde_json::json!(42)));
assert_eq!(
store.get(ns, "name").unwrap(),
Some(serde_json::json!("algocline"))
);
assert_eq!(store.get(ns, "missing").unwrap(), None);
let k = store.keys(ns).unwrap();
assert!(k.contains(&"count".to_string()));
assert!(k.contains(&"name".to_string()));
assert!(store.delete(ns, "count").unwrap());
assert!(!store.delete(ns, "count").unwrap());
assert_eq!(store.get(ns, "count").unwrap(), None);
}
#[test]
fn invalid_namespace() {
let (store, _tmp) = new_store();
assert!(store.state_path("../evil").is_err());
assert!(store.state_path("foo/bar").is_err());
assert!(store.state_path("foo\\bar").is_err());
assert!(store.state_path("").is_err());
assert!(store.state_path("foo\0bar").is_err());
}
#[test]
fn get_nonexistent_namespace_returns_empty() {
let (store, _tmp) = new_store();
let result = store.get("ghost_ns", "any_key").unwrap();
assert_eq!(result, None);
}
#[test]
fn keys_nonexistent_namespace_returns_empty() {
let (store, _tmp) = new_store();
let result = store.keys("ghost_ns").unwrap();
assert!(result.is_empty());
}
#[test]
fn delete_nonexistent_key_returns_false() {
let (store, _tmp) = new_store();
assert!(!store.delete("delns", "nope").unwrap());
}
#[test]
fn set_overwrites_existing_value() {
let (store, _tmp) = new_store();
let ns = "ow";
store.set(ns, "k", serde_json::json!(1)).unwrap();
store.set(ns, "k", serde_json::json!(2)).unwrap();
assert_eq!(store.get(ns, "k").unwrap(), Some(serde_json::json!(2)));
}
#[test]
fn state_path_valid_namespaces() {
let (store, _tmp) = new_store();
assert!(store.state_path("default").is_ok());
assert!(store.state_path("my-app").is_ok());
assert!(store.state_path("test_123").is_ok());
}
#[test]
fn has_returns_existence() {
let (store, _tmp) = new_store();
let ns = "hasns";
assert!(!store.has(ns, "x").unwrap());
store.set(ns, "x", serde_json::json!(1)).unwrap();
assert!(store.has(ns, "x").unwrap());
}
#[test]
fn set_nx_only_sets_if_absent() {
let (store, _tmp) = new_store();
let ns = "snx";
assert!(store.set_nx(ns, "k", serde_json::json!("first")).unwrap());
assert!(!store.set_nx(ns, "k", serde_json::json!("second")).unwrap());
assert_eq!(
store.get(ns, "k").unwrap(),
Some(serde_json::json!("first")),
"set_nx should not overwrite"
);
}
#[test]
fn incr_initialises_and_increments() {
let (store, _tmp) = new_store();
let ns = "inc";
let v = store.incr(ns, "counter", 1.0, 0.0).unwrap();
assert!((v - 1.0).abs() < f64::EPSILON);
let v = store.incr(ns, "counter", 5.0, 0.0).unwrap();
assert!((v - 6.0).abs() < f64::EPSILON);
let v = store.incr(ns, "counter", -2.0, 0.0).unwrap();
assert!((v - 4.0).abs() < f64::EPSILON);
}
#[test]
fn incr_rejects_non_numeric() {
let (store, _tmp) = new_store();
let ns = "incerr";
store.set(ns, "s", serde_json::json!("hello")).unwrap();
let err = store.incr(ns, "s", 1.0, 0.0).unwrap_err();
assert!(err.contains("not a number"), "got: {err}");
}
#[test]
fn incr_custom_default() {
let (store, _tmp) = new_store();
let ns = "incdef";
let v = store.incr(ns, "score", 10.0, 100.0).unwrap();
assert!((v - 110.0).abs() < f64::EPSILON, "100 + 10 = 110");
}
}
#[cfg(test)]
mod proptests {
use super::*;
use proptest::prelude::*;
fn new_store() -> (JsonFileStore, tempfile::TempDir) {
let tmp = tempfile::tempdir().unwrap();
let store = JsonFileStore::new(tmp.path().to_path_buf());
(store, tmp)
}
proptest! {
#[test]
fn roundtrip_arbitrary_values(
key in "[a-z]{1,20}",
val in any::<i64>(),
) {
let (store, _tmp) = new_store();
let ns = "rt";
let json_val = serde_json::json!(val);
store.set(ns, &key, json_val.clone()).unwrap();
let got = store.get(ns, &key).unwrap();
prop_assert_eq!(got, Some(json_val));
let _ = store.delete(ns, &key);
}
#[test]
fn traversal_always_rejected(
prefix in "[a-z]{0,5}",
suffix in "[a-z]{0,5}",
) {
let (store, _tmp) = new_store();
let evil = format!("{prefix}/../{suffix}");
prop_assert!(store.state_path(&evil).is_err());
}
#[test]
fn nul_byte_always_rejected(
prefix in "[a-z]{0,10}",
suffix in "[a-z]{0,10}",
) {
let (store, _tmp) = new_store();
let evil = format!("{prefix}\0{suffix}");
prop_assert!(store.state_path(&evil).is_err());
}
}
}