Repository: firecracker-microvm/firecracker
File: src/vmm/src/vstate/memory.rs (Rust)

Code before:
        &self,
        writer: &mut T,
        kvm_bitmap: &[u64],
        page_size: usize,
    ) -> Result<(), MemoryError> {
        let firecracker_bitmap = self.slice.bitmap();
        let mut write_size = 0;
        let mut skip_size = 0;
        let mut dirty_batch_start = 0;

        'outer: for (i, v) in kvm_bitmap.iter().enumerate() {
            for j in 0..64 {
                let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
                let page_offset = ((i * 64) + j) * page_size;
                let is_firecracker_page_dirty = firecracker_bitmap.dirty_at(page_offset);

                // We process 64 pages at a time, however the number of pages
                // in the slot might not be a multiple of 64. We need to break
                // once we go past the last page that is actually part of the
                // region.
                if page_offset >= self.slice.len() {
                    // Ensure there are no more dirty bits after this point
                    if (v >> j) != 0 || i + 1 != kvm_bitmap.len() {
                        return Err(MemoryError::DirtyBitmapTooLarge);
                    }
                    break 'outer;
                }

                if is_kvm_page_dirty || is_firecracker_page_dirty {
                    // We are at the start of a new batch of dirty pages.
                    if skip_size > 0 {
                        // Seek forward over the unmodified pages.
                        let offset = skip_size
                            .try_into()
                            .map_err(|_| MemoryError::SlotSizeTooLarge)?;
                        writer
                            .seek(SeekFrom::Current(offset))
                            .map_err(MemoryError::IOError)?;
                        dirty_batch_start = page_offset;
                        skip_size = 0;
                    }
                    write_size += page_size;
                } else {
                    // We are at the end of a batch of dirty pages.
                    if write_size > 0 {
                        // Dump the dirty pages.
                        let slice = &self.slice.subslice(dirty_batch_start, write_size)?;
                        writer.write_all_volatile(slice)?;
                        write_size = 0;
                    }
                    skip_size += page_size;

Review comment (bug):
what I had in mind was simply checking that `kvm_bitmap.len() == self.slice.len() / self.page_size / 64` to catch the case where it's a wrong size (it could also be too short), as it can be too short).
I think we could keep the trailing 0s check, however I'm not too sure about the logic, shouldn't it be `j+1`?
I also think that the outer break makes this complex code even more complicated. A simple break will work just fine as it will break for a few times anyways.

Code after:
    /// Dumps the dirty pages in this slot onto the writer
    pub(crate) fn dump_dirty<T: WriteVolatile + std::io::Seek>(
        &self,
        writer: &mut T,
        kvm_bitmap: &[u64],
        page_size: usize,
    ) -> Result<(), MemoryError> {
        let firecracker_bitmap = self.slice.bitmap();
        let mut write_size = 0;
        let mut skip_size = 0;
        let mut dirty_batch_start = 0;

        let expected_bitmap_array_len = (self.slice.len() / page_size).div_ceil(64);
        if kvm_bitmap.len() > expected_bitmap_array_len {
            return Err(MemoryError::DirtyBitmapTooLarge);
        } else if kvm_bitmap.len() < expected_bitmap_array_len {
            return Err(MemoryError::DirtyBitmapTooSmall);
        }

        for (i, v) in kvm_bitmap.iter().enumerate() {
            for j in 0..64 {
                let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
                let page_offset = ((i * 64) + j) * page_size;
                let is_firecracker_page_dirty = firecracker_bitmap.dirty_at(page_offset);

                // We process 64 pages at a time, however the number of pages
                // in the slot might not be a multiple of 64. We need to break
                // once we go past the last page that is actually part of the
                // region.
                if page_offset >= self.slice.len() {
                    // Ensure there are no more dirty bits after this point
                    if (v >> j) != 0 {
                        return Err(MemoryError::DirtyBitmapTooLarge);
                    }
                    break;
                }

                if is_kvm_page_dirty || is_firecracker_page_dirty {
                    // We are at the start of a new batch of dirty pages.
                    if skip_size > 0 {
                        // Seek forward over the unmodified pages.
                        let offset = skip_size
                            .try_into()
                            .map_err(|_| MemoryError::SlotSizeTooLarge)?;
                        writer
                            .seek(SeekFrom::Current(offset))
                            .map_err(MemoryError::SeekError)?;
                        dirty_batch_start = page_offset;
                        skip_size = 0;
                    }
                    write_size += page_size;