From b910a7922d2ce3d694d561e9816494f7dbd1b623 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Fri, 24 Sep 2021 15:29:13 +0200
Subject: [PATCH] vmm: Fix migration when writing/reading big chunks of data

Both read_exact_from() and write_all_to() functions from the GuestMemory
trait implementation in vm-memory are buggy. They should retry until
they wrote or read the amount of data that was expected, but instead
they simply return an error when this happens. This causes the migration
to fail when trying to send important amount of data through the
migration socket, due to large memory regions.

This should be eventually fixed in vm-memory, and here is the link to
follow up on the issue: https://github.com/rust-vmm/vm-memory/issues/174

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/vm.rs | 65 ++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 51 insertions(+), 14 deletions(-)

diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index cae848978..d08da9818 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2115,14 +2115,33 @@ impl Vm {
         let mem = guest_memory.memory();
 
         for range in ranges.regions() {
-            mem.read_exact_from(GuestAddress(range.gpa), fd, range.length as usize)
-                .map_err(|e| {
-                    MigratableError::MigrateReceive(anyhow!(
-                        "Error transferring memory to socket: {}",
-                        e
-                    ))
-                })?;
+            let mut offset: u64 = 0;
+            // Here we are manually handling the retry in case we can't the
+            // whole region at once because we can't use the implementation
+            // from vm-memory::GuestMemory of read_exact_from() as it is not
+            // following the correct behavior. For more info about this issue
+            // see: https://github.com/rust-vmm/vm-memory/issues/174
+            loop {
+                let bytes_read = mem
+                    .read_from(
+                        GuestAddress(range.gpa + offset),
+                        fd,
+                        (range.length - offset) as usize,
+                    )
+                    .map_err(|e| {
+                        MigratableError::MigrateReceive(anyhow!(
+                            "Error receiving memory from socket: {}",
+                            e
+                        ))
+                    })?;
+                offset += bytes_read as u64;
+
+                if offset == range.length {
+                    break;
+                }
+            }
         }
+
         Ok(())
     }
 
@@ -2138,13 +2157,31 @@ impl Vm {
         let mem = guest_memory.memory();
 
         for range in ranges.regions() {
-            mem.write_all_to(GuestAddress(range.gpa), fd, range.length as usize)
-                .map_err(|e| {
-                    MigratableError::MigrateSend(anyhow!(
-                        "Error transferring memory to socket: {}",
-                        e
-                    ))
-                })?;
+            let mut offset: u64 = 0;
+            // Here we are manually handling the retry in case we can't the
+            // whole region at once because we can't use the implementation
+            // from vm-memory::GuestMemory of write_all_to() as it is not
+            // following the correct behavior. For more info about this issue
+            // see: https://github.com/rust-vmm/vm-memory/issues/174
+            loop {
+                let bytes_written = mem
+                    .write_to(
+                        GuestAddress(range.gpa + offset),
+                        fd,
+                        (range.length - offset) as usize,
+                    )
+                    .map_err(|e| {
+                        MigratableError::MigrateSend(anyhow!(
+                            "Error transferring memory to socket: {}",
+                            e
+                        ))
+                    })?;
+                offset += bytes_written as u64;
+
+                if offset == range.length {
+                    break;
+                }
+            }
         }
 
         Ok(())