vhost_user: Add Inflight I/O tracking support

The inflight I/O tracking feature is useful for handling crashes and
disconnections from the backend. It allows the backend to rely on a
buffer that was shared earlier with the VMM to restore to the previous
state it was before the crash.

This feature depends on the availability of the protocol feature
VHOST_USER_PROTOCOL_F_INFLIGHT_SHMFD, and it implements both
VHOST_USER_GET_INFLIGHT_FD and VHOST_USER_SET_INFLIGHT_FD messages.

Fixes #43

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

7 files changed, 201 insertions, 3 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 0cd15f7..a9bfe83 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,4 +26,5 @@ vmm-sys-util = ">=0.3.1"
 vm-memory = { version = "0.2.0", optional = true }
 
 [dev-dependencies]
+tempfile = ">=3.2.0"
 vm-memory = { version = "0.2.0", features=["backend-mmap"] }
diff --git a/coverage_config_x86_64.json b/coverage_config_x86_64.json
index f9164a8..73a5c25 100644
--- a/coverage_config_x86_64.json
+++ b/coverage_config_x86_64.json
@@ -1 +1 @@
-{"coverage_score": 81.0, "exclude_path": "src/vhost_kern/", "crate_features": "vhost-user-master,vhost-user-slave"}
+{"coverage_score": 80.9, "exclude_path": "src/vhost_kern/", "crate_features": "vhost-user-master,vhost-user-slave"}
diff --git a/src/vhost_user/dummy_slave.rs b/src/vhost_user/dummy_slave.rs
index b2b83d2..dc9eed5 100644
--- a/src/vhost_user/dummy_slave.rs
+++ b/src/vhost_user/dummy_slave.rs
@@ -1,7 +1,8 @@
 // Copyright (C) 2019 Alibaba Cloud Computing. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::os::unix::io::RawFd;
+use std::fs::File;
+use std::os::unix::io::{AsRawFd, RawFd};
 
 use super::message::*;
 use super::*;
@@ -25,6 +26,7 @@ pub struct DummySlaveReqHandler {
     pub err_fd: [Option<RawFd>; MAX_QUEUE_NUM],
     pub vring_started: [bool; MAX_QUEUE_NUM],
     pub vring_enabled: [bool; MAX_QUEUE_NUM],
+    pub inflight_file: Option<File>,
 }
 
 impl DummySlaveReqHandler {
@@ -245,6 +247,28 @@ impl VhostUserSlaveReqHandlerMut for DummySlaveReqHandler {
         Ok(())
     }
 
+    fn get_inflight_fd(
+        &mut self,
+        inflight: &VhostUserInflight,
+    ) -> Result<(VhostUserInflight, RawFd)> {
+        let file = tempfile::tempfile().unwrap();
+        let fd = file.as_raw_fd();
+        self.inflight_file = Some(file);
+        Ok((
+            VhostUserInflight {
+                mmap_size: 0x1000,
+                mmap_offset: 0,
+                num_queues: inflight.num_queues,
+                queue_size: inflight.queue_size,
+            },
+            fd,
+        ))
+    }
+
+    fn set_inflight_fd(&mut self, _inflight: &VhostUserInflight, _file: File) -> Result<()> {
+        Ok(())
+    }
+
     fn get_max_mem_slots(&mut self) -> Result<u64> {
         Ok(MAX_MEM_SLOTS as u64)
     }
diff --git a/src/vhost_user/master.rs b/src/vhost_user/master.rs
index b8ba0af..2d98987 100644
--- a/src/vhost_user/master.rs
+++ b/src/vhost_user/master.rs
@@ -3,8 +3,9 @@
 
 //! Traits and Struct for vhost-user master.
 
+use std::fs::File;
 use std::mem;
-use std::os::unix::io::{AsRawFd, RawFd};
+use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use std::os::unix::net::UnixStream;
 use std::path::Path;
 use std::sync::{Arc, Mutex, MutexGuard};
@@ -51,6 +52,15 @@ pub trait VhostUserMaster: VhostBackend {
     /// Setup slave communication channel.
     fn set_slave_request_fd(&mut self, fd: RawFd) -> Result<()>;
 
+    /// Retrieve shared buffer for inflight I/O tracking.
+    fn get_inflight_fd(
+        &mut self,
+        inflight: &VhostUserInflight,
+    ) -> Result<(VhostUserInflight, File)>;
+
+    /// Set shared buffer for inflight I/O tracking.
+    fn set_inflight_fd(&mut self, inflight: &VhostUserInflight, fd: RawFd) -> Result<()>;
+
     /// Query the maximum amount of memory slots supported by the backend.
     fn get_max_mem_slots(&mut self) -> Result<u64>;
 
@@ -452,6 +462,47 @@ impl VhostUserMaster for Master {
         Ok(())
     }
 
+    fn get_inflight_fd(
+        &mut self,
+        inflight: &VhostUserInflight,
+    ) -> Result<(VhostUserInflight, File)> {
+        let mut node = self.node();
+        if node.acked_protocol_features & VhostUserProtocolFeatures::INFLIGHT_SHMFD.bits() == 0 {
+            return error_code(VhostUserError::InvalidOperation);
+        }
+
+        let hdr = node.send_request_with_body(MasterReq::GET_INFLIGHT_FD, inflight, None)?;
+        let (inflight, fds) = node.recv_reply_with_fds::<VhostUserInflight>(&hdr)?;
+
+        if let Some(fds) = &fds {
+            if fds.len() == 1 && fds[0] >= 0 {
+                // Safe because we know the fd is valid.
+                let file = unsafe { File::from_raw_fd(fds[0]) };
+                return Ok((inflight, file));
+            }
+        }
+
+        // Make sure to close the fds before returning the error.
+        Endpoint::<MasterReq>::close_rfds(fds);
+
+        error_code(VhostUserError::IncorrectFds)
+    }
+
+    fn set_inflight_fd(&mut self, inflight: &VhostUserInflight, fd: RawFd) -> Result<()> {
+        let mut node = self.node();
+        if node.acked_protocol_features & VhostUserProtocolFeatures::INFLIGHT_SHMFD.bits() == 0 {
+            return error_code(VhostUserError::InvalidOperation);
+        }
+
+        if inflight.mmap_size == 0 || inflight.num_queues == 0 || inflight.queue_size == 0 || fd < 0
+        {
+            return error_code(VhostUserError::InvalidParam);
+        }
+
+        let hdr = node.send_request_with_body(MasterReq::SET_INFLIGHT_FD, inflight, Some(&[fd]))?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
+    }
+
     fn get_max_mem_slots(&mut self) -> Result<u64> {
         let mut node = self.node();
         if node.acked_protocol_features & VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS.bits() == 0
@@ -646,6 +697,23 @@ impl MasterInternal {
         Ok(body)
     }
 
+    fn recv_reply_with_fds<T: Sized + Default + VhostUserMsgValidator>(
+        &mut self,
+        hdr: &VhostUserMsgHeader<MasterReq>,
+    ) -> VhostUserResult<(T, Option<Vec<RawFd>>)> {
+        if mem::size_of::<T>() > MAX_MSG_SIZE || hdr.is_reply() {
+            return Err(VhostUserError::InvalidParam);
+        }
+        self.check_state()?;
+
+        let (reply, body, rfds) = self.main_sock.recv_body::<T>()?;
+        if !reply.is_reply_for(&hdr) || rfds.is_none() || !body.is_valid() {
+            Endpoint::<MasterReq>::close_rfds(rfds);
+            return Err(VhostUserError::InvalidMessage);
+        }
+        Ok((body, rfds))
+    }
+
     fn recv_reply_with_payload<T: Sized + Default + VhostUserMsgValidator>(
         &mut self,
         hdr: &VhostUserMsgHeader<MasterReq>,
diff --git a/src/vhost_user/message.rs b/src/vhost_user/message.rs
index e906fb1..df9d720 100644
--- a/src/vhost_user/message.rs
+++ b/src/vhost_user/message.rs
@@ -673,6 +673,42 @@ impl VhostUserMsgValidator for VhostUserConfig {
 /// Payload for the VhostUserConfig message.
 pub type VhostUserConfigPayload = Vec<u8>;
 
+/// Single memory region descriptor as payload for ADD_MEM_REG and REM_MEM_REG
+/// requests.
+#[repr(C)]
+#[derive(Default, Clone)]
+pub struct VhostUserInflight {
+    /// Size of the area to track inflight I/O.
+    pub mmap_size: u64,
+    /// Offset of this area from the start of the supplied file descriptor.
+    pub mmap_offset: u64,
+    /// Number of virtqueues.
+    pub num_queues: u16,
+    /// Size of virtqueues.
+    pub queue_size: u16,
+}
+
+impl VhostUserInflight {
+    /// Create a new instance.
+    pub fn new(mmap_size: u64, mmap_offset: u64, num_queues: u16, queue_size: u16) -> Self {
+        VhostUserInflight {
+            mmap_size,
+            mmap_offset,
+            num_queues,
+            queue_size,
+        }
+    }
+}
+
+impl VhostUserMsgValidator for VhostUserInflight {
+    fn is_valid(&self) -> bool {
+        if self.num_queues == 0 || self.queue_size == 0 {
+            return false;
+        }
+        true
+    }
+}
+
 /*
  * TODO: support dirty log, live migration and IOTLB operations.
 #[repr(packed)]
diff --git a/src/vhost_user/mod.rs b/src/vhost_user/mod.rs
index 4e1b529..52d97f7 100644
--- a/src/vhost_user/mod.rs
+++ b/src/vhost_user/mod.rs
@@ -307,6 +307,11 @@ mod tests {
                 VhostUserProtocolFeatures::all().bits()
             );
 
+            // get_inflight_fd()
+            slave.handle_request().unwrap();
+            // set_inflight_fd()
+            slave.handle_request().unwrap();
+
             // get_queue_num()
             slave.handle_request().unwrap();
 
@@ -359,6 +364,19 @@ mod tests {
         assert_eq!(features.bits(), VhostUserProtocolFeatures::all().bits());
         master.set_protocol_features(features).unwrap();
 
+        // Retrieve inflight I/O tracking information
+        let (inflight_info, inflight_file) = master
+            .get_inflight_fd(&VhostUserInflight {
+                num_queues: 2,
+                queue_size: 256,
+                ..Default::default()
+            })
+            .unwrap();
+        // Set the buffer back to the backend
+        master
+            .set_inflight_fd(&inflight_info, inflight_file.as_raw_fd())
+            .unwrap();
+
         let num = master.get_queue_num().unwrap();
         assert_eq!(num, 2);
 
diff --git a/src/vhost_user/slave_req_handler.rs b/src/vhost_user/slave_req_handler.rs
index 9d7ea10..08b9ca3 100644
--- a/src/vhost_user/slave_req_handler.rs
+++ b/src/vhost_user/slave_req_handler.rs
@@ -1,6 +1,7 @@
 // Copyright (C) 2019 Alibaba Cloud Computing. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::fs::File;
 use std::mem;
 use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use std::os::unix::net::UnixStream;
@@ -62,6 +63,8 @@ pub trait VhostUserSlaveReqHandler {
     fn get_config(&self, offset: u32, size: u32, flags: VhostUserConfigFlags) -> Result<Vec<u8>>;
     fn set_config(&self, offset: u32, buf: &[u8], flags: VhostUserConfigFlags) -> Result<()>;
     fn set_slave_req_fd(&self, _vu_req: SlaveFsCacheReq) {}
+    fn get_inflight_fd(&self, inflight: &VhostUserInflight) -> Result<(VhostUserInflight, RawFd)>;
+    fn set_inflight_fd(&self, inflight: &VhostUserInflight, file: File) -> Result<()>;
     fn get_max_mem_slots(&self) -> Result<u64>;
     fn add_mem_region(&self, region: &VhostUserSingleMemoryRegion, fd: RawFd) -> Result<()>;
     fn remove_mem_region(&self, region: &VhostUserSingleMemoryRegion) -> Result<()>;
@@ -105,6 +108,11 @@ pub trait VhostUserSlaveReqHandlerMut {
     ) -> Result<Vec<u8>>;
     fn set_config(&mut self, offset: u32, buf: &[u8], flags: VhostUserConfigFlags) -> Result<()>;
     fn set_slave_req_fd(&mut self, _vu_req: SlaveFsCacheReq) {}
+    fn get_inflight_fd(
+        &mut self,
+        inflight: &VhostUserInflight,
+    ) -> Result<(VhostUserInflight, RawFd)>;
+    fn set_inflight_fd(&mut self, inflight: &VhostUserInflight, file: File) -> Result<()>;
     fn get_max_mem_slots(&mut self) -> Result<u64>;
     fn add_mem_region(&mut self, region: &VhostUserSingleMemoryRegion, fd: RawFd) -> Result<()>;
     fn remove_mem_region(&mut self, region: &VhostUserSingleMemoryRegion) -> Result<()>;
@@ -197,6 +205,14 @@ impl<T: VhostUserSlaveReqHandlerMut> VhostUserSlaveReqHandler for Mutex<T> {
         self.lock().unwrap().set_slave_req_fd(vu_req)
     }
 
+    fn get_inflight_fd(&self, inflight: &VhostUserInflight) -> Result<(VhostUserInflight, RawFd)> {
+        self.lock().unwrap().get_inflight_fd(inflight)
+    }
+
+    fn set_inflight_fd(&self, inflight: &VhostUserInflight, file: File) -> Result<()> {
+        self.lock().unwrap().set_inflight_fd(inflight, file)
+    }
+
     fn get_max_mem_slots(&self) -> Result<u64> {
         self.lock().unwrap().get_max_mem_slots()
     }
@@ -435,6 +451,41 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
                 self.check_request_size(&hdr, size, hdr.get_size() as usize)?;
                 self.set_slave_req_fd(&hdr, rfds)?;
             }
+            MasterReq::GET_INFLIGHT_FD => {
+                if self.acked_protocol_features & VhostUserProtocolFeatures::INFLIGHT_SHMFD.bits()
+                    == 0
+                {
+                    return Err(Error::InvalidOperation);
+                }
+
+                let msg = self.extract_request_body::<VhostUserInflight>(&hdr, size, &buf)?;
+                let (inflight, fd) = self.backend.get_inflight_fd(&msg)?;
+                let reply_hdr = self.new_reply_header::<VhostUserInflight>(&hdr, 0)?;
+                self.main_sock
+                    .send_message(&reply_hdr, &inflight, Some(&[fd]))?;
+            }
+            MasterReq::SET_INFLIGHT_FD => {
+                if self.acked_protocol_features & VhostUserProtocolFeatures::INFLIGHT_SHMFD.bits()
+                    == 0
+                {
+                    return Err(Error::InvalidOperation);
+                }
+                let file = if let Some(fds) = rfds {
+                    if fds.len() != 1 || fds[0] < 0 {
+                        Endpoint::<MasterReq>::close_rfds(Some(fds));
+                        return Err(Error::IncorrectFds);
+                    }
+
+                    // Safe because we know the fd is valid.
+                    unsafe { File::from_raw_fd(fds[0]) }
+                } else {
+                    return Err(Error::IncorrectFds);
+                };
+
+                let msg = self.extract_request_body::<VhostUserInflight>(&hdr, size, &buf)?;
+                let res = self.backend.set_inflight_fd(&msg, file);
+                self.send_ack_message(&hdr, res)?;
+            }
             MasterReq::GET_MAX_MEM_SLOTS => {
                 if self.acked_protocol_features
                     & VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS.bits()