Extracts BalloonAdjust from VmRequest.

Extracts BalloonAdjust from VmRequest into BalloonControlCommand.

BUG=None
TEST=cargo test
TEST=cargo test --package msg_socket
TEST=cargo test --package devices
TEST=cargo test --package vm_control
TEST=tast -verbose run ${IP} vm.CrostiniStartEverything

Change-Id: Ia9f5778c37c8fd4fa560df413134d1b441142f64
Reviewed-on: https://chromium-review.googlesource.com/1565298
Commit-Ready: Jakub Staroń <jstaron@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Stephen Barber <smbarber@chromium.org>

6 files changed, 84 insertions, 53 deletions

diff --git a/devices/src/virtio/balloon.rs b/devices/src/virtio/balloon.rs
index baff7ab..2187381 100644
--- a/devices/src/virtio/balloon.rs
+++ b/devices/src/virtio/balloon.rs
@@ -6,17 +6,17 @@ use std;
 use std::cmp;
 use std::fmt::{self, Display};
 use std::io::Write;
-use std::mem;
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::thread;
 
-use byteorder::{ByteOrder, LittleEndian, ReadBytesExt, WriteBytesExt};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use msg_socket::MsgReceiver;
 use sys_util::{
-    self, error, info, net::UnixSeqpacket, warn, EventFd, GuestAddress, GuestMemory, PollContext,
-    PollToken,
+    self, error, info, warn, EventFd, GuestAddress, GuestMemory, PollContext, PollToken,
 };
+use vm_control::{BalloonControlCommand, BalloonControlResponseSocket};
 
 use super::{
     DescriptorChain, Queue, VirtioDevice, INTERRUPT_STATUS_CONFIG_CHANGED,
@@ -69,7 +69,7 @@ struct Worker {
     interrupt_evt: EventFd,
     interrupt_resample_evt: EventFd,
     config: Arc<BalloonConfig>,
-    command_socket: UnixSeqpacket,
+    command_socket: BalloonControlResponseSocket,
 }
 
 fn valid_inflate_desc(desc: &DescriptorChain) -> bool {
@@ -194,17 +194,17 @@ impl Worker {
                         needs_interrupt |= self.process_inflate_deflate(false);
                     }
                     Token::CommandSocket => {
-                        let mut buf = [0u8; mem::size_of::<u64>()];
-                        if let Ok(count) = self.command_socket.recv(&mut buf) {
-                            // Ignore any malformed messages that are not exactly 8 bytes long.
-                            if count == mem::size_of::<u64>() {
-                                let num_bytes = LittleEndian::read_u64(&buf);
-                                let num_pages = (num_bytes >> VIRTIO_BALLOON_PFN_SHIFT) as usize;
-                                info!("ballon config changed to consume {} pages", num_pages);
-
-                                self.config.num_pages.store(num_pages, Ordering::Relaxed);
-                                self.signal_config_changed();
-                            }
+                        if let Ok(req) = self.command_socket.recv() {
+                            match req {
+                                BalloonControlCommand::Adjust { num_bytes } => {
+                                    let num_pages =
+                                        (num_bytes >> VIRTIO_BALLOON_PFN_SHIFT) as usize;
+                                    info!("ballon config changed to consume {} pages", num_pages);
+
+                                    self.config.num_pages.store(num_pages, Ordering::Relaxed);
+                                    self.signal_config_changed();
+                                }
+                            };
                         }
                     }
                     Token::InterruptResample => {
@@ -232,7 +232,7 @@ impl Worker {
 
 /// Virtio device for memory balloon inflation/deflation.
 pub struct Balloon {
-    command_socket: Option<UnixSeqpacket>,
+    command_socket: Option<BalloonControlResponseSocket>,
     config: Arc<BalloonConfig>,
     features: u64,
     kill_evt: Option<EventFd>,
@@ -240,7 +240,7 @@ pub struct Balloon {
 
 impl Balloon {
     /// Create a new virtio balloon device.
-    pub fn new(command_socket: UnixSeqpacket) -> Result<Balloon> {
+    pub fn new(command_socket: BalloonControlResponseSocket) -> Result<Balloon> {
         Ok(Balloon {
             command_socket: Some(command_socket),
             config: Arc::new(BalloonConfig {
diff --git a/msg_socket/src/msg_on_socket.rs b/msg_socket/src/msg_on_socket.rs
index 635820d..0a9ab37 100644
--- a/msg_socket/src/msg_on_socket.rs
+++ b/msg_socket/src/msg_on_socket.rs
@@ -172,6 +172,24 @@ impl<T: MsgOnSocket> MsgOnSocket for Option<T> {
     }
 }
 
+impl MsgOnSocket for () {
+    fn msg_size() -> usize {
+        0
+    }
+
+    fn max_fd_count() -> usize {
+        0
+    }
+
+    unsafe fn read_from_buffer(_buffer: &[u8], _fds: &[RawFd]) -> MsgResult<(Self, usize)> {
+        Ok(((), 0))
+    }
+
+    fn write_to_buffer(&self, _buffer: &mut [u8], _fds: &mut [RawFd]) -> MsgResult<usize> {
+        Ok(0)
+    }
+}
+
 macro_rules! rawfd_impl {
     ($type:ident) => {
         impl MsgOnSocket for $type {
diff --git a/msg_socket/tests/unit.rs b/msg_socket/tests/unit.rs
new file mode 100644
index 0000000..9855752
--- /dev/null
+++ b/msg_socket/tests/unit.rs
@@ -0,0 +1,12 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use msg_socket::*;
+
+#[test]
+fn sock_send_recv_unit() {
+    let (req, res) = pair::<(), ()>().unwrap();
+    req.send(&()).unwrap();
+    let _ = res.recv().unwrap();
+}
diff --git a/src/linux.rs b/src/linux.rs
index 6e2c2c9..07cb5a9 100644
--- a/src/linux.rs
+++ b/src/linux.rs
@@ -9,7 +9,6 @@ use std::ffi::CStr;
 use std::fmt::{self, Display};
 use std::fs::{File, OpenOptions};
 use std::io::{self, stdin, Read};
-use std::mem;
 use std::net::Ipv4Addr;
 use std::os::unix::io::{FromRawFd, RawFd};
 use std::os::unix::net::UnixStream;
@@ -23,7 +22,6 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use libc::{self, c_int, gid_t, uid_t};
 
 use audio_streams::DummyStreamSource;
-use byteorder::{ByteOrder, LittleEndian};
 use devices::virtio::{self, VirtioDevice};
 use devices::{self, HostBackendDeviceProvider, PciDevice, VirtioPciDevice, XhciController};
 use io_jail::{self, Minijail};
@@ -37,7 +35,7 @@ use remain::sorted;
 #[cfg(feature = "gpu-forward")]
 use resources::Alloc;
 use sync::{Condvar, Mutex};
-use sys_util::net::{UnixSeqpacket, UnixSeqpacketListener, UnlinkUnixSeqpacketListener};
+use sys_util::net::{UnixSeqpacketListener, UnlinkUnixSeqpacketListener};
 use sys_util::{
     self, block_signal, clear_signal, drop_capabilities, error, flock, get_blocked_signals,
     get_group_id, get_user_id, getegid, geteuid, info, register_signal_handler, set_cpu_affinity,
@@ -48,6 +46,7 @@ use sys_util::{
 use sys_util::{GuestAddress, MemoryMapping, Protection};
 use vhost;
 use vm_control::{
+    BalloonControlCommand, BalloonControlRequestSocket, BalloonControlResponseSocket,
     DiskControlCommand, DiskControlRequestSocket, DiskControlResponseSocket, DiskControlResult,
     UsbControlSocket, VmControlRequestSocket, VmControlResponseSocket, VmRequest, VmResponse,
     VmRunMode,
@@ -457,7 +456,7 @@ fn create_vinput_device(cfg: &Config, dev_path: &Path) -> DeviceResult {
     })
 }
 
-fn create_balloon_device(cfg: &Config, socket: UnixSeqpacket) -> DeviceResult {
+fn create_balloon_device(cfg: &Config, socket: BalloonControlResponseSocket) -> DeviceResult {
     let dev = virtio::Balloon::new(socket).map_err(Error::BalloonDeviceNew)?;
 
     Ok(VirtioDeviceStub {
@@ -673,7 +672,7 @@ fn create_virtio_devices(
     mem: &GuestMemory,
     _exit_evt: &EventFd,
     wayland_device_socket: VmControlRequestSocket,
-    balloon_device_socket: UnixSeqpacket,
+    balloon_device_socket: BalloonControlResponseSocket,
     disk_device_sockets: &mut Vec<DiskControlResponseSocket>,
 ) -> DeviceResult<Vec<VirtioDeviceStub>> {
     let mut devs = Vec::new();
@@ -773,7 +772,7 @@ fn create_devices(
     mem: &GuestMemory,
     exit_evt: &EventFd,
     wayland_device_socket: VmControlRequestSocket,
-    balloon_device_socket: UnixSeqpacket,
+    balloon_device_socket: BalloonControlResponseSocket,
     disk_device_sockets: &mut Vec<DiskControlResponseSocket>,
     usb_provider: HostBackendDeviceProvider,
 ) -> DeviceResult<Vec<(Box<dyn PciDevice>, Option<Minijail>)>> {
@@ -1131,7 +1130,7 @@ pub fn run_config(cfg: Config) -> Result<()> {
     control_sockets.push(wayland_host_socket);
     // Balloon gets a special socket so balloon requests can be forwarded from the main process.
     let (balloon_host_socket, balloon_device_socket) =
-        UnixSeqpacket::pair().map_err(Error::CreateSocket)?;
+        msg_socket::pair::<BalloonControlCommand, ()>().map_err(Error::CreateSocket)?;
 
     // Create one control socket per disk.
     let mut disk_device_sockets = Vec::new();
@@ -1215,7 +1214,7 @@ fn run_control(
     mut linux: RunnableLinuxVm,
     control_server_socket: Option<UnlinkUnixSeqpacketListener>,
     mut control_sockets: Vec<VmControlResponseSocket>,
-    balloon_host_socket: UnixSeqpacket,
+    balloon_host_socket: BalloonControlRequestSocket,
     disk_host_sockets: &[DiskControlRequestSocket],
     usb_control_socket: UsbControlSocket,
     sigchld_fd: SignalFd,
@@ -1416,9 +1415,10 @@ fn run_control(
                             } else {
                                 0
                             };
-                        let mut buf = [0u8; mem::size_of::<u64>()];
-                        LittleEndian::write_u64(&mut buf, current_balloon_memory);
-                        if let Err(e) = balloon_host_socket.send(&buf) {
+                        let command = BalloonControlCommand::Adjust {
+                            num_bytes: current_balloon_memory,
+                        };
+                        if let Err(e) = balloon_host_socket.send(&command) {
                             warn!("failed to send memory value to balloon device: {}", e);
                         }
                     }
@@ -1431,9 +1431,10 @@ fn run_control(
                             max_balloon_memory,
                         );
                         if current_balloon_memory != old_balloon_memory {
-                            let mut buf = [0u8; mem::size_of::<u64>()];
-                            LittleEndian::write_u64(&mut buf, current_balloon_memory);
-                            if let Err(e) = balloon_host_socket.send(&buf) {
+                            let command = BalloonControlCommand::Adjust {
+                                num_bytes: current_balloon_memory,
+                            };
+                            if let Err(e) = balloon_host_socket.send(&command) {
                                 warn!("failed to send memory value to balloon device: {}", e);
                             }
                         }
diff --git a/src/main.rs b/src/main.rs
index 01c096b..b5ba7e7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -27,8 +27,8 @@ use sys_util::{
     validate_raw_fd, warn,
 };
 use vm_control::{
-    DiskControlCommand, MaybeOwnedFd, UsbControlCommand, UsbControlResult, VmControlRequestSocket,
-    VmRequest, VmResponse,
+    BalloonControlCommand, DiskControlCommand, MaybeOwnedFd, UsbControlCommand, UsbControlResult,
+    VmControlRequestSocket, VmRequest, VmResponse,
 };
 
 use crate::argument::{print_help, set_arguments, Argument};
@@ -891,7 +891,8 @@ fn balloon_vms(mut args: std::env::Args) -> std::result::Result<(), ()> {
         }
     };
 
-    vms_request(&VmRequest::BalloonAdjust(num_bytes), args)
+    let command = BalloonControlCommand::Adjust { num_bytes };
+    vms_request(&VmRequest::BalloonCommand(command), args)
 }
 
 fn create_qcow2(mut args: std::env::Args) -> std::result::Result<(), ()> {
diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index 474380c..9d68f52 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -17,13 +17,10 @@ use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 
 use libc::{EINVAL, EIO, ENODEV};
 
-use byteorder::{LittleEndian, WriteBytesExt};
 use kvm::Vm;
 use msg_socket::{MsgOnSocket, MsgReceiver, MsgResult, MsgSender, MsgSocket};
 use resources::{GpuMemoryDesc, SystemAllocator};
-use sys_util::{
-    error, net::UnixSeqpacket, Error as SysError, GuestAddress, MemoryMapping, MmapError, Result,
-};
+use sys_util::{error, Error as SysError, GuestAddress, MemoryMapping, MmapError, Result};
 
 /// A file descriptor either borrowed or owned by this.
 #[derive(Debug)]
@@ -92,6 +89,12 @@ impl Default for VmRunMode {
 }
 
 #[derive(MsgOnSocket, Debug)]
+pub enum BalloonControlCommand {
+    /// Set the size of the VM's balloon.
+    Adjust { num_bytes: u64 },
+}
+
+#[derive(MsgOnSocket, Debug)]
 pub enum DiskControlCommand {
     /// Resize a disk to `new_size` in bytes.
     Resize { new_size: u64 },
@@ -155,6 +158,9 @@ impl Display for UsbControlResult {
     }
 }
 
+pub type BalloonControlRequestSocket = MsgSocket<BalloonControlCommand, ()>;
+pub type BalloonControlResponseSocket = MsgSocket<(), BalloonControlCommand>;
+
 pub type DiskControlRequestSocket = MsgSocket<DiskControlCommand, DiskControlResult>;
 pub type DiskControlResponseSocket = MsgSocket<DiskControlResult, DiskControlCommand>;
 
@@ -168,8 +174,6 @@ pub type VmControlResponseSocket = MsgSocket<VmResponse, VmRequest>;
 /// Unless otherwise noted, each request should expect a `VmResponse::Ok` to be received on success.
 #[derive(MsgOnSocket, Debug)]
 pub enum VmRequest {
-    /// Set the size of the VM's balloon in bytes.
-    BalloonAdjust(u64),
     /// Break the VM's run loop and exit.
     Exit,
     /// Suspend the VM's VCPUs until resume.
@@ -188,6 +192,8 @@ pub enum VmRequest {
         height: u32,
         format: u32,
     },
+    /// Command for balloon driver.
+    BalloonCommand(BalloonControlCommand),
     /// Send a command to a disk chosen by `disk_index`.
     /// `disk_index` is a 0-based count of `--disk`, `--rwdisk`, and `-r` command-line options.
     DiskCommand {
@@ -242,7 +248,7 @@ impl VmRequest {
         vm: &mut Vm,
         sys_allocator: &mut SystemAllocator,
         run_mode: &mut Option<VmRunMode>,
-        balloon_host_socket: &UnixSeqpacket,
+        balloon_host_socket: &BalloonControlRequestSocket,
         disk_host_sockets: &[DiskControlRequestSocket],
         usb_control_socket: &UsbControlSocket,
     ) -> VmResponse {
@@ -269,17 +275,10 @@ impl VmRequest {
                 Ok(_) => VmResponse::Ok,
                 Err(e) => VmResponse::Err(e),
             },
-            VmRequest::BalloonAdjust(num_pages) => {
-                let mut buf = [0u8; 8];
-                // write_u64 can't fail as the buffer is 8 bytes long.
-                (&mut buf[0..])
-                    .write_u64::<LittleEndian>(num_pages)
-                    .unwrap();
-                match balloon_host_socket.send(&buf) {
-                    Ok(_) => VmResponse::Ok,
-                    Err(_) => VmResponse::Err(SysError::last()),
-                }
-            }
+            VmRequest::BalloonCommand(ref command) => match balloon_host_socket.send(command) {
+                Ok(_) => VmResponse::Ok,
+                Err(_) => VmResponse::Err(SysError::last()),
+            },
             VmRequest::AllocateAndRegisterGpuMemory {
                 width,
                 height,