affine_importer 0.1.2

AFFiNE import planning and batching core.
Documentation
#[cfg(test)]
use std::io::Cursor;
use std::{
  fs::{self, File},
  io::Read,
  path::{Path, PathBuf},
};

use zip::ZipArchive;

use crate::{ImportError, ImportResult, vfs::normalize_import_path};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ArchiveEntryMeta {
  pub index: usize,
  pub path: String,
  pub compressed_size: u64,
  pub uncompressed_size: u64,
}

pub trait ArchiveSource: Send {
  fn entries(&self) -> ImportResult<Vec<ArchiveEntryMeta>>;
  fn read_entry(&self, index: usize) -> ImportResult<Vec<u8>>;
}

#[derive(Debug, Clone)]
pub struct DirectoryPathSource {
  path: PathBuf,
}

impl DirectoryPathSource {
  pub fn new(path: impl Into<PathBuf>) -> Self {
    Self { path: path.into() }
  }

  pub fn path(&self) -> &Path {
    &self.path
  }

  fn files(&self) -> ImportResult<Vec<DirectoryEntry>> {
    let mut entries = Vec::new();
    collect_directory_entries(&self.path, &self.path, &mut entries)?;
    entries.sort_by(|left, right| left.path.cmp(&right.path));
    Ok(entries)
  }
}

impl ArchiveSource for DirectoryPathSource {
  fn entries(&self) -> ImportResult<Vec<ArchiveEntryMeta>> {
    let entries = self.files()?;
    if entries.is_empty() {
      return Err(ImportError::InvalidSource(
        "directory contains no importable files".to_string(),
      ));
    }
    Ok(
      entries
        .into_iter()
        .enumerate()
        .map(|(index, entry)| ArchiveEntryMeta {
          index,
          path: entry.path,
          compressed_size: entry.size,
          uncompressed_size: entry.size,
        })
        .collect(),
    )
  }

  fn read_entry(&self, index: usize) -> ImportResult<Vec<u8>> {
    let entries = self.files()?;
    let entry = entries
      .get(index)
      .ok_or_else(|| ImportError::InvalidSource(format!("missing directory entry: {index}")))?;
    Ok(fs::read(&entry.fs_path)?)
  }
}

#[derive(Debug, Clone)]
struct DirectoryEntry {
  path: String,
  fs_path: PathBuf,
  size: u64,
}

fn collect_directory_entries(root: &Path, current: &Path, entries: &mut Vec<DirectoryEntry>) -> ImportResult<()> {
  for entry in fs::read_dir(current)? {
    let entry = entry?;
    let path = entry.path();
    let metadata = entry.metadata()?;
    if metadata.is_dir() {
      collect_directory_entries(root, &path, entries)?;
    } else if metadata.is_file() {
      let relative = path
        .strip_prefix(root)
        .map_err(|_| ImportError::InvalidSource(format!("path escapes import root: {}", path.to_string_lossy())))?;
      let import_path = normalize_import_path(&relative.to_string_lossy());
      if import_path.is_empty() || is_system_path(&import_path) {
        continue;
      }
      entries.push(DirectoryEntry {
        path: import_path,
        fs_path: path,
        size: metadata.len(),
      });
    }
  }
  Ok(())
}

#[derive(Debug, Clone)]
pub struct ZipPathSource {
  path: PathBuf,
}

impl ZipPathSource {
  pub fn new(path: impl Into<PathBuf>) -> Self {
    Self { path: path.into() }
  }

  pub fn path(&self) -> &Path {
    &self.path
  }

  fn open_archive(&self) -> ImportResult<ZipArchive<File>> {
    Ok(ZipArchive::new(File::open(&self.path)?)?)
  }
}

impl ArchiveSource for ZipPathSource {
  fn entries(&self) -> ImportResult<Vec<ArchiveEntryMeta>> {
    let mut archive = self.open_archive()?;
    let mut entries = Vec::new();
    for index in 0..archive.len() {
      let file = archive.by_index(index)?;
      if file.is_dir() {
        continue;
      }
      let Some(path) = file
        .enclosed_name()
        .map(|path| normalize_import_path(&path.to_string_lossy()))
      else {
        continue;
      };
      if path.is_empty() || is_system_path(&path) {
        continue;
      }
      entries.push(ArchiveEntryMeta {
        index,
        path,
        compressed_size: file.compressed_size(),
        uncompressed_size: file.size(),
      });
    }
    if entries.is_empty() {
      return Err(ImportError::InvalidSource(
        "zip contains no importable files".to_string(),
      ));
    }
    Ok(entries)
  }

  fn read_entry(&self, index: usize) -> ImportResult<Vec<u8>> {
    let mut archive = self.open_archive()?;
    let mut file = archive.by_index(index)?;
    let mut bytes = Vec::with_capacity(file.size() as usize);
    file.read_to_end(&mut bytes)?;
    Ok(bytes)
  }
}

#[cfg(test)]
#[derive(Debug, Clone)]
pub struct ZipBytesSource {
  bytes: Vec<u8>,
}

#[cfg(test)]
impl ZipBytesSource {
  pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
    Self { bytes: bytes.into() }
  }

  fn open_archive(&self) -> ImportResult<ZipArchive<Cursor<&[u8]>>> {
    Ok(ZipArchive::new(Cursor::new(self.bytes.as_slice()))?)
  }
}

#[cfg(test)]
impl ArchiveSource for ZipBytesSource {
  fn entries(&self) -> ImportResult<Vec<ArchiveEntryMeta>> {
    let mut archive = self.open_archive()?;
    let mut entries = Vec::new();
    for index in 0..archive.len() {
      let file = archive.by_index(index)?;
      if file.is_dir() {
        continue;
      }
      let Some(path) = file
        .enclosed_name()
        .map(|path| normalize_import_path(&path.to_string_lossy()))
      else {
        continue;
      };
      if path.is_empty() || is_system_path(&path) {
        continue;
      }
      entries.push(ArchiveEntryMeta {
        index,
        path,
        compressed_size: file.compressed_size(),
        uncompressed_size: file.size(),
      });
    }
    if entries.is_empty() {
      return Err(ImportError::InvalidSource(
        "zip contains no importable files".to_string(),
      ));
    }
    Ok(entries)
  }

  fn read_entry(&self, index: usize) -> ImportResult<Vec<u8>> {
    let mut archive = self.open_archive()?;
    let mut file = archive.by_index(index)?;
    let mut bytes = Vec::with_capacity(file.size() as usize);
    file.read_to_end(&mut bytes)?;
    Ok(bytes)
  }
}

pub(crate) fn is_system_path(path: &str) -> bool {
  path.split('/').any(|part| part == "__MACOSX" || part == ".DS_Store")
}