1use std::{
41 fs::File,
42 io::{BufReader, Read, Seek, Write},
43 path::{Component, Path, PathBuf},
44};
45
46pub type Result<M> = std::result::Result<M, Error>;
48
49#[derive(thiserror::Error, Debug)]
51pub enum Error {
52 #[error("{source}")]
53 Io { source: std::io::Error },
54 #[error("{source}")]
55 Zip { source: zip::result::ZipError },
56 #[error("file not found")]
57 FileNotFound,
58 #[error("filename not utf8")]
59 FilenameNotUtf8,
60 #[error("unexpected zip file content")]
61 UnexpectedZipContent,
62 #[error("unexpected zip file name")]
63 UnexpectedZipName,
64 #[error("directory does not exist: {0}")]
65 NotDirectory(String),
66}
67
68impl From<zip::result::ZipError> for Error {
69 fn from(source: zip::result::ZipError) -> Self {
70 match source {
71 zip::result::ZipError::FileNotFound => Error::FileNotFound,
72 source => Error::Zip { source },
73 }
74 }
75}
76
77impl From<std::io::Error> for Error {
78 fn from(source: std::io::Error) -> Self {
79 match source.kind() {
80 std::io::ErrorKind::NotFound => Error::FileNotFound,
81 _ => Error::Io { source },
82 }
83 }
84}
85
86pub struct ZipDirArchive<R: Read + Seek> {
92 path: PathBuf,
94 zip_archive: Option<zip::ZipArchive<R>>,
96}
97
98impl<R: Read + Seek> std::fmt::Debug for ZipDirArchive<R> {
99 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
100 f.debug_struct("KalmanEstimatesInfo")
101 .field("path", &self.path)
102 .finish_non_exhaustive()
103 }
104}
105
106impl ZipDirArchive<BufReader<File>> {
107 pub fn auto_from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
114 if path.as_ref().exists() {
115 if path.as_ref().is_dir() {
116 Self::from_dir(path.as_ref().to_path_buf())
117 } else {
118 let reader = BufReader::new(File::open(&path)?);
119 Self::from_zip(
120 reader,
121 path.as_ref().as_os_str().to_str().unwrap().to_string(),
122 )
123 }
124 } else {
125 Err(Error::FileNotFound)
126 }
127 }
128
129 pub fn from_dir(path: PathBuf) -> Result<Self> {
131 Ok(ZipDirArchive {
132 path,
133 zip_archive: None,
134 })
135 }
136}
137
138impl<R: Read + Seek> ZipDirArchive<R> {
139 pub fn from_zip(reader: R, display_name: String) -> Result<Self> {
141 let zip_archive = Some(zip::ZipArchive::new(reader)?);
142 Ok(ZipDirArchive {
143 path: display_name.into(),
144 zip_archive,
145 })
146 }
147 pub fn path_starter(&mut self) -> PathLike<'_, R> {
148 let parent = self;
149 PathLike {
150 parent,
151 relname: "".into(),
152 }
153 }
154 pub fn display(&self) -> std::path::Display<'_> {
155 self.path.display()
156 }
157 pub fn path(&self) -> &std::path::Path {
158 &self.path
159 }
160
161 fn rel(&self, relname: &Path) -> PathBuf {
163 match self.zip_archive {
164 Some(_) => self.path.slash_join(relname),
167 None => self.path.join(relname),
168 }
169 }
170 pub fn exists(&mut self, relname: &Path) -> bool {
172 match &mut self.zip_archive {
173 Some(zip_archive) => {
174 let relname_str = relname.as_os_str().to_str().unwrap();
175 zip_archive.by_name(relname_str).is_ok()
176 }
177 None => self.rel(relname).exists(),
178 }
179 }
180
181 pub fn open<P: AsRef<Path>>(&mut self, relname: P) -> Result<FileReader<'_>> {
183 let dirpath = self.rel(relname.as_ref());
184 match &mut self.zip_archive {
185 Some(zip_archive) => {
186 let relname_str = relname.as_ref().as_os_str().to_str().unwrap();
187 let zipfile = zip_archive.by_name(relname_str)?;
188 Ok(FileReader::from_zip(zipfile)?)
189 }
190 None => Ok(FileReader::open_file(dirpath)?),
191 }
192 }
193
194 pub fn is_file<P: AsRef<Path>>(&mut self, relname: P) -> bool {
195 let dirpath = self.rel(relname.as_ref());
196 match &mut self.zip_archive {
197 Some(zip_archive) => {
198 let relname_str = relname.as_ref().as_os_str().to_str().unwrap();
199 zip_archive.by_name(relname_str).is_ok()
200 }
201 None => dirpath.is_file(),
202 }
203 }
204
205 pub fn list_paths<P: AsRef<Path> + std::fmt::Debug>(
211 &self,
212 relname: Option<P>,
213 ) -> Result<Vec<PathBuf>> {
214 let dirpath = match &relname {
216 Some(rn) => self.rel(rn.as_ref()),
217 None => self.path.clone(),
218 };
219 let mut result = vec![];
220 let mut unique_single_components = std::collections::BTreeSet::new();
221
222 match &self.zip_archive {
223 Some(zip_archive) => {
224 let mut suffixes: Vec<PathBuf> = vec![];
226
227 let mut found_any = false;
228 for deep_path in zip_archive.file_names().map(PathBuf::from) {
229 if let Some(prefix) = &relname {
230 let (has_match, suffix) = remove_shared_prefix(&deep_path, prefix);
231 if has_match {
232 found_any = true;
233 match suffix {
234 None => {}
235 Some(trailing) => {
236 suffixes.push(trailing);
237 }
238 }
239 }
240 } else {
241 suffixes.push(deep_path);
243 }
244
245 for p1 in suffixes.iter() {
247 match p1.components().next() {
248 Some(Component::Normal(next)) => {
249 unique_single_components.insert(PathBuf::from(next));
250 }
251 Some(_) | None => {
252 return Err(Error::UnexpectedZipName);
253 }
254 }
255 }
256 }
257 result = unique_single_components.into_iter().collect();
258 if let Some(relname) = relname
259 && result.is_empty()
260 && !found_any
261 {
262 return Err(not_dir_error(relname));
263 }
264 }
265 None => {
266 let dir_result = std::fs::read_dir(&dirpath);
267 let dir_result = match dir_result {
268 Ok(d) => d,
269 Err(e) => {
270 match e.kind() {
271 std::io::ErrorKind::NotFound => {
272 return Err(not_dir_error(dirpath));
274 }
275 _ => {
276 return Err(e.into());
277 }
278 }
279 }
280 };
281 for entry in dir_result {
283 let entry = entry?;
284 if let (has_match, Some(suffix)) = remove_shared_prefix(entry.path(), &dirpath)
286 {
287 debug_assert!(has_match);
288 result.push(suffix);
289 }
290 }
291 }
292 }
293 Ok(result)
294 }
295
296 pub fn open_raw_or_gz(&mut self, src_fname: &str) -> Result<MaybeGzReader<'_>> {
300 let gz_fname = format!("{src_fname}.gz");
301 let gz_exists = self.path_starter().join(&gz_fname).exists();
302
303 if gz_exists {
304 let gz_fd = self.path_starter().join(gz_fname).open()?;
305 let decoder = libflate::gzip::Decoder::new(gz_fd)?;
306 Ok(MaybeGzReader::Gz(decoder))
307 } else {
308 let fd = self.path_starter().join(src_fname).open()?;
309 Ok(MaybeGzReader::Raw(fd))
310 }
311 }
312}
313
314pub enum MaybeGzReader<'a> {
316 Raw(FileReader<'a>),
317 Gz(libflate::gzip::Decoder<FileReader<'a>>),
318}
319
320impl<'a> Read for MaybeGzReader<'a> {
321 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
322 match self {
323 Self::Raw(fd) => fd.read(buf),
324 Self::Gz(gz_fd) => gz_fd.read(buf),
325 }
326 }
327}
328
329#[derive(Debug)]
331pub struct FileReader<'a> {
332 inner: FileReaderInner<'a>,
333}
334
335enum FileReaderInner<'a> {
336 File(BufReader<File>),
337 ZipFile(Box<BufReader<zip::read::ZipFile<'a>>>),
338}
339
340impl<'a> std::fmt::Debug for FileReaderInner<'a> {
341 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
342 match self {
343 FileReaderInner::File(_) => write!(f, "FileReaderInner::File"),
344 FileReaderInner::ZipFile(_) => write!(f, "FileReaderInner::ZipFile"),
345 }
346 }
347}
348
349impl<'a> FileReader<'a> {
350 fn from_inner(inner: FileReaderInner<'a>) -> Result<FileReader<'a>> {
351 Ok(FileReader { inner })
352 }
353 fn open_file<P: AsRef<std::path::Path>>(path: P) -> Result<FileReader<'a>> {
354 let f = File::open(path)?;
355 Self::from_inner(FileReaderInner::File(BufReader::new(f)))
356 }
357 fn from_zip(zipfile: zip::read::ZipFile<'a>) -> Result<FileReader<'a>> {
358 Self::from_inner(FileReaderInner::ZipFile(Box::new(BufReader::new(zipfile))))
359 }
360}
361
362impl<'a> Read for FileReader<'a> {
363 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
364 let n_bytes = match &mut self.inner {
365 FileReaderInner::File(f) => f.read(buf)?,
366 FileReaderInner::ZipFile(zf) => zf.read(buf)?,
367 };
368 Ok(n_bytes)
369 }
370}
371
372fn remove_shared_prefix<P1: AsRef<Path>, P2: AsRef<Path>>(
382 src: P1,
383 prefix: P2,
384) -> (bool, Option<PathBuf>) {
385 let mut src_components = src.as_ref().components();
387 let prefix_components = prefix.as_ref().components();
388
389 for c in prefix_components {
391 let sc = src_components.next();
392 match sc {
393 Some(scc) => {
394 if c != scc {
395 return (false, None);
397 }
398 }
399 None => return (false, None),
401 }
402 }
403
404 let result: PathBuf = src_components.as_path().into();
405
406 let suffix = if result.as_os_str() == std::ffi::OsStr::new("") {
407 None
408 } else {
409 Some(result)
410 };
411 (true, suffix)
412}
413
414#[test]
416fn test_implements_send() {
417 fn implements_send<F: Send>() {}
418 implements_send::<ZipDirArchive<File>>();
419 implements_send::<PathLike<File>>();
420}
421
422#[derive(Debug)]
428pub struct PathLike<'a, R: Read + Seek> {
429 parent: &'a mut ZipDirArchive<R>,
430 relname: PathBuf,
431}
432
433impl<'a, R: Read + Seek> PathLike<'a, R> {
434 pub fn push<P: AsRef<std::path::Path>>(&mut self, p: P) {
435 match self.parent.zip_archive {
436 Some(_) => self.relname.slash_push(p),
437 None => self.relname.push(p),
438 }
439 }
440 pub fn join<P: AsRef<std::path::Path>>(mut self, p: P) -> Self {
441 self.push(p);
442 self
443 }
444 pub fn path(&mut self) -> &std::path::Path {
445 &self.relname
446 }
447 pub fn replace(&mut self, relname: PathBuf) -> PathBuf {
448 std::mem::replace(&mut self.relname, relname)
449 }
450 pub fn extension(&mut self) -> Option<&std::ffi::OsStr> {
451 self.relname.extension()
452 }
453 pub fn set_extension(&mut self, e: &str) -> bool {
454 self.relname.set_extension(e)
455 }
456 pub fn display(&self) -> std::path::Display<'_> {
457 Path::display(&self.relname)
458 }
459 pub fn exists(&mut self) -> bool {
460 self.parent.exists(&self.relname)
461 }
462 pub fn open(self) -> Result<FileReader<'a>> {
463 self.parent.open(&self.relname)
464 }
465 pub fn is_file(&mut self) -> bool {
466 self.parent.is_file(&self.relname)
467 }
468 pub fn list_paths(&self) -> Result<Vec<PathBuf>> {
470 self.parent.list_paths(Some(&self.relname))
471 }
472}
473
474trait SlashJoin {
475 fn slash_join<P: AsRef<Path>>(&self, p: P) -> PathBuf;
476 fn slash_push<P: AsRef<Path>>(&mut self, path: P);
477}
478
479impl SlashJoin for PathBuf {
480 fn slash_join<P: AsRef<Path>>(&self, p: P) -> PathBuf {
481 let mut buf = self.to_path_buf();
482 buf.slash_push(p);
483 buf
484 }
485 fn slash_push<P: AsRef<Path>>(&mut self, path: P) {
486 let self_str = format!("{}", self.display());
488 let new_str = if self_str.is_empty() {
489 std::path::PathBuf::from(path.as_ref())
490 } else {
491 let my_str = format!("{}/{}", self_str, path.as_ref().display());
492 let remove_double = my_str.replace("//", "/");
493 std::path::PathBuf::from(remove_double)
494 };
495 *self = new_str;
496 }
497}
498
499pub fn copy_to_zip<P1: AsRef<Path>, P2: AsRef<Path>>(src: P1, dest: P2) -> Result<()> {
506 let mut src_archive = ZipDirArchive::auto_from_path(src.as_ref()).unwrap();
507 let mut zipfile = File::create(dest)?;
508 copy_archive_to_zipfile(&mut src_archive, &mut zipfile)
509}
510
511pub fn copy_archive_to_zipfile<R: Read + Seek>(
526 src: &mut ZipDirArchive<R>,
527 dest: &mut File,
528) -> Result<()> {
529 let mut zip_writer = zip::ZipWriter::new(dest);
530 copy_dir(src, None, &mut zip_writer)?;
531 zip_writer.finish()?;
532 Ok(())
533}
534
535fn copy_dir<R: Read + Seek>(
537 src: &mut ZipDirArchive<R>,
538 relname: Option<&Path>,
539 zip_writer: &mut zip::ZipWriter<&mut File>,
540) -> zip::result::ZipResult<()> {
541 let parent = match relname {
542 None => PathBuf::new(),
543 Some(parent) => PathBuf::from(parent),
544 };
545
546 let paths = src
548 .list_paths::<PathBuf>(relname.map(PathBuf::from))
549 .unwrap();
550
551 for entry in paths.iter() {
553 let full_entry = parent.join(entry);
554 let mut ep = src.path_starter();
556 ep.push(full_entry.as_os_str().to_str().unwrap());
557
558 if ep.is_file() {
560 let mut fd = ep.open().unwrap();
561 let mut buf = vec![];
562 fd.read_to_end(&mut buf).unwrap();
563
564 let mut options = zip::write::SimpleFileOptions::default();
565 if buf.len() as u64 >= 0xFFFFFFFF {
568 println!("setting large file to true");
569 options = options.large_file(true);
570 }
571
572 zip_writer
573 .start_file(full_entry.to_str().unwrap(), options)
574 .unwrap();
575 zip_writer.write_all(&buf).unwrap();
576 } else {
577 let subpath: PathBuf = match relname {
579 None => full_entry,
580 Some(_) => full_entry, };
583
584 copy_dir(src, Some(&subpath), zip_writer)?;
585 }
586 }
587
588 Ok(())
589}
590
591fn not_dir_error<P: AsRef<Path>>(relname: P) -> Error {
592 Error::NotDirectory(format!("{}", relname.as_ref().display()))
593}
594
595#[cfg(test)]
596mod tests {
597 use crate::*;
598
599 fn create_files(fnames: &[&str], basepath: &Path) -> std::io::Result<()> {
600 for fname in fnames {
601 let f = basepath.join(fname);
602 let mut fd = File::create(&f).unwrap();
603 fd.write_all(fname.as_bytes())?;
604 }
605 Ok(())
606 }
607
608 #[test]
609 fn test_remove_shared_prefix() {
610 let prefix = "root";
611 let a = "root/hello";
612 let (has_match, actual) = remove_shared_prefix(a, prefix);
613 assert!(has_match);
614 assert_eq!(actual.unwrap().as_os_str().to_str().unwrap(), "hello");
615
616 let prefix = "root/b";
617 let a = "root/b/hello";
618 let (has_match, actual) = remove_shared_prefix(a, prefix);
619 assert!(has_match);
620 assert_eq!(actual.unwrap().as_os_str().to_str().unwrap(), "hello");
621
622 let prefix = "root/a";
623 let a = "root/b/hello";
624 let (has_match, actual) = remove_shared_prefix(a, prefix);
625 assert!(!has_match);
626 assert_eq!(actual, None);
627
628 let prefix = "root/b";
629 let a = "root/b";
630 let (has_match, actual) = remove_shared_prefix(a, prefix);
631 assert!(has_match);
632 assert_eq!(actual, None);
633 }
634
635 #[test]
645 fn it_works() {
646 let tempdir = tempfile::tempdir().unwrap();
674 let root = tempdir.keep(); create_files(&["1", "2", "3"], &root).unwrap();
682
683 let subdir1 = root.join("subdir1");
685 std::fs::create_dir(&subdir1).unwrap();
686
687 create_files(&["4", "5", "6"], &subdir1).unwrap();
689
690 let subdir2 = root.join("subdir2");
692 std::fs::create_dir(&subdir2).unwrap();
693
694 create_files(&["7", "8", "9"], &subdir2).unwrap();
696
697 let subsub = subdir2.join("subsub");
699 std::fs::create_dir(&subsub).unwrap();
700
701 create_files(&["subsub1", "subsub2", "subsub2"], &subsub).unwrap();
703
704 let mut dirarchive = ZipDirArchive::from_dir(root.clone()).unwrap();
705
706 let mut zipfile = tempfile::tempfile().unwrap();
712
713 copy_archive_to_zipfile(&mut dirarchive, &mut zipfile).unwrap();
718 zipfile.seek(std::io::SeekFrom::Start(0)).unwrap();
719
720 let mut ziparchive = ZipDirArchive::from_zip(&mut zipfile, "archive.zip".into()).unwrap();
721
722 println!("checking zip");
723 check_archive(&mut ziparchive).unwrap();
724
725 println!("checking dirs");
726 check_archive(&mut dirarchive).unwrap();
727
728 std::fs::remove_dir_all(root).unwrap();
729 }
730
731 fn check_archive<R: Read + Seek>(archive: &mut ZipDirArchive<R>) -> Result<()> {
732 let paths = archive.list_paths::<PathBuf>(None)?;
733 assert_eq!(paths.len(), 5);
734 assert!(paths.contains(&PathBuf::from("1")));
735 assert!(paths.contains(&PathBuf::from("2")));
736 assert!(paths.contains(&PathBuf::from("3")));
737 assert!(paths.contains(&PathBuf::from("subdir1")));
738 assert!(paths.contains(&PathBuf::from("subdir2")));
739
740 let subs = PathBuf::from("subdir2").slash_join("subsub");
741
742 let subpaths = archive.list_paths::<PathBuf>(Some(subs))?;
743 assert_eq!(subpaths.len(), 2);
744 assert!(subpaths.contains(&PathBuf::from("subsub1")));
745 assert!(subpaths.contains(&PathBuf::from("subsub2")));
746
747 for not_exist_dir in &["not-exist", "abc/def", "abc\\def"] {
748 match archive
749 .list_paths::<PathBuf>(Some(PathBuf::from(not_exist_dir)))
750 .unwrap_err()
751 {
752 Error::NotDirectory(_) => {}
753 _ => {
754 panic!("returned wrong error. Should return NotDirectory");
755 }
756 }
757 }
758
759 {
760 let mut buf = String::new();
761 archive
762 .open_raw_or_gz("subdir2/8")?
763 .read_to_string(&mut buf)?;
764 assert_eq!(&buf, "8");
765 }
766
767 Ok(())
768 }
769}