feat(vhost-user-backend): added processing of POSTCOPY messages

Added processing of POSTCOPY messages for vhost-user-backend crate.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

vhost-user-backend/Cargo.toml

@@ -14,6 +14,7 @@ xen = ["vm-memory/xen", "vhost/xen"]
 [dependencies]
 libc = "0.2.39"
 log = "0.4.17"
+userfaultfd = "0.7.0"
 vhost = { path = "../vhost", version = "0.9", features = ["vhost-user-backend"] }
 virtio-bindings = "0.2.1"
 virtio-queue = "0.10.0"

vhost-user-backend/src/handler.rs

@@ -6,10 +6,12 @@
 use std::error;
 use std::fs::File;
 use std::io;
+use std::os::fd::FromRawFd;
 use std::os::unix::io::AsRawFd;
 use std::sync::Arc;
 use std::thread;
 
+use userfaultfd::{Uffd, UffdBuilder};
 use vhost::vhost_user::message::{
     VhostUserConfigFlags, VhostUserMemoryRegion, VhostUserProtocolFeatures,
     VhostUserSingleMemoryRegion, VhostUserVirtioFeatures, VhostUserVringAddrFlags,
@@ -68,6 +70,7 @@ impl error::Error for VhostUserHandlerError {}
 pub type VhostUserHandlerResult<T> = std::result::Result<T, VhostUserHandlerError>;
 
 struct AddrMapping {
+    local_addr: u64,
     vmm_addr: u64,
     size: u64,
     gpa_base: u64,
@@ -86,6 +89,7 @@ pub struct VhostUserHandler<T: VhostUserBackend> {
     mappings: Vec<AddrMapping>,
     atomic_mem: GM<T::Bitmap>,
     vrings: Vec<T::Vring>,
+    uffd: Option<Uffd>,
     worker_threads: Vec<thread::JoinHandle<VringEpollResult<()>>>,
 }
 
@@ -145,6 +149,7 @@ where
             mappings: Vec::new(),
             atomic_mem,
             vrings,
+            uffd: None,
             worker_threads,
         })
     }
@@ -278,20 +283,20 @@ where
         let mut mappings: Vec<AddrMapping> = Vec::new();
 
         for (region, file) in ctx.iter().zip(files) {
-            regions.push(
-                GuestRegionMmap::new(
-                    region.mmap_region(file)?,
-                    GuestAddress(region.guest_phys_addr),
-                )
-                .map_err(|e| {
-                    VhostUserError::ReqHandlerError(io::Error::new(io::ErrorKind::Other, e))
-                })?,
-            );
+            let guest_region = GuestRegionMmap::new(
+                region.mmap_region(file)?,
+                GuestAddress(region.guest_phys_addr),
+            )
+            .map_err(|e| {
+                VhostUserError::ReqHandlerError(io::Error::new(io::ErrorKind::Other, e))
+            })?;
             mappings.push(AddrMapping {
+                local_addr: guest_region.as_ptr() as u64,
                 vmm_addr: region.user_addr,
                 size: region.memory_size,
                 gpa_base: region.guest_phys_addr,
             });
+            regions.push(guest_region);
         }
 
         let mem = GuestMemoryMmap::from_regions(regions).map_err(|e| {
@@ -568,6 +573,13 @@ where
             })?,
         );
 
+        let addr_mapping = AddrMapping {
+            local_addr: guest_region.as_ptr() as u64,
+            vmm_addr: region.user_addr,
+            size: region.memory_size,
+            gpa_base: region.guest_phys_addr,
+        };
+
         let mem = self
             .atomic_mem
             .memory()
@@ -584,11 +596,7 @@ where
                 VhostUserError::ReqHandlerError(io::Error::new(io::ErrorKind::Other, e))
             })?;
 
-        self.mappings.push(AddrMapping {
-            vmm_addr: region.user_addr,
-            size: region.memory_size,
-            gpa_base: region.guest_phys_addr,
-        });
+        self.mappings.push(addr_mapping);
 
         Ok(())
     }
@@ -615,6 +623,53 @@ where
 
         Ok(())
     }
+
+    fn postcopy_advice(&mut self) -> VhostUserResult<File> {
+        let mut uffd_builder = UffdBuilder::new();
+
+        let uffd = uffd_builder
+            .close_on_exec(true)
+            .non_blocking(true)
+            .user_mode_only(false)
+            .create()
+            .map_err(|e| {
+                VhostUserError::ReqHandlerError(io::Error::new(io::ErrorKind::Other, e))
+            })?;
+
+        // SAFETY:
+        // We know that uffd is a valid fd.
+        let uffd_file = unsafe { File::from_raw_fd(uffd.as_raw_fd()) };
+
+        self.uffd = Some(uffd);
+
+        Ok(uffd_file)
+    }
+
+    fn postcopy_listen(&mut self) -> VhostUserResult<()> {
+        let Some(ref uffd) = self.uffd else {
+            return Err(VhostUserError::ReqHandlerError(io::Error::new(
+                io::ErrorKind::Other,
+                "No registered UFFD handler",
+            )));
+        };
+
+        for mapping in self.mappings.iter() {
+            uffd.register(
+                mapping.local_addr as *mut libc::c_void,
+                mapping.size as usize,
+            )
+            .map_err(|e| {
+                VhostUserError::ReqHandlerError(io::Error::new(io::ErrorKind::Other, e))
+            })?;
+        }
+
+        Ok(())
+    }
+
+    fn postcopy_end(&mut self) -> VhostUserResult<()> {
+        _ = self.uffd.take();
+        Ok(())
+    }
 }
 
 impl<T: VhostUserBackend> Drop for VhostUserHandler<T> {