1use crate::error::FaucetError;
17use async_trait::async_trait;
18use serde_json::Value;
19use std::collections::HashMap;
20use std::path::{Path, PathBuf};
21use tokio::io::AsyncWriteExt;
22use tokio::sync::Mutex;
23
24#[async_trait]
30pub trait StateStore: Send + Sync {
31 async fn get(&self, key: &str) -> Result<Option<Value>, FaucetError>;
33
34 async fn put(&self, key: &str, value: &Value) -> Result<(), FaucetError>;
39
40 async fn delete(&self, key: &str) -> Result<(), FaucetError>;
42
43 async fn check(
50 &self,
51 _ctx: &crate::check::CheckContext,
52 ) -> Result<crate::check::CheckReport, FaucetError> {
53 Ok(crate::check::CheckReport::not_implemented())
54 }
55}
56
57pub const DOCTOR_SENTINEL_KEY: &str = "faucet_doctor_probe";
60
61pub fn validate_state_key(key: &str) -> Result<(), FaucetError> {
65 if key.is_empty() {
66 return Err(FaucetError::State("state key must not be empty".into()));
67 }
68 if key.len() > 256 {
69 return Err(FaucetError::State(format!(
70 "state key '{key}' exceeds 256 characters"
71 )));
72 }
73 for (i, c) in key.char_indices() {
74 let ok = c.is_ascii_alphanumeric() || matches!(c, '_' | '-' | ':' | '.');
75 if !ok {
76 return Err(FaucetError::State(format!(
77 "state key '{key}' contains illegal character {c:?} at byte {i}"
78 )));
79 }
80 }
81 if key == "." || key == ".." || key.starts_with('.') {
82 return Err(FaucetError::State(format!(
83 "state key '{key}' must not begin with a dot"
84 )));
85 }
86 Ok(())
87}
88
89#[derive(Default)]
93pub struct MemoryStateStore {
94 inner: Mutex<HashMap<String, Value>>,
95}
96
97impl MemoryStateStore {
98 pub fn new() -> Self {
100 Self::default()
101 }
102}
103
104#[async_trait]
105impl StateStore for MemoryStateStore {
106 async fn get(&self, key: &str) -> Result<Option<Value>, FaucetError> {
107 validate_state_key(key)?;
108 Ok(self.inner.lock().await.get(key).cloned())
109 }
110
111 async fn put(&self, key: &str, value: &Value) -> Result<(), FaucetError> {
112 validate_state_key(key)?;
113 self.inner
114 .lock()
115 .await
116 .insert(key.to_owned(), value.clone());
117 Ok(())
118 }
119
120 async fn delete(&self, key: &str) -> Result<(), FaucetError> {
121 validate_state_key(key)?;
122 self.inner.lock().await.remove(key);
123 Ok(())
124 }
125
126 async fn check(
127 &self,
128 _ctx: &crate::check::CheckContext,
129 ) -> Result<crate::check::CheckReport, FaucetError> {
130 Ok(crate::check::CheckReport::single(
132 crate::check::Probe::pass("sentinel", std::time::Duration::ZERO),
133 ))
134 }
135}
136
137fn safe_filename(key: &str) -> String {
145 key.replace(':', "%3A")
146}
147
148pub struct FileStateStore {
163 root: PathBuf,
164 write_lock: Mutex<()>,
165}
166
167impl FileStateStore {
168 pub fn new(root: impl Into<PathBuf>) -> Self {
170 Self {
171 root: root.into(),
172 write_lock: Mutex::new(()),
173 }
174 }
175
176 fn entry_path(&self, key: &str) -> PathBuf {
177 self.root.join(format!("{}.json", safe_filename(key)))
178 }
179
180 fn temp_path(&self, key: &str) -> PathBuf {
181 use std::sync::OnceLock;
198 use std::sync::atomic::{AtomicU64, Ordering};
199 static PROC_TOKEN: OnceLock<String> = OnceLock::new();
200 static SEQ: AtomicU64 = AtomicU64::new(0);
201 let token = PROC_TOKEN.get_or_init(|| uuid::Uuid::new_v4().simple().to_string());
202 let seq = SEQ.fetch_add(1, Ordering::Relaxed);
203 self.root
204 .join(format!("{}.{}.{}.json.tmp", safe_filename(key), token, seq))
205 }
206
207 async fn ensure_root(&self) -> Result<(), FaucetError> {
208 tokio::fs::create_dir_all(&self.root).await.map_err(|e| {
209 FaucetError::State(format!(
210 "failed to create state dir {}: {e}",
211 self.root.display()
212 ))
213 })
214 }
215
216 pub fn root(&self) -> &Path {
218 &self.root
219 }
220}
221
222#[async_trait]
223impl StateStore for FileStateStore {
224 async fn get(&self, key: &str) -> Result<Option<Value>, FaucetError> {
225 validate_state_key(key)?;
226 let path = self.entry_path(key);
227 match tokio::fs::read(&path).await {
228 Ok(bytes) => {
229 let value: Value = serde_json::from_slice(&bytes).map_err(|e| {
230 FaucetError::State(format!(
231 "failed to parse state file {}: {e}",
232 path.display()
233 ))
234 })?;
235 Ok(Some(value))
236 }
237 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
238 Err(e) => Err(FaucetError::State(format!(
239 "failed to read state file {}: {e}",
240 path.display()
241 ))),
242 }
243 }
244
245 async fn put(&self, key: &str, value: &Value) -> Result<(), FaucetError> {
246 validate_state_key(key)?;
247 let _guard = self.write_lock.lock().await;
248 self.ensure_root().await?;
249 let bytes = serde_json::to_vec(value).map_err(|e| {
250 FaucetError::State(format!("failed to serialize state for key '{key}': {e}"))
251 })?;
252 let final_path = self.entry_path(key);
253 let tmp_path = self.temp_path(key);
254
255 {
261 let mut file = tokio::fs::File::create(&tmp_path).await.map_err(|e| {
262 FaucetError::State(format!(
263 "failed to create temp state file {}: {e}",
264 tmp_path.display()
265 ))
266 })?;
267 file.write_all(&bytes).await.map_err(|e| {
268 FaucetError::State(format!(
269 "failed to write temp state file {}: {e}",
270 tmp_path.display()
271 ))
272 })?;
273 file.sync_all().await.map_err(|e| {
274 FaucetError::State(format!(
275 "failed to fsync temp state file {}: {e}",
276 tmp_path.display()
277 ))
278 })?;
279 }
280
281 tokio::fs::rename(&tmp_path, &final_path)
282 .await
283 .map_err(|e| {
284 FaucetError::State(format!(
285 "failed to commit state file {}: {e}",
286 final_path.display()
287 ))
288 })?;
289
290 #[cfg(unix)]
296 {
297 let dir = tokio::fs::File::open(&self.root).await.map_err(|e| {
298 FaucetError::State(format!(
299 "failed to open state dir {} for fsync: {e}",
300 self.root.display()
301 ))
302 })?;
303 dir.sync_all().await.map_err(|e| {
304 FaucetError::State(format!(
305 "failed to fsync state dir {}: {e}",
306 self.root.display()
307 ))
308 })?;
309 }
310
311 tracing::debug!(
312 key,
313 path = %final_path.display(),
314 "state file written"
315 );
316 Ok(())
317 }
318
319 async fn delete(&self, key: &str) -> Result<(), FaucetError> {
320 validate_state_key(key)?;
321 let path = self.entry_path(key);
322 match tokio::fs::remove_file(&path).await {
323 Ok(()) => Ok(()),
324 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
325 Err(e) => Err(FaucetError::State(format!(
326 "failed to delete state file {}: {e}",
327 path.display()
328 ))),
329 }
330 }
331
332 async fn check(
333 &self,
334 _ctx: &crate::check::CheckContext,
335 ) -> Result<crate::check::CheckReport, FaucetError> {
336 use crate::check::{CheckReport, Probe};
337 let start = std::time::Instant::now();
341 let probe = match self.sentinel_roundtrip().await {
342 Ok(()) => Probe::pass("sentinel", start.elapsed()),
343 Err(e) => Probe::fail_hint(
344 "sentinel",
345 start.elapsed(),
346 e.to_string(),
347 format!("ensure {} exists and is writable", self.root.display()),
348 ),
349 };
350 Ok(CheckReport::single(probe))
351 }
352}
353
354impl FileStateStore {
355 async fn sentinel_roundtrip(&self) -> Result<(), FaucetError> {
358 let probe = serde_json::json!({ "faucet_doctor": true });
359 self.put(DOCTOR_SENTINEL_KEY, &probe).await?;
360 let got = self.get(DOCTOR_SENTINEL_KEY).await?;
361 let _ = self.delete(DOCTOR_SENTINEL_KEY).await;
363 match got {
364 Some(v) if v == probe => Ok(()),
365 _ => Err(FaucetError::State(
366 "sentinel readback did not match what was written".into(),
367 )),
368 }
369 }
370}
371
372#[cfg(test)]
373mod tests {
374 use super::*;
375 use serde_json::json;
376 use std::sync::Arc;
377 use tempfile::TempDir;
378
379 #[test]
382 fn rejects_empty_key() {
383 let err = validate_state_key("").unwrap_err();
384 assert!(matches!(err, FaucetError::State(_)));
385 }
386
387 #[test]
388 fn rejects_path_traversal_segments() {
389 for k in ["../etc/passwd", "a/b", "a\\b", "..", "."] {
390 assert!(validate_state_key(k).is_err(), "expected reject for {k:?}");
391 }
392 }
393
394 #[test]
395 fn rejects_leading_dot() {
396 assert!(validate_state_key(".hidden").is_err());
397 }
398
399 #[test]
400 fn rejects_over_long_key() {
401 let k = "a".repeat(257);
402 assert!(validate_state_key(&k).is_err());
403 }
404
405 #[test]
406 fn accepts_typical_keys() {
407 for k in [
408 "github_issues",
409 "pipeline:rest:issues",
410 "with.dot",
411 "with-dash_and_underscore",
412 "lower-Case_99",
413 ] {
414 validate_state_key(k).unwrap_or_else(|e| panic!("expected ok for {k:?}: {e}"));
415 }
416 }
417
418 #[tokio::test]
421 async fn memory_get_returns_none_for_missing_key() {
422 let s = MemoryStateStore::new();
423 assert!(s.get("nope").await.unwrap().is_none());
424 }
425
426 #[tokio::test]
427 async fn memory_put_then_get_round_trips() {
428 let s = MemoryStateStore::new();
429 s.put("k", &json!({"cursor": "abc", "n": 7})).await.unwrap();
430 let got = s.get("k").await.unwrap().unwrap();
431 assert_eq!(got["cursor"], "abc");
432 assert_eq!(got["n"], 7);
433 }
434
435 #[tokio::test]
436 async fn memory_put_overwrites_previous_value() {
437 let s = MemoryStateStore::new();
438 s.put("k", &json!(1)).await.unwrap();
439 s.put("k", &json!(2)).await.unwrap();
440 assert_eq!(s.get("k").await.unwrap().unwrap(), json!(2));
441 }
442
443 #[tokio::test]
444 async fn memory_delete_makes_get_return_none() {
445 let s = MemoryStateStore::new();
446 s.put("k", &json!("v")).await.unwrap();
447 s.delete("k").await.unwrap();
448 assert!(s.get("k").await.unwrap().is_none());
449 }
450
451 #[tokio::test]
452 async fn memory_delete_missing_key_is_ok() {
453 let s = MemoryStateStore::new();
454 s.delete("absent").await.unwrap();
455 }
456
457 #[tokio::test]
458 async fn memory_rejects_invalid_keys() {
459 let s = MemoryStateStore::new();
460 assert!(s.get("a/b").await.is_err());
461 assert!(s.put("a/b", &json!(1)).await.is_err());
462 assert!(s.delete("a/b").await.is_err());
463 }
464
465 #[tokio::test]
468 async fn file_get_returns_none_for_missing_key() {
469 let dir = TempDir::new().unwrap();
470 let s = FileStateStore::new(dir.path());
471 assert!(s.get("nope").await.unwrap().is_none());
472 }
473
474 #[tokio::test]
475 async fn file_put_creates_root_directory_lazily() {
476 let dir = TempDir::new().unwrap();
477 let root = dir.path().join("nested/state");
478 let s = FileStateStore::new(&root);
479 s.put("k", &json!("v")).await.unwrap();
480 assert!(root.is_dir(), "root dir should be created on first put");
481 }
482
483 #[tokio::test]
484 async fn file_put_then_get_round_trips() {
485 let dir = TempDir::new().unwrap();
486 let s = FileStateStore::new(dir.path());
487 let value = json!({"cursor": "abc", "n": 42, "nested": {"flag": true}});
488 s.put("github_issues", &value).await.unwrap();
489 let got = s.get("github_issues").await.unwrap().unwrap();
490 assert_eq!(got, value);
491 }
492
493 #[test]
494 fn temp_path_is_unique_and_not_pid_derived() {
495 let dir = TempDir::new().unwrap();
501 let s = FileStateStore::new(dir.path());
502
503 let a = s.temp_path("k");
504 let b = s.temp_path("k");
505 assert_ne!(a, b);
507
508 let name_a = a.file_name().unwrap().to_str().unwrap();
509 let pid = std::process::id().to_string();
510 assert!(
512 !name_a.split('.').any(|seg| seg == pid),
513 "temp filename {name_a} must not embed the process id ({pid})"
514 );
515 assert!(name_a.ends_with(".json.tmp"));
516 }
517
518 #[test]
519 fn safe_filename_percent_encodes_colon() {
520 assert_eq!(
521 safe_filename("pipeline:rest:issues"),
522 "pipeline%3Arest%3Aissues"
523 );
524 assert_eq!(safe_filename("plain_key-1.v2"), "plain_key-1.v2");
525 }
526
527 #[tokio::test]
528 async fn file_round_trips_colon_keys_with_safe_filename() {
529 let dir = TempDir::new().unwrap();
533 let s = FileStateStore::new(dir.path());
534 let value = json!({"cursor": "z"});
535 s.put("pipeline:rest:issues", &value).await.unwrap();
536 assert_eq!(s.get("pipeline:rest:issues").await.unwrap().unwrap(), value);
537 assert!(dir.path().join("pipeline%3Arest%3Aissues.json").exists());
539 let mut has_colon = false;
540 for entry in std::fs::read_dir(dir.path()).unwrap() {
541 if entry.unwrap().file_name().to_string_lossy().contains(':') {
542 has_colon = true;
543 }
544 }
545 assert!(!has_colon, "no state filename may contain ':'");
546 }
547
548 fn has_tmp_residue(dir: &std::path::Path) -> bool {
552 std::fs::read_dir(dir)
553 .unwrap()
554 .filter_map(|e| e.ok())
555 .any(|e| e.file_name().to_string_lossy().ends_with(".json.tmp"))
556 }
557
558 #[tokio::test]
559 async fn file_put_overwrites_previous_value_atomically() {
560 let dir = TempDir::new().unwrap();
561 let s = FileStateStore::new(dir.path());
562 s.put("k", &json!({"v": 1})).await.unwrap();
563 s.put("k", &json!({"v": 2})).await.unwrap();
564 assert_eq!(s.get("k").await.unwrap().unwrap(), json!({"v": 2}));
565 assert!(!has_tmp_residue(dir.path()), "no temp residue after put");
567 }
568
569 #[test]
570 fn file_temp_paths_are_unique_per_write() {
571 let dir = TempDir::new().unwrap();
576 let s = FileStateStore::new(dir.path());
577 let a = s.temp_path("k");
578 let b = s.temp_path("k");
579 assert_ne!(a, b, "each write must get a distinct temp path");
580 assert_eq!(s.entry_path("k"), s.entry_path("k"));
582 }
583
584 #[tokio::test]
585 async fn file_put_writes_complete_durable_file_with_no_temp_residue() {
586 let dir = TempDir::new().unwrap();
593 let s = FileStateStore::new(dir.path());
594 let big: Vec<Value> = (0..1_000)
595 .map(|i| json!({"i": i, "s": "x".repeat(20)}))
596 .collect();
597 let value = json!({"cursor": "abc", "rows": big});
598
599 s.put("github_issues", &value).await.unwrap();
600
601 let raw = tokio::fs::read(dir.path().join("github_issues.json"))
603 .await
604 .expect("state file must exist after put");
605 assert!(!raw.is_empty(), "state file must not be zero-length");
606 let parsed: Value = serde_json::from_slice(&raw).expect("state file must be valid JSON");
607 assert_eq!(parsed, value);
608
609 assert!(!has_tmp_residue(dir.path()), "no temp residue after put");
611 }
612
613 #[tokio::test]
614 async fn file_delete_removes_file() {
615 let dir = TempDir::new().unwrap();
616 let s = FileStateStore::new(dir.path());
617 s.put("k", &json!("v")).await.unwrap();
618 s.delete("k").await.unwrap();
619 assert!(s.get("k").await.unwrap().is_none());
620 assert!(!dir.path().join("k.json").exists());
621 }
622
623 #[tokio::test]
624 async fn file_delete_missing_key_is_ok() {
625 let dir = TempDir::new().unwrap();
626 let s = FileStateStore::new(dir.path());
627 s.delete("absent").await.unwrap();
628 }
629
630 #[tokio::test]
631 async fn file_get_returns_error_for_corrupt_json() {
632 let dir = TempDir::new().unwrap();
633 let s = FileStateStore::new(dir.path());
634 tokio::fs::create_dir_all(dir.path()).await.unwrap();
635 tokio::fs::write(dir.path().join("bad.json"), b"not json")
636 .await
637 .unwrap();
638 let err = s.get("bad").await.unwrap_err();
639 match err {
640 FaucetError::State(msg) => assert!(msg.contains("bad.json")),
641 other => panic!("expected State error, got {other:?}"),
642 }
643 }
644
645 #[tokio::test]
646 async fn file_concurrent_puts_do_not_corrupt_or_leak_temp() {
647 let dir = TempDir::new().unwrap();
648 let s = Arc::new(FileStateStore::new(dir.path()));
649 let mut handles = vec![];
650 for i in 0..50 {
651 let s = Arc::clone(&s);
652 handles.push(tokio::spawn(async move {
653 s.put("k", &json!({"i": i})).await.unwrap();
654 }));
655 }
656 for h in handles {
657 h.await.unwrap();
658 }
659 let got = s.get("k").await.unwrap().unwrap();
661 let i = got["i"].as_i64().unwrap();
662 assert!((0..50).contains(&i));
663 assert!(
665 !has_tmp_residue(dir.path()),
666 "no temp residue after concurrent puts"
667 );
668 }
669
670 #[tokio::test]
671 async fn file_store_works_through_trait_object() {
672 let dir = TempDir::new().unwrap();
673 let s: Box<dyn StateStore> = Box::new(FileStateStore::new(dir.path()));
674 s.put("k", &json!(1)).await.unwrap();
675 assert_eq!(s.get("k").await.unwrap().unwrap(), json!(1));
676 }
677
678 #[tokio::test]
681 async fn memory_check_passes() {
682 let s = MemoryStateStore::new();
683 let report = s
684 .check(&crate::check::CheckContext::default())
685 .await
686 .unwrap();
687 assert_eq!(report.failed_count(), 0);
688 assert!(
689 report
690 .probes
691 .iter()
692 .all(|p| matches!(p.status, crate::check::ProbeStatus::Pass))
693 );
694 }
695
696 #[tokio::test]
697 async fn file_check_passes_for_writable_root() {
698 let dir = TempDir::new().unwrap();
699 let s = FileStateStore::new(dir.path());
700 let report = s
701 .check(&crate::check::CheckContext::default())
702 .await
703 .unwrap();
704 assert_eq!(report.failed_count(), 0, "writable root should pass");
705 let leftovers: Vec<_> = std::fs::read_dir(dir.path()).unwrap().collect();
707 assert!(leftovers.is_empty(), "check() must not leave files behind");
708 }
709
710 #[tokio::test]
711 async fn file_store_root_returns_configured_directory() {
712 let dir = TempDir::new().unwrap();
713 let s = FileStateStore::new(dir.path());
714 assert_eq!(s.root(), dir.path());
715 }
716
717 #[tokio::test]
718 async fn default_check_reports_not_implemented() {
719 struct BareStore;
722 #[async_trait]
723 impl StateStore for BareStore {
724 async fn get(&self, _key: &str) -> Result<Option<Value>, FaucetError> {
725 Ok(None)
726 }
727 async fn put(&self, _key: &str, _value: &Value) -> Result<(), FaucetError> {
728 Ok(())
729 }
730 async fn delete(&self, _key: &str) -> Result<(), FaucetError> {
731 Ok(())
732 }
733 }
734 let s = BareStore;
735 let report = s
736 .check(&crate::check::CheckContext::default())
737 .await
738 .unwrap();
739 assert_eq!(report.failed_count(), 0);
741 assert!(
742 report
743 .probes
744 .iter()
745 .any(|p| matches!(p.status, crate::check::ProbeStatus::Skip { .. })),
746 "default check must surface a skipped (not-implemented) probe"
747 );
748 }
749
750 #[tokio::test]
751 async fn file_check_fails_when_root_unusable() {
752 let dir = TempDir::new().unwrap();
754 let file = dir.path().join("not_a_dir");
755 std::fs::write(&file, b"x").unwrap();
756 let s = FileStateStore::new(file.join("state"));
757 let report = s
758 .check(&crate::check::CheckContext::default())
759 .await
760 .unwrap();
761 assert_eq!(report.failed_count(), 1, "unusable root should fail");
762 }
763}