path: root/src/main.rs

mod block_device;
mod cached_inode;
mod disk;

use and_then_some::BoolExt;
use clap::Parser;
use fuser::{
    Filesystem, KernelConfig, MountOption, ReplyAttr, ReplyData, ReplyDirectory, ReplyEmpty,
    ReplyEntry, ReplyLseek, ReplyWrite, Request, TimeOrNow,
};
use log::debug;
use lru::LruCache;
use std::ffi::OsStr;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::cached_inode::{convert, convert_mut, CachedBlock, CachedInode};
use crate::disk::data_block::{
    Block, DataBlock, DoubleIndirectBlock, IndirectBlock, InodeBlock, TripleIndirectBlock,
};
use crate::disk::inode::InodeMode;
use block_device::{memory_disk::MemoryDisk, BlockDevice, BLOCK_SIZE};
use disk::bitmap::Bitmap;
use disk::inode::{Inode, INODE_SIZE};
use libc::{c_int, ENOENT, ENOSPC, ENOSYS};
use users::{get_current_gid, get_current_uid};

#[derive(Parser, Debug)]
#[command(author, version, about)]
struct Args {
    mount_point: Option<String>,
    #[arg(long)]
    auto_unmount: bool,
    #[arg(long)]
    allow_root: bool,
}

const TTL: Duration = Duration::from_secs(1);
const INODE_PER_BLOCK: usize = BLOCK_SIZE / INODE_SIZE;

/// The design of MyFS is rather simple:
/// +-------------------+
/// | Super Block       |
/// +-------------------+
/// | Inode Bitmap      |
/// +-------------------+
/// | ...               |
/// +-------------------+
/// | Data Block Bitmap |
/// +-------------------+
/// | ...               |
/// +-------------------+
/// | Inode Block       |
/// +-------------------+
/// | ...               |
/// +-------------------+
/// | Data Block        |
/// +-------------------+
/// With each block 4KiB, each Inode entry 128B

#[repr(C)]
struct MyFS {
    device: Arc<dyn BlockDevice>,
    data_bitmap: Bitmap,
    inode_bitmap: Bitmap,
    inode_start_block: usize,
    data_start_block: usize,

    cached_inodes: LruCache<usize, CachedInode>,
    cached_blocks: LruCache<usize, CachedBlock<DataBlock>>,
    // cached_block_cleanness: LruCache<usize, bool>,
}

impl MyFS {
    fn new(device: Arc<dyn BlockDevice>, total_block_number: usize) -> Self {
        let max_inode_number: usize = 16384; // TODO: remove hard-coded magic number
        let inode_block_number = max_inode_number / INODE_PER_BLOCK; // == 128
        let inode_bitmap_block_number = (inode_block_number + BLOCK_SIZE - 1) / BLOCK_SIZE;

        let blocks_remaining =
            total_block_number - inode_block_number - inode_bitmap_block_number - 1;
        // let number of data blocks be x, the remaining block number be C,
        // the corresponding data bitmap length should be ceil(x / BLK_SIZE),
        // thus we have BLK_SIZE * (C-1) / (BLK_SIZE+1) <= x <= BLK_SIZE * C / (BLK_SIZE+1)
        // the difference of the two bounds is less than 1,
        // meaning only 1 integer could be in between.
        // Thus we have x = floor(BLK_SIZE * C / (BLK_SIZE + 1))
        let data_block_number = BLOCK_SIZE * blocks_remaining / (BLOCK_SIZE + 1);
        let data_bitmap_block_number = blocks_remaining - data_block_number;
        debug!("data_bitmap_block_number: {}", data_bitmap_block_number);
        debug!("inode_bitmap_block_number: {}", inode_bitmap_block_number);
        debug!("inode_block_number:  {}", inode_block_number);
        debug!("data_block_number:  {}", data_block_number);
        debug!(
            "sum:  {}",
            1 + data_bitmap_block_number
                + inode_bitmap_block_number
                + inode_block_number
                + data_block_number
        );

        let mut data_bitmap = Bitmap::new(1, data_bitmap_block_number, device.clone());
        let _ = data_bitmap.allocate().unwrap(); // data block 0 is not usable
        let mut inode_bitmap = Bitmap::new(
            data_bitmap_block_number + 1,
            inode_bitmap_block_number,
            device.clone(),
        );
        let _ = inode_bitmap.allocate().unwrap(); // inode block 0 is not usable

        let mut fs = Self {
            device,
            data_bitmap,
            inode_bitmap,
            inode_start_block: data_bitmap_block_number + inode_bitmap_block_number + 1,
            data_start_block: data_bitmap_block_number
                + inode_bitmap_block_number
                + inode_block_number
                + 1,
            cached_inodes: LruCache::new(NonZeroUsize::new(256).unwrap()),
            cached_blocks: LruCache::new(NonZeroUsize::new(256).unwrap()),
        };

        fs.create_inode_2(
            0o755,
            InodeMode::IFDIR,
            get_current_uid(),
            get_current_gid(),
            0,
        );

        fs
    }

    fn time_now() -> u32 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("How can current time be earlier than UNIX_EPOCH?")
            .as_secs() as u32
    }

    pub fn create_inode(
        &mut self,
        permissions: u16,
        mode: InodeMode,
        uid: u32,
        gid: u32,
        flags: u32,
    ) -> Option<usize> {
        self.inode_bitmap.allocate().map(|inode_index| {
            let inode = CachedInode {
                inode: Inode::make_inode(
                    permissions,
                    mode,
                    uid,
                    gid,
                    Self::time_now(),
                    flags,
                    0,
                    0,
                    0,
                ),
                index: inode_index,
                dirty: false,
            };
            self.cached_inodes.put(inode_index, inode); // TODO write back evicted inode
            inode_index
        })
    }

    pub fn create_inode_2(
        &mut self,
        permissions: u16,
        mode: InodeMode,
        uid: u32,
        gid: u32,
        flags: u32,
    ) -> Option<usize> {
        self.inode_bitmap.allocate().map(|inode_index| {
            self.get_inode_mut_2(inode_index).map(|inode| {
                *inode = Inode::make_inode(
                    permissions,
                    mode,
                    uid,
                    gid,
                    Self::time_now(),
                    flags,
                    0,
                    0,
                    0,
                );
            });
            inode_index
        })
    }

    fn update_block<T: Block>(&mut self, block: CachedBlock<T>) -> bool {
        if self.cached_blocks.contains(&block.index) {
            let data_block = convert::<T, DataBlock>(&block).clone();
            self.cached_blocks.push(block.index, data_block);
            true
        } else {
            false
        }
    }

    fn get_block<T: Block>(&mut self, index: usize) -> Option<&CachedBlock<T>> {
        self.load_block(index)
            .and_then(|| self.cached_blocks.get(&index).map(convert::<DataBlock, T>))
    }

    fn get_block_mut<T: Block>(&mut self, index: usize) -> Option<&mut CachedBlock<T>> {
        self.load_block(index).and_then(|| {
            self.cached_blocks
                .get_mut(&index)
                .map(convert_mut::<DataBlock, T>)
        })
    }

    fn load_block(&mut self, index: usize) -> bool {
        // 反正我自己保证所有实现了 Block trait 的数据结构都是 4K bytes 长, 来回 cast 没问题
        // 如果 block cache 里没有这个 block
        if self.cached_blocks.contains(&index) == false {
            if self.data_bitmap.query(index) == false {
                return false;
            } // 先看这个 block 是不是 valid, 不 valid 直接返回 None
            let block = DataBlock::default();
            let buffer = unsafe {
                std::slice::from_raw_parts_mut(&block as *const DataBlock as *mut u8, BLOCK_SIZE)
            };
            self.device.read(index, buffer);
            let cached_block = CachedBlock {
                block,
                index,
                dirty: false,
            };
            if let Some((old_index, old_block)) = self.cached_blocks.push(index, cached_block) {
                assert_ne!(old_index, index); // 只有 block 不在 cache 里的时候才会插入
                if old_block.dirty {
                    let old_block_ptr = &old_block.block as *const DataBlock as *mut u8;
                    let old_block_buffer =
                        unsafe { std::slice::from_raw_parts(old_block_ptr, BLOCK_SIZE) };
                    self.device.write(old_index, old_block_buffer);
                }
            } // 如果 valid 就放到 block cache 里, 同时将（可能）被驱逐的 block 依据其是否为脏块进行写回
        }

        true
    }

    fn get_inode_2(&mut self, inode_index: usize) -> Option<&Inode> {
        self.inode_bitmap.query(inode_index).and_then(|| {
            let (block_index, offset) = self.locate_inode(inode_index);
            self.get_block::<InodeBlock>(block_index)
                .map(|cached_block| &cached_block.block.inodes[offset / INODE_SIZE])
        })
    }

    fn get_inode_mut_2(&mut self, inode_index: usize) -> Option<&mut Inode> {
        self.inode_bitmap.query(inode_index).and_then(|| {
            let (block_index, offset) = self.locate_inode(inode_index);
            self.get_block_mut::<InodeBlock>(block_index)
                .map(|cached_block| {
                    cached_block.dirty = true; // 保守一些, 只要返回了 &mut Inode 这个页一定标记为脏
                    &mut cached_block.block.inodes[offset / INODE_SIZE]
                })
        })
    }

    /// 不使用 inode cache, 只用 block cache 的方案
    fn load_inode_2(&mut self, inode_index: usize) -> bool {
        // 首先保证 inode 有效, 无效返回 false
        if !self.inode_bitmap.query(inode_index) {
            return false;
        }
        // 找到 inode
        let (block_index, _) = self.locate_inode(inode_index);
        if self.cached_blocks.contains(&block_index) {
            return true;
        }

        return self.load_block(block_index);
    }

    fn get_inode(&mut self, inode_index: usize) -> Option<&CachedInode> {
        self.load_inode(inode_index)
            .and_then(|| self.cached_inodes.get(&inode_index))
    }

    fn get_inode_mut(&mut self, inode_index: usize) -> Option<&mut CachedInode> {
        self.load_inode(inode_index)
            .and_then(|| self.cached_inodes.get_mut(&inode_index))
    }

    fn load_inode(&mut self, inode_index: usize) -> bool {
        // 首先保证 inode 有效, 无效返回 false
        if !self.inode_bitmap.query(inode_index) {
            return false;
        }
        // 如果已经缓存了就返回 true
        if self.cached_inodes.contains(&inode_index) {
            return true;
        }
        // inode 有效但没有在 cache 里
        let (block_index, offset) = self.locate_inode(inode_index);
        // 找到 inode 所在的 block
        if let Some(inode_block) = self.get_block::<DataBlock>(block_index) {
            // 创建新的 inode 并且从 block 中加载.
            let inode = unsafe { std::mem::zeroed::<Inode>() };
            let inode_slice = unsafe {
                std::slice::from_raw_parts_mut(&inode as *const Inode as *mut u8, INODE_SIZE)
            };
            inode_slice.copy_from_slice(&inode_block.block.0[offset..offset + INODE_SIZE]);

            let cached_inode = CachedInode {
                inode,
                index: inode_index,
                dirty: false,
            };

            // 将加载好的 inode 放进 inode cache
            if let Some((old_index, old_inode)) = self.cached_inodes.push(inode_index, cached_inode)
            {
                assert_ne!(old_index, inode_index); // 如果 old_index == inode index 说明它原本就在 inode cache 里了.
                if old_inode.dirty {
                    // 如果旧的 inode 被修改过了, 将它写回 block cache 并且标记对应 block 为脏
                    let (old_block_index, old_offset) = self.locate_inode(old_index);
                    self.get_block_mut::<DataBlock>(old_block_index)
                        .map(|cached_block| {
                            let old_inode_ptr = &old_inode.inode as *const Inode as *const u8;
                            let old_inode_slice =
                                unsafe { std::slice::from_raw_parts(old_inode_ptr, INODE_SIZE) };
                            cached_block.dirty = true;
                            cached_block.block.0[old_offset..old_offset + INODE_SIZE]
                                .copy_from_slice(old_inode_slice);
                        })
                        .expect("Writing inode back to a invalid block!"); // Debug use
                }
            }
        } else {
            panic!("Getting inode from a invalid block!"); // Debug use
        }

        true
    }

    /// 输入 inode 编号, 返回它对应的 block number 和 block 内 offset
    pub fn locate_inode(&self, inode_index: usize) -> (usize, usize) {
        let block_number =
            inode_index / INODE_PER_BLOCK + 1 + self.inode_bitmap.length + self.data_bitmap.length;
        let block_offset = inode_index % INODE_PER_BLOCK * INODE_SIZE;
        (block_number, block_offset)
    }

    /// 为 Inode 分配新 block, 返回 block 的编号
    pub fn allocate_block(&mut self, inode: &mut Inode) -> Option<u32> {
        // 先看这个 inode 的 direct block 有没有空闲
        for index in inode.direct.iter_mut() {
            if !self.data_bitmap.query(*index as usize) {
                let block_index = self
                    .data_bitmap
                    .allocate()
                    .unwrap() as u32;
                *index = block_index;
                inode.n_blocks += 1;
                // 当调用 get_inode_mut 拿出 &mut Inode 的时候对应的 block 在 cache 里已经脏了
                return Some(block_index);
            }
        }

        // direct block 全部分配完了, 先检查 indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.single_indirect as usize) {
            inode.single_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
        }
        // 在 indirect block 里尝试分配
        if let Some(block_index) = self.allocate_in_indirect(inode.single_indirect) {
            inode.n_blocks += 1;
            return Some(block_index);
        }

        // direct & indirect block 全部分配完了, 先检查 double indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.double_indirect as usize) {
            inode.double_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
        }
        // 在 double indirect block 里尝试分配
        if let Some(block_index) = self.alloc_in_double_indirect(inode.double_indirect) {
            inode.n_blocks += 1;
            return Some(block_index);
        }

        // direct, indirect & double indirect block 全部分配完了, 先检查 triple indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.triple_indirect as usize) {
            inode.triple_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
        }
        // 在 double indirect block 里尝试分配
        if let Some(block_index) = self.alloc_in_triple_indirect(inode.triple_indirect) {
            inode.n_blocks += 1;
            return Some(block_index);
        }
        None
    }

    fn allocate_in_indirect(&mut self, indirect_entry: u32) -> Option<u32> {
        // 取出 single indirect block, 尝试在里面分配
        let indirect_entry = indirect_entry as usize;

        if let Some(block) = self
            .get_block(indirect_entry)
            .map(convert::<DataBlock, IndirectBlock>)
        {
            let mut indirect_block = block.clone();
            for entry in indirect_block.block.entries.iter_mut() {
                if self.data_bitmap.query(*entry as usize) == false {
                    indirect_block.dirty = true; // 把这个块标记为 dirty
                    let block_index = self.data_bitmap.allocate().expect("No free space") as u32;
                    *entry = block_index;
                    self.update_block(indirect_block);
                    return Some(block_index);
                }
            }
        }
        None
    }

    fn alloc_in_double_indirect(&mut self, double_indirect_entry: u32) -> Option<u32> {
        let double_indirect_entry = double_indirect_entry as usize;

        if let Some(block) = self
            .get_block(double_indirect_entry)
            .map(convert::<DataBlock, DoubleIndirectBlock>)
        {
            let mut double_indirect_block = block.clone();
            for indirect_entry in double_indirect_block.block.indirect.iter_mut() {
                if self.data_bitmap.query(*indirect_entry as usize) == false {
                    double_indirect_block.dirty = true;
                    *indirect_entry = self.data_bitmap.allocate().expect("No free space") as u32;
                }

                if let Some(block_index) = self.allocate_in_indirect(*indirect_entry) {
                    if double_indirect_block.dirty {
                        self.update_block(double_indirect_block);
                    }
                    return Some(block_index);
                }
            }
        }
        None
    }

    fn alloc_in_triple_indirect(&mut self, triple_indirect_entry: u32) -> Option<u32> {
        let triple_indirect_entry = triple_indirect_entry as usize;

        if let Some(block) = self
            .get_block(triple_indirect_entry)
            .map(convert::<DataBlock, TripleIndirectBlock>)
        {
            let mut triple_indirect_block = block.clone();
            for double_indirect_entry in triple_indirect_block.block.double_indirect.iter_mut() {
                if self.data_bitmap.query(*double_indirect_entry as usize) == false {
                    triple_indirect_block.dirty = true;
                    *double_indirect_entry =
                        self.data_bitmap.allocate().expect("No free space") as u32;
                }
                if let Some(block_index) = self.alloc_in_double_indirect(*double_indirect_entry) {
                    if triple_indirect_block.dirty {
                        self.update_block(triple_indirect_block);
                    }
                    return Some(block_index);
                }
            }
        }
        None
    }
}

impl Filesystem for MyFS {
    fn init(&mut self, _req: &Request<'_>, _config: &mut KernelConfig) -> Result<(), c_int> {
        debug!("Filesystem::init called.");
        Ok(())
    }

    fn destroy(&mut self) {
        debug!("Filesystem::destroy()");
    }

    fn lookup(&mut self, _req: &Request<'_>, parent: u64, name: &OsStr, reply: ReplyEntry) {
        debug!(
            "Filesystem::lookup called with parent {} name {}",
            parent,
            name.to_str().unwrap()
        );
        let parent = parent as usize;
        if let Some(inode) = self.get_inode_2(parent) {
            // debug!("{:?}", inode);
        }
        // if self.inode_active(parent) {
        //     let (block, offset) = self.locate_inode(parent);
        //     let inode = self.get_inode_2(block, offset);
        //     debug!("{:?}", inode);
        // }
        reply.error(ENOENT);
    }

    fn forget(&mut self, _req: &Request<'_>, _ino: u64, _nlookup: u64) {
        debug!("Filesystem::forget()");
        todo!("This is a dumb implementation")
    }

    fn getattr(&mut self, _req: &Request<'_>, ino: u64, reply: ReplyAttr) {
        debug!("Filesystem::getattr(ino: {})", ino);
        let ino = ino as usize;
        if let Some(inode) = self.get_inode_2(ino) {
            // debug!("{:?}", inode);
        }
        reply.error(ENOENT);
    }

    fn setattr(
        &mut self,
        _req: &Request<'_>,
        ino: u64,
        mode: Option<u32>,
        uid: Option<u32>,
        gid: Option<u32>,
        size: Option<u64>,
        _atime: Option<TimeOrNow>,
        _mtime: Option<TimeOrNow>,
        _ctime: Option<SystemTime>,
        fh: Option<u64>,
        _crtime: Option<SystemTime>,
        _chgtime: Option<SystemTime>,
        _bkuptime: Option<SystemTime>,
        flags: Option<u32>,
        reply: ReplyAttr,
    ) {
        debug!(
            "Filesystem::setattr(ino: {:#x?}, mode: {:?}, uid: {:?}, \
            gid: {:?}, size: {:?}, fh: {:?}, flags: {:?})",
            ino, mode, uid, gid, size, fh, flags
        );
        reply.error(ENOSYS);
    }

    fn readlink(&mut self, _req: &Request<'_>, ino: u64, reply: ReplyData) {
        debug!("[Not Implemented] readlink(ino: {})", ino);
        reply.error(ENOSYS);
    }

    fn mknod(
        &mut self,
        _req: &Request<'_>,
        parent: u64,
        name: &OsStr,
        mode: u32,
        umask: u32,
        rdev: u32,
        reply: ReplyEntry,
    ) {
        debug!(
            "Filesystem::mknod(parent: {}, name: {:?}, mode: {}, umask: {}, rdev: {})",
            parent, name, mode, umask, rdev
        );
        reply.error(ENOSPC);
    }

    fn mkdir(
        &mut self,
        _req: &Request<'_>,
        parent: u64,
        name: &OsStr,
        mode: u32,
        umask: u32,
        reply: ReplyEntry,
    ) {
        debug!(
            "Filesystem::mkdir(parent: {}, name: {:?}, mode: {}, umask: {})",
            parent, name, mode, umask
        );
        if let Some(inode) = self.get_inode_2(parent as usize) {
        } else {
            reply.error(ENOENT);
        }
        // reply.error(ENOSPC);
    }

    fn read(
        &mut self,
        _req: &Request<'_>,
        ino: u64,
        _fh: u64,
        offset: i64,
        _size: u32,
        _flags: i32,
        _lock_owner: Option<u64>,
        reply: ReplyData,
    ) {
        todo!()
    }

    fn readdir(
        &mut self,
        _req: &Request<'_>,
        ino: u64,
        _fh: u64,
        offset: i64,
        mut reply: ReplyDirectory,
    ) {
        todo!()
    }

    fn access(&mut self, _req: &Request<'_>, ino: u64, mask: i32, reply: ReplyEmpty) {
        debug!("Filesystem::getattr(ino: {}, mask: {})", ino, mask);
        if let Some(inode) = self.get_inode_2(ino as usize) {
            reply.ok()
        } else {
            reply.error(ENOENT)
        }
    }

    fn lseek(
        &mut self,
        _req: &Request<'_>,
        ino: u64,
        fh: u64,
        offset: i64,
        whence: i32,
        reply: ReplyLseek,
    ) {
        debug!(
            "lseek(ino: {:#x?}, fh: {}, offset: {}, whence: {})",
            ino, fh, offset, whence
        );
        reply.error(ENOSYS);
    }

    fn copy_file_range(
        &mut self,
        _req: &Request<'_>,
        ino_in: u64,
        fh_in: u64,
        offset_in: i64,
        ino_out: u64,
        fh_out: u64,
        offset_out: i64,
        len: u64,
        flags: u32,
        reply: ReplyWrite,
    ) {
        debug!(
            "copy_file_range(ino_in: {:#x?}, fh_in: {}, \
            offset_in: {}, ino_out: {:#x?}, fh_out: {}, offset_out: {}, \
            len: {}, flags: {})",
            ino_in, fh_in, offset_in, ino_out, fh_out, offset_out, len, flags
        );
        reply.error(ENOSYS);
    }
}

fn main() {
    env_logger::init();
    let args = Args::parse();
    let mount_point = args.mount_point.unwrap();
    let options = vec![
        // MountOption::RO,
        MountOption::FSName("hello".to_string()),
        MountOption::AutoUnmount,
        MountOption::AllowRoot,
    ];
    let mem_disk = Arc::new(MemoryDisk::new());
    let filesystem = MyFS::new(mem_disk, 16384);

    fuser::mount2(filesystem, mount_point, &options).unwrap();
}