nix-nar 0.5.0

use std::{
    cmp::min,
    fs::File,
    io::{self, Read},
    path::Path,
};

use camino::{Utf8Path, Utf8PathBuf};

use crate::{coder, NarError};

mod filesystem;
pub use filesystem::{FileSystem, FileSystemMetadata, NativeFileSystem};

/// Encoder which can archive a given path as a NAR file.
///
/// The encoder will read the data to archive using its [`FileSystem`]
/// implementation. By default, it reads the host filesystem using
/// [`std::fs`].
pub struct Encoder<FS: FileSystem = NativeFileSystem> {
    stack: Vec<CurrentActivity<FS>>,
    internal_buffer_size: usize,
    fs: FS,
}

/// Builder for [`Encoder`].
pub struct EncoderBuilder<P: AsRef<Path>, FS: FileSystem = NativeFileSystem> {
    path: P,
    internal_buffer_size: usize,
    fs: FS,
}

#[derive(Debug)]
enum CurrentActivity<FS: FileSystem> {
    StartArchive,
    StartEntry,
    Toplevel {
        path: Utf8PathBuf,
    },
    WalkingDir {
        dir_path: Utf8PathBuf,
        files_rev: Vec<String>,
    },
    EncodingFile {
        file: FS::File,
        length_remaining: u64,
    },
    WritePadding {
        padding: u64,
    },
    WriteMoreBytes {
        bytes: Vec<u8>,
    },
    CloseDirEntry,
    CloseEntry,
}

impl Encoder<NativeFileSystem> {
    /// Create a new encoder for file hierarchy at the given path
    /// using default filesystem implementation.
    ///
    /// # Errors
    ///
    /// Returns an error if the path is not valid UTF-8.
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, NarError> {
        Self::new_with_filesystem(path, NativeFileSystem {})
    }

    /// Create a builder for this encoder that can take additional
    /// options.
    ///
    /// ```
    /// # use std::{io, fs::File};
    /// use nix_nar::{Encoder, NativeFileSystem};
    /// let mut enc = Encoder::builder("./test-data/")
    ///     .internal_buffer_size(2048)
    ///     .filesystem(NativeFileSystem {})
    ///     .build().unwrap();
    /// let mut nar = File::create("output2.nar")?;
    /// io::copy(&mut enc, &mut nar)?;
    /// # std::fs::remove_file("output2.nar")?;
    /// # Ok::<(), nix_nar::NarError>(())
    /// ```
    pub fn builder<P: AsRef<Path>>(path: P) -> EncoderBuilder<P> {
        EncoderBuilder {
            path,
            internal_buffer_size: 1024,
            fs: NativeFileSystem {},
        }
    }
}

impl<FS: FileSystem> Encoder<FS> {
    /// Create a new encoder for file hierarchy at the given path and
    /// provided filesystem implementation.
    ///
    /// # Errors
    ///
    /// Returns an error if the path is not valid UTF-8.
    pub fn new_with_filesystem<P: AsRef<Path>>(
        path: P,
        fs: FS,
    ) -> Result<Self, NarError> {
        let path = to_utf8_path(path)?;
        Ok(Self {
            stack: vec![
                CurrentActivity::CloseEntry,
                CurrentActivity::Toplevel { path },
                CurrentActivity::StartEntry,
                CurrentActivity::StartArchive,
            ],
            internal_buffer_size: 1024,
            fs,
        })
    }

    /// Archive to the given path.
    ///
    /// This is equivalent to creating the file, and [`io::copy`]ing
    /// to it.
    ///
    /// # Errors
    ///
    /// Returns an error if the path already exists or an I/O error occurs.
    pub fn pack<P: AsRef<Path>>(&mut self, dst: P) -> Result<(), NarError> {
        let dst = to_utf8_path(dst)?;
        if dst.symlink_metadata().is_ok() {
            return Err(NarError::PackError(format!(
                "Destination {dst} already exists. Delete it first."
            )));
        }
        let mut nar = File::create(&dst)?;
        io::copy(self, &mut nar)?;
        Ok(())
    }
}

impl<P: AsRef<Path>, FS: FileSystem> EncoderBuilder<P, FS> {
    /// Build the encoder.
    ///
    /// # Errors
    ///
    /// Returns an error if the path is not valid UTF-8.
    pub fn build(self) -> Result<Encoder<FS>, NarError> {
        let Self {
            path,
            internal_buffer_size,
            fs,
        } = self;
        let mut enc = Encoder::new_with_filesystem(path, fs)?;
        enc.internal_buffer_size = internal_buffer_size;
        Ok(enc)
    }

    /// Configure the internal buffer size.  This should be at least
    /// 200 bytes larger than the longest filename.
    ///
    /// # Panics
    ///
    /// Panics if the given number is smaller than 200.
    #[must_use]
    pub fn internal_buffer_size(mut self, x: usize) -> Self {
        assert!(
            x >= 200,
            "internal_buffer_size should be at least 200 bytes larger than the longest filename you have"
        );
        self.internal_buffer_size = x;
        self
    }

    /// Configure filesystem implementation.
    #[must_use]
    pub fn filesystem<FS2: FileSystem>(self, fs: FS2) -> EncoderBuilder<P, FS2> {
        let Self {
            path,
            internal_buffer_size,
            fs: _,
        } = self;
        EncoderBuilder {
            path,
            internal_buffer_size,
            fs,
        }
    }
}

impl<FS: FileSystem> Encoder<FS> {
    fn start_encoding<P: AsRef<Utf8Path>>(
        &mut self,
        buf: &mut [u8],
        path: P,
        metadata: FileSystemMetadata,
        dir_entry: Option<String>,
    ) -> Result<usize, io::Error> {
        match metadata {
            FileSystemMetadata::Directory => {
                self.start_encoding_dir(buf, path, dir_entry)
            }
            FileSystemMetadata::File { executable, length } => {
                self.start_encoding_file(buf, path, executable, length, dir_entry)
            }
            FileSystemMetadata::Symlink { target_path } => {
                self.start_encoding_symlink(buf, &target_path, dir_entry)
            }
        }
    }

    fn start_encoding_file<P: AsRef<Utf8Path>>(
        &mut self,
        buf: &mut [u8],
        path: P,
        executable: bool,
        length: u64,
        dir_entry: Option<String>,
    ) -> Result<usize, io::Error> {
        let path = path.as_ref();
        let file_handle = self.fs.open(path).map_err(annotate_err_with_path(&path))?;
        let file_len_rounded_up = (length + 7) & !7;
        if file_len_rounded_up > length {
            self.stack.push(CurrentActivity::WritePadding {
                padding: file_len_rounded_up - length,
            });
        }
        self.stack.push(CurrentActivity::EncodingFile {
            file: file_handle,
            length_remaining: length,
        });
        self.write_with_buffer(buf, move |buf| {
            let mut len = 0;
            if let Some(ref file) = dir_entry {
                len += coder::start_dir_entry(&mut buf[len..], file)?;
            }
            len += coder::write_file_regular(&mut buf[len..], executable)?;
            len += coder::write_u64_le(&mut buf[len..], length)?;
            Ok(len)
        })
    }

    fn start_encoding_dir<P: AsRef<Utf8Path>>(
        &mut self,
        buf: &mut [u8],
        path: P,
        dir_entry: Option<String>,
    ) -> Result<usize, io::Error> {
        let path = path.as_ref();
        let mut files_rev = self
            .fs
            .read_dir(path)
            .map_err(annotate_err_with_path(path))?;
        files_rev.sort_by(|a, b| b.cmp(a));

        self.stack.push(CurrentActivity::WalkingDir {
            files_rev,
            dir_path: path.into(),
        });
        self.write_with_buffer(buf, move |buf| {
            let mut len = 0;
            if let Some(ref file) = dir_entry {
                len += coder::start_dir_entry(&mut buf[len..], file)?;
            }
            len += coder::start_dir(&mut buf[len..])?;
            Ok(len)
        })
    }

    fn start_encoding_symlink<P: AsRef<Utf8Path>>(
        &mut self,
        buf: &mut [u8],
        target_path: P,
        dir_entry: Option<String>,
    ) -> Result<usize, io::Error> {
        self.write_with_buffer(buf, move |buf| {
            let mut len = 0;
            if let Some(ref file) = dir_entry {
                len += coder::start_dir_entry(buf, file)?;
            }
            len += coder::write_symlink(&mut buf[len..], target_path)?;
            Ok(len)
        })
    }

    /// Execute a write-into-buffer operation.  If the given buffer is
    /// big enough, then we just write into that.  Otherwise, we
    /// create a temporary buffer, we write into that, we copy as much
    /// data as we can to the given buffer, and store the remainer
    /// into a `CurrentActivity::WriteMoreBytes` on the `self.stack`.
    fn write_with_buffer<F>(&mut self, dst_buf: &mut [u8], f: F) -> io::Result<usize>
    where
        F: FnOnce(&mut [u8]) -> io::Result<usize>,
    {
        if dst_buf.len() >= 1024 {
            f(dst_buf)
        } else {
            let mut buf = vec![0; self.internal_buffer_size];
            let len = f(&mut buf)?;
            let to_write_len = min(len, dst_buf.len());
            dst_buf[..to_write_len].copy_from_slice(&buf[..to_write_len]);
            if len > to_write_len {
                buf.truncate(len);
                self.stack.push(CurrentActivity::WriteMoreBytes {
                    bytes: buf.split_off(to_write_len),
                });
            }
            Ok(to_write_len)
        }
    }
}

/// Read the encoded NAR file.
///
/// # Errors
///
/// [`Self::read`] will return underlying I/O errors returned by [`FileSystem`].
///
/// It will also return an error if the data read from an archived
/// file is not the same length as declared by
/// [`FileSystem::metadata`]. The length from the metadata was already
/// written to the NAR output at this point, so it's not possible to
/// recover and produce a valid archive.
impl<FS: FileSystem> Read for Encoder<FS> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.stack.pop() {
            None => Ok(0),
            Some(CurrentActivity::StartArchive) => {
                self.write_with_buffer(buf, coder::start_archive)
            }
            Some(CurrentActivity::StartEntry) => {
                self.write_with_buffer(buf, coder::start_entry)
            }
            Some(CurrentActivity::CloseDirEntry) => {
                self.write_with_buffer(buf, coder::close_dir_entry)
            }
            Some(CurrentActivity::CloseEntry) => {
                self.write_with_buffer(buf, coder::close_entry)
            }
            Some(CurrentActivity::Toplevel { path }) => {
                let metadata = self
                    .fs
                    .metadata(&path)
                    .map_err(annotate_err_with_path(&path))?;
                self.start_encoding(buf, path, metadata, None)
            }
            Some(CurrentActivity::WalkingDir {
                dir_path,
                mut files_rev,
            }) => match files_rev.pop() {
                None => self.read(buf),
                Some(file) => {
                    let path = dir_path.join(&file);

                    self.stack.push(CurrentActivity::WalkingDir {
                        dir_path,
                        files_rev,
                    });

                    self.stack.push(CurrentActivity::CloseDirEntry);

                    let metadata = self
                        .fs
                        .metadata(&path)
                        .map_err(annotate_err_with_path(&path))?;
                    self.start_encoding(buf, path, metadata, Some(file))
                }
            },
            Some(CurrentActivity::EncodingFile {
                mut file,
                length_remaining,
            }) => {
                let len = file.read(buf)?;
                if len != 0 {
                    self.stack.push(CurrentActivity::EncodingFile {
                        file,
                        length_remaining: length_remaining
                            .checked_sub(len.try_into().map_err(io::Error::other)?)
                            .ok_or_else(|| {
                                other_io_error("File was longer than declared length")
                            })?,
                    });
                    Ok(len)
                } else if length_remaining > 0 {
                    Err(other_io_error("File was shorter than declared length"))
                } else {
                    self.read(buf)
                }
            }
            Some(CurrentActivity::WritePadding { padding }) => {
                #[allow(clippy::cast_possible_truncation)]
                let len = min(padding, buf.len() as u64) as usize;
                buf.fill(0);
                if (len as u64) < padding {
                    self.stack.push(CurrentActivity::WritePadding {
                        padding: padding - len as u64,
                    });
                }
                Ok(len)
            }
            Some(CurrentActivity::WriteMoreBytes { bytes }) => {
                let len = min(bytes.len(), buf.len());
                buf[..len].copy_from_slice(&bytes[..len]);
                if len < bytes.len() {
                    self.stack.push(CurrentActivity::WriteMoreBytes {
                        bytes: bytes[len..].to_vec(),
                    });
                }
                Ok(len)
            }
        }
    }
}

fn annotate_err_with_path<P: AsRef<Utf8Path>>(
    path: P,
) -> impl FnOnce(io::Error) -> io::Error {
    let path = path.as_ref().to_path_buf();
    move |err: io::Error| other_io_error(format!("IO error on {path}: {err}"))
}

fn other_io_error<S: AsRef<str>>(message: S) -> io::Error {
    io::Error::other(message.as_ref())
}

fn to_utf8_path<P: AsRef<Path>>(path: P) -> Result<Utf8PathBuf, NarError> {
    let path = path.as_ref();
    path.try_into()
        .map(|x: &Utf8Path| x.to_path_buf())
        .map_err(|err| {
            NarError::Utf8PathError(format!(
                "Failed to convert '{}' to UTF-8: {err}",
                path.display()
            ))
        })
}