| 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 |