cramfs 0.1.0

use binrw::BinRead;
use flate2::read::ZlibDecoder;
use log::trace;

use crate::PAGE_SIZE;
use crate::constant::{
    BLK_DIRECT_PTR_SHIFT, BLK_FLAG_DIRECT_PTR, BLK_FLAG_UNCOMPRESSED, BLK_FLAGS,
    FLAG_EXT_BLOCK_POINTERS, S_ISBLK, S_ISCHR, S_ISFIFO, S_ISLNK, S_ISREG, S_ISSOCK,
};
use crate::error::Result;
use crate::sblk::print_node;
use crate::{Cramfs, constant::S_ISDIR, error::Error, sblk::INode};
use std::io::{Cursor, Write};
use std::path::Path;
use std::{
    io::{Read, Seek},
    path::PathBuf,
};

/// Iterator over each [`DirEntry`] in a directory inode.
pub struct ReadDir<R: Read + Seek> {
    fs: Cramfs<R>,
    count: u32,
    offset: u32,
    start_dir: u32,
    end_dir: u32,
}

impl<R: Read + Seek> ReadDir<R> {
    pub(crate) fn new(fs: Cramfs<R>, inode: INode) -> Result<Self> {
        let count = inode.size();
        let offset = inode.offset() << 2;
        let mut start_dir = 0;
        let end_dir = 0;

        if offset == 0 && count != 0 {
            return Err(("directory inode has zero offset and non-zero size").into());
        }
        if offset != 0 && offset < start_dir {
            start_dir = offset;
        }

        Ok(Self {
            fs,
            count,
            offset,
            start_dir,
            end_dir,
        })
    }
}

impl<R: Read + Seek> Iterator for ReadDir<R> {
    type Item = Result<DirEntry<R>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.count == 0 {
            return None;
        }

        Some(self._next())
    }
}

impl<R: Read + Seek> ReadDir<R> {
    fn _next(&mut self) -> Result<DirEntry<R>> {
        let mut child = Cursor::new(self.fs.romfs_read(self.offset as u64)?);
        let child = INode::read(&mut child)?;
        let new_len = child.namelen() << 2;
        if new_len == 0 {
            return Err("filename length is zero".into());
        }

        let size = INode::size_of_no_padding() as u32 + new_len;
        self.count -= size;

        self.offset += INode::size_of_no_padding() as u32;

        let newpath = &self.fs.romfs_read(self.offset as u64)?[..new_len as usize];
        let newpath = String::from_utf8_lossy(newpath).to_string();

        self.offset += new_len;

        if self.offset <= self.start_dir {
            return Err("bad inode offset".into());
        }
        if self.offset > self.end_dir {
            self.end_dir = self.offset;
        }

        Ok(DirEntry::new(
            self.fs.clone(),
            newpath.trim_end_matches('\0'),
            child,
        ))
    }
}

/// Directory entry
pub struct DirEntry<R: Read + Seek> {
    inode: INode,
    fs: Cramfs<R>,
    path: PathBuf,
    start_data: u64,
    end_data: u64,
    out_buffer: Vec<u8>,
}

macro_rules! try_into_num {
    ($n:ty, $t:expr) => {
        <$n>::from_le_bytes(($t)[..std::mem::size_of::<$n>()].try_into()?)
    };
}

macro_rules! to_u32 {
    ($t:expr) => {
        try_into_num!(u32, $t)
    };
}

macro_rules! to_u16 {
    ($t:expr) => {
        try_into_num!(u16, $t)
    };
}

impl<R: Read + Seek> DirEntry<R> {
    pub(crate) fn new(fs: Cramfs<R>, path: impl Into<PathBuf>, inode: INode) -> Self {
        Self {
            inode,
            path: path.into(),
            fs,
            start_data: !0,
            end_data: 0,
            out_buffer: Vec::with_capacity(PAGE_SIZE as usize * 2),
        }
    }

    /// Get the entry's path
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Check if the entry is a special inode
    pub fn is_special_inode(&self) -> bool {
        !(self.is_dir() || self.is_file() || self.is_symlink())
    }

    /// Check if the entry is a file
    pub fn is_file(&self) -> bool {
        S_ISREG(self.inode.mode())
    }

    /// Check if the entry is a symbolic link
    pub fn is_symlink(&self) -> bool {
        S_ISLNK(self.inode.mode())
    }

    /// Check if the entry is a directory
    pub fn is_dir(&self) -> bool {
        S_ISDIR(self.inode.mode())
    }

    /// Read the symbolic link target name
    pub fn read_symlink(&mut self) -> Result<String> {
        if !self.is_symlink() {
            return Err("inode is not symbolic link".into());
        }

        let offset = self.inode.offset() << 2;
        if offset == 0 {
            return Err("symbolic link has zero offset".into());
        }
        if self.inode.size() == 0 {
            return Err("symbolic link has zero size".into());
        }

        let size = self.read_block(offset as u64, 0, self.inode.size())?;
        if size != self.inode.size() {
            return Err(format!("size error in symlink: {}", self.path.display()).into());
        }
        let target = String::from_utf8_lossy(&self.out_buffer[..size as usize]);
        print_node(
            'l',
            &self.inode,
            format!("{} -> {}", self.path.display(), target),
        );

        Ok(target.to_string())
    }

    /// Read the dev_type for `mknod` function
    pub fn read_special_inode(&mut self) -> Result<u32> {
        if !self.is_special_inode() {
            return Err("inode is not special inode".into());
        }

        if self.inode.offset() != 0 {
            return Err(
                format!("special inode has non-zero offset: {}", self.path.display()).into(),
            );
        }

        let mode = self.inode.mode();
        let size = self.inode.size();

        let mut dev_type: u32 = 0;
        let r#type: char;
        if S_ISCHR(mode) {
            dev_type = size;
            r#type = 'c';
        } else if S_ISBLK(mode) {
            dev_type = size;
            r#type = 'b';
        } else if S_ISFIFO(mode) {
            if size != 0 {
                return Err(format!("fifo has non-zero size: {}", self.path.display()).into());
            }
            r#type = 'p';
        } else if S_ISSOCK(mode) {
            if size != 0 {
                return Err(format!("socket has non-zero size: {}", self.path.display()).into());
            }
            r#type = 's';
        } else {
            return Err(format!("bogus mode: {} ({:o})", self.path.display(), mode).into());
        }
        print_node(r#type, &self.inode, self.path.to_string_lossy());

        Ok(dev_type)
    }

    /// Get an iterator over the directory entry's children
    pub fn read_dir(&self) -> Result<ReadDir<R>> {
        if !self.is_dir() {
            return Err("inode is not directory".into());
        }

        print_node('d', &self.inode, self.path.to_string_lossy());
        ReadDir::new(self.fs.clone(), self.inode)
    }

    /// Extract the file contents
    pub fn read_file(&mut self, writer: &mut impl Write) -> Result<()> {
        if !self.is_file() {
            return Err("inode is not file".into());
        }

        let offset = (self.inode.offset() << 2) as u64;

        if offset == 0 && self.inode.size() != 0 {
            return Err("file inode has zero offset and non-zero size".into());
        }
        if self.inode.size() == 0 && offset != 0 {
            return Err("file inode has zero size and non-zero offset".into());
        }
        if offset != 0 && offset < self.start_data {
            self.start_data = offset;
        }

        print_node('f', &self.inode, self.path.to_string_lossy());
        if self.inode.size() > 0 {
            self.extract(offset, writer)?;
        }

        Ok(())
    }

    fn extract(&mut self, offset: u64, writer: &mut impl Write) -> Result<()> {
        let size = self.inode.size();
        let mut left = size as u64;
        let mut block_nr = 0u32;

        loop {
            let out = self.read_block(offset, block_nr, size)? as u64;
            if left >= PAGE_SIZE as u64 {
                if out != PAGE_SIZE as u64 {
                    return Err(format!("non-block ({}) bytes", out).into());
                }
            } else if out != left {
                return Err(format!("non-size ({} vs {}) bytes", out, left).into());
            }
            left = left.saturating_sub(out);
            writer
                .write_all(&self.out_buffer[..out as usize])
                .map_err(|e| -> Error {
                    format!("write failed, {}: {:?}", self.path.display(), e).into()
                })?;

            block_nr += 1;
            if left == 0 {
                break;
            }
        }

        Ok(())
    }

    fn read_block(&mut self, offset: u64, block_nr: u32, size: u32) -> Result<u32> {
        let blkptr_offset = offset + block_nr as u64 * 4;
        let max_block = size.div_ceil(PAGE_SIZE);

        if offset < self.start_data {
            self.start_data = offset;
        }

        let mut block_ptr = to_u32!(self.fs.romfs_read(blkptr_offset)?);
        if (block_ptr & BLK_FLAGS != 0)
            && (self.fs.super_block.flags & FLAG_EXT_BLOCK_POINTERS == 0)
        {
            return Err("block pointer extension usage not in super block".into());
        }
        let uncompressed = block_ptr & BLK_FLAG_UNCOMPRESSED;
        let direct = block_ptr & BLK_FLAG_DIRECT_PTR;
        block_ptr &= !BLK_FLAGS;

        let mut block_len: u32;
        let mut block_start: u32;
        if direct != 0 {
            // The block pointer is an absolute start pointer.
            // shifted by 2 bits. The size is included in the
            // first 2 bytes of the data block when compressed,
            // or PAGE_SIZE otherwise.
            block_start = block_ptr << BLK_DIRECT_PTR_SHIFT;
            if (block_start as u64) < self.start_data {
                self.start_data = block_start as u64;
            }

            if uncompressed != 0 {
                block_len = PAGE_SIZE;
                // if last block: cap to file length
                if block_nr == max_block - 1 {
                    block_len = size & PAGE_SIZE;
                }
            } else {
                block_len = to_u16!(self.fs.romfs_read(block_start as u64)?) as u32;
                block_start += 2;
            }
        } else {
            // The block pointer indicates one past the end of
            // the current block (start of next block). If this
            // is the first block then it starts where then block
            // pointer table ends, otherwise its start comes
            // from the previous block's poitner
            block_start = offset as u32 + max_block * 4;
            if block_nr != 0 {
                block_start = to_u32!(self.fs.romfs_read(blkptr_offset - 4)?);
            }
            // Beware... previous ptr might be a direct ptr
            if block_start & BLK_FLAG_DIRECT_PTR != 0 {
                let prev_start = block_start;
                block_start = prev_start & !BLK_FLAGS;
                block_start <<= BLK_DIRECT_PTR_SHIFT;
                if (block_start as u64) < self.start_data {
                    self.start_data = block_start as u64;
                }
                if prev_start & BLK_FLAG_UNCOMPRESSED != 0 {
                    block_start += PAGE_SIZE;
                } else {
                    block_len = to_u16!(self.fs.romfs_read(block_start as u64)?) as u32;
                    block_start += 2 + block_len;
                }
            }

            block_start &= !BLK_FLAGS;
            block_len = block_ptr - block_start;
        }

        if block_len > 2 * PAGE_SIZE || (uncompressed != 0 && block_len > PAGE_SIZE) {
            return Err(format!("block too large ({} bytes)", block_len).into());
        }
        if (block_start + block_len) as u64 > self.end_data {
            self.end_data = (block_start + block_len) as u64;
        }

        let mut out;
        if block_len == 0 {
            // hole
            out = PAGE_SIZE;
            // if last block: cap to file length
            if block_nr == max_block - 1 {
                out = size % PAGE_SIZE;
            }
            trace!("  hole at {} ({})", block_start, out);
            self.out_buffer.fill(0);
        } else if uncompressed != 0 {
            trace!(
                "  non-compressed {}block at {} to {} ({})",
                if direct != 0 { "direct " } else { "" },
                block_start,
                block_start + block_len,
                block_len
            );

            self.out_buffer
                .copy_from_slice(&self.fs.romfs_read(block_start as u64)?[..block_len as usize]);
            out = block_len;
        } else {
            trace!(
                "  uncompressing {}block at {} to {} ({})",
                if direct != 0 { "direct " } else { "" },
                block_start,
                block_start + block_len,
                block_len
            );
            out = self.uncompress_block(&self.fs.romfs_read(block_start as u64)?, block_len)?
        }

        Ok(out)
    }

    fn uncompress_block(&mut self, src: &[u8], len: u32) -> Result<u32> {
        if len > PAGE_SIZE * 2 {
            return Err("data block too large".into());
        }

        let mut decoder = ZlibDecoder::new(&src[..len as usize]);
        self.out_buffer.clear();
        decoder
            .read_to_end(&mut self.out_buffer)
            .map_err(|e| format!("decompression error ({}): {:?}", len, e).into())
            .map(|u| u as u32)
    }
}