haggis-rs/src/lib.rs

#![warn(clippy::all, clippy::pedantic)]
#![doc = include_str!("../README.md")]
use std::{
    collections::HashMap,
    fs,
    io::{self, BufWriter, Read, Seek, Write},
    sync::Mutex,
};
#[cfg(feature = "parallel")]
use {
    rayon::prelude::{IntoParallelRefIterator, ParallelIterator},
    std::sync::mpsc::Sender,
};

mod checksum;
mod error;
mod file;
mod filetype;
mod listing;
mod listing_stream;
pub(crate) mod nix;
mod node;
mod special;
mod stream;

pub use {
    checksum::{Algorithm, Checksum},
    error::Error,
    file::File,
    filetype::FileType,
    listing::Kind as ListingKind,
    listing::Listing,
    listing_stream::ListingStream,
    node::Node,
    special::Special,
    stream::Stream as NodeStream,
};

#[cfg(feature = "parallel")]
pub use stream::Message as StreamMessage;

/// The *magic* number for a Haggis archive
pub static MAGIC: [u8; 7] = [0x89, b'h', b'a', b'g', b'g', b'i', b's'];

static ZSTD_MAGIC: [u8; 4] = [0x28, 0xb5, 0x2f, 0xfd];

/// Tries to detect if the bytes in *reader* are zstd compressed
/// # Errors
/// Returns `Error` if io fails
pub fn detect_zstd<R: Read + Seek>(reader: &mut R) -> Result<bool, Error> {
    let mut buf = [0; 4];
    reader.read_exact(&mut buf)?;
    reader.rewind()?;
    Ok(buf == ZSTD_MAGIC)
}

pub(crate) fn load_string<R: Read>(reader: &mut R) -> Result<String, Error> {
    let mut len = [0; 2];
    reader.read_exact(&mut len)?;
    let len = u16::from_le_bytes(len);
    let mut buf = Vec::with_capacity(len.into());
    let mut handle = reader.take(len.into());
    handle.read_to_end(&mut buf)?;
    Ok(String::from_utf8(buf)?)
}

#[allow(clippy::similar_names)]
/// Creates a haggis archive from a list of files
/// # Errors
/// Returns `crate::Error` if io fails or several other error conditions
pub fn create_archive(
    path: &str,
    files: &[String],
    algorithm: Algorithm,
    uid: Option<u32>,
    gid: Option<u32>,
) -> Result<(), Error> {
    let fd = fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .open(path)?;
    let mut writer = BufWriter::new(fd);
    stream_archive(&mut writer, files, algorithm, uid, gid)?;
    Ok(())
}

#[allow(clippy::similar_names)]
/// Creates a haggis archive and writes it to stdout
/// # Errors
/// Returns `crate::Error` if io fails
pub fn create_archive_stdout(
    files: &[String],
    algorithm: Algorithm,
    uid: Option<u32>,
    gid: Option<u32>,
) -> Result<(), Error> {
    let stdout = io::stdout().lock();
    let mut writer = BufWriter::new(stdout);
    stream_archive(&mut writer, files, algorithm, uid, gid)?;
    Ok(())
}

#[allow(clippy::similar_names)]
/// Creates and streams a haggis archive over something which implements `Write`
/// # Errors
/// Returns `crate::Error` if io fails or several other error conditions
pub fn stream_archive<W: Write>(
    mut writer: W,
    files: &[String],
    algorithm: Algorithm,
    uid: Option<u32>,
    gid: Option<u32>,
) -> Result<(), Error> {
    writer.write_all(&MAGIC)?;
    let len = u32::try_from(files.len())?;
    writer.write_all(&len.to_le_bytes())?;
    let links = Mutex::new(HashMap::new());
    for f in files {
        let mut node = Node::from_path(f, algorithm, &links)?;
        if let Some(n) = uid {
            node.uid = n;
        }
        if let Some(n) = gid {
            node.gid = n;
        }
        node.write(&mut writer)?;
    }
    writer.write_all(&[0; 8])?;
    Ok(())
}

/// A message to be passed to another thread when a `Node` is created or saved.
/// Useful when creating a `Node` in a background thread and displaying progress
/// in the main thread
#[cfg(feature = "parallel")]
#[derive(Debug)]
pub enum Message {
    /// A `Node` has successfully been created from this path
    NodeCreated(
        /// The path of the created `Node`
        String,
    ),
    /// The `Node` has successfully been written into the writer
    NodeSaved {
        /// The path of the saved `Node`
        name: String,
        /// The size in bytes of the created `Node`
        size: u64,
    },
    /// An error occurred creating or writing out the node
    Err {
        /// The pathname of the node
        name: String,
        /// The error which occurred
        error: Error,
    },
    Eof,
}

/// Creates a Haggis archive from a list of files, processing each file in parallel
/// # Errors
/// Returns `crate::Error` if io fails or several other error conditions
#[cfg(feature = "parallel")]
pub fn par_create_archive(
    path: &str,
    files: &[String],
    algorithm: Algorithm,
    sender: &Sender<Message>,
    uid: Option<u32>,
    gid: Option<u32>,
) -> Result<(), Error> {
    let fd = fs::File::create(path)?;
    let writer = BufWriter::new(fd);
    par_stream_archive(writer, files, algorithm, sender, uid, gid)?;
    Ok(())
}

/// Creates a Haggis archive from a list of file, processing each file in parallel
/// and writing the resulting archive to stdout
/// # Errors
/// Returns `crate::Error` if io fails of one of several other conditions
#[cfg(feature = "parallel")]
pub fn par_create_archive_stdout(
    files: &[String],
    algorithm: Algorithm,
    sender: &Sender<Message>,
    uid: Option<u32>,
    gid: Option<u32>,
) -> Result<(), Error> {
    let stdout = io::stdout();
    let writer = BufWriter::new(stdout);
    par_stream_archive(writer, files, algorithm, sender, uid, gid)?;
    Ok(())
}

/// Creates and streams a Haggis archive from a list of files, processing each
/// file in parallel
/// # Errors
/// Returns `crate::Error` if io fails or several other error conditions
#[cfg(feature = "parallel")]
pub fn par_stream_archive<W: Write + Send>(
    mut writer: W,
    files: &[String],
    algorithm: Algorithm,
    sender: &Sender<Message>,
    uid: Option<u32>,
    gid: Option<u32>,
) -> Result<(), Error> {
    writer.write_all(&MAGIC)?;
    let len = u32::try_from(files.len())?;
    writer.write_all(&len.to_le_bytes())?;
    let links = Mutex::new(HashMap::<u64, String>::new());
    let writer = Mutex::new(writer);
    let s = sender.clone();
    files.par_iter().try_for_each_with(s, |s, f| {
        let mut node = match Node::from_path(f, algorithm, &links) {
            Ok(n) => n,
            Err(error) => {
                s.send(Message::Err {
                    name: f.to_string(),
                    error,
                })
                .map_err(|_| Error::SenderError)?;
                return Ok(());
            }
        };
        if let Some(n) = uid {
            node.uid = n;
        }
        if let Some(n) = gid {
            node.gid = n;
        }
        if let Ok(mut writer) = writer.lock() {
            let mut writer = &mut *writer;
            if let Err(error) = node.write(&mut writer) {
                s.send(Message::Err {
                    name: node.name.clone(),
                    error,
                })
                .map_err(|_| Error::SenderError)?;
            }
            match node.filetype {
                FileType::Normal(n) => s
                    .send(Message::NodeSaved {
                        name: node.name.clone(),
                        size: n.len,
                    })
                    .map_err(|_| Error::SenderError)?,
                _ => s
                    .send(Message::NodeCreated(node.name.clone()))
                    .map_err(|_| Error::SenderError)?,
            }
            Ok(())
        } else {
            Err(Error::MutexError)
        }
    })?;
    let mut writer = writer.into_inner().map_err(|_| Error::MutexError)?;
    writer.write_all(&[0; 8])?;
    sender.send(Message::Eof).map_err(|_| Error::SenderError)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn check_zstd() {
        let mut fd = fs::File::open("test/li.txt.zst").unwrap();
        assert!(detect_zstd(&mut fd).unwrap());
    }
}