返回 CodeWhale
persistence.rs
根目录 / crates / tui / src / tui / pet_watch / persistence.rs
1 //! Session-owned pet sidecar, using the existing confined artifact I/O. A
2 //! writer lock plus content revision rejects concurrent or external edits.
3 use std::fs::File;
4 use std::io::{self, Read};
5 use std::path::{Path, PathBuf};
6
7 use sha2::{Digest, Sha256};
8
9 #[cfg(test)]
10 use crate::artifacts::open_session_relative;
11 use crate::artifacts::write_session_relative_immutable;
12 use crate::fleet::files::{WorkspaceFile, same_file};
13
14 const MAX_BYTES: usize = 8 * 1024 * 1024;
15 pub(super) const MAX_EXPORT_BYTES: usize = 64 * 1024 * 1024;
16 #[cfg(test)]
17 const HABITAT: &str = "artifacts/pet/habitat.json";
18
19 pub struct Store {
20 data: WorkspaceFile,
21 lock: WorkspaceFile,
22 original_lock: File,
23 expected: Option<[u8; 32]>,
24 }
25
26 impl Store {
27 /// Shared presentation-owner state. Legacy session habitats stay in place.
28 pub fn at(root: &Path) -> io::Result<Self> {
29 let data = WorkspaceFile::open(root, Path::new("habitat.json"), true)?;
30 let lock = WorkspaceFile::open(root, Path::new("habitat.lock"), true)?;
31 let original_lock = lock.open_update(true, false)?;
32 Ok(Self {
33 data,
34 lock,
35 original_lock,
36 expected: None,
37 })
38 }
39 #[cfg(test)]
40 pub fn open(session: &str) -> io::Result<Self> {
41 let data = open_session_relative(session, Path::new(HABITAT), true)?;
42 let lock = open_session_relative(session, Path::new("artifacts/pet/habitat.lock"), true)?;
43 let original_lock = lock.open_update(true, false)?;
44 Ok(Self {
45 data,
46 lock,
47 original_lock,
48 expected: None,
49 })
50 }
51
52 fn with_lock<T>(&self, action: impl FnOnce() -> io::Result<T>) -> io::Result<T> {
53 let file = self.lock.open_update(false, false)?;
54 if !same_file(&file, &self.original_lock)? {
55 return Err(io::Error::other("Pet habitat lock was replaced"));
56 }
57 let mut lock = fd_lock::RwLock::new(file);
58 // Never wait behind another process in the world worker.
59 let _guard = lock.try_write()?;
60 action()
61 }
62
63 fn read(&self) -> io::Result<Option<Vec<u8>>> {
64 let file = match self.data.open_file() {
65 Ok(file) => file,
66 Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None),
67 Err(error) => return Err(error),
68 };
69 let mut bytes = Vec::new();
70 file.take(MAX_BYTES as u64 + 1).read_to_end(&mut bytes)?;
71 if bytes.len() > MAX_BYTES {
72 return Err(io::Error::other("Pet habitat exceeds 8 MiB"));
73 }
74 Ok(Some(bytes))
75 }
76
77 pub fn load(&mut self) -> io::Result<Option<String>> {
78 let bytes = self.with_lock(|| self.read())?;
79 let text = bytes
80 .as_ref()
81 .map(|b| {
82 String::from_utf8(b.clone())
83 .map_err(|_| io::Error::other("Pet habitat is not UTF-8"))
84 })
85 .transpose()?;
86 self.expected = bytes.map(|b| Sha256::digest(b).into());
87 Ok(text)
88 }
89
90 #[cfg(test)]
91 pub fn save(&mut self, text: &str) -> io::Result<()> {
92 self.save_archived(text, None)
93 }
94
95 pub fn save_archived(&mut self, text: &str, archive: Option<(&[u8], u64)>) -> io::Result<()> {
96 if text.len() > MAX_BYTES {
97 return Err(io::Error::other("Pet habitat exceeds 8 MiB"));
98 }
99 self.with_lock(|| {
100 let current = self.read()?.map(|b| <[u8; 32]>::from(Sha256::digest(b)));
101 if current != self.expected {
102 return Err(io::Error::other("Another writer changed the pet habitat"));
103 }
104 if let Some((bytes, tick)) = archive {
105 if bytes.len() > MAX_EXPORT_BYTES {
106 return Err(io::Error::other("Pet recording exceeds 64 MiB"));
107 }
108 let hash: String = Sha256::digest(bytes)
109 .iter()
110 .map(|b| format!("{b:02x}"))
111 .collect();
112 let archived = self
113 .data
114 .sibling(&format!("segment-{tick:012}-{hash}.json"))?;
115 if let Err(error) = archived.publish(bytes) {
116 if error.kind() != io::ErrorKind::AlreadyExists {
117 return Err(error);
118 }
119 let mut existing = Vec::new();
120 archived
121 .open_file()?
122 .take(MAX_EXPORT_BYTES as u64 + 1)
123 .read_to_end(&mut existing)?;
124 if existing != bytes {
125 return Err(io::Error::other("An archived recording was changed"));
126 }
127 }
128 }
129 self.data.replace(text.as_bytes())
130 })?;
131 self.expected = Some(Sha256::digest(text.as_bytes()).into());
132 Ok(())
133 }
134 }
135
136 pub fn export(session: &str, bytes: &[u8]) -> io::Result<PathBuf> {
137 if bytes.len() > MAX_EXPORT_BYTES {
138 return Err(io::Error::other("Pet recording exceeds 64 MiB"));
139 }
140 let relative = PathBuf::from(format!(
141 "artifacts/pet/replay-{}.json",
142 uuid::Uuid::new_v4()
143 ));
144 write_session_relative_immutable(session, &relative, bytes)
145 }
146
147 /// One command owns the world until export completes. Keep the large buffer in
148 /// the host, outside QuickJS's 64 MiB heap, and retain the exact checkpoint.
149 pub(super) fn export_recording(
150 ctx: &rquickjs::Ctx<'_>,
151 completed: bool,
152 ) -> rquickjs::Result<Vec<u8>> {
153 let mut bytes = Vec::new();
154 let mut index = 0usize;
155 loop {
156 let chunk: Option<String> = ctx.eval(format!("pet.recordingChunk({index},{completed})"))?;
157 let Some(chunk) = chunk else { return Ok(bytes) };
158 if bytes.len() + chunk.len() > MAX_EXPORT_BYTES {
159 return Err(rquickjs::Error::Unknown);
160 }
161 bytes.extend_from_slice(chunk.as_bytes());
162 index += 1;
163 }
164 }
165
166 #[cfg(test)]
167 mod tests {
168 use super::*;
169
170 fn store(root: &Path) -> Store {
171 let data = WorkspaceFile::open(root, Path::new("habitat.json"), true).unwrap();
172 let lock = WorkspaceFile::open(root, Path::new("habitat.lock"), true).unwrap();
173 let original_lock = lock.open_update(true, false).unwrap();
174 Store {
175 data,
176 lock,
177 original_lock,
178 expected: None,
179 }
180 }
181
182 #[test]
183 fn stale_writer_cannot_replace_a_newer_recording_or_external_edit() {
184 let root = tempfile::tempdir().unwrap();
185 let mut first = store(root.path());
186 let mut second = store(root.path());
187 assert!(first.load().unwrap().is_none());
188 assert!(second.load().unwrap().is_none());
189 first.save("first recording").unwrap();
190 assert!(second.save("stale recording").is_err());
191 assert_eq!(
192 std::fs::read_to_string(root.path().join("habitat.json")).unwrap(),
193 "first recording"
194 );
195 std::fs::write(root.path().join("habitat.json"), "external edit").unwrap();
196 assert!(first.save("lost edit").is_err());
197 assert_eq!(first.load().unwrap().as_deref(), Some("external edit"));
198 }
199
200 #[test]
201 fn invalid_and_oversized_habitats_remain_intact() {
202 let root = tempfile::tempdir().unwrap();
203 let mut files = store(root.path());
204 let path = root.path().join("habitat.json");
205 std::fs::write(&path, [0xff]).unwrap();
206 assert!(files.load().is_err());
207 assert!(files.save("replacement").is_err());
208 assert_eq!(std::fs::read(&path).unwrap(), [0xff]);
209 std::fs::write(&path, vec![b' '; MAX_BYTES + 1]).unwrap();
210 assert!(files.load().is_err());
211 assert!(files.save("replacement").is_err());
212 assert_eq!(
213 std::fs::metadata(&path).unwrap().len(),
214 MAX_BYTES as u64 + 1
215 );
216 }
217
218 #[test]
219 fn segments_are_immutable_and_must_publish_before_the_habitat_advances() {
220 let root = tempfile::tempdir().unwrap();
221 let mut files = store(root.path());
222 files.load().unwrap();
223 files.save("before").unwrap();
224 files
225 .save_archived("after", Some((b"complete history", 123)))
226 .unwrap();
227 let name = "segment-000000000123-42fcd454bac01f468e693701bf88157cd7a540556f85d26be0009392e43ecbd4.json";
228 let archive = root.path().join(name);
229 assert_eq!(std::fs::read(&archive).unwrap(), b"complete history");
230 std::fs::write(&archive, "damaged archive").unwrap();
231 assert!(
232 files
233 .save_archived("lost", Some((b"complete history", 123)))
234 .is_err()
235 );
236 assert_eq!(
237 std::fs::read_to_string(root.path().join("habitat.json")).unwrap(),
238 "after"
239 );
240 assert_eq!(
241 std::fs::read_to_string(&archive).unwrap(),
242 "damaged archive"
243 );
244 std::fs::write(root.path().join("habitat.json"), "another writer").unwrap();
245 assert!(
246 files
247 .save_archived("lost", Some((b"new history", 124)))
248 .is_err()
249 );
250 assert_eq!(std::fs::read_dir(root.path()).unwrap().count(), 3);
251 }
252
253 #[cfg(unix)]
254 #[test]
255 fn private_atomic_files_reject_links_and_replaced_locks() {
256 use std::os::unix::fs::{PermissionsExt, symlink};
257 let root = tempfile::tempdir().unwrap();
258 let mut files = store(root.path());
259 files.save("private").unwrap();
260 let path = root.path().join("habitat.json");
261 assert_eq!(
262 std::fs::metadata(&path).unwrap().permissions().mode() & 0o777,
263 0o600
264 );
265 let outside = root.path().join("outside");
266 std::fs::rename(&path, &outside).unwrap();
267 symlink(&outside, &path).unwrap();
268 assert!(files.load().is_err());
269 assert!(files.save("replacement").is_err());
270 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "private");
271 std::fs::remove_file(&path).unwrap();
272 std::fs::hard_link(&outside, &path).unwrap();
273 assert!(files.load().is_err());
274 std::fs::remove_file(root.path().join("habitat.lock")).unwrap();
275 let _new_owner = store(root.path());
276 assert!(files.save("split lock").is_err());
277 }
278 }
279
279 lines RUST