dropbox_content_hash.rs

const BLOCK_SIZE: usize = 4 * 1024 * 1024;

/// Implementation of Dropbox's Content Hash.
/// See https://www.dropbox.com/developers/reference/content-hash
pub struct ContentHash {
    entire_hasher: ring::digest::Context,
    block_hasher: ring::digest::Context,
    block_position: usize,
}

impl ContentHash {
    pub fn new() -> Self {
        Self {
            entire_hasher: ring::digest::Context::new(&ring::digest::SHA256),
            block_hasher: ring::digest::Context::new(&ring::digest::SHA256),
            block_position: 0,
        }
    }

    pub fn update(&mut self, data: &[u8]) {
        let remain_bytes = BLOCK_SIZE - self.block_position;
        if data.len() > remain_bytes {
            self.update(&data[..remain_bytes]);
            self.update(&data[remain_bytes..]);
            return;
        }
        self.block_position += data.len();
        assert!(self.block_position <= BLOCK_SIZE);
        self.block_hasher.update(data);
        if self.block_position == BLOCK_SIZE {
            let mut block_hasher = ring::digest::Context::new(&ring::digest::SHA256);
            std::mem::swap(&mut block_hasher, &mut self.block_hasher);
            let block_hash = block_hasher.finish();
            self.entire_hasher.update(block_hash.as_ref());
            self.block_position = 0;
        }
    }

    pub fn finalize(mut self) -> [u8; 32] {
        if self.block_position > 0 {
            let block_hash = self.block_hasher.finish();
            self.entire_hasher.update(block_hash.as_ref());
        }
        let entire_hash = self.entire_hasher.finish();
        let mut result = [0; 32];
        result.copy_from_slice(entire_hash.as_ref());
        result
    }
}

#[cfg(test)]
mod tests {
    const DROPBOX_EXAMPLE_URL: &str = "https://www.dropbox.com/static/images/developers/milky-way-nasa.jpg";
    const DROPBOX_EXAMPLE_SIZE: usize = 9_711_423;
    const DROPBOX_EXAMPLE_CHECKSUM: [u8; 32] = [
        0x48, 0x52, 0x91, 0xfa, 0x0e, 0xe5, 0x0c, 0x01,
        0x69, 0x82, 0xab, 0xbf, 0xa9, 0x43, 0x95, 0x7b,
        0xcd, 0x23, 0x1a, 0xae, 0x04, 0x92, 0xcc, 0xba,
        0xa2, 0x2c, 0x58, 0xe3, 0x99, 0x7b, 0x35, 0xe0,
    ];

    use crate::content_hash::BLOCK_SIZE;

    #[tokio::test(flavor = "multi_thread")]
    async fn test_with_dropbox_example() {
        let dropbox_binary = reqwest::Client::new().get(DROPBOX_EXAMPLE_URL)
            .send().await.unwrap().bytes().await.unwrap();
        if dropbox_binary.len() != DROPBOX_EXAMPLE_SIZE {
            panic!("Dropbox binary size changed");
        }
        let mut hasher = super::ContentHash::new();
        hasher.update(&dropbox_binary);
        let result = hasher.finalize();
        assert_eq!(result, DROPBOX_EXAMPLE_CHECKSUM);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_with_dropbox_example_weird_chunk() {
        let dropbox_binary = reqwest::Client::new().get(DROPBOX_EXAMPLE_URL)
            .send().await.unwrap().bytes().await.unwrap();
        if dropbox_binary.len() != DROPBOX_EXAMPLE_SIZE {
            panic!("Dropbox binary size changed");
        }
        let mut hasher = super::ContentHash::new();
        hasher.update(&dropbox_binary[0..BLOCK_SIZE+1]);
        hasher.update(&dropbox_binary[BLOCK_SIZE+1..BLOCK_SIZE*2-1]);
        hasher.update(&dropbox_binary[BLOCK_SIZE*2-1..]);
        let result = hasher.finalize();
        assert_eq!(result, DROPBOX_EXAMPLE_CHECKSUM);
    }
}