android-sparse 0.3.1

An implementation of Android's sparse file format.
Documentation
use std::fs::File as StdFile;
use std::io::{ErrorKind, Read, Seek, SeekFrom};

use byteorder::{LittleEndian, ReadBytesExt};
use crc::crc32;
use crc::crc32::Hasher32;

use constants::BLOCK_SIZE;
use file::{Chunk, File};
use headers::{ChunkHeader, ChunkType, FileHeader};
use result::Result;

pub struct Reader {
    src: StdFile,
    crc: Option<crc32::Digest>,
}

impl Reader {
    pub fn new(src: StdFile) -> Self {
        Self {
            src: src,
            crc: None,
        }
    }

    pub fn with_crc(mut self) -> Self {
        self.crc = Some(crc32::Digest::new(crc32::IEEE));
        self
    }

    pub fn read(mut self, mut sparse_file: &mut File) -> Result<()> {
        let header = FileHeader::deserialize(&mut self.src)?;
        for _ in 0..header.total_chunks {
            self.read_chunk(&mut sparse_file)?;
        }

        Ok(())
    }

    fn read_chunk(&mut self, mut spf: &mut File) -> Result<()> {
        let header = ChunkHeader::deserialize(&mut self.src)?;
        let num_blocks = header.chunk_size;

        match header.chunk_type {
            ChunkType::Raw => self.read_raw_chunk(&mut spf, num_blocks),
            ChunkType::Fill => self.read_fill_chunk(&mut spf, num_blocks),
            ChunkType::DontCare => Ok(self.read_dont_care_chunk(&mut spf, num_blocks)),
            ChunkType::Crc32 => self.read_crc32_chunk(&mut spf),
        }
    }

    fn read_raw_chunk(&mut self, spf: &mut File, num_blocks: u32) -> Result<()> {
        let off = self.src.seek(SeekFrom::Current(0))?;

        if let Some(ref mut digest) = self.crc {
            let mut block = [0; BLOCK_SIZE as usize];
            for _ in 0..num_blocks {
                self.src.read_exact(&mut block)?;
                digest.write(&block);
            }
        } else {
            let size = i64::from(num_blocks * BLOCK_SIZE);
            self.src.seek(SeekFrom::Current(size))?;
        }

        let chunk = Chunk::Raw {
            file: self.src.try_clone()?,
            offset: off,
            num_blocks: num_blocks,
        };
        spf.add_chunk(chunk);

        Ok(())
    }

    fn read_fill_chunk(&mut self, spf: &mut File, num_blocks: u32) -> Result<()> {
        let fill = read4(&mut self.src)?;

        if let Some(ref mut digest) = self.crc {
            for _ in 0..(num_blocks * BLOCK_SIZE / 4) {
                digest.write(&fill);
            }
        }

        let chunk = Chunk::Fill { fill, num_blocks };
        spf.add_chunk(chunk);

        Ok(())
    }

    fn read_dont_care_chunk(&mut self, spf: &mut File, num_blocks: u32) {
        if let Some(ref mut digest) = self.crc {
            let block = [0; BLOCK_SIZE as usize];
            for _ in 0..num_blocks {
                digest.write(&block);
            }
        }

        let chunk = Chunk::DontCare { num_blocks };
        spf.add_chunk(chunk);
    }

    fn read_crc32_chunk(&mut self, spf: &mut File) -> Result<()> {
        let crc = self.src.read_u32::<LittleEndian>()?;
        self.check_crc(crc)?;

        let chunk = Chunk::Crc32 { crc };
        spf.add_chunk(chunk);

        Ok(())
    }

    fn check_crc(&self, crc: u32) -> Result<()> {
        if let Some(ref digest) = self.crc {
            if digest.sum32() != crc {
                return Err("Checksum does not match".into());
            }
        }
        Ok(())
    }
}

pub struct Encoder {
    src: StdFile,
    chunk: Option<Chunk>,
}

impl Encoder {
    pub fn new(src: StdFile) -> Self {
        Self {
            src: src,
            chunk: None,
        }
    }

    pub fn read(mut self, mut sparse_file: &mut File) -> Result<()> {
        let block_size = BLOCK_SIZE as usize;
        let mut block = vec![0; block_size];
        loop {
            let bytes_read = read_all(&mut self.src, &mut block)?;
            self.read_block(&mut sparse_file, &block[..bytes_read])?;
            if bytes_read != block_size {
                break;
            }
        }

        if let Some(last_chunk) = self.chunk.take() {
            sparse_file.add_chunk(last_chunk);
        }

        Ok(())
    }

    fn read_block(&mut self, spf: &mut File, block: &[u8]) -> Result<()> {
        if block.is_empty() {
            return Ok(());
        }

        if let Some(chunk) = self.merge_block(block)? {
            spf.add_chunk(chunk);
        }

        Ok(())
    }

    fn merge_block(&mut self, block: &[u8]) -> Result<Option<Chunk>> {
        if is_sparse_block(block) {
            let fill = read4(block)?;
            if fill == [0; 4] {
                Ok(self.merge_don_care_block())
            } else {
                Ok(self.merge_fill_block(fill))
            }
        } else {
            self.merge_raw_block()
        }
    }

    fn merge_raw_block(&mut self) -> Result<Option<Chunk>> {
        let (old, new) = match self.chunk.take() {
            Some(Chunk::Raw {
                file,
                offset,
                num_blocks,
            }) => (
                None,
                Chunk::Raw {
                    file: file,
                    offset: offset,
                    num_blocks: num_blocks + 1,
                },
            ),
            old_chunk => {
                let mut file = self.src.try_clone()?;
                let curr_off = file.seek(SeekFrom::Current(0))?;
                (
                    old_chunk,
                    Chunk::Raw {
                        file: file,
                        offset: curr_off - u64::from(BLOCK_SIZE),
                        num_blocks: 1,
                    },
                )
            }
        };
        self.chunk = Some(new);
        Ok(old)
    }

    fn merge_fill_block(&mut self, fill: [u8; 4]) -> Option<Chunk> {
        let new_fill = fill;
        let (old, new) = match self.chunk.take() {
            Some(Chunk::Fill { fill, num_blocks }) if fill == new_fill => (
                None,
                Chunk::Fill {
                    fill: fill,
                    num_blocks: num_blocks + 1,
                },
            ),
            old_chunk => (
                old_chunk,
                Chunk::Fill {
                    fill: new_fill,
                    num_blocks: 1,
                },
            ),
        };
        self.chunk = Some(new);
        old
    }

    fn merge_don_care_block(&mut self) -> Option<Chunk> {
        let (old, new) = match self.chunk.take() {
            Some(Chunk::DontCare { num_blocks }) => (
                None,
                Chunk::DontCare {
                    num_blocks: num_blocks + 1,
                },
            ),
            old_chunk => (old_chunk, Chunk::DontCare { num_blocks: 1 }),
        };
        self.chunk = Some(new);
        old
    }
}

fn is_sparse_block(block: &[u8]) -> bool {
    if block.len() != BLOCK_SIZE as usize {
        return false;
    }

    let mut words = block.chunks(4);
    let first = words.next().unwrap();
    for word in words {
        if word != first {
            return false;
        }
    }

    true
}

fn read4<R: Read>(mut r: R) -> Result<[u8; 4]> {
    let mut buf = [0; 4];
    r.read_exact(&mut buf)?;
    Ok(buf)
}

fn read_all<R: Read>(mut r: R, mut buf: &mut [u8]) -> Result<usize> {
    let buf_size = buf.len();

    while !buf.is_empty() {
        match r.read(buf) {
            Ok(0) => break,
            Ok(n) => {
                let tmp = buf;
                buf = &mut tmp[n..]
            }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => (),
            Err(e) => return Err(e.into()),
        };
    }

    Ok(buf_size - buf.len())
}