1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use crate::backup::BackupStore;
5use crate::error::AftError;
6
7#[derive(Debug, Clone)]
9pub struct CheckpointInfo {
10 pub name: String,
11 pub file_count: usize,
12 pub created_at: u64,
13 pub skipped: Vec<(PathBuf, String)>,
18}
19
20#[derive(Debug, Clone)]
22struct Checkpoint {
23 name: String,
24 file_contents: HashMap<PathBuf, String>,
25 created_at: u64,
26}
27
28#[derive(Debug)]
37pub struct CheckpointStore {
38 checkpoints: HashMap<String, HashMap<String, Checkpoint>>,
40}
41
42impl CheckpointStore {
43 pub fn new() -> Self {
44 CheckpointStore {
45 checkpoints: HashMap::new(),
46 }
47 }
48
49 pub fn create(
63 &mut self,
64 session: &str,
65 name: &str,
66 files: Vec<PathBuf>,
67 backup_store: &BackupStore,
68 ) -> Result<CheckpointInfo, AftError> {
69 let explicit_request = !files.is_empty();
70 let file_list = if files.is_empty() {
71 backup_store.tracked_files(session)
72 } else {
73 files
74 };
75
76 let mut file_contents = HashMap::new();
77 let mut skipped: Vec<(PathBuf, String)> = Vec::new();
78 for path in &file_list {
79 match std::fs::read_to_string(path) {
80 Ok(content) => {
81 file_contents.insert(path.clone(), content);
82 }
83 Err(e) => {
84 crate::slog_warn!(
85 "checkpoint {}: skipping unreadable file {}: {}",
86 name,
87 path.display(),
88 e
89 );
90 skipped.push((path.clone(), e.to_string()));
91 }
92 }
93 }
94
95 if explicit_request && file_contents.is_empty() && !skipped.is_empty() {
101 let (path, err) = &skipped[0];
102 return Err(AftError::FileNotFound {
103 path: format!("{}: {}", path.display(), err),
104 });
105 }
106
107 let created_at = current_timestamp();
108 let file_count = file_contents.len();
109
110 let checkpoint = Checkpoint {
111 name: name.to_string(),
112 file_contents,
113 created_at,
114 };
115
116 self.checkpoints
117 .entry(session.to_string())
118 .or_default()
119 .insert(name.to_string(), checkpoint);
120
121 if skipped.is_empty() {
122 crate::slog_info!("checkpoint created: {} ({} files)", name, file_count);
123 } else {
124 crate::slog_info!(
125 "checkpoint created: {} ({} files, {} skipped)",
126 name,
127 file_count,
128 skipped.len()
129 );
130 }
131
132 Ok(CheckpointInfo {
133 name: name.to_string(),
134 file_count,
135 created_at,
136 skipped,
137 })
138 }
139
140 pub fn restore(&self, session: &str, name: &str) -> Result<CheckpointInfo, AftError> {
142 let checkpoint = self.get(session, name)?;
143 let mut paths = checkpoint.file_contents.keys().cloned().collect::<Vec<_>>();
144 paths.sort();
145
146 restore_paths_atomically(checkpoint, &paths)?;
147
148 crate::slog_info!("checkpoint restored: {}", name);
149
150 Ok(CheckpointInfo {
151 name: checkpoint.name.clone(),
152 file_count: checkpoint.file_contents.len(),
153 created_at: checkpoint.created_at,
154 skipped: Vec::new(),
155 })
156 }
157
158 pub fn restore_validated(
160 &self,
161 session: &str,
162 name: &str,
163 validated_paths: &[PathBuf],
164 ) -> Result<CheckpointInfo, AftError> {
165 let checkpoint = self.get(session, name)?;
166
167 for path in validated_paths {
168 checkpoint
169 .file_contents
170 .get(path)
171 .ok_or_else(|| AftError::FileNotFound {
172 path: path.display().to_string(),
173 })?;
174 }
175 restore_paths_atomically(checkpoint, validated_paths)?;
176
177 crate::slog_info!("checkpoint restored: {}", name);
178
179 Ok(CheckpointInfo {
180 name: checkpoint.name.clone(),
181 file_count: checkpoint.file_contents.len(),
182 created_at: checkpoint.created_at,
183 skipped: Vec::new(),
184 })
185 }
186
187 pub fn file_paths(&self, session: &str, name: &str) -> Result<Vec<PathBuf>, AftError> {
189 let checkpoint = self.get(session, name)?;
190 Ok(checkpoint.file_contents.keys().cloned().collect())
191 }
192
193 pub fn delete(&mut self, session: &str, name: &str) -> bool {
195 let Some(session_checkpoints) = self.checkpoints.get_mut(session) else {
196 return false;
197 };
198 let removed = session_checkpoints.remove(name).is_some();
199 if session_checkpoints.is_empty() {
200 self.checkpoints.remove(session);
201 }
202 removed
203 }
204
205 pub fn list(&self, session: &str) -> Vec<CheckpointInfo> {
207 self.checkpoints
208 .get(session)
209 .map(|s| {
210 s.values()
211 .map(|cp| CheckpointInfo {
212 name: cp.name.clone(),
213 file_count: cp.file_contents.len(),
214 created_at: cp.created_at,
215 skipped: Vec::new(),
216 })
217 .collect()
218 })
219 .unwrap_or_default()
220 }
221
222 pub fn total_count(&self) -> usize {
224 self.checkpoints.values().map(|s| s.len()).sum()
225 }
226
227 pub fn cleanup(&mut self, ttl_hours: u32) {
230 let now = current_timestamp();
231 let ttl_secs = ttl_hours as u64 * 3600;
232 self.checkpoints.retain(|_, session_cps| {
233 session_cps.retain(|_, cp| now.saturating_sub(cp.created_at) < ttl_secs);
234 !session_cps.is_empty()
235 });
236 }
237
238 fn get(&self, session: &str, name: &str) -> Result<&Checkpoint, AftError> {
239 self.checkpoints
240 .get(session)
241 .and_then(|s| s.get(name))
242 .ok_or_else(|| AftError::CheckpointNotFound {
243 name: name.to_string(),
244 })
245 }
246}
247
248fn restore_paths_atomically(checkpoint: &Checkpoint, paths: &[PathBuf]) -> Result<(), AftError> {
249 let mut pre_restore_snapshot: HashMap<PathBuf, Option<String>> = HashMap::new();
250 for path in paths {
251 let current = if path.exists() {
252 Some(
253 std::fs::read_to_string(path).map_err(|_| AftError::FileNotFound {
254 path: path.display().to_string(),
255 })?,
256 )
257 } else {
258 None
259 };
260 pre_restore_snapshot.insert(path.clone(), current);
261 }
262
263 let mut restored_paths: Vec<PathBuf> = Vec::new();
264 let mut created_dirs: Vec<PathBuf> = Vec::new();
265 for path in paths {
266 let content = checkpoint
267 .file_contents
268 .get(path)
269 .ok_or_else(|| AftError::FileNotFound {
270 path: path.display().to_string(),
271 })?;
272 if let Err(e) = write_restored_file(path, content, &mut created_dirs) {
273 let mut rollback_errors = Vec::new();
274 if let Some(snapshot) = pre_restore_snapshot.get(path) {
275 if let Err(rollback_error) = restore_snapshot_file(path, snapshot.as_deref()) {
276 rollback_errors.push(format!("{}: {}", path.display(), rollback_error));
277 }
278 }
279 for restored_path in restored_paths.iter().rev() {
280 if let Some(snapshot) = pre_restore_snapshot.get(restored_path) {
281 if let Err(rollback_error) =
282 restore_snapshot_file(restored_path, snapshot.as_deref())
283 {
284 rollback_errors.push(format!(
285 "{}: {}",
286 restored_path.display(),
287 rollback_error
288 ));
289 }
290 }
291 }
292 let dirs_rollback_ok = rollback_created_dirs(&created_dirs);
293 if rollback_errors.is_empty() && dirs_rollback_ok {
294 return Err(e);
295 }
296 return Err(AftError::IoError {
297 path: path.display().to_string(),
298 message: format!(
299 "{}; restore_checkpoint rollback_succeeded: {}; rollback_errors: {}",
300 e,
301 rollback_errors.is_empty() && dirs_rollback_ok,
302 if rollback_errors.is_empty() {
303 "none".to_string()
304 } else {
305 rollback_errors.join("; ")
306 }
307 ),
308 });
309 }
310 restored_paths.push(path.clone());
311 }
312
313 Ok(())
314}
315
316fn restore_snapshot_file(path: &Path, content: Option<&str>) -> Result<(), AftError> {
317 match content {
318 Some(content) => write_restored_file(path, content, &mut Vec::new()),
319 None => match std::fs::remove_file(path) {
320 Ok(()) => Ok(()),
321 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
322 Err(_) => Err(AftError::FileNotFound {
323 path: path.display().to_string(),
324 }),
325 },
326 }
327}
328
329fn write_restored_file(
330 path: &Path,
331 content: &str,
332 created_dirs: &mut Vec<PathBuf>,
333) -> Result<(), AftError> {
334 if let Some(parent) = path.parent() {
335 let missing_dirs = missing_parent_dirs(parent);
336 std::fs::create_dir_all(parent).map_err(|_| AftError::FileNotFound {
337 path: path.display().to_string(),
338 })?;
339 created_dirs.extend(missing_dirs);
340 }
341 std::fs::write(path, content).map_err(|_| AftError::FileNotFound {
342 path: path.display().to_string(),
343 })
344}
345
346fn missing_parent_dirs(parent: &Path) -> Vec<PathBuf> {
347 let mut dirs = Vec::new();
348 let mut current = Some(parent);
349
350 while let Some(dir) = current {
351 if dir.as_os_str().is_empty() || dir.exists() {
352 break;
353 }
354 dirs.push(dir.to_path_buf());
355 current = dir.parent();
356 }
357
358 dirs
359}
360
361fn rollback_created_dirs(dirs: &[PathBuf]) -> bool {
362 let mut dirs = dirs.to_vec();
363 dirs.sort_by_key(|dir| std::cmp::Reverse(dir.components().count()));
364 dirs.dedup();
365
366 let mut ok = true;
367 for dir in dirs {
368 match std::fs::remove_dir(&dir) {
369 Ok(()) => {}
370 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
371 Err(_) => ok = false,
372 }
373 }
374 ok
375}
376
377fn current_timestamp() -> u64 {
378 std::time::SystemTime::now()
379 .duration_since(std::time::UNIX_EPOCH)
380 .unwrap_or_default()
381 .as_secs()
382}
383
384#[cfg(test)]
385mod tests {
386 use super::*;
387 use crate::protocol::DEFAULT_SESSION_ID;
388 use std::fs;
389
390 fn temp_file(name: &str, content: &str) -> PathBuf {
391 let dir = std::env::temp_dir().join("aft_checkpoint_tests");
392 fs::create_dir_all(&dir).unwrap();
393 let path = dir.join(name);
394 fs::write(&path, content).unwrap();
395 path
396 }
397
398 #[test]
399 fn create_and_restore_round_trip() {
400 let path1 = temp_file("cp_rt1.txt", "hello");
401 let path2 = temp_file("cp_rt2.txt", "world");
402
403 let backup_store = BackupStore::new();
404 let mut store = CheckpointStore::new();
405
406 let info = store
407 .create(
408 DEFAULT_SESSION_ID,
409 "snap1",
410 vec![path1.clone(), path2.clone()],
411 &backup_store,
412 )
413 .unwrap();
414 assert_eq!(info.name, "snap1");
415 assert_eq!(info.file_count, 2);
416
417 fs::write(&path1, "changed1").unwrap();
419 fs::write(&path2, "changed2").unwrap();
420
421 let info = store.restore(DEFAULT_SESSION_ID, "snap1").unwrap();
423 assert_eq!(info.file_count, 2);
424 assert_eq!(fs::read_to_string(&path1).unwrap(), "hello");
425 assert_eq!(fs::read_to_string(&path2).unwrap(), "world");
426 }
427
428 #[test]
429 fn overwrite_existing_name() {
430 let path = temp_file("cp_overwrite.txt", "v1");
431 let backup_store = BackupStore::new();
432 let mut store = CheckpointStore::new();
433
434 store
435 .create(DEFAULT_SESSION_ID, "dup", vec![path.clone()], &backup_store)
436 .unwrap();
437 fs::write(&path, "v2").unwrap();
438 store
439 .create(DEFAULT_SESSION_ID, "dup", vec![path.clone()], &backup_store)
440 .unwrap();
441
442 fs::write(&path, "v3").unwrap();
444 store.restore(DEFAULT_SESSION_ID, "dup").unwrap();
445 assert_eq!(fs::read_to_string(&path).unwrap(), "v2");
446 }
447
448 #[test]
449 fn list_returns_metadata_scoped_to_session() {
450 let path = temp_file("cp_list.txt", "data");
451 let backup_store = BackupStore::new();
452 let mut store = CheckpointStore::new();
453
454 store
455 .create(DEFAULT_SESSION_ID, "a", vec![path.clone()], &backup_store)
456 .unwrap();
457 store
458 .create(DEFAULT_SESSION_ID, "b", vec![path.clone()], &backup_store)
459 .unwrap();
460 store
461 .create("other_session", "c", vec![path.clone()], &backup_store)
462 .unwrap();
463
464 let default_list = store.list(DEFAULT_SESSION_ID);
465 assert_eq!(default_list.len(), 2);
466 let names: Vec<&str> = default_list.iter().map(|i| i.name.as_str()).collect();
467 assert!(names.contains(&"a"));
468 assert!(names.contains(&"b"));
469
470 let other_list = store.list("other_session");
471 assert_eq!(other_list.len(), 1);
472 assert_eq!(other_list[0].name, "c");
473 }
474
475 #[test]
476 fn sessions_isolate_checkpoint_names() {
477 let path_a = temp_file("cp_isolated_a.txt", "a-original");
479 let path_b = temp_file("cp_isolated_b.txt", "b-original");
480 let backup_store = BackupStore::new();
481 let mut store = CheckpointStore::new();
482
483 store
485 .create("session_a", "snap", vec![path_a.clone()], &backup_store)
486 .unwrap();
487 store
488 .create("session_b", "snap", vec![path_b.clone()], &backup_store)
489 .unwrap();
490
491 fs::write(&path_a, "a-modified").unwrap();
492 fs::write(&path_b, "b-modified").unwrap();
493
494 store.restore("session_a", "snap").unwrap();
496 assert_eq!(fs::read_to_string(&path_a).unwrap(), "a-original");
497 assert_eq!(fs::read_to_string(&path_b).unwrap(), "b-modified");
498
499 fs::write(&path_a, "a-modified").unwrap();
501 store.restore("session_b", "snap").unwrap();
502 assert_eq!(fs::read_to_string(&path_a).unwrap(), "a-modified");
503 assert_eq!(fs::read_to_string(&path_b).unwrap(), "b-original");
504 }
505
506 #[test]
507 fn cleanup_removes_expired_across_sessions() {
508 let path = temp_file("cp_cleanup.txt", "data");
509 let backup_store = BackupStore::new();
510 let mut store = CheckpointStore::new();
511
512 store
513 .create(
514 DEFAULT_SESSION_ID,
515 "recent",
516 vec![path.clone()],
517 &backup_store,
518 )
519 .unwrap();
520
521 store
523 .checkpoints
524 .entry("other".to_string())
525 .or_default()
526 .insert(
527 "old".to_string(),
528 Checkpoint {
529 name: "old".to_string(),
530 file_contents: HashMap::new(),
531 created_at: 1000, },
533 );
534
535 assert_eq!(store.total_count(), 2);
536 store.cleanup(24); assert_eq!(store.total_count(), 1);
538 assert_eq!(store.list(DEFAULT_SESSION_ID)[0].name, "recent");
539 assert!(store.list("other").is_empty());
540 }
541
542 #[test]
543 fn restore_nonexistent_returns_error() {
544 let store = CheckpointStore::new();
545 let result = store.restore(DEFAULT_SESSION_ID, "nope");
546 assert!(result.is_err());
547 match result.unwrap_err() {
548 AftError::CheckpointNotFound { name } => {
549 assert_eq!(name, "nope");
550 }
551 other => panic!("expected CheckpointNotFound, got: {:?}", other),
552 }
553 }
554
555 #[test]
556 fn restore_nonexistent_in_other_session_returns_error() {
557 let path = temp_file("cp_cross_session.txt", "data");
559 let backup_store = BackupStore::new();
560 let mut store = CheckpointStore::new();
561 store
562 .create("session_a", "only_a", vec![path], &backup_store)
563 .unwrap();
564 assert!(store.restore("session_b", "only_a").is_err());
565 }
566
567 #[test]
568 fn create_skips_missing_files_from_backup_tracked_set() {
569 let readable = temp_file("cp_skip_readable.txt", "still_here");
575 let deleted = temp_file("cp_skip_deleted.txt", "about_to_vanish");
576
577 let deleted_canonical = fs::canonicalize(&deleted).unwrap();
580
581 let mut backup_store = BackupStore::new();
582 backup_store
583 .snapshot(DEFAULT_SESSION_ID, &readable, "auto")
584 .unwrap();
585 backup_store
586 .snapshot(DEFAULT_SESSION_ID, &deleted, "auto")
587 .unwrap();
588
589 fs::remove_file(&deleted).unwrap();
590
591 let mut store = CheckpointStore::new();
592 let info = store
593 .create(DEFAULT_SESSION_ID, "partial", vec![], &backup_store)
594 .expect("checkpoint should succeed despite one missing file");
595 assert_eq!(info.file_count, 1);
596 assert_eq!(info.skipped.len(), 1);
597 assert_eq!(info.skipped[0].0, deleted_canonical);
598 assert!(!info.skipped[0].1.is_empty());
599 }
600
601 #[test]
602 fn create_with_explicit_single_missing_file_errors() {
603 let missing = std::env::temp_dir()
606 .join("aft_checkpoint_tests/cp_explicit_missing_does_not_exist.txt");
607 let _ = fs::remove_file(&missing);
608
609 let backup_store = BackupStore::new();
610 let mut store = CheckpointStore::new();
611 let result = store.create(
612 DEFAULT_SESSION_ID,
613 "explicit",
614 vec![missing.clone()],
615 &backup_store,
616 );
617
618 assert!(result.is_err());
619 match result.unwrap_err() {
620 AftError::FileNotFound { path } => {
621 assert!(path.contains(&missing.display().to_string()));
622 }
623 other => panic!("expected FileNotFound, got: {:?}", other),
624 }
625 }
626
627 #[test]
628 fn create_with_explicit_mixed_files_keeps_readable_and_reports_skipped() {
629 let good = temp_file("cp_mixed_good.txt", "ok");
633 let missing = std::env::temp_dir().join("aft_checkpoint_tests/cp_mixed_missing.txt");
634 let _ = fs::remove_file(&missing);
635
636 let backup_store = BackupStore::new();
637 let mut store = CheckpointStore::new();
638 let info = store
639 .create(
640 DEFAULT_SESSION_ID,
641 "mixed",
642 vec![good.clone(), missing.clone()],
643 &backup_store,
644 )
645 .expect("mixed checkpoint should succeed when any file is readable");
646 assert_eq!(info.file_count, 1);
647 assert_eq!(info.skipped.len(), 1);
648 assert_eq!(info.skipped[0].0, missing);
649 }
650
651 #[test]
652 fn create_with_empty_files_uses_backup_tracked() {
653 let path = temp_file("cp_tracked.txt", "tracked_content");
654 let mut backup_store = BackupStore::new();
655 backup_store
656 .snapshot(DEFAULT_SESSION_ID, &path, "auto")
657 .unwrap();
658
659 let mut store = CheckpointStore::new();
660 let info = store
661 .create(DEFAULT_SESSION_ID, "from_tracked", vec![], &backup_store)
662 .unwrap();
663 assert!(info.file_count >= 1);
664
665 fs::write(&path, "modified").unwrap();
667 store.restore(DEFAULT_SESSION_ID, "from_tracked").unwrap();
668 assert_eq!(fs::read_to_string(&path).unwrap(), "tracked_content");
669 }
670
671 #[test]
672 fn restore_recreates_missing_parent_directories() {
673 let dir = tempfile::tempdir().unwrap();
674 let path = dir.path().join("nested").join("deeper").join("file.txt");
675 fs::create_dir_all(path.parent().unwrap()).unwrap();
676 fs::write(&path, "original nested content").unwrap();
677
678 let backup_store = BackupStore::new();
679 let mut store = CheckpointStore::new();
680 store
681 .create(
682 DEFAULT_SESSION_ID,
683 "nested",
684 vec![path.clone()],
685 &backup_store,
686 )
687 .unwrap();
688
689 fs::remove_dir_all(dir.path().join("nested")).unwrap();
690
691 store.restore(DEFAULT_SESSION_ID, "nested").unwrap();
692 assert_eq!(
693 fs::read_to_string(&path).unwrap(),
694 "original nested content"
695 );
696 }
697
698 #[cfg(unix)]
699 #[test]
700 fn checkpoint_restore_rolls_back_on_partial_failure() {
701 use std::os::unix::fs::PermissionsExt;
702
703 let dir = tempfile::tempdir().unwrap();
704 let path_a = dir.path().join("a.txt");
705 let path_b = dir.path().join("b.txt");
706 fs::write(&path_a, "checkpoint-a").unwrap();
707 fs::write(&path_b, "checkpoint-b").unwrap();
708
709 let backup_store = BackupStore::new();
710 let mut store = CheckpointStore::new();
711 store
712 .create(
713 DEFAULT_SESSION_ID,
714 "partial_failure",
715 vec![path_a.clone(), path_b.clone()],
716 &backup_store,
717 )
718 .unwrap();
719
720 fs::write(&path_a, "pre-restore-a").unwrap();
721 fs::write(&path_b, "pre-restore-b").unwrap();
722 let mut readonly = fs::metadata(&path_b).unwrap().permissions();
723 readonly.set_mode(0o444);
724 fs::set_permissions(&path_b, readonly).unwrap();
725
726 let result = store.restore(DEFAULT_SESSION_ID, "partial_failure");
727 let mut writable = fs::metadata(&path_b).unwrap().permissions();
728 writable.set_mode(0o644);
729 fs::set_permissions(&path_b, writable).unwrap();
730
731 assert!(result.is_err(), "restore should surface write failure");
732 assert_eq!(fs::read_to_string(&path_a).unwrap(), "pre-restore-a");
733 assert_eq!(fs::read_to_string(&path_b).unwrap(), "pre-restore-b");
734 }
735
736 #[test]
737 fn checkpoint_restore_failure_removes_created_parent_dirs() {
738 let dir = tempfile::tempdir().unwrap();
739 let missing_root = dir.path().join("created");
740 let path_a = missing_root.join("nested").join("a.txt");
741 let path_b = dir.path().join("blocking-dir");
742 fs::create_dir(&path_b).unwrap();
743
744 let checkpoint = Checkpoint {
745 name: "dir-cleanup".to_string(),
746 file_contents: HashMap::from([
747 (path_a.clone(), "checkpoint-a".to_string()),
748 (path_b.clone(), "checkpoint-b".to_string()),
749 ]),
750 created_at: current_timestamp(),
751 };
752
753 let result = restore_paths_atomically(&checkpoint, &[path_a.clone(), path_b.clone()]);
754
755 assert!(
756 result.is_err(),
757 "second restore write should fail on directory"
758 );
759 assert!(!path_a.exists(), "restored file should be rolled back");
760 assert!(
761 !missing_root.exists(),
762 "new parent directories should be removed on rollback"
763 );
764 assert!(path_b.is_dir(), "pre-existing blocking directory remains");
765 }
766}