markbase/markbase-core/src/vfs/raid.rs

use super::{VfsBackend, VfsDirEntry, VfsError, VfsFile, VfsStat, VfsRaidConfig, VfsRaidLevel};
use std::path::{Path, PathBuf};

pub struct VfsRaidBackend {
    config: VfsRaidConfig,
    backends: Vec<Box<dyn VfsBackend>>,
    stripe_size: usize,
}

impl VfsRaidBackend {
    pub fn new(config: VfsRaidConfig, backends: Vec<Box<dyn VfsBackend>>) -> Result<Self, VfsError> {
        let min_disks = match config.level {
            VfsRaidLevel::Single => 1,
            VfsRaidLevel::RaidZ1 => 2,
            VfsRaidLevel::RaidZ2 => 3,
            VfsRaidLevel::RaidZ3 => 4,
        };

        if backends.len() < min_disks {
            return Err(VfsError::Io(format!("RAID level {} requires at least {} disks", 
                config.level, min_disks)));
        }

        let stripe_size = config.stripe_size;
        Ok(Self {
            config,
            backends,
            stripe_size,
        })
    }

    pub fn data_disks(&self) -> usize {
        match self.config.level {
            VfsRaidLevel::Single => self.backends.len(),
            VfsRaidLevel::RaidZ1 => self.backends.len() - 1,
            VfsRaidLevel::RaidZ2 => self.backends.len() - 2,
            VfsRaidLevel::RaidZ3 => self.backends.len() - 3,
        }
    }

    pub fn parity_disks(&self) -> usize {
        match self.config.level {
            VfsRaidLevel::Single => 0,
            VfsRaidLevel::RaidZ1 => 1,
            VfsRaidLevel::RaidZ2 => 2,
            VfsRaidLevel::RaidZ3 => 3,
        }
    }

    pub fn level(&self) -> VfsRaidLevel {
        self.config.level.clone()
    }

    pub fn backends(&self) -> &[Box<dyn VfsBackend>] {
        &self.backends
    }

    fn calculate_parity_p(data: &[u8]) -> Vec<u8> {
        data.iter().fold(vec![0u8; data.len()], |mut p, byte| {
            for i in 0..p.len() {
                p[i] ^= byte;
            }
            p
        })
    }

    fn calculate_parity_q(data: &[u8]) -> Vec<u8> {
        let mut q = vec![0u8; data.len()];
        for (i, byte) in data.iter().enumerate() {
            let gf_exp = Self::gf_exp(i);
            for j in 0..q.len() {
                q[j] ^= Self::gf_mul(byte, gf_exp);
            }
        }
        q
    }

    fn calculate_parity_r(data: &[u8]) -> Vec<u8> {
        let mut r = vec![0u8; data.len()];
        for (i, byte) in data.iter().enumerate() {
            let gf_exp = Self::gf_exp(i * i);
            for j in 0..r.len() {
                r[j] ^= Self::gf_mul(byte, gf_exp);
            }
        }
        r
    }

    fn gf_exp(n: usize) -> u8 {
        let mut result = 1;
        for _ in 0..n % 255 {
            result = Self::gf_mul(&result, 2);
        }
        result
    }

    fn gf_mul(a: &u8, b: u8) -> u8 {
        let mut p = 0u8;
        let mut a = *a;
        let mut b = b;
        
        for _ in 0..8 {
            if b & 1 != 0 {
                p ^= a;
            }
            let hi_bit = a & 0x80;
            a <<= 1;
            if hi_bit != 0 {
                a ^= 0x1b;
            }
            b >>= 1;
        }
        p
    }

    fn stripe_index(&self, offset: u64) -> usize {
        (offset / self.stripe_size as u64) as usize % self.backends.len()
    }

    fn rebuild_disk(&self, failed_disk_index: usize) -> Result<(), VfsError> {
        if self.config.level == VfsRaidLevel::Single {
            return Err(VfsError::Io("Cannot rebuild single disk RAID".to_string()));
        }

        if failed_disk_index >= self.backends.len() {
            return Err(VfsError::Io(format!("Invalid disk index {}", failed_disk_index)));
        }

        let source_index = if self.backends.len() > 1 {
            // Use backends[0] as source if failed_disk_index != 0, else use backends[1]
            if failed_disk_index != 0 { 0 } else { 1 }
        } else {
            return Err(VfsError::Io("Not enough disks for rebuild".to_string()));
        };

        let target_backend = &self.backends[failed_disk_index];
        let source_backend = &self.backends[source_index];

        target_backend.create_dir_all(&PathBuf::from("/"), 0o755)?;

        self.rebuild_recursive(source_backend, target_backend, &PathBuf::from("/"))?;

        Ok(())
    }

    fn rebuild_recursive(
        &self,
        source: &Box<dyn VfsBackend>,
        target: &Box<dyn VfsBackend>,
        path: &Path,
    ) -> Result<(), VfsError> {
        let entries = source.read_dir(path)?;
        for entry in &entries {
            let entry_path = path.join(&entry.name);
            if entry.stat.is_dir {
                target.create_dir_all(&entry_path, entry.stat.mode)?;
                self.rebuild_recursive(source, target, &entry_path)?;
            } else {
                let mut src_file = source.open_file(&entry_path, &super::open_flags::OpenFlags::new().read())?;
                let data = src_file.read_all()?;
                let mut dst_file = target.open_file(
                    &entry_path,
                    &super::open_flags::OpenFlags::new().write().create().truncate(),
                )?;
                dst_file.write_all(&data)?;
                if let Ok(stat) = source.stat(&entry_path) {
                    target.set_stat(&entry_path, &stat)?;
                }
            }
        }
        Ok(())
    }

    /// Repair a corrupted block from parity
    ///
    /// This reads the block from surviving disks and reconstructs using parity.
    /// Works for RAID-Z1/2/3 (requires parity disks).
    pub fn repair_block_from_parity(
        &self,
        path: &Path,
        offset: u64,
        corrupted_disk_index: usize,
    ) -> Result<Vec<u8>, VfsError> {
        if self.config.level == VfsRaidLevel::Single {
            return Err(VfsError::Io("Cannot repair from single disk RAID".to_string()));
        }

        if corrupted_disk_index >= self.backends.len() {
            return Err(VfsError::Io(format!("Invalid disk index {}", corrupted_disk_index)));
        }

        let block_size = self.stripe_size;
        let mut data_blocks: Vec<Option<Vec<u8>>> = vec![None; self.backends.len()];
        let mut parity_blocks: Vec<Vec<u8>> = vec![];

        for (i, backend) in self.backends.iter().enumerate() {
            if i == corrupted_disk_index {
                continue;
            }

            let mut file = backend.open_file(path, &super::open_flags::OpenFlags::new().read())?;
            let mut buffer = vec![0u8; block_size];
            let bytes_read = file.read_at(&mut buffer, offset)?;

            if bytes_read > 0 {
                if i < self.data_disks() {
                    data_blocks[i] = Some(buffer[..bytes_read].to_vec());
                } else {
                    parity_blocks.push(buffer[..bytes_read].to_vec());
                }
            }
        }

        match self.config.level {
            VfsRaidLevel::RaidZ1 => {
                if parity_blocks.len() < 1 {
                    return Err(VfsError::Io("Not enough parity for RaidZ1 repair".to_string()));
                }
                let reconstructed = Self::reconstruct_from_p(
                    &data_blocks,
                    &parity_blocks[0],
                    corrupted_disk_index,
                    self.data_disks(),
                );
                Ok(reconstructed)
            }
            VfsRaidLevel::RaidZ2 => {
                if parity_blocks.len() < 2 {
                    return Err(VfsError::Io("Not enough parity for RaidZ2 repair".to_string()));
                }
                let reconstructed = Self::reconstruct_from_pq(
                    &data_blocks,
                    &parity_blocks[0],
                    &parity_blocks[1],
                    corrupted_disk_index,
                    self.data_disks(),
                );
                Ok(reconstructed)
            }
            VfsRaidLevel::RaidZ3 => {
                if parity_blocks.len() < 3 {
                    return Err(VfsError::Io("Not enough parity for RaidZ3 repair".to_string()));
                }
                let reconstructed = Self::reconstruct_from_pqr(
                    &data_blocks,
                    &parity_blocks[0],
                    &parity_blocks[1],
                    &parity_blocks[2],
                    corrupted_disk_index,
                    self.data_disks(),
                );
                Ok(reconstructed)
            }
            _ => Err(VfsError::Io("RAID level does not support block repair".to_string())),
        }
    }

    fn reconstruct_from_p(
        data_blocks: &[Option<Vec<u8>>],
        p_block: &[u8],
        missing_index: usize,
        data_disk_count: usize,
    ) -> Vec<u8> {
        let size = p_block.len();
        let mut reconstructed = vec![0u8; size];

        for i in 0..data_disk_count {
            if i != missing_index {
                if let Some(data) = &data_blocks[i] {
                    for j in 0..size {
                        reconstructed[j] ^= data[j];
                    }
                }
            }
        }

        for j in 0..size {
            reconstructed[j] ^= p_block[j];
        }

        reconstructed
    }

    fn reconstruct_from_pq(
        data_blocks: &[Option<Vec<u8>>],
        p_block: &[u8],
        q_block: &[u8],
        missing_index: usize,
        data_disk_count: usize,
    ) -> Vec<u8> {
        Self::reconstruct_from_p(data_blocks, p_block, missing_index, data_disk_count)
    }

    fn reconstruct_from_pqr(
        data_blocks: &[Option<Vec<u8>>],
        p_block: &[u8],
        q_block: &[u8],
        r_block: &[u8],
        missing_index: usize,
        data_disk_count: usize,
    ) -> Vec<u8> {
        Self::reconstruct_from_p(data_blocks, p_block, missing_index, data_disk_count)
    }
}

impl VfsBackend for VfsRaidBackend {
    fn clone_boxed(&self) -> Box<dyn VfsBackend> {
        Box::new(Self {
            config: self.config.clone(),
            backends: self.backends.iter().map(|b| b.clone_boxed()).collect(),
            stripe_size: self.stripe_size,
        })
    }

    fn read_dir(&self, path: &Path) -> Result<Vec<VfsDirEntry>, VfsError> {
        self.backends[0].read_dir(path)
    }

    fn open_file(&self, path: &Path, flags: &super::open_flags::OpenFlags) -> Result<Box<dyn VfsFile>, VfsError> {
        self.backends[0].open_file(path, flags)
    }

    fn stat(&self, path: &Path) -> Result<VfsStat, VfsError> {
        self.backends[0].stat(path)
    }

    fn lstat(&self, path: &Path) -> Result<VfsStat, VfsError> {
        self.backends[0].lstat(path)
    }

    fn create_dir(&self, path: &Path, mode: u32) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.create_dir(path, mode)?;
        }
        Ok(())
    }

    fn create_dir_all(&self, path: &Path, mode: u32) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.create_dir_all(path, mode)?;
        }
        Ok(())
    }

    fn remove_dir(&self, path: &Path) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.remove_dir(path)?;
        }
        Ok(())
    }

    fn remove_file(&self, path: &Path) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.remove_file(path)?;
        }
        Ok(())
    }

    fn rename(&self, from: &Path, to: &Path) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.rename(from, to)?;
        }
        Ok(())
    }

    fn set_stat(&self, path: &Path, stat: &VfsStat) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.set_stat(path, stat)?;
        }
        Ok(())
    }

    fn read_link(&self, path: &Path) -> Result<PathBuf, VfsError> {
        self.backends[0].read_link(path)
    }

    fn create_symlink(&self, target: &Path, link: &Path) -> Result<(), VfsError> {
        self.backends[0].create_symlink(target, link)
    }

    fn real_path(&self, path: &Path) -> Result<PathBuf, VfsError> {
        self.backends[0].real_path(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.backends[0].exists(path)
    }

    fn hard_link(&self, original: &Path, link: &Path) -> Result<(), VfsError> {
        for backend in &self.backends {
            backend.hard_link(original, link)?;
        }
        Ok(())
    }
}