ddup_bak/

lib.rs

use flate2::{read::GzDecoder, write::GzEncoder};
use positioned_io::ReadAt;
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
use std::{
    fmt::{Debug, Formatter},
    fs::{DirEntry, File, Permissions},
    io::{Read, Seek, SeekFrom, Write},
};

pub const FILE_SIGNATURE: &[u8] = b"DDUPBAK\0";

#[derive(Debug, Clone, Copy)]
pub enum CompressionFormat {
    None,
    Gzip,
}

impl CompressionFormat {
    pub fn encode(&self) -> u8 {
        match self {
            CompressionFormat::None => 0,
            CompressionFormat::Gzip => 1,
        }
    }

    pub fn decode(value: u8) -> Self {
        match value {
            0 => CompressionFormat::None,
            1 => CompressionFormat::Gzip,
            _ => panic!("Invalid compression format"),
        }
    }
}

pub struct VarInt(u64);

impl VarInt {
    pub fn new(value: u64) -> Self {
        VarInt(value)
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut value = self.0;
        let mut buffer = Vec::new();

        while value > 0x7F {
            buffer.push((value & 0x7F) as u8 | 0x80);
            value >>= 7;
        }

        buffer.push(value as u8);

        buffer
    }

    pub fn decode<S: Read>(stream: &mut S) -> Result<u64, std::io::Error> {
        let mut value = 0;
        let mut shift = 0;

        loop {
            let mut byte = [0; 1];
            stream.read_exact(&mut byte)?;

            value |= ((byte[0] & 0x7F) as u64) << shift;
            if byte[0] & 0x80 == 0 {
                break;
            }
            shift += 7;
        }

        Ok(value)
    }
}

fn encode_file_permissions(permissions: Permissions) -> u32 {
    #[cfg(unix)]
    {
        permissions.mode()
    }
    #[cfg(windows)]
    {
        if permissions.readonly() { 1 } else { 0 }
    }
}

fn decode_file_permissions(mode: u32) -> Permissions {
    #[cfg(unix)]
    {
        Permissions::from_mode(mode)
    }
    #[cfg(windows)]
    {
        let mut permissions = unsafe { std::mem::zeroed::<Permissions>() };
        if mode == 1 {
            permissions.set_readonly(true);
        } else {
            permissions.set_readonly(false);
        }

        permissions
    }
}

pub struct FileEntry {
    pub name: String,
    pub mode: Permissions,

    file: File,
    decoder: Option<GzDecoder<File>>,
    size: u64,
    offset: u64,
    consumed: u64,
    compression: CompressionFormat,
}

impl Debug for FileEntry {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FileEntry")
            .field("name", &self.name)
            .field("mode", &self.mode)
            .field("size", &self.size)
            .field("offset", &self.offset)
            .field("compression", &self.compression)
            .finish()
    }
}

impl Read for FileEntry {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.consumed >= self.size {
            return Ok(0);
        }

        let remaining = self.size - self.consumed;

        match self.compression {
            CompressionFormat::Gzip => {
                if self.decoder.is_none() {
                    self.file
                        .seek(SeekFrom::Start(self.offset + self.consumed))?;
                    let decoder = GzDecoder::new(self.file.try_clone()?);
                    self.decoder = Some(decoder);
                }

                let decoder = self.decoder.as_mut().unwrap();
                let bytes_read = decoder.read(buf)?;

                if bytes_read > remaining as usize {
                    self.consumed += remaining;
                    return Ok(remaining as usize);
                }

                self.consumed += bytes_read as u64;
                Ok(bytes_read)
            }
            CompressionFormat::None => {
                let bytes_read = self.file.read_at(self.offset + self.consumed, buf)?;

                if bytes_read > remaining as usize {
                    self.consumed += remaining;
                    return Ok(remaining as usize);
                }

                self.consumed += bytes_read as u64;
                Ok(bytes_read)
            }
        }
    }
}

#[derive(Debug)]
pub struct DirectoryEntry {
    pub name: String,
    pub mode: Permissions,
    pub entries: Vec<Entry>,
}

#[derive(Debug)]
pub struct SymlinkEntry {
    pub name: String,
    pub mode: Permissions,
    pub target: String,
    pub target_dir: bool,
}

#[derive(Debug)]
pub enum Entry {
    File(Box<FileEntry>),
    Directory(DirectoryEntry),
    Symlink(SymlinkEntry),
}

impl Entry {
    /// Returns the name of the entry.
    /// This is the name of the file or directory, not the full path.
    /// For example, if the entry is under `path/to/file.txt`, this will return `file.txt`.
    pub fn name(&self) -> &str {
        match self {
            Entry::File(entry) => &entry.name,
            Entry::Directory(entry) => &entry.name,
            Entry::Symlink(entry) => &entry.name,
        }
    }

    /// Returns the mode of the entry.
    /// This is the file permissions of the entry.
    pub fn mode(&self) -> Permissions {
        match self {
            Entry::File(entry) => entry.mode.clone(),
            Entry::Directory(entry) => entry.mode.clone(),
            Entry::Symlink(entry) => entry.mode.clone(),
        }
    }
}

pub struct Archive {
    file: File,

    entries: Vec<Entry>,
    entries_offset: u64,
}

impl Debug for Archive {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Archive")
            .field("entries", &self.entries)
            .finish()
    }
}

type ProgressCallback = Option<fn(&std::path::PathBuf)>;

impl Archive {
    /// Creates a new archive file.
    /// The file signature is written to the beginning of the file.
    /// The file is truncated to 0 bytes.
    pub fn new(mut file: File) -> Self {
        file.set_len(0).unwrap();
        file.write_all(FILE_SIGNATURE).unwrap();
        file.sync_all().unwrap();

        Self {
            file,
            entries: Vec::new(),
            entries_offset: 8,
        }
    }

    /// Opens an existing archive file for reading and writing.
    /// This will not overwrite the file, but append to it.
    pub fn open(path: &str) -> Result<Self, std::io::Error> {
        let mut file = File::open(path)?;
        let len = file.metadata()?.len();

        let mut buffer = [0; 8];
        file.read_exact(&mut buffer)?;
        if buffer != FILE_SIGNATURE {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Invalid file signature",
            ));
        }

        file.read_exact_at(len - 16, &mut buffer)?;
        let entries_count = u64::from_le_bytes(buffer);
        file.read_exact_at(len - 8, &mut buffer)?;
        let entries_offset = u64::from_le_bytes(buffer);

        println!("Entries count: {}", entries_count);
        println!("Entries offset: {}", entries_offset);

        let mut entries = Vec::with_capacity(entries_count as usize);
        file.seek(SeekFrom::Start(entries_offset))?;

        let mut decoder = GzDecoder::new(file.try_clone()?);
        for _ in 0..entries_count {
            let file_clone = file.try_clone()?;
            let entry = Self::decode_entry(&mut decoder, file_clone)?;
            entries.push(entry);
        }

        Ok(Self {
            file,
            entries,
            entries_offset,
        })
    }

    /// Adds all files in the given directory to the archive. (including subdirectories)
    /// This will append the directory to the end of the archive, if this directory already exists, it will not be replaced.
    ///
    /// After this function is called, the existing header will be trimmed to the end of the archive, then readded upon completion.
    ///
    /// # Panics
    /// This function will panic if any filename is not valid UTF-8 or longer than 255 bytes.
    pub fn add_directory(
        &mut self,
        path: &str,
        progress: ProgressCallback,
    ) -> Result<(), std::io::Error> {
        self.trim_end_header()?;

        for entry in std::fs::read_dir(path)?.flatten() {
            self.encode_entry(None, entry, progress)?;
        }

        self.write_end_header()?;

        Ok(())
    }

    /// Returns the entries in the archive.
    pub fn entries(&self) -> &[Entry] {
        &self.entries
    }

    /// Consumes the archive and returns the entries.
    pub fn into_entries(self) -> Vec<Entry> {
        self.entries
    }

    /// Adds a single file entry to the archive. (including subdirectories)
    /// This will append the entry to the end of the archive, if this entry already exists, it will not be replaced.
    ///
    /// After this function is called, the existing header will be trimmed to the end of the archive, then readded upon completion.
    ///
    /// # Panics
    /// This function will panic if any filename is not valid UTF-8 or longer than 255 bytes.
    pub fn add_entry(
        &mut self,
        entry: DirEntry,
        progress: ProgressCallback,
    ) -> Result<(), std::io::Error> {
        self.trim_end_header()?;
        self.encode_entry(None, entry, progress)?;

        self.write_end_header()?;

        Ok(())
    }

    fn trim_end_header(&mut self) -> Result<(), std::io::Error> {
        if self.entries_offset == 0 {
            return Ok(());
        }

        self.file.set_len(self.entries_offset)?;

        Ok(())
    }

    fn write_end_header(&mut self) -> Result<(), std::io::Error> {
        let mut encoder = GzEncoder::new(&mut self.file, flate2::Compression::default());
        for entry in &self.entries {
            Self::encode_entry_metadata(&mut encoder, entry)?;
        }

        encoder.finish()?;

        self.file
            .write_all(&(self.entries.len() as u64).to_le_bytes())?;
        self.file.write_all(&self.entries_offset.to_le_bytes())?;
        self.file.sync_all()?;

        Ok(())
    }

    fn encode_entry_metadata<S: Write>(
        writer: &mut S,
        entry: &Entry,
    ) -> Result<(), std::io::Error> {
        let name = entry.name();
        let name_length = name.len() as u8;

        let mut buffer = Vec::with_capacity(1 + name.len() + 4);

        buffer.push(name_length);
        buffer.extend_from_slice(name.as_bytes());

        let mode = encode_file_permissions(entry.mode());
        let compression = match entry {
            Entry::File(file_entry) => file_entry.compression,
            _ => CompressionFormat::None,
        };
        let entry_type = match entry {
            Entry::File(_) => 0,
            Entry::Directory(_) => 1,
            Entry::Symlink(_) => 2,
        };

        let type_compression_mode =
            (entry_type << 30) | ((compression.encode() as u32) << 26) | (mode & 0x3FFFFFFF);
        buffer.extend_from_slice(&type_compression_mode.to_le_bytes()[..4]);

        writer.write_all(&buffer)?;

        match entry {
            Entry::File(file_entry) => {
                writer.write_all(VarInt::new(file_entry.size).encode().as_slice())?;
                writer.write_all(VarInt::new(file_entry.offset).encode().as_slice())?;
            }
            Entry::Directory(dir_entry) => {
                writer.write_all(
                    VarInt::new(dir_entry.entries.len() as u64)
                        .encode()
                        .as_slice(),
                )?;

                for sub_entry in &dir_entry.entries {
                    Self::encode_entry_metadata(writer, sub_entry)?;
                }
            }
            Entry::Symlink(link_entry) => {
                writer.write_all(
                    VarInt::new(link_entry.target.len() as u64)
                        .encode()
                        .as_slice(),
                )?;
                writer.write_all(link_entry.target.as_bytes())?;
                writer.write_all(&[link_entry.target_dir as u8])?;
            }
        }

        Ok(())
    }

    fn encode_entry(
        &mut self,
        entries: Option<&mut Vec<Entry>>,
        fs_entry: DirEntry,
        progress: ProgressCallback,
    ) -> Result<(), std::io::Error> {
        let path = fs_entry.path();
        if let Some(f) = progress {
            f(&path)
        }

        let file_name = path.file_name().unwrap().to_string_lossy();
        let metadata = path.symlink_metadata()?;

        if metadata.is_file() {
            let mut file = File::open(&path)?;
            let mut buffer = vec![0; 1024 * 1024];
            let mut bytes_read = file.read(&mut buffer)?;

            let compression = if metadata.len() > 16 {
                CompressionFormat::Gzip
            } else {
                CompressionFormat::None
            };
            let entry = FileEntry {
                name: file_name.to_string(),
                mode: metadata.permissions(),
                file: self.file.try_clone()?,
                decoder: None,
                size: metadata.len(),
                offset: self.entries_offset,
                consumed: 0,
                compression,
            };

            if let Some(entries) = entries {
                entries.push(Entry::File(Box::new(entry)));
            } else {
                self.entries.push(Entry::File(Box::new(entry)));
            }

            match compression {
                CompressionFormat::Gzip => {
                    let mut encoder =
                        GzEncoder::new(&mut self.file, flate2::Compression::default());
                    loop {
                        encoder.write_all(&buffer[..bytes_read])?;

                        bytes_read = file.read(&mut buffer)?;
                        if bytes_read == 0 {
                            break;
                        }
                    }
                    encoder.finish()?;

                    self.entries_offset = self.file.stream_position()?;
                }
                CompressionFormat::None => loop {
                    self.file.write_all(&buffer[..bytes_read])?;
                    self.entries_offset += bytes_read as u64;

                    bytes_read = file.read(&mut buffer)?;
                    if bytes_read == 0 {
                        break;
                    }
                },
            }
        } else if metadata.is_dir() {
            let mut dir_entries = Vec::new();
            for entry in std::fs::read_dir(&path)?.flatten() {
                self.encode_entry(Some(&mut dir_entries), entry, progress)?;
            }

            let dir_entry = DirectoryEntry {
                name: file_name.to_string(),
                mode: metadata.permissions(),
                entries: dir_entries,
            };

            if let Some(entries) = entries {
                entries.push(Entry::Directory(dir_entry));
            } else {
                self.entries.push(Entry::Directory(dir_entry));
            }
        } else if metadata.is_symlink() {
            let target = std::fs::read_link(&path)?;
            let target = target.to_string_lossy().to_string();

            let link_entry = SymlinkEntry {
                name: file_name.to_string(),
                mode: metadata.permissions(),
                target,
                target_dir: std::fs::metadata(&path)?.is_dir(),
            };

            if let Some(entries) = entries {
                entries.push(Entry::Symlink(link_entry));
            } else {
                self.entries.push(Entry::Symlink(link_entry));
            }
        }

        Ok(())
    }

    fn decode_entry<S: Read>(decoder: &mut S, file: File) -> Result<Entry, std::io::Error> {
        let mut name_length = [0; 1];
        decoder.read_exact(&mut name_length)?;
        let name_length = name_length[0] as usize;

        let mut name_bytes = vec![0; name_length];
        decoder.read_exact(&mut name_bytes)?;
        let name = String::from_utf8(name_bytes).unwrap();

        let mut type_mode_bytes = [0; 4];
        decoder.read_exact(&mut type_mode_bytes)?;
        let type_compression_mode = u32::from_le_bytes(type_mode_bytes);

        let entry_type = (type_compression_mode >> 30) & 0b11;
        let compression = CompressionFormat::decode(((type_compression_mode >> 26) & 0b1111) as u8);
        let mode = decode_file_permissions(type_compression_mode & 0x3FFFFFFF);
        let size = VarInt::decode(decoder)?;

        match entry_type {
            0 => {
                let offset = VarInt::decode(decoder)?;

                Ok(Entry::File(Box::new(FileEntry {
                    name,
                    mode,
                    file,
                    decoder: None,
                    size,
                    offset,
                    consumed: 0,
                    compression,
                })))
            }
            1 => {
                let mut entries: Vec<Entry> = Vec::with_capacity(size as usize);
                for _ in 0..size {
                    let entry = Self::decode_entry(decoder, file.try_clone()?)?;
                    entries.push(entry);
                }

                Ok(Entry::Directory(DirectoryEntry {
                    name,
                    mode,
                    entries,
                }))
            }
            2 => {
                let mut target_bytes = vec![0; size as usize];
                decoder.read_exact(&mut target_bytes)?;

                let target = String::from_utf8(target_bytes).unwrap();
                let target = std::path::PathBuf::from(target);

                let target = target
                    .canonicalize()
                    .unwrap_or_else(|_| target.clone())
                    .to_string_lossy()
                    .to_string();
                let metadata = std::fs::metadata(&target)?;

                Ok(Entry::Symlink(SymlinkEntry {
                    name,
                    mode,
                    target,
                    target_dir: metadata.is_dir(),
                }))
            }
            _ => panic!("Unsupported entry type"),
        }
    }
}