backhand 0.25.5

//! File Data

use no_std_io2::io::{Read, Seek, SeekFrom, Write};
use std::collections::HashMap;
use std::collections::hash_map::Entry;

use deku::prelude::*;
use solana_nohash_hasher::IntMap;
use xxhash_rust::xxh64::Xxh64;

use crate::error::BackhandError;
use crate::v4::filesystem::writer::FilesystemCompressor;
use crate::v4::fragment::Fragment;
use crate::v4::reader::WriteSeek;

use crate::v4::filesystem::reader::SquashfsRawData;

// bitflag for data size field in inode for signifying that the data is uncompressed
const DATA_STORED_UNCOMPRESSED: u32 = 1 << 24;

#[derive(Copy, Clone, Debug, PartialEq, Eq, DekuRead, DekuWrite, DekuSize)]
#[deku(endian = "endian", ctx = "endian: deku::ctx::Endian")]
pub struct DataSize(u32);
impl DataSize {
    #[inline]
    pub fn new(size: u32, uncompressed: bool) -> Self {
        let mut value: u32 = size;
        if value > DATA_STORED_UNCOMPRESSED {
            panic!("value is too big");
        }
        if uncompressed {
            value |= DATA_STORED_UNCOMPRESSED;
        }
        Self(value)
    }

    #[inline]
    pub fn new_compressed(size: u32) -> Self {
        Self::new(size, false)
    }

    #[inline]
    pub fn new_uncompressed(size: u32) -> Self {
        Self::new(size, true)
    }

    /// A block of all zeros, with no bytes stored in the image
    #[inline]
    pub const fn hole() -> Self {
        Self(0)
    }

    #[inline]
    pub fn uncompressed(&self) -> bool {
        self.0 & DATA_STORED_UNCOMPRESSED != 0
    }

    #[inline]
    pub fn set_uncompressed(&mut self) {
        self.0 |= DATA_STORED_UNCOMPRESSED
    }

    #[inline]
    pub fn set_compressed(&mut self) {
        self.0 &= !DATA_STORED_UNCOMPRESSED
    }

    #[inline]
    pub fn size(&self) -> u32 {
        self.0 & !DATA_STORED_UNCOMPRESSED
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Added {
    // Only Data was added. `sparse` is the number of bytes stored as holes.
    Data { blocks_start: u64, block_sizes: Vec<DataSize>, sparse: u64 },
    // Only Fragment was added
    Fragment { frag_index: u32, block_offset: u32 },
}

struct DataWriterChunkReader<R: std::io::Read> {
    chunk: Vec<u8>,
    file_len: usize,
    reader: R,
}
impl<R: std::io::Read> DataWriterChunkReader<R> {
    pub fn read_chunk(&mut self) -> std::io::Result<&[u8]> {
        use std::io::ErrorKind;
        let mut buf: &mut [u8] = &mut self.chunk;
        let mut read_len = 0;
        while !buf.is_empty() {
            match self.reader.read(buf) {
                Ok(0) => break,
                Ok(n) => {
                    read_len += n;
                    let tmp = buf;
                    buf = &mut tmp[n..];
                }
                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        self.file_len += read_len;
        Ok(&self.chunk[..read_len])
    }
}

pub(crate) struct DataWriter<'a> {
    compressor: &'a (
            dyn crate::traits::CompressionAction<
        Compressor = super::compressor::Compressor,
        FilesystemCompressor = FilesystemCompressor,
        SuperBlock = super::squashfs::SuperBlock,
        Error = crate::BackhandError,
    > + Send
                + Sync
        ),
    block_size: u32,
    fs_compressor: FilesystemCompressor,
    /// If some, cache of HashMap<file_len, HashMap<hash of whole file, (file_len, Added)>>
    #[allow(clippy::type_complexity)]
    dup_cache: Option<IntMap<u64, IntMap<u64, (usize, Added)>>>,
    /// Un-written fragment_bytes
    pub(crate) fragment_bytes: Vec<u8>,
    pub(crate) fragment_table: Vec<Fragment>,
}

impl<'a> DataWriter<'a> {
    pub fn new(
        compressor: &'a (
                dyn crate::traits::CompressionAction<
            Compressor = super::compressor::Compressor,
            FilesystemCompressor = FilesystemCompressor,
            SuperBlock = super::squashfs::SuperBlock,
            Error = crate::BackhandError,
        > + Send
                    + Sync
            ),
        fs_compressor: FilesystemCompressor,
        block_size: u32,
        no_duplicate_files: bool,
    ) -> Self {
        Self {
            compressor,
            block_size,
            fs_compressor,
            dup_cache: no_duplicate_files.then_some(HashMap::default()),
            fragment_bytes: Vec::with_capacity(block_size as usize),
            fragment_table: vec![],
        }
    }

    /// Add to data writer, either a pre-compressed Data or Fragment
    // TODO: support tail-end fragments (off by default in squashfs-tools/mksquashfs)
    pub(crate) fn just_copy_it<W: WriteSeek>(
        &mut self,
        mut reader: SquashfsRawData,
        mut writer: W,
    ) -> Result<(usize, Added), BackhandError> {
        //just clone it, because block sizes where never modified, just copy it
        let source = reader.file;
        let mut block_sizes = source.block_sizes().to_vec();
        let sparse = source.sparse();
        let mut read_buf = vec![];
        let mut decompress_buf = vec![];

        // if the first block is not full (fragment), store only a fragment
        // otherwise processed to store blocks
        let blocks_start = writer.stream_position()?;
        // Older backhand versions gave an empty file a fragment, do not copy that
        if source.file_len() == 0 {
            return Ok((0, Added::Data { blocks_start, block_sizes: vec![], sparse: 0 }));
        }
        let first_block = match reader.next_block(&mut read_buf) {
            Some(Ok(first_block)) => first_block,
            Some(Err(x)) => return Err(x),
            None => return Ok((0, Added::Data { blocks_start, block_sizes, sparse })),
        };

        // write and early return if fragment
        if first_block.fragment {
            reader.decompress(first_block, &mut read_buf, &mut decompress_buf)?;
            // if this doesn't fit in the current fragment bytes
            // compress the current fragment bytes and add to data_bytes
            if (decompress_buf.len() + self.fragment_bytes.len()) > self.block_size as usize {
                self.finalize(writer)?;
            }
            // add to fragment bytes
            let frag_index = self.fragment_table.len() as u32;
            let block_offset = self.fragment_bytes.len() as u32;
            self.fragment_bytes.write_all(&decompress_buf)?;

            return Ok((decompress_buf.len(), Added::Fragment { frag_index, block_offset }));
        }

        // The reader gives a hole as a block of zeros, but a hole has no bytes in the image
        let mut is_hole = source.block_sizes().iter().map(|size| size.size() == 0);
        let mut copy_block = |writer: &mut W, raw: &[u8]| -> Result<(), BackhandError> {
            if is_hole.next() != Some(true) {
                writer.write_all(raw)?;
            }
            Ok(())
        };

        copy_block(&mut writer, &read_buf)?;
        while let Some(block) = reader.next_block(&mut read_buf) {
            let block = block?;
            if block.fragment {
                reader.decompress(block, &mut read_buf, &mut decompress_buf)?;
                // TODO: support tail-end fragments, for now just treat it like a block
                let cb = self.compressor.compress(
                    &decompress_buf,
                    self.fs_compressor,
                    self.block_size,
                )?;
                // compression didn't reduce size
                if cb.len() > decompress_buf.len() {
                    // store uncompressed
                    block_sizes.push(DataSize::new_uncompressed(decompress_buf.len() as u32));
                    writer.write_all(&decompress_buf)?;
                } else {
                    // store compressed
                    block_sizes.push(DataSize::new_compressed(cb.len() as u32));
                    writer.write_all(&cb)?;
                }
            } else {
                copy_block(&mut writer, &read_buf)?;
            }
        }
        let file_size = source.file_len();
        Ok((file_size, Added::Data { blocks_start, block_sizes, sparse }))
    }

    /// Add to data writer, either a Data or Fragment
    ///
    /// If `self.dup_cache` is on, return alrady added `(usize, Added)` if duplicate
    /// is found
    // TODO: support tail-end fragments (off by default in squashfs-tools/mksquashfs)
    pub(crate) fn add_bytes<W: WriteSeek>(
        &mut self,
        reader: impl Read,
        mut writer: W,
    ) -> Result<(usize, Added), BackhandError> {
        let mut chunk_reader = DataWriterChunkReader {
            chunk: vec![0u8; self.block_size as usize],
            file_len: 0,
            reader,
        };

        // read entire chunk (file)
        let mut chunk = chunk_reader.read_chunk()?;

        // an empty file must carry no fragment reference: a zero-byte
        // fragment entry makes the kernel squashfs driver reject the inode
        // with EINVAL on stat/open (Added::Data with no blocks encodes
        // frag_index 0xffffffff, matching mksquashfs)
        if chunk.is_empty() {
            let blocks_start = writer.stream_position()?;
            return Ok((0, Added::Data { blocks_start, block_sizes: vec![], sparse: 0 }));
        }

        // chunk size not exactly the size of the block
        if chunk.len() != self.block_size as usize {
            // if this doesn't fit in the current fragment bytes
            // compress the current fragment bytes and add to data_bytes
            if (chunk.len() + self.fragment_bytes.len()) > self.block_size as usize {
                self.finalize(writer)?;
            }

            // add to fragment bytes
            let frag_index = self.fragment_table.len() as u32;
            let block_offset = self.fragment_bytes.len() as u32;
            self.fragment_bytes.write_all(chunk)?;

            return Ok((chunk_reader.file_len, Added::Fragment { frag_index, block_offset }));
        }

        // Add to data bytes
        let blocks_start = writer.stream_position()?;
        let mut block_sizes = vec![];

        let mut hasher = Xxh64::new(0);
        let mut sparse = 0;
        while !chunk.is_empty() {
            if chunk.iter().all(|byte| *byte == 0) {
                // like mksquashfs, store a block of zeros as a hole
                block_sizes.push(DataSize::hole());
                sparse += chunk.len() as u64;
            } else {
                let cb = self.compressor.compress(chunk, self.fs_compressor, self.block_size)?;

                // compression didn't reduce size
                if cb.len() > chunk.len() {
                    // store uncompressed
                    block_sizes.push(DataSize::new_uncompressed(chunk.len() as u32));
                    writer.write_all(chunk)?;
                } else {
                    // store compressed
                    block_sizes.push(DataSize::new_compressed(cb.len() as u32));
                    writer.write_all(&cb)?;
                }
            }
            hasher.update(chunk);
            chunk = chunk_reader.read_chunk()?;
        }

        let added = (chunk_reader.file_len, Added::Data { blocks_start, block_sizes, sparse });
        let Some(dup_cache) = &mut self.dup_cache else {
            return Ok(added);
        };

        // The hash is known only after the whole file is read. Thus, like mksquashfs, write the
        // data first, then seek back so that the next data writes over a duplicate.
        let files_with_len = dup_cache.entry(chunk_reader.file_len as u64).or_default();
        match files_with_len.entry(hasher.digest()) {
            Entry::Occupied(original) => {
                trace!("duplicate file data found");
                writer.seek(SeekFrom::Start(blocks_start))?;
                Ok(original.get().clone())
            }
            Entry::Vacant(entry) => Ok(entry.insert(added).clone()),
        }
    }

    /// Compress the fragments that were under length, write to data, add to fragment table, clear
    /// current fragment_bytes
    pub fn finalize<W: Write + Seek>(&mut self, mut writer: W) -> Result<(), BackhandError> {
        let start = writer.stream_position()?;
        let cb =
            self.compressor.compress(&self.fragment_bytes, self.fs_compressor, self.block_size)?;

        // compression didn't reduce size
        let size = if cb.len() > self.fragment_bytes.len() {
            // store uncompressed
            writer.write_all(&self.fragment_bytes)?;
            DataSize::new_uncompressed(self.fragment_bytes.len() as u32)
        } else {
            // store compressed
            writer.write_all(&cb)?;
            DataSize::new_compressed(cb.len() as u32)
        };
        self.fragment_table.push(Fragment::new(start, size, 0));
        self.fragment_bytes.clear();
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use super::*;
    use crate::{
        DEFAULT_BLOCK_SIZE,
        compression::{Compressor, DefaultCompressor},
    };

    #[test]
    #[cfg(feature = "gzip")]
    fn test_duplicate_check() {
        let mut data_writer = DataWriter::new(
            &DefaultCompressor,
            FilesystemCompressor::new(Compressor::Gzip, None).unwrap(),
            DEFAULT_BLOCK_SIZE,
            true,
        );
        let bytes = [0xff_u8; DEFAULT_BLOCK_SIZE as usize * 2];
        let mut writer = Cursor::new(vec![]);
        let added_1 = data_writer.add_bytes(&bytes[..], &mut writer).unwrap();
        let added_2 = data_writer.add_bytes(&bytes[..], &mut writer).unwrap();
        assert_eq!(added_1, added_2);
    }

    #[test]
    #[cfg(feature = "gzip")]
    fn test_no_duplicate_check() {
        let mut data_writer = DataWriter::new(
            &DefaultCompressor,
            FilesystemCompressor::new(Compressor::Gzip, None).unwrap(),
            DEFAULT_BLOCK_SIZE,
            false,
        );
        let bytes = [0xff_u8; DEFAULT_BLOCK_SIZE as usize * 2];
        let mut writer = Cursor::new(vec![]);
        let added_1 = data_writer.add_bytes(&bytes[..], &mut writer).unwrap();
        let added_2 = data_writer.add_bytes(&bytes[..], &mut writer).unwrap();
        assert_ne!(added_1, added_2);
    }
}