vhost_user: Fix NEED_REPLY for all messages

Except the list of commands from which we can expect a reply, all the
other commands should be able to reply with an acknowledgment if being
asked for.

For the commands such as SET_PROTOCOL_FEATURES, which usually happens
before the protocol features negotiation is complete, we should also be
able to reply with an ACK if needed. Indeed, the protocol features
negotiation could have happened earlier, which means we might expect a
reply ACK from GET_FEATURES command if this happens later on.

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

2 files changed, 35 insertions, 46 deletions

diff --git a/src/vhost_user/master.rs b/src/vhost_user/master.rs
index 2d98987..0e8b22c 100644
--- a/src/vhost_user/master.rs
+++ b/src/vhost_user/master.rs
@@ -158,11 +158,9 @@ impl VhostBackend for Master {
     fn set_features(&self, features: u64) -> Result<()> {
         let mut node = self.node();
         let val = VhostUserU64::new(features);
-        let _ = node.send_request_with_body(MasterReq::SET_FEATURES, &val, None)?;
-        // Don't wait for ACK here because the protocol feature negotiation process hasn't been
-        // completed yet.
+        let hdr = node.send_request_with_body(MasterReq::SET_FEATURES, &val, None)?;
         node.acked_virtio_features = features & node.virtio_features;
-        Ok(())
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     /// Set the current Master as an owner of the session.
@@ -170,18 +168,14 @@ impl VhostBackend for Master {
         // We unwrap() the return value to assert that we are not expecting threads to ever fail
         // while holding the lock.
         let mut node = self.node();
-        let _ = node.send_request_header(MasterReq::SET_OWNER, None)?;
-        // Don't wait for ACK here because the protocol feature negotiation process hasn't been
-        // completed yet.
-        Ok(())
+        let hdr = node.send_request_header(MasterReq::SET_OWNER, None)?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     fn reset_owner(&self) -> Result<()> {
         let mut node = self.node();
-        let _ = node.send_request_header(MasterReq::RESET_OWNER, None)?;
-        // Don't wait for ACK here because the protocol feature negotiation process hasn't been
-        // completed yet.
-        Ok(())
+        let hdr = node.send_request_header(MasterReq::RESET_OWNER, None)?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     /// Set the memory map regions on the slave so it can translate the vring
@@ -237,8 +231,8 @@ impl VhostBackend for Master {
     fn set_log_fd(&self, fd: RawFd) -> Result<()> {
         let mut node = self.node();
         let fds = [fd];
-        node.send_request_header(MasterReq::SET_LOG_FD, Some(&fds))?;
-        Ok(())
+        let hdr = node.send_request_header(MasterReq::SET_LOG_FD, Some(&fds))?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     /// Set the size of the queue.
@@ -300,8 +294,8 @@ impl VhostBackend for Master {
         if queue_index as u64 >= node.max_queue_num {
             return error_code(VhostUserError::InvalidParam);
         }
-        node.send_fd_for_vring(MasterReq::SET_VRING_CALL, queue_index, fd.as_raw_fd())?;
-        Ok(())
+        let hdr = node.send_fd_for_vring(MasterReq::SET_VRING_CALL, queue_index, fd.as_raw_fd())?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     /// Set the event file descriptor for adding buffers to the vring.
@@ -313,8 +307,8 @@ impl VhostBackend for Master {
         if queue_index as u64 >= node.max_queue_num {
             return error_code(VhostUserError::InvalidParam);
         }
-        node.send_fd_for_vring(MasterReq::SET_VRING_KICK, queue_index, fd.as_raw_fd())?;
-        Ok(())
+        let hdr = node.send_fd_for_vring(MasterReq::SET_VRING_KICK, queue_index, fd.as_raw_fd())?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     /// Set the event file descriptor to signal when error occurs.
@@ -325,8 +319,8 @@ impl VhostBackend for Master {
         if queue_index as u64 >= node.max_queue_num {
             return error_code(VhostUserError::InvalidParam);
         }
-        node.send_fd_for_vring(MasterReq::SET_VRING_ERR, queue_index, fd.as_raw_fd())?;
-        Ok(())
+        let hdr = node.send_fd_for_vring(MasterReq::SET_VRING_ERR, queue_index, fd.as_raw_fd())?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 }
 
@@ -355,12 +349,12 @@ impl VhostUserMaster for Master {
             return error_code(VhostUserError::InvalidOperation);
         }
         let val = VhostUserU64::new(features.bits());
-        let _ = node.send_request_with_body(MasterReq::SET_PROTOCOL_FEATURES, &val, None)?;
+        let hdr = node.send_request_with_body(MasterReq::SET_PROTOCOL_FEATURES, &val, None)?;
         // Don't wait for ACK here because the protocol feature negotiation process hasn't been
         // completed yet.
         node.acked_protocol_features = features.bits();
         node.protocol_features_ready = true;
-        Ok(())
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     fn get_queue_num(&mut self) -> Result<u64> {
@@ -458,8 +452,8 @@ impl VhostUserMaster for Master {
         }
 
         let fds = [fd];
-        node.send_request_header(MasterReq::SET_SLAVE_REQ_FD, Some(&fds))?;
-        Ok(())
+        let hdr = node.send_request_header(MasterReq::SET_SLAVE_REQ_FD, Some(&fds))?;
+        node.wait_for_ack(&hdr).map_err(|e| e.into())
     }
 
     fn get_inflight_fd(
diff --git a/src/vhost_user/slave_req_handler.rs b/src/vhost_user/slave_req_handler.rs
index 08b9ca3..f2130b7 100644
--- a/src/vhost_user/slave_req_handler.rs
+++ b/src/vhost_user/slave_req_handler.rs
@@ -318,11 +318,13 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
         match hdr.get_code() {
             MasterReq::SET_OWNER => {
                 self.check_request_size(&hdr, size, 0)?;
-                self.backend.set_owner()?;
+                let res = self.backend.set_owner();
+                self.send_ack_message(&hdr, res)?;
             }
             MasterReq::RESET_OWNER => {
                 self.check_request_size(&hdr, size, 0)?;
-                self.backend.reset_owner()?;
+                let res = self.backend.reset_owner();
+                self.send_ack_message(&hdr, res)?;
             }
             MasterReq::GET_FEATURES => {
                 self.check_request_size(&hdr, size, 0)?;
@@ -334,9 +336,10 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
             }
             MasterReq::SET_FEATURES => {
                 let msg = self.extract_request_body::<VhostUserU64>(&hdr, size, &buf)?;
-                self.backend.set_features(msg.value)?;
+                let res = self.backend.set_features(msg.value);
                 self.acked_virtio_features = msg.value;
                 self.update_reply_ack_flag();
+                self.send_ack_message(&hdr, res)?;
             }
             MasterReq::SET_MEM_TABLE => {
                 let res = self.set_mem_table(&hdr, size, &buf, rfds);
@@ -401,9 +404,10 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
             }
             MasterReq::SET_PROTOCOL_FEATURES => {
                 let msg = self.extract_request_body::<VhostUserU64>(&hdr, size, &buf)?;
-                self.backend.set_protocol_features(msg.value)?;
+                let res = self.backend.set_protocol_features(msg.value);
                 self.acked_protocol_features = msg.value;
                 self.update_reply_ack_flag();
+                self.send_ack_message(&hdr, res)?;
             }
             MasterReq::GET_QUEUE_NUM => {
                 if self.acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() == 0 {
@@ -442,14 +446,16 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
                     return Err(Error::InvalidOperation);
                 }
                 self.check_request_size(&hdr, size, hdr.get_size() as usize)?;
-                self.set_config(&hdr, size, &buf)?;
+                let res = self.set_config(size, &buf);
+                self.send_ack_message(&hdr, res)?;
             }
             MasterReq::SET_SLAVE_REQ_FD => {
                 if self.acked_protocol_features & VhostUserProtocolFeatures::SLAVE_REQ.bits() == 0 {
                     return Err(Error::InvalidOperation);
                 }
                 self.check_request_size(&hdr, size, hdr.get_size() as usize)?;
-                self.set_slave_req_fd(&hdr, rfds)?;
+                let res = self.set_slave_req_fd(rfds);
+                self.send_ack_message(&hdr, res)?;
             }
             MasterReq::GET_INFLIGHT_FD => {
                 if self.acked_protocol_features & VhostUserProtocolFeatures::INFLIGHT_SHMFD.bits()
@@ -631,12 +637,7 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
         Ok(())
     }
 
-    fn set_config(
-        &mut self,
-        hdr: &VhostUserMsgHeader<MasterReq>,
-        size: usize,
-        buf: &[u8],
-    ) -> Result<()> {
+    fn set_config(&mut self, size: usize, buf: &[u8]) -> Result<()> {
         if size > MAX_MSG_SIZE || size < mem::size_of::<VhostUserConfig>() {
             return Err(Error::InvalidMessage);
         }
@@ -653,22 +654,16 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
             None => return Err(Error::InvalidMessage),
         }
 
-        let res = self.backend.set_config(msg.offset, buf, flags);
-        self.send_ack_message(&hdr, res)?;
-        Ok(())
+        self.backend.set_config(msg.offset, buf, flags)
     }
 
-    fn set_slave_req_fd(
-        &mut self,
-        hdr: &VhostUserMsgHeader<MasterReq>,
-        rfds: Option<Vec<RawFd>>,
-    ) -> Result<()> {
+    fn set_slave_req_fd(&mut self, rfds: Option<Vec<RawFd>>) -> Result<()> {
         if let Some(fds) = rfds {
             if fds.len() == 1 {
                 let sock = unsafe { UnixStream::from_raw_fd(fds[0]) };
                 let vu_req = SlaveFsCacheReq::from_stream(sock);
                 self.backend.set_slave_req_fd(vu_req);
-                self.send_ack_message(&hdr, Ok(()))
+                Ok(())
             } else {
                 Err(Error::InvalidMessage)
             }
@@ -825,7 +820,7 @@ impl<S: VhostUserSlaveReqHandler> SlaveReqHandler<S> {
             let msg = VhostUserU64::new(val);
             self.main_sock.send_message(&hdr, &msg, None)?;
         }
-        Ok(())
+        res
     }
 
     fn send_reply_message<T>(