k-os/crates/servers/storage-server/src/fs/fat16.rs

use crate::filesystem::{DirEntry, FileHandle, FileName, FileSystem};
use port_block_device::{BlockDevice, SectorId};

const SECTOR_SIZE: usize = 512;
const MAX_OPEN_FILES: usize = 8;

struct Bpb {
    bytes_per_sector: u16,
    sectors_per_cluster: u8,
    reserved_sectors: u16,
    num_fats: u8,
    root_entry_count: u16,
    total_sectors: u16,
    sectors_per_fat: u16,
}

impl Bpb {
    /// Reads the BPB, refusing geometry the rest of this module cannot
    /// work with: the disk is untrusted input like any client request.
    fn parse(sector: &[u8]) -> Option<Self> {
        if sector.len() < SECTOR_SIZE {
            return None;
        }
        if sector[510] != 0x55 || sector[511] != 0xAA {
            return None;
        }
        let bpb = Self {
            bytes_per_sector: u16::from_le_bytes([sector[11], sector[12]]),
            sectors_per_cluster: sector[13],
            reserved_sectors: u16::from_le_bytes([sector[14], sector[15]]),
            num_fats: sector[16],
            root_entry_count: u16::from_le_bytes([sector[17], sector[18]]),
            total_sectors: u16::from_le_bytes([sector[19], sector[20]]),
            sectors_per_fat: u16::from_le_bytes([sector[22], sector[23]]),
        };
        let usable = bpb.bytes_per_sector as usize == SECTOR_SIZE
            && bpb.sectors_per_cluster.is_power_of_two()
            && bpb.num_fats > 0
            && bpb.sectors_per_fat > 0;
        usable.then_some(bpb)
    }

    fn fat_start(&self) -> u32 {
        self.reserved_sectors as u32
    }

    fn root_dir_start(&self) -> u32 {
        self.fat_start() + (self.num_fats as u32) * (self.sectors_per_fat as u32)
    }

    fn root_dir_sectors(&self) -> u32 {
        ((self.root_entry_count as u32) * 32 + (self.bytes_per_sector as u32) - 1)
            / (self.bytes_per_sector as u32)
    }

    fn data_start(&self) -> u32 {
        self.root_dir_start() + self.root_dir_sectors()
    }

    /// The first sector of data cluster `cluster`. Clusters 0 and 1 are
    /// reserved and hold no data.
    fn cluster_to_sector(&self, cluster: u16) -> Option<u32> {
        let index = (cluster as u32).checked_sub(2)?;
        Some(self.data_start() + index * (self.sectors_per_cluster as u32))
    }

    fn cluster_size(&self) -> usize {
        self.sectors_per_cluster as usize * SECTOR_SIZE
    }
}

#[derive(Clone, Copy)]
struct OpenFile {
    first_cluster: u16,
    size: usize,
    offset: usize,
    name: FileName,
}

pub struct Fat16<B: BlockDevice> {
    device: B,
    bpb: Bpb,
    open_files: [Option<OpenFile>; MAX_OPEN_FILES],
}

impl<B: BlockDevice> Fat16<B> {
    pub fn mount(mut device: B) -> Option<Self> {
        let mut buf = [0u8; SECTOR_SIZE];
        if !device.read_sector(SectorId::new(0), &mut buf) {
            return None;
        }
        let bpb = Bpb::parse(&buf)?;
        Some(Self {
            device,
            bpb,
            open_files: [const { None }; MAX_OPEN_FILES],
        })
    }

    fn read_fat_entry(&mut self, cluster: u16) -> u16 {
        let fat_offset = (cluster as u32) * 2;
        let sector = self.bpb.fat_start() + fat_offset / (SECTOR_SIZE as u32);
        let offset = (fat_offset % (SECTOR_SIZE as u32)) as usize;
        let mut buf = [0u8; SECTOR_SIZE];
        if !self.device.read_sector(SectorId::new(sector), &mut buf) {
            return 0xFFFF;
        }
        u16::from_le_bytes([buf[offset], buf[offset + 1]])
    }

    fn write_fat_entry(&mut self, cluster: u16, value: u16) {
        let fat_offset = (cluster as u32) * 2;
        let sector = self.bpb.fat_start() + fat_offset / (SECTOR_SIZE as u32);
        let offset = (fat_offset % (SECTOR_SIZE as u32)) as usize;
        let mut buf = [0u8; SECTOR_SIZE];
        self.device.read_sector(SectorId::new(sector), &mut buf);
        buf[offset] = value as u8;
        buf[offset + 1] = (value >> 8) as u8;
        self.device.write_sector(SectorId::new(sector), &buf);
    }

    fn alloc_cluster(&mut self) -> Option<u16> {
        let max = self.bpb.total_sectors as u16 / self.bpb.sectors_per_cluster as u16;
        for i in 2..max {
            if self.read_fat_entry(i) == 0 {
                self.write_fat_entry(i, 0xFFFF);
                return Some(i);
            }
        }
        None
    }

    fn free_chain(&mut self, start: u16) {
        let mut cluster = start;
        loop {
            if cluster < 2 || cluster >= 0xFFF8 {
                break;
            }
            let next = self.read_fat_entry(cluster);
            self.write_fat_entry(cluster, 0);
            cluster = next;
        }
    }

    fn find_in_root_dir(&mut self, name: &[u8]) -> Option<(u16, usize)> {
        let mut buf = [0u8; SECTOR_SIZE];
        let fat_name = to_fat_name(name);
        for s in 0..self.bpb.root_dir_sectors() {
            if !self
                .device
                .read_sector(SectorId::new(self.bpb.root_dir_start() + s), &mut buf)
            {
                continue;
            }
            for i in 0..(SECTOR_SIZE / 32) {
                let off = i * 32;
                if buf[off] == 0x00 || buf[off] == 0xE5 {
                    continue;
                }
                if buf[off..off + 11] == fat_name {
                    let cluster = u16::from_le_bytes([buf[off + 26], buf[off + 27]]);
                    let size = u32::from_le_bytes([
                        buf[off + 28],
                        buf[off + 29],
                        buf[off + 30],
                        buf[off + 31],
                    ]);
                    return Some((cluster, size as usize));
                }
            }
        }
        None
    }

    fn create_dir_entry(&mut self, name: &[u8], cluster: u16) -> bool {
        let mut buf = [0u8; SECTOR_SIZE];
        let fat_name = to_fat_name(name);
        for s in 0..self.bpb.root_dir_sectors() {
            let sector = self.bpb.root_dir_start() + s;
            if !self.device.read_sector(SectorId::new(sector), &mut buf) {
                continue;
            }
            for i in 0..(SECTOR_SIZE / 32) {
                let off = i * 32;
                if buf[off] == 0x00 || buf[off] == 0xE5 {
                    buf[off..off + 11].copy_from_slice(&fat_name);
                    buf[off + 11] = 0x20;
                    for b in &mut buf[off + 12..off + 26] {
                        *b = 0;
                    }
                    buf[off + 26] = cluster as u8;
                    buf[off + 27] = (cluster >> 8) as u8;
                    for b in &mut buf[off + 28..off + 32] {
                        *b = 0;
                    }
                    self.device.write_sector(SectorId::new(sector), &buf);
                    return true;
                }
            }
        }
        false
    }

    fn delete_dir_entry(&mut self, name: &[u8]) -> Option<u16> {
        let mut buf = [0u8; SECTOR_SIZE];
        let fat_name = to_fat_name(name);
        for s in 0..self.bpb.root_dir_sectors() {
            let sector = self.bpb.root_dir_start() + s;
            if !self.device.read_sector(SectorId::new(sector), &mut buf) {
                continue;
            }
            for i in 0..(SECTOR_SIZE / 32) {
                let off = i * 32;
                if buf[off..off + 11] == fat_name {
                    let cluster = u16::from_le_bytes([buf[off + 26], buf[off + 27]]);
                    buf[off] = 0xE5;
                    self.device.write_sector(SectorId::new(sector), &buf);
                    return Some(cluster);
                }
            }
        }
        None
    }

    fn update_dir_size(&mut self, name: &[u8], size: u32, cluster: u16) {
        let mut buf = [0u8; SECTOR_SIZE];
        let fat_name = to_fat_name(name);
        for s in 0..self.bpb.root_dir_sectors() {
            let sector = self.bpb.root_dir_start() + s;
            if !self.device.read_sector(SectorId::new(sector), &mut buf) {
                continue;
            }
            for i in 0..(SECTOR_SIZE / 32) {
                let off = i * 32;
                if buf[off..off + 11] == fat_name {
                    buf[off + 28..off + 32].copy_from_slice(&size.to_le_bytes());
                    buf[off + 26] = cluster as u8;
                    buf[off + 27] = (cluster >> 8) as u8;
                    self.device.write_sector(SectorId::new(sector), &buf);
                    return;
                }
            }
        }
    }

    fn follow_chain(&mut self, start: u16, skip_clusters: usize) -> u16 {
        let mut cluster = start;
        for _ in 0..skip_clusters {
            let next = self.read_fat_entry(cluster);
            if next >= 0xFFF8 {
                return cluster;
            }
            cluster = next;
        }
        cluster
    }

    /// The open-file slot `handle` names; `None` for handles past the
    /// table, which come straight from clients.
    fn slot(&mut self, handle: FileHandle) -> Option<&mut Option<OpenFile>> {
        self.open_files.get_mut(handle.raw() as usize)
    }

    fn alloc_handle(&self) -> Option<FileHandle> {
        for i in 0..MAX_OPEN_FILES {
            if self.open_files[i].is_none() {
                return Some(FileHandle::new(i as u16));
            }
        }
        None
    }
}

impl<B: BlockDevice> FileSystem for Fat16<B> {
    fn open(&mut self, path: &[u8]) -> Option<FileHandle> {
        let (cluster, size) = self.find_in_root_dir(path)?;
        let handle = self.alloc_handle()?;
        *self.slot(handle)? = Some(OpenFile {
            first_cluster: cluster,
            size,
            offset: 0,
            name: FileName::from_bytes(path),
        });
        Some(handle)
    }

    fn read(&mut self, handle: FileHandle, buf: &mut [u8]) -> Option<usize> {
        let file = (*self.slot(handle)?)?;
        if file.offset >= file.size {
            return Some(0);
        }

        let remaining = file.size - file.offset;
        let to_read = buf.len().min(remaining);
        let cs = self.bpb.cluster_size();

        let skip_clusters = file.offset / cs;
        let mut cluster = self.follow_chain(file.first_cluster, skip_clusters);
        let mut skip = file.offset % cs;
        let mut read = 0;
        let mut sector_buf = [0u8; SECTOR_SIZE];

        while read < to_read {
            let sector_in_cluster = skip / SECTOR_SIZE;
            let byte_in_sector = skip % SECTOR_SIZE;
            let Some(base) = self.bpb.cluster_to_sector(cluster) else {
                break;
            };
            let sector = base + sector_in_cluster as u32;
            if !self
                .device
                .read_sector(SectorId::new(sector), &mut sector_buf)
            {
                break;
            }

            let chunk = (SECTOR_SIZE - byte_in_sector).min(to_read - read);
            buf[read..read + chunk]
                .copy_from_slice(&sector_buf[byte_in_sector..byte_in_sector + chunk]);
            read += chunk;
            skip += chunk;

            if skip >= cs {
                let next = self.read_fat_entry(cluster);
                if next >= 0xFFF8 {
                    break;
                }
                cluster = next;
                skip = 0;
            }
        }

        if let Some(Some(f)) = self.slot(handle) {
            f.offset += read;
        }
        Some(read)
    }

    fn write(&mut self, handle: FileHandle, buf: &[u8]) -> Option<usize> {
        let mut file = (*self.slot(handle)?)?;
        let cs = self.bpb.cluster_size();

        if file.first_cluster < 2 {
            let c = self.alloc_cluster()?;
            file.first_cluster = c;
        }

        let skip_clusters = file.offset / cs;
        let mut cluster = self.follow_chain(file.first_cluster, skip_clusters);
        let mut pos = file.offset % cs;
        let mut written = 0;
        let mut sector_buf = [0u8; SECTOR_SIZE];

        while written < buf.len() {
            let sector_in_cluster = pos / SECTOR_SIZE;
            let byte_in_sector = pos % SECTOR_SIZE;
            let Some(base) = self.bpb.cluster_to_sector(cluster) else {
                break;
            };
            let sector = base + sector_in_cluster as u32;
            if !self.device.read_sector(SectorId::new(sector), &mut sector_buf) {
                break;
            }

            let chunk = (SECTOR_SIZE - byte_in_sector).min(buf.len() - written);
            sector_buf[byte_in_sector..byte_in_sector + chunk]
                .copy_from_slice(&buf[written..written + chunk]);
            // Bytes count as written only once they reached the disk.
            if !self.device.write_sector(SectorId::new(sector), &sector_buf) {
                break;
            }

            written += chunk;
            pos += chunk;

            if pos >= cs {
                let next = self.read_fat_entry(cluster);
                if next >= 0xFFF8 {
                    if written < buf.len() {
                        if let Some(new_c) = self.alloc_cluster() {
                            self.write_fat_entry(cluster, new_c);
                            cluster = new_c;
                        } else {
                            break;
                        }
                    }
                } else {
                    cluster = next;
                }
                pos = 0;
            }
        }

        file.offset += written;
        if file.offset > file.size {
            file.size = file.offset;
        }
        *self.slot(handle)? = Some(file);
        Some(written)
    }

    fn close(&mut self, handle: FileHandle) {
        let Some(file) = self.slot(handle).and_then(Option::take) else {
            return;
        };
        let mut name_buf = [0u8; 32];
        let name_bytes = file.name.as_bytes();
        let len = name_bytes.len().min(32);
        name_buf[..len].copy_from_slice(&name_bytes[..len]);
        self.update_dir_size(&name_buf[..len], file.size as u32, file.first_cluster);
    }

    fn list_dir(&mut self, _path: &[u8], entries: &mut [DirEntry]) -> Option<usize> {
        let mut buf = [0u8; SECTOR_SIZE];
        let mut count = 0;
        for s in 0..self.bpb.root_dir_sectors() {
            if !self
                .device
                .read_sector(SectorId::new(self.bpb.root_dir_start() + s), &mut buf)
            {
                continue;
            }
            for i in 0..(SECTOR_SIZE / 32) {
                if count >= entries.len() {
                    return Some(count);
                }
                let off = i * 32;
                if buf[off] == 0x00 {
                    return Some(count);
                }
                if buf[off] == 0xE5 || buf[off + 11] & 0x08 != 0 {
                    continue;
                }
                let size = u32::from_le_bytes([
                    buf[off + 28],
                    buf[off + 29],
                    buf[off + 30],
                    buf[off + 31],
                ]);
                entries[count] = DirEntry {
                    name: from_fat_name(&buf[off..off + 11]),
                    size: size as usize,
                    is_dir: buf[off + 11] & 0x10 != 0,
                };
                count += 1;
            }
        }
        Some(count)
    }

    fn create(&mut self, path: &[u8]) -> Option<FileHandle> {
        if self.find_in_root_dir(path).is_some() {
            return None;
        }
        let cluster = self.alloc_cluster()?;
        if !self.create_dir_entry(path, cluster) {
            self.write_fat_entry(cluster, 0);
            return None;
        }
        let handle = self.alloc_handle()?;
        *self.slot(handle)? = Some(OpenFile {
            first_cluster: cluster,
            size: 0,
            offset: 0,
            name: FileName::from_bytes(path),
        });
        Some(handle)
    }

    fn delete(&mut self, path: &[u8]) -> bool {
        if let Some(cluster) = self.delete_dir_entry(path) {
            self.free_chain(cluster);
            true
        } else {
            false
        }
    }

    fn file_size(&mut self, handle: FileHandle) -> Option<usize> {
        self.slot(handle)?.as_ref().map(|f| f.size)
    }

    fn mkdir(&mut self, _path: &[u8]) -> bool {
        false
    }

    fn writable(&self) -> bool {
        self.device.writable()
    }
}

fn to_fat_name(name: &[u8]) -> [u8; 11] {
    let mut fat = [b' '; 11];
    let dot_pos = name.iter().position(|&b| b == b'.').unwrap_or(name.len());
    for i in 0..dot_pos.min(8) {
        fat[i] = name[i].to_ascii_uppercase();
    }
    if dot_pos < name.len() {
        for i in 0..(name.len() - dot_pos - 1).min(3) {
            fat[8 + i] = name[dot_pos + 1 + i].to_ascii_uppercase();
        }
    }
    fat
}

fn from_fat_name(raw: &[u8]) -> FileName {
    let mut buf = [0u8; 32];
    let mut pos = 0;
    let name_end = raw[..8]
        .iter()
        .rposition(|&b| b != b' ')
        .map(|i| i + 1)
        .unwrap_or(0);
    for i in 0..name_end {
        buf[pos] = raw[i].to_ascii_lowercase();
        pos += 1;
    }
    let ext_end = raw[8..11]
        .iter()
        .rposition(|&b| b != b' ')
        .map(|i| i + 1)
        .unwrap_or(0);
    if ext_end > 0 {
        buf[pos] = b'.';
        pos += 1;
        for i in 0..ext_end {
            buf[pos] = raw[8 + i].to_ascii_lowercase();
            pos += 1;
        }
    }
    FileName::from_bytes(&buf[..pos])
}