acadia/rust/lib/fs/ext2/src/ext2_driver.rs

use core::cmp::min;

use alloc::{collections::BTreeMap, string::String, vec::Vec};
use denali_client::{DenaliClient, DiskBlock, DiskReader, ReadRequest};
use mammoth::{cap::Capability, debug, mem::MemoryRegion, zion::ZError};
use yellowstone_yunq::DenaliInfo;

use crate::types::{BlockGroupDescriptor, DirEntry, Inode, Superblock};

pub struct FileInfo {
    pub inode: u32,
    pub name: String,
}

/// Ext2 Driver with the ability to read files and directories from the given disk.
///
/// Implementation based on the information available at
/// https://www.nongnu.org/ext2-doc/ext2.html
pub struct Ext2Driver {
    reader: DiskReader,
    superblock_region: MemoryRegion,
    bgdt_region: MemoryRegion,

    /// Cache of the memory regions for the inode tables available indexed by
    /// the block_group number.
    inode_table_map: Vec<Option<MemoryRegion>>,

    /// Cache of inode_num to memory capability.
    /// This is particularly important for directories so we
    /// don't iterate over the disk each time.
    inode_cache: BTreeMap<u32, Capability>,
}

impl Ext2Driver {
    pub fn new(denali_info: DenaliInfo) -> Self {
        let mut client = DenaliClient::new(Capability::take(denali_info.denali_endpoint));

        // Calculate the absolute offset and size of the superblock. It is located at
        // offset 1024 of the partition and is 1024 bytes long. (Mostly extra
        // reserved space).
        // Ref: https://www.nongnu.org/ext2-doc/ext2.html#def-superblock
        let abs_superblock_start = denali_info.lba_offset + 2;
        let abs_superblock_size = 2; // TODO: This assumes 512 bytes sectors.
        let superblock_region = MemoryRegion::from_cap(Capability::take(
            client
                .read(&ReadRequest {
                    device_id: denali_info.device_id,
                    block: DiskBlock {
                        lba: abs_superblock_start,
                        size: abs_superblock_size,
                    },
                })
                .unwrap()
                .memory,
        ))
        .unwrap();
        let superblock: &Superblock = superblock_region.as_ref();
        assert!(superblock.is_valid());

        let mut reader = DiskReader::new(
            client,
            denali_info.device_id,
            denali_info.lba_offset,
            superblock.sectors_per_block(),
        );

        let bgdt_region = MemoryRegion::from_cap(
            reader
                .read(superblock.bgdt_block_num(), superblock.bgdt_block_size())
                .unwrap(),
        )
        .unwrap();

        let mut inode_table_map = Vec::new();
        inode_table_map.resize_with(superblock.num_block_groups() as usize, || None);

        Self {
            reader,
            superblock_region,
            bgdt_region,
            inode_table_map,
            inode_cache: BTreeMap::new(),
        }
    }

    fn superblock(&self) -> &Superblock {
        self.superblock_region.as_ref()
    }

    fn bgdt(&self) -> &[BlockGroupDescriptor] {
        self.bgdt_region.slice()
    }

    /// Updates the cached inode tables to contain the inode table for
    /// a specific group.
    fn populate_inode_table_if_none(&mut self, block_group_num: usize) {
        if self.inode_table_map[block_group_num].is_none() {
            debug!(
                "Cache MISS on inode table for block_group {}",
                block_group_num
            );
            let inode_table = self.bgdt()[block_group_num].inode_table;
            self.inode_table_map[block_group_num] = Some(
                MemoryRegion::from_cap(
                    self.reader
                        .read(
                            inode_table as u64,
                            self.superblock().inode_table_block_size(),
                        )
                        .unwrap(),
                )
                .unwrap(),
            );
        } else {
            debug!(
                "Cache HIT  on inode table for block_group {}",
                block_group_num
            );
        }
    }

    pub fn get_inode(&mut self, inode_num: u32) -> Inode {
        // See the following for a description of finding an inode.
        // https://www.nongnu.org/ext2-doc/ext2.html#idm140660447281728
        let block_group_num = (inode_num - 1) / self.superblock().inodes_per_group;
        self.populate_inode_table_if_none(block_group_num as usize);
        let region = self.inode_table_map[block_group_num as usize]
            .as_ref()
            .unwrap();

        let local_index = (inode_num - 1) % self.superblock().inodes_per_group;
        let offset = self.superblock().inode_size() * local_index as u64;
        unsafe { region.raw_ptr_at_offset::<Inode>(offset).read().clone() }
    }

    fn get_blocks_from_single_indirect(&mut self, block_num: u64, num_blocks: usize) -> Vec<u32> {
        assert!(num_blocks <= 256);
        let single_indr_block_mem =
            MemoryRegion::from_cap(self.reader.read(block_num, 1).unwrap()).unwrap();

        single_indr_block_mem.slice()[..num_blocks].to_vec()
    }

    fn get_blocks_from_double_indirect(&mut self, block_num: u64, num_blocks: usize) -> Vec<u32> {
        assert!(num_blocks > 0 && num_blocks <= (256 * 256));
        let num_dbl_indr = ((num_blocks - 1) / 256) + 1;

        let dbl_indr_block_mem =
            MemoryRegion::from_cap(self.reader.read(block_num, 1).unwrap()).unwrap();

        let dbl_indr_blocks: &[u32] = &dbl_indr_block_mem.slice()[0..num_dbl_indr];

        let mut blocks_to_read = Vec::new();

        for (i, dbl_indr_block) in dbl_indr_blocks.iter().enumerate() {
            let num_blocks_in_single = min(num_blocks - (256 * i), 256);
            blocks_to_read.append(
                &mut self
                    .get_blocks_from_single_indirect(*dbl_indr_block as u64, num_blocks_in_single),
            );
        }

        blocks_to_read
    }

    fn run_len_compress_blocks(&self, blocks: Vec<u32>) -> Vec<DiskBlock> {
        let mut curr_block = DiskBlock {
            lba: blocks[0] as u64,
            size: 1,
        };

        let mut iter = blocks.into_iter();
        iter.next();

        let mut blocks = Vec::new();

        for block in iter {
            if block as u64 == (curr_block.lba + curr_block.size) {
                curr_block.size += 1;
            } else {
                blocks.push(curr_block.clone());

                curr_block.lba = block as u64;
                curr_block.size = 1;
            }
        }

        blocks.push(curr_block);

        blocks
    }

    fn read_inode(&mut self, _inode_num: u32, inode: Inode) -> Result<Capability, ZError> {
        // TODO: Cache this method using _inode_num
        // TODO: This assumes 512 byte sectors.
        let real_block_cnt = (inode.blocks as u64 - 1) / (self.superblock().block_size() / 512) + 1;
        if inode.block[14] != 0 {
            debug!("Can't handle triply indirect inodes yet.");
            return Err(ZError::UNIMPLEMENTED);
        }

        let mut blocks_to_read = Vec::new();

        for i in 0..min(12, real_block_cnt) {
            blocks_to_read.push(inode.block[i as usize])
        }

        // Singly indirect block.
        if inode.block[12] != 0 {
            let num_blocks = min(256, real_block_cnt - 12) as usize;
            blocks_to_read.append(
                &mut self.get_blocks_from_single_indirect(inode.block[12] as u64, num_blocks),
            );
        }

        // Doubly indirect block.
        if inode.block[13] != 0 {
            let num_blocks = min(256 * 256, real_block_cnt - 268) as usize;
            blocks_to_read.append(
                &mut self.get_blocks_from_double_indirect(inode.block[13] as u64, num_blocks),
            );
        };

        self.reader
            .read_many(&self.run_len_compress_blocks(blocks_to_read))
    }

    fn read_inode_into_mem(
        &mut self,
        inode_num: u32,
        inode: Inode,
    ) -> Result<MemoryRegion, ZError> {
        if !self.inode_cache.contains_key(&inode_num) {
            debug!("Cache MISS for inode_num: {}", inode_num);
            let inode_cap = self.read_inode(inode_num, inode)?;
            self.inode_cache.insert(inode_num, inode_cap);
        } else {
            debug!("Cache HIT  for inode_num: {}", inode_num);
        }

        MemoryRegion::from_cap(self.inode_cache[&inode_num].duplicate(Capability::PERMS_ALL)?)
    }

    pub fn read_file(&mut self, inode_num: u32) -> Result<Capability, ZError> {
        let inode = self.get_inode(inode_num);
        if (inode.mode & 0x8000) == 0 {
            debug!("Reading non file.");
            return Err(ZError::INVALID_ARGUMENT);
        }
        self.read_inode(inode_num, inode)
    }

    pub fn read_directory(&mut self, inode_num: u32) -> Result<Vec<FileInfo>, ZError> {
        let inode = self.get_inode(inode_num);
        if (inode.mode & 0x4000) == 0 {
            let mode = inode.mode;
            debug!("Reading non directory. Inode {:?}, Mode {}", inode, mode);
            return Err(ZError::INVALID_ARGUMENT);
        }

        let dir = self.read_inode_into_mem(inode_num, inode)?;

        let mut file_names = Vec::new();

        let mut offset = 0;
        while offset < dir.size() {
            let dir_ptr: DirEntry = unsafe { dir.raw_ptr_at_offset::<DirEntry>(offset).read() };

            let name = dir_ptr.name;
            let file_name: String =
                String::from_utf8(name[..dir_ptr.name_len as usize].to_vec()).unwrap();
            file_names.push(FileInfo {
                inode: dir_ptr.inode,
                name: file_name,
            });

            offset += dir_ptr.record_length as u64;
        }

        Ok(file_names)
    }
}