12345678module snapshot910import os11import src.storage1213141516const restoring_key = 'restoring'1718pub const restore_unfinished = 'an earlier restore stopped before all its files were written — run koh revert to finish it (it records the tree as it is in a checkpoint first)'1920pub fn restore_pending(store &storage.Store) bool {21 return store.state(restoring_key) != none22}23242526pub fn mark_restore(store &storage.Store) ! {27 store.set_state(restoring_key, 'yes')!28}2930pub fn clear_restore_mark(store &storage.Store) ! {31 store.remove_state(restoring_key)!32}3334pub struct RestoreReport {35pub:36 written int37 removed int38}3940414243444546pub fn restore_files(store &storage.Store, root string, ig Ignore, target []FileEntry, remove []string) !RestoreReport {47 48 in_repo := root == store.root49 report := write_tree(store, root, ig, target, remove) or {50 51 if in_repo {52 return error('${err.msg()} — the restore stopped part way; fix that, then run koh revert to finish it')53 }54 return err55 }56 if in_repo {57 clear_restore_mark(store)!58 }59 return report60}6162636465fn write_tree(store &storage.Store, root string, ig Ignore, target []FileEntry, remove []string) !RestoreReport {66 check_containment(root, target, remove)!67 mut wanted := map[string]bool{}68 for entry in target {69 wanted[entry.path] = true70 }71 mut removed := 072 for path in remove {73 full := os.join_path(root, path)74 if !os.is_file(full) || ig.matches(path, false) || path in wanted {75 continue76 }77 os.rm(full)!78 removed++79 prune_empty_dirs(root, os.dir(full))80 }81 82 83 84 for entry in target {85 full := os.join_path(root, entry.path)86 if os.is_dir(full) {87 os.rmdir(full) or {88 return error('${entry.path} is a folder of ignored files where a saved file belongs — move it')89 }90 }91 }92 mut written := 093 cache := store.load_cache()!94 for entry in target {95 if ig.matches(entry.path, false) && os.exists(os.join_path(root, entry.path)) {96 continue97 }98 99 100 if already_there(root, entry, cache) {101 continue102 }103 data := store.get_object(entry.hash)!104 write_file(os.join_path(root, entry.path), data, entry.executable)!105 written++106 }107 return RestoreReport{108 written: written109 removed: removed110 }111}112113114fn check_containment(root string, target []FileEntry, remove []string) ! {115 base := os.abs_path(root).trim_right(os.path_separator) + os.path_separator116 mut paths := target.map(it.path)117 paths << remove118 for path in paths {119 full := os.abs_path(os.join_path(root, path))120 if !legal_path(path) || !full.starts_with(base) {121 return error('unsafe path ${escape_path(path)} is outside the repo')122 }123 }124}125126127128129fn already_there(root string, entry FileEntry, cache map[string]storage.CacheRow) bool {130 full := os.join_path(root, entry.path)131 st := lstat(full) or { return false }132 if st.kind != .file || st.size != entry.size || is_executable(st.mode) != entry.executable {133 return false134 }135 if row := cache[entry.path] {136 if row.size == st.size && row.mtime_ns == st.mtime_ns && row.ctime_ns == st.ctime_ns {137 return row.hash == entry.hash138 }139 }140 data := os.read_bytes(full) or { return false }141 return storage.hash_hex(data) == entry.hash142}143144145146147148pub fn check_links(root string, files []FileEntry) ! {149 for file in files {150 mut folder := ''151 for part in file.path.split('/')#[..-1] {152 folder = if folder == '' { part } else { '${folder}/${part}' }153 st := lstat(os.join_path(root, folder)) or { break }154 if st.kind == .symlink {155 return error('${folder} is a symbolic link where a saved folder belongs — koh will not write through it; move it, then run the command again')156 }157 }158 }159}160161fn write_file(path string, data []u8, executable bool) ! {162 os.mkdir_all(os.dir(path), os.MkdirParams{})!163 tmp := '${path}.koh-tmp-${os.getpid()}'164 os.write_file_array(tmp, data)!165 mode := if executable { 0o755 } else { 0o644 }166 os.chmod(tmp, mode)!167 os.rename(tmp, path)!168}169170171fn prune_empty_dirs(root string, dir string) {172 mut current := dir173 inside := root.trim_right(os.path_separator) + os.path_separator174 for current.starts_with(inside) {175 if (os.ls(current) or { return }).len > 0 {176 return177 }178 os.rmdir(current) or { return }179 current = os.dir(current)180 }181}