1 files changed, 41 insertions, 16 deletions

diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index c63fbfe..d4922ab 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -19,8 +19,9 @@ use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use libc::{EINVAL, EIO, ENODEV};
 
 use kvm::{IrqRoute, IrqSource, Vm};
-use msg_socket::{MsgError, MsgOnSocket, MsgReceiver, MsgResult, MsgSender, MsgSocket};
+use msg_socket::{MsgError, MsgOnSocket, MsgResult, MsgSocket, UnixSeqpacketExt};
 use resources::{Alloc, GpuMemoryDesc, MmioType, SystemAllocator};
+use sys_util::net::UnixSeqpacket;
 use sys_util::{error, Error as SysError, EventFd, GuestAddress, MemoryMapping, MmapError, Result};
 
 /// A data structure that either owns or borrows a file descriptor.
@@ -189,6 +190,7 @@ impl Display for DiskControlCommand {
 
 #[derive(MsgOnSocket, Debug)]
 pub enum DiskControlResult {
+    Ready,
     Ok,
     Err(SysError),
 }
@@ -563,17 +565,34 @@ fn register_memory(
 }
 
 impl VmRequest {
+    pub fn socket<'s>(
+        &self,
+        balloon_host_socket: &'s BalloonControlRequestSocket,
+        disk_host_sockets: &'s [DiskControlRequestSocket],
+        usb_control_socket: &'s UsbControlSocket,
+    ) -> Option<&'s UnixSeqpacket> {
+        use VmRequest::*;
+        match *self {
+            Exit => None,
+            Suspend => None,
+            Resume => None,
+            BalloonCommand(_) => Some(balloon_host_socket.as_ref()),
+            DiskCommand { disk_index, .. } => disk_host_sockets.get(disk_index).map(AsRef::as_ref),
+            UsbCommand(_) => Some(usb_control_socket.as_ref()),
+        }
+    }
+
     /// Executes this request on the given Vm and other mutable state.
     ///
     /// This does not return a result, instead encapsulating the success or failure in a
     /// `VmResponse` with the intended purpose of sending the response back over the  socket that
     /// received this `VmRequest`.
+    ///
+    /// The `socket` parameter must be the value that was obtained by calling `VmRequest::socket`.
     pub fn execute(
         &self,
         run_mode: &mut Option<VmRunMode>,
-        balloon_host_socket: &BalloonControlRequestSocket,
-        disk_host_sockets: &[DiskControlRequestSocket],
-        usb_control_socket: &UsbControlSocket,
+        socket: Option<&UnixSeqpacket>,
     ) -> VmResponse {
         match *self {
             VmRequest::Exit => {
@@ -589,14 +608,18 @@ impl VmRequest {
                 VmResponse::Ok
             }
             VmRequest::BalloonCommand(BalloonControlCommand::Adjust { num_bytes }) => {
-                match balloon_host_socket.send(&BalloonControlCommand::Adjust { num_bytes }) {
+                match socket
+                    .unwrap()
+                    .send_msg_on_socket(&BalloonControlCommand::Adjust { num_bytes })
+                {
                     Ok(_) => VmResponse::Ok,
                     Err(_) => VmResponse::Err(SysError::last()),
                 }
             }
             VmRequest::BalloonCommand(BalloonControlCommand::Stats) => {
-                match balloon_host_socket.send(&BalloonControlCommand::Stats {}) {
-                    Ok(_) => match balloon_host_socket.recv() {
+                let socket = socket.unwrap();
+                match socket.send_msg_on_socket(&BalloonControlCommand::Stats {}) {
+                    Ok(_) => match socket.recv_msg_on_socket() {
                         Ok(BalloonControlResult::Stats {
                             stats,
                             balloon_actual,
@@ -612,19 +635,20 @@ impl VmRequest {
                     Err(_) => VmResponse::Err(SysError::last()),
                 }
             }
-            VmRequest::DiskCommand {
-                disk_index,
-                ref command,
-            } => {
+            VmRequest::DiskCommand { ref command, .. } => {
                 // Forward the request to the block device process via its control socket.
-                if let Some(sock) = disk_host_sockets.get(disk_index) {
-                    if let Err(e) = sock.send(command) {
+                if let Some(sock) = socket {
+                    if let Err(e) = sock.send_msg_on_socket(command) {
                         error!("disk socket send failed: {}", e);
                         VmResponse::Err(SysError::new(EINVAL))
                     } else {
-                        match sock.recv() {
+                        match sock.recv_msg_on_socket() {
                             Ok(DiskControlResult::Ok) => VmResponse::Ok,
                             Ok(DiskControlResult::Err(e)) => VmResponse::Err(e),
+                            Ok(resp) => {
+                                error!("unexpected disk socket result: {:?}", resp);
+                                VmResponse::Err(SysError::new(EINVAL))
+                            }
                             Err(e) => {
                                 error!("disk socket recv failed: {}", e);
                                 VmResponse::Err(SysError::new(EINVAL))
@@ -636,12 +660,13 @@ impl VmRequest {
                 }
             }
             VmRequest::UsbCommand(ref cmd) => {
-                let res = usb_control_socket.send(cmd);
+                let socket = socket.unwrap();
+                let res = socket.send_msg_on_socket(cmd);
                 if let Err(e) = res {
                     error!("fail to send command to usb control socket: {}", e);
                     return VmResponse::Err(SysError::new(EIO));
                 }
-                match usb_control_socket.recv() {
+                match socket.recv_msg_on_socket() {
                     Ok(response) => VmResponse::UsbResponse(response),
                     Err(e) => {
                         error!("fail to recv command from usb control socket: {}", e);