1use std::collections::HashMap;
49use std::fs;
50use std::path::{Path, PathBuf};
51use std::sync::{Arc, Mutex};
52
53use serde_json::Value;
54
55pub trait StateStore: Send + Sync {
64 fn get(&self, ns: &str, key: &str) -> Result<Option<Value>, String>;
66
67 fn set(&self, ns: &str, key: &str, value: Value) -> Result<(), String>;
69
70 fn delete(&self, ns: &str, key: &str) -> Result<bool, String>;
72
73 fn keys(&self, ns: &str) -> Result<Vec<String>, String>;
75
76 fn has(&self, ns: &str, key: &str) -> Result<bool, String>;
82
83 fn set_nx(&self, ns: &str, key: &str, value: Value) -> Result<bool, String>;
93
94 fn incr(&self, ns: &str, key: &str, delta: f64, default: f64) -> Result<f64, String>;
110}
111
112pub struct JsonFileStore {
142 root: PathBuf,
143 locks: Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>,
146}
147
148impl JsonFileStore {
149 pub fn new(root: PathBuf) -> Self {
154 Self {
155 root,
156 locks: Mutex::new(HashMap::new()),
157 }
158 }
159
160 fn ns_lock(&self, path: &Path) -> Result<Arc<Mutex<()>>, String> {
167 let mut map = self
168 .locks
169 .lock()
170 .map_err(|_| "state: locks map poisoned".to_string())?;
171 Ok(Arc::clone(
172 map.entry(path.to_path_buf())
173 .or_insert_with(|| Arc::new(Mutex::new(()))),
174 ))
175 }
176
177 pub fn root(&self) -> &Path {
179 &self.root
180 }
181
182 fn ensure_root(&self) -> Result<&Path, String> {
184 if !self.root.exists() {
185 fs::create_dir_all(&self.root)
186 .map_err(|e| format!("Failed to create state dir: {e}"))?;
187 }
188 Ok(&self.root)
189 }
190
191 pub fn state_path(&self, ns: &str) -> Result<PathBuf, String> {
194 if ns.contains('/')
195 || ns.contains('\\')
196 || ns.contains("..")
197 || ns.contains('\0')
198 || ns.is_empty()
199 {
200 return Err(format!("Invalid namespace: '{ns}'"));
201 }
202 let dir = self.ensure_root()?;
203 Ok(dir.join(format!("{ns}.json")))
204 }
205
206 fn load(&self, ns: &str) -> Result<HashMap<String, Value>, String> {
207 let path = self.state_path(ns)?;
208 if !path.exists() {
209 return Ok(HashMap::new());
210 }
211 let content =
212 fs::read_to_string(&path).map_err(|e| format!("Failed to read state '{ns}': {e}"))?;
213 serde_json::from_str(&content).map_err(|e| format!("Failed to parse state '{ns}': {e}"))
214 }
215
216 fn save(&self, ns: &str, data: &HashMap<String, Value>) -> Result<(), String> {
217 let path = self.state_path(ns)?;
218 let tmp = path.with_extension("json.tmp");
219 let content = serde_json::to_string_pretty(data)
220 .map_err(|e| format!("Failed to serialize state: {e}"))?;
221 fs::write(&tmp, &content).map_err(|e| format!("Failed to write state tmp: {e}"))?;
222 fs::rename(&tmp, &path).map_err(|e| format!("Failed to rename state file: {e}"))?;
223 Ok(())
224 }
225}
226
227impl StateStore for JsonFileStore {
228 fn get(&self, ns: &str, key: &str) -> Result<Option<Value>, String> {
229 let path = self.state_path(ns)?;
230 let lock = self.ns_lock(&path)?;
231 let _guard = lock
232 .lock()
233 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
234 let state = self.load(ns)?;
235 Ok(state.get(key).cloned())
236 }
237
238 fn set(&self, ns: &str, key: &str, value: Value) -> Result<(), String> {
239 let path = self.state_path(ns)?;
240 let lock = self.ns_lock(&path)?;
241 let _guard = lock
242 .lock()
243 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
244 let mut state = self.load(ns)?;
245 state.insert(key.to_string(), value);
246 self.save(ns, &state)
247 }
248
249 fn delete(&self, ns: &str, key: &str) -> Result<bool, String> {
250 let path = self.state_path(ns)?;
251 let lock = self.ns_lock(&path)?;
252 let _guard = lock
253 .lock()
254 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
255 let mut state = self.load(ns)?;
256 let existed = state.remove(key).is_some();
257 if existed {
258 self.save(ns, &state)?;
259 }
260 Ok(existed)
261 }
262
263 fn keys(&self, ns: &str) -> Result<Vec<String>, String> {
264 let path = self.state_path(ns)?;
265 let lock = self.ns_lock(&path)?;
266 let _guard = lock
267 .lock()
268 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
269 let state = self.load(ns)?;
270 Ok(state.keys().cloned().collect())
271 }
272
273 fn has(&self, ns: &str, key: &str) -> Result<bool, String> {
274 let path = self.state_path(ns)?;
275 let lock = self.ns_lock(&path)?;
276 let _guard = lock
277 .lock()
278 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
279 let state = self.load(ns)?;
280 Ok(state.contains_key(key))
281 }
282
283 fn set_nx(&self, ns: &str, key: &str, value: Value) -> Result<bool, String> {
284 let path = self.state_path(ns)?;
285 let lock = self.ns_lock(&path)?;
286 let _guard = lock
287 .lock()
288 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
289 let mut state = self.load(ns)?;
290 if state.contains_key(key) {
291 return Ok(false);
292 }
293 state.insert(key.to_string(), value);
294 self.save(ns, &state)?;
295 Ok(true)
296 }
297
298 fn incr(&self, ns: &str, key: &str, delta: f64, default: f64) -> Result<f64, String> {
299 let path = self.state_path(ns)?;
300 let lock = self.ns_lock(&path)?;
301 let _guard = lock
302 .lock()
303 .map_err(|_| format!("state: lock poisoned for ns '{ns}'"))?;
304 let mut state = self.load(ns)?;
305 let current = match state.get(key) {
306 Some(v) => v
307 .as_f64()
308 .ok_or_else(|| format!("incr: value at '{key}' is not a number"))?,
309 None => default,
310 };
311 let new_val = current + delta;
312 state.insert(key.to_string(), serde_json::json!(new_val));
313 self.save(ns, &state)?;
314 Ok(new_val)
315 }
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321 use tempfile::TempDir;
322
323 fn new_store() -> (JsonFileStore, TempDir) {
326 let tmp = tempfile::tempdir().unwrap();
327 let store = JsonFileStore::new(tmp.path().to_path_buf());
328 (store, tmp)
329 }
330
331 #[test]
332 fn roundtrip() {
333 let (store, _tmp) = new_store();
334 let ns = "rt";
335
336 store.set(ns, "count", serde_json::json!(42)).unwrap();
337 store
338 .set(ns, "name", serde_json::json!("algocline"))
339 .unwrap();
340
341 assert_eq!(store.get(ns, "count").unwrap(), Some(serde_json::json!(42)));
342 assert_eq!(
343 store.get(ns, "name").unwrap(),
344 Some(serde_json::json!("algocline"))
345 );
346 assert_eq!(store.get(ns, "missing").unwrap(), None);
347
348 let k = store.keys(ns).unwrap();
349 assert!(k.contains(&"count".to_string()));
350 assert!(k.contains(&"name".to_string()));
351
352 assert!(store.delete(ns, "count").unwrap());
353 assert!(!store.delete(ns, "count").unwrap());
354 assert_eq!(store.get(ns, "count").unwrap(), None);
355 }
356
357 #[test]
358 fn invalid_namespace() {
359 let (store, _tmp) = new_store();
360 assert!(store.state_path("../evil").is_err());
361 assert!(store.state_path("foo/bar").is_err());
362 assert!(store.state_path("foo\\bar").is_err());
363 assert!(store.state_path("").is_err());
364 assert!(store.state_path("foo\0bar").is_err());
365 }
366
367 #[test]
368 fn get_nonexistent_namespace_returns_empty() {
369 let (store, _tmp) = new_store();
370 let result = store.get("ghost_ns", "any_key").unwrap();
371 assert_eq!(result, None);
372 }
373
374 #[test]
375 fn keys_nonexistent_namespace_returns_empty() {
376 let (store, _tmp) = new_store();
377 let result = store.keys("ghost_ns").unwrap();
378 assert!(result.is_empty());
379 }
380
381 #[test]
382 fn delete_nonexistent_key_returns_false() {
383 let (store, _tmp) = new_store();
384 assert!(!store.delete("delns", "nope").unwrap());
385 }
386
387 #[test]
388 fn set_overwrites_existing_value() {
389 let (store, _tmp) = new_store();
390 let ns = "ow";
391
392 store.set(ns, "k", serde_json::json!(1)).unwrap();
393 store.set(ns, "k", serde_json::json!(2)).unwrap();
394 assert_eq!(store.get(ns, "k").unwrap(), Some(serde_json::json!(2)));
395 }
396
397 #[test]
398 fn state_path_valid_namespaces() {
399 let (store, _tmp) = new_store();
400 assert!(store.state_path("default").is_ok());
401 assert!(store.state_path("my-app").is_ok());
402 assert!(store.state_path("test_123").is_ok());
403 }
404
405 #[test]
408 fn has_returns_existence() {
409 let (store, _tmp) = new_store();
410 let ns = "hasns";
411
412 assert!(!store.has(ns, "x").unwrap());
413 store.set(ns, "x", serde_json::json!(1)).unwrap();
414 assert!(store.has(ns, "x").unwrap());
415 }
416
417 #[test]
418 fn set_nx_only_sets_if_absent() {
419 let (store, _tmp) = new_store();
420 let ns = "snx";
421
422 assert!(store.set_nx(ns, "k", serde_json::json!("first")).unwrap());
423 assert!(!store.set_nx(ns, "k", serde_json::json!("second")).unwrap());
424 assert_eq!(
425 store.get(ns, "k").unwrap(),
426 Some(serde_json::json!("first")),
427 "set_nx should not overwrite"
428 );
429 }
430
431 #[test]
432 fn incr_initialises_and_increments() {
433 let (store, _tmp) = new_store();
434 let ns = "inc";
435
436 let v = store.incr(ns, "counter", 1.0, 0.0).unwrap();
438 assert!((v - 1.0).abs() < f64::EPSILON);
439
440 let v = store.incr(ns, "counter", 5.0, 0.0).unwrap();
442 assert!((v - 6.0).abs() < f64::EPSILON);
443
444 let v = store.incr(ns, "counter", -2.0, 0.0).unwrap();
446 assert!((v - 4.0).abs() < f64::EPSILON);
447 }
448
449 #[test]
450 fn incr_rejects_non_numeric() {
451 let (store, _tmp) = new_store();
452 let ns = "incerr";
453
454 store.set(ns, "s", serde_json::json!("hello")).unwrap();
455 let err = store.incr(ns, "s", 1.0, 0.0).unwrap_err();
456 assert!(err.contains("not a number"), "got: {err}");
457 }
458
459 #[test]
460 fn incr_custom_default() {
461 let (store, _tmp) = new_store();
462 let ns = "incdef";
463
464 let v = store.incr(ns, "score", 10.0, 100.0).unwrap();
465 assert!((v - 110.0).abs() < f64::EPSILON, "100 + 10 = 110");
466 }
467}
468
469#[cfg(test)]
470mod proptests {
471 use super::*;
472 use proptest::prelude::*;
473
474 fn new_store() -> (JsonFileStore, tempfile::TempDir) {
475 let tmp = tempfile::tempdir().unwrap();
476 let store = JsonFileStore::new(tmp.path().to_path_buf());
477 (store, tmp)
478 }
479
480 proptest! {
481 #[test]
483 fn roundtrip_arbitrary_values(
484 key in "[a-z]{1,20}",
485 val in any::<i64>(),
486 ) {
487 let (store, _tmp) = new_store();
488 let ns = "rt";
489 let json_val = serde_json::json!(val);
490 store.set(ns, &key, json_val.clone()).unwrap();
491 let got = store.get(ns, &key).unwrap();
492 prop_assert_eq!(got, Some(json_val));
493 let _ = store.delete(ns, &key);
494 }
495
496 #[test]
498 fn traversal_always_rejected(
499 prefix in "[a-z]{0,5}",
500 suffix in "[a-z]{0,5}",
501 ) {
502 let (store, _tmp) = new_store();
503 let evil = format!("{prefix}/../{suffix}");
504 prop_assert!(store.state_path(&evil).is_err());
505 }
506
507 #[test]
509 fn nul_byte_always_rejected(
510 prefix in "[a-z]{0,10}",
511 suffix in "[a-z]{0,10}",
512 ) {
513 let (store, _tmp) = new_store();
514 let evil = format!("{prefix}\0{suffix}");
515 prop_assert!(store.state_path(&evil).is_err());
516 }
517 }
518}