markbase/markbase-core/src/rsync/checksum.rs

use anyhow::Result;
use md5::compute;

pub struct RollingChecksum {
    a: u16,
    b: u16,
}

impl RollingChecksum {
    pub fn new(data: &[u8]) -> Self {
        let mut a = 1u16;
        let mut b = 0u16;

        for byte in data {
            a = (a + *byte as u16) % 65521;
            b = (b + a) % 65521;
        }

        Self { a, b }
    }

    pub fn sum(&self) -> u32 {
        ((self.b as u32) << 16) | (self.a as u32)
    }

    pub fn update(&mut self, remove: u8, add: u8, len: usize) {
        self.a = (self.a - remove as u16 + add as u16) % 65521;
        self.b = (self.b - (len as u16 * remove as u16) + self.a) % 65521;
    }

    pub fn reset(&mut self) {
        self.a = 1;
        self.b = 0;
    }
}

pub fn adler32(data: &[u8]) -> u32 {
    let rolling = RollingChecksum::new(data);
    rolling.sum()
}

pub fn md5_checksum(data: &[u8]) -> [u8; 16] {
    let result = compute(data);
    let mut checksum = [0u8; 16];
    checksum.copy_from_slice(&result.0);
    checksum
}

pub fn md5_checksum_with_seed(data: &[u8], seed: u32) -> [u8; 16] {
    let mut input = Vec::with_capacity(data.len() + 4);
    input.extend_from_slice(data);
    input.extend_from_slice(&seed.to_le_bytes());

    let result = compute(&input);
    let mut checksum = [0u8; 16];
    checksum.copy_from_slice(&result.0);
    checksum
}

#[derive(Debug, Clone)]
pub struct BlockChecksum {
    pub rolling: u32,
    pub strong: [u8; 16],
    pub offset: usize,
    pub length: usize,
}

impl BlockChecksum {
    pub fn new(data: &[u8], offset: usize) -> Self {
        let rolling = adler32(data);
        let strong = md5_checksum(data);

        Self {
            rolling,
            strong,
            offset,
            length: data.len(),
        }
    }

    pub fn new_with_seed(data: &[u8], offset: usize, seed: u32) -> Self {
        let rolling = adler32(data);
        let strong = md5_checksum_with_seed(data, seed);

        Self {
            rolling,
            strong,
            offset,
            length: data.len(),
        }
    }

    pub fn verify(&self, data: &[u8]) -> bool {
        let rolling = adler32(data);
        let strong = md5_checksum(data);

        rolling == self.rolling && strong == self.strong
    }

    pub fn verify_with_seed(&self, data: &[u8], seed: u32) -> bool {
        let rolling = adler32(data);
        let strong = md5_checksum_with_seed(data, seed);

        rolling == self.rolling && strong == self.strong
    }
}

pub fn compute_block_checksums(data: &[u8], block_size: usize) -> Vec<BlockChecksum> {
    let mut checksums = Vec::new();
    let mut offset = 0;

    while offset < data.len() {
        let end = std::cmp::min(offset + block_size, data.len());
        let block_data = &data[offset..end];

        checksums.push(BlockChecksum::new(block_data, offset));
        offset += block_size;
    }

    checksums
}

pub fn compute_block_checksums_with_seed(
    data: &[u8],
    block_size: usize,
    seed: u32,
) -> Vec<BlockChecksum> {
    let mut checksums = Vec::new();
    let mut offset = 0;

    while offset < data.len() {
        let end = std::cmp::min(offset + block_size, data.len());
        let block_data = &data[offset..end];

        checksums.push(BlockChecksum::new_with_seed(block_data, offset, seed));
        offset += block_size;
    }

    checksums
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rolling_checksum() {
        let data = b"hello world";
        let rolling = RollingChecksum::new(data);

        let sum = rolling.sum();
        assert!(sum > 0);
    }

    #[test]
    fn test_rolling_update() {
        let data1 = b"hello";
        let data2 = b"ello "; // Shifted by 1, replace 'h' with ' '

        let rolling1 = RollingChecksum::new(data1);
        let rolling2 = RollingChecksum::new(data2);

        // Test rolling checksum basic functionality
        assert!(rolling1.sum() > 0);
        assert!(rolling2.sum() > 0);
    }

    #[test]
    fn test_md5_checksum() {
        let data = b"hello world";
        let checksum = md5_checksum(data);

        assert_eq!(checksum.len(), 16);
        assert!(checksum.iter().any(|&b| b != 0));
    }

    #[test]
    fn test_block_checksum() {
        let data = b"hello world test data";
        let block_checksum = BlockChecksum::new(data, 0);

        assert!(block_checksum.verify(data));
    }

    #[test]
    fn test_compute_block_checksums() {
        let data = b"hello world test data long string";
        let checksums = compute_block_checksums(data, 10);

        // Test basic functionality
        assert!(checksums.len() > 0);
        assert_eq!(checksums[0].length, 10);
    }
}