1use std::path::{Path, PathBuf};
10
11use super::loader::LoadError;
12use crate::validator::{BoundedZipRead, ValidatorLimits, read_zip_entry_bounded};
13
14pub enum EntitySource {
16 Directory { root: PathBuf },
19 ZipArchive(PathBuf),
23}
24
25#[derive(Debug, Clone)]
27pub struct SourceEntry {
28 pub relative_path: String,
32 pub source_path: PathBuf,
37 pub content: String,
39}
40
41#[derive(Debug)]
45pub struct SourceReadError {
46 pub source_path: PathBuf,
47 pub error: std::io::Error,
48}
49
50impl EntitySource {
51 pub fn read_all(&self) -> Result<(Vec<SourceEntry>, Vec<SourceReadError>), LoadError> {
58 match self {
59 EntitySource::Directory { root } => read_directory(root),
60 EntitySource::ZipArchive(archive) => read_zip_archive(archive),
61 }
62 }
63}
64
65fn read_directory(root: &Path) -> Result<(Vec<SourceEntry>, Vec<SourceReadError>), LoadError> {
66 if !root.exists() {
67 return Err(LoadError::DirNotFound(root.display().to_string()));
68 }
69
70 let mut files = Vec::new();
71 find_markdown_files(root, &mut files)?;
72 files.sort();
73
74 let mut entries = Vec::with_capacity(files.len());
75 let mut errors = Vec::new();
76
77 for file in &files {
78 match std::fs::read_to_string(file) {
79 Ok(content) => {
80 let relative_path = file
81 .strip_prefix(root)
82 .unwrap_or(file)
83 .to_string_lossy()
84 .to_string();
85 entries.push(SourceEntry {
86 relative_path,
87 source_path: file.clone(),
88 content,
89 });
90 }
91 Err(error) => errors.push(SourceReadError {
92 source_path: file.clone(),
93 error,
94 }),
95 }
96 }
97
98 Ok((entries, errors))
99}
100
101fn read_zip_archive(
102 archive_path: &Path,
103) -> Result<(Vec<SourceEntry>, Vec<SourceReadError>), LoadError> {
104 if !archive_path.is_file() {
105 return Err(LoadError::ArchiveNotFound(
106 archive_path.display().to_string(),
107 ));
108 }
109
110 let file = std::fs::File::open(archive_path)?;
111 let mut archive = zip::ZipArchive::new(file)?;
112
113 let limits = ValidatorLimits::DEFAULT;
114 if archive.len() as u32 > limits.max_file_count {
115 return Err(LoadError::InvalidArchive(format!(
116 "archive contains {} entries, exceeding the {}-entry cap",
117 archive.len(),
118 limits.max_file_count
119 )));
120 }
121
122 let mut entries = Vec::new();
123 let mut errors = Vec::new();
124 let mut uncompressed_total: u64 = 0;
125
126 for i in 0..archive.len() {
127 let mut entry = archive.by_index(i)?;
128 let raw_name = entry.name().to_string();
129
130 if entry.is_symlink() {
134 return Err(LoadError::InvalidArchive(format!(
135 "entry '{raw_name}': symlinks are not allowed in sealed mem archives"
136 )));
137 }
138
139 let safe_path = match entry.enclosed_name() {
144 Some(p) => p,
145 None => {
146 return Err(LoadError::InvalidArchive(format!(
147 "entry '{raw_name}': path escapes archive root \
148 (absolute, '..'-components, or otherwise unsafe)"
149 )));
150 }
151 };
152
153 if entry.is_dir() {
154 continue;
155 }
156
157 let relative_path = safe_path.to_string_lossy().to_string();
158 if !relative_path.ends_with(".md") {
159 continue;
164 }
165
166 let bytes = match read_zip_entry_bounded(&mut entry, limits.max_uncompressed_entry)? {
170 BoundedZipRead::Within(bytes) => bytes,
171 BoundedZipRead::ExceedsCap => {
172 return Err(LoadError::InvalidArchive(format!(
173 "entry '{relative_path}' exceeds the {}-byte uncompressed cap",
174 limits.max_uncompressed_entry
175 )));
176 }
177 };
178 uncompressed_total = uncompressed_total.saturating_add(bytes.len() as u64);
179 if uncompressed_total > limits.max_uncompressed_archive {
180 return Err(LoadError::InvalidArchive(format!(
181 "archive exceeds the {}-byte total uncompressed cap",
182 limits.max_uncompressed_archive
183 )));
184 }
185 match String::from_utf8(bytes) {
186 Ok(content) => entries.push(SourceEntry {
187 relative_path: relative_path.clone(),
188 source_path: PathBuf::from(&relative_path),
189 content,
190 }),
191 Err(error) => errors.push(SourceReadError {
192 source_path: PathBuf::from(&relative_path),
193 error: std::io::Error::new(std::io::ErrorKind::InvalidData, error),
194 }),
195 }
196 }
197
198 entries.sort_by(|a, b| a.relative_path.cmp(&b.relative_path));
202
203 Ok((entries, errors))
204}
205
206fn find_markdown_files(dir: &Path, files: &mut Vec<PathBuf>) -> Result<(), LoadError> {
216 let entries = std::fs::read_dir(dir)?;
217
218 for entry in entries {
219 let entry = entry?;
220 let path = entry.path();
221 let file_name = entry.file_name();
222 let name = file_name.to_string_lossy();
223
224 if path.is_dir() {
225 if name.as_ref() == ".git" || name.as_ref() == crate::mem::MEM_META_DIR {
226 continue;
227 }
228 find_markdown_files(&path, files)?;
229 } else if name.ends_with(".md") {
230 files.push(path);
231 }
232 }
233
234 Ok(())
235}
236
237#[cfg(test)]
238mod tests {
239 use super::*;
240 use std::fs;
241 use tempfile::TempDir;
242
243 #[test]
244 fn directory_reads_markdown_in_sorted_order() {
245 let dir = TempDir::new().unwrap();
246 fs::write(dir.path().join("b.md"), "b").unwrap();
247 fs::write(dir.path().join("a.md"), "a").unwrap();
248
249 let (entries, errors) = EntitySource::Directory {
250 root: dir.path().to_path_buf(),
251 }
252 .read_all()
253 .unwrap();
254 assert!(errors.is_empty());
255 let paths: Vec<_> = entries.iter().map(|e| e.relative_path.as_str()).collect();
256 assert_eq!(paths, vec!["a.md", "b.md"]);
257 }
258
259 #[test]
260 fn directory_skips_engine_internal_dirs_and_non_md() {
261 let dir = TempDir::new().unwrap();
264 fs::write(dir.path().join("keep.md"), "k").unwrap();
265 fs::write(dir.path().join("ignore.txt"), "i").unwrap();
266 fs::create_dir_all(dir.path().join(".git")).unwrap();
267 fs::write(dir.path().join(".git/secret.md"), "s").unwrap();
268 fs::create_dir_all(dir.path().join(".memstead")).unwrap();
269 fs::write(dir.path().join(".memstead/note.md"), "n").unwrap();
270 fs::create_dir_all(dir.path().join(".obsidian")).unwrap();
271 fs::write(dir.path().join(".obsidian/vis.md"), "v").unwrap();
272
273 let (entries, _) = EntitySource::Directory {
274 root: dir.path().to_path_buf(),
275 }
276 .read_all()
277 .unwrap();
278 let paths: Vec<_> = entries.iter().map(|e| e.relative_path.as_str()).collect();
279 assert!(paths.contains(&"keep.md"), "keep.md must load: {paths:?}");
280 assert!(
281 paths.iter().any(|p| p.ends_with("vis.md")),
282 ".obsidian/vis.md must load: {paths:?}"
283 );
284 assert!(
285 !paths.iter().any(|p| p.contains(".git")),
286 ".git/* must be skipped: {paths:?}"
287 );
288 assert!(
289 !paths.iter().any(|p| p.contains(".memstead")),
290 ".memstead/* must be skipped: {paths:?}"
291 );
292 }
293
294 #[test]
295 fn directory_missing_root_returns_error() {
296 let err = EntitySource::Directory {
297 root: PathBuf::from("/nonexistent/path/xyz"),
298 }
299 .read_all()
300 .unwrap_err();
301 assert!(matches!(err, LoadError::DirNotFound(_)));
302 }
303
304 use std::io::Write;
307 use zip::CompressionMethod;
308 use zip::write::SimpleFileOptions;
309
310 fn write_zip(path: &Path, entries: &[(&str, &str)]) {
313 let file = fs::File::create(path).unwrap();
314 let mut zip = zip::ZipWriter::new(file);
315 let opts = SimpleFileOptions::default().compression_method(CompressionMethod::Stored);
316 for (name, content) in entries {
317 zip.start_file(*name, opts).unwrap();
318 zip.write_all(content.as_bytes()).unwrap();
319 }
320 zip.finish().unwrap();
321 }
322
323 #[test]
324 fn zip_archive_reads_markdown_in_sorted_order() {
325 let dir = TempDir::new().unwrap();
326 let archive = dir.path().join("pkg.mem");
327 write_zip(
328 &archive,
329 &[
330 ("b.md", "b"),
331 ("a.md", "a"),
332 ("meta.json", "{\"name\":\"pkg\"}"),
333 (".memstead/config.json", "{}"),
334 ("readme.txt", "ignored"),
335 ("nested/c.md", "c"),
336 ],
337 );
338
339 let (entries, errors) = EntitySource::ZipArchive(archive).read_all().unwrap();
340 assert!(errors.is_empty());
341 let paths: Vec<_> = entries.iter().map(|e| e.relative_path.as_str()).collect();
342 assert_eq!(paths, vec!["a.md", "b.md", "nested/c.md"]);
343 let contents: Vec<_> = entries.iter().map(|e| e.content.as_str()).collect();
344 assert_eq!(contents, vec!["a", "b", "c"]);
345 }
346
347 #[test]
348 fn zip_archive_missing_file_returns_error() {
349 let err = EntitySource::ZipArchive(PathBuf::from("/nonexistent/pkg.mem"))
350 .read_all()
351 .unwrap_err();
352 assert!(matches!(err, LoadError::ArchiveNotFound(_)));
353 }
354
355 #[test]
356 fn zip_archive_corrupt_file_returns_zip_error() {
357 let dir = TempDir::new().unwrap();
358 let archive = dir.path().join("bad.mem");
359 fs::write(&archive, b"not a zip file at all, just bytes").unwrap();
360
361 let err = EntitySource::ZipArchive(archive).read_all().unwrap_err();
362 assert!(
363 matches!(err, LoadError::Zip(_)),
364 "corrupt archive should surface as LoadError::Zip, got {err:?}"
365 );
366 }
367
368 #[test]
369 fn zip_archive_rejects_parent_dir_escape() {
370 let dir = TempDir::new().unwrap();
371 let archive = dir.path().join("evil.mem");
372 write_zip(&archive, &[("../escape.md", "bad")]);
373
374 let err = EntitySource::ZipArchive(archive).read_all().unwrap_err();
375 let msg = format!("{err}");
376 assert!(matches!(err, LoadError::InvalidArchive(_)));
377 assert!(
378 msg.contains("escape") || msg.contains("..") || msg.contains("unsafe"),
379 "zip-slip error should explain the rejection: {msg}"
380 );
381 }
382
383 #[test]
384 fn zip_archive_rejects_nested_parent_dir_escape() {
385 let dir = TempDir::new().unwrap();
390 let archive = dir.path().join("evil.mem");
391 write_zip(&archive, &[("subdir/../../outside.md", "bad")]);
392
393 let err = EntitySource::ZipArchive(archive).read_all().unwrap_err();
394 assert!(matches!(err, LoadError::InvalidArchive(_)));
395 }
396
397 #[test]
398 fn zip_archive_rejects_oversized_entry() {
399 let dir = TempDir::new().unwrap();
403 let archive = dir.path().join("bomb.mem");
404 let big = "a".repeat((ValidatorLimits::DEFAULT.max_uncompressed_entry + 1) as usize);
405 let file = fs::File::create(&archive).unwrap();
406 let mut zip = zip::ZipWriter::new(file);
407 let opts = SimpleFileOptions::default().compression_method(CompressionMethod::Deflated);
408 zip.start_file("bomb.md", opts).unwrap();
409 zip.write_all(big.as_bytes()).unwrap();
410 zip.finish().unwrap();
411
412 let err = EntitySource::ZipArchive(archive).read_all().unwrap_err();
413 let msg = format!("{err}");
414 assert!(matches!(err, LoadError::InvalidArchive(_)), "got {err:?}");
415 assert!(msg.contains("cap"), "error should name the cap: {msg}");
416 }
417
418 }