crosvm: Use Reader/Writer interfaces in various virtio devices.

Switching the devices to the new interface reduces code duplication and
will ease fuzzing the devices as they now have a common input and output
interface for descriptors.

BUG=chromium:966258
TEST=vm.CrostiniStartEverything

Change-Id: I823c04dfc24e017433f8e8ab167bbd5dfafd338b
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/1647371
Reviewed-by: Stephen Barber <smbarber@chromium.org>
Reviewed-by: Zach Reizner <zachr@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Daniel Verkamp <dverkamp@chromium.org>

3 files changed, 51 insertions, 76 deletions

diff --git a/devices/src/virtio/balloon.rs b/devices/src/virtio/balloon.rs
index 8109395..633b3fc 100644
--- a/devices/src/virtio/balloon.rs
+++ b/devices/src/virtio/balloon.rs
@@ -17,7 +17,7 @@ use sys_util::{
 use vm_control::{BalloonControlCommand, BalloonControlResponseSocket};
 
 use super::{
-    copy_config, DescriptorChain, Queue, VirtioDevice, INTERRUPT_STATUS_CONFIG_CHANGED,
+    copy_config, Queue, Reader, VirtioDevice, INTERRUPT_STATUS_CONFIG_CHANGED,
     INTERRUPT_STATUS_USED_RING, TYPE_BALLOON, VIRTIO_F_VERSION_1,
 };
 
@@ -81,10 +81,6 @@ struct Worker {
     command_socket: BalloonControlResponseSocket,
 }
 
-fn valid_inflate_desc(desc: &DescriptorChain) -> bool {
-    !desc.is_write_only() && desc.len % 4 == 0
-}
-
 impl Worker {
     fn process_inflate_deflate(&mut self, inflate: bool) -> bool {
         let queue = if inflate {
@@ -95,19 +91,28 @@ impl Worker {
 
         let mut needs_interrupt = false;
         while let Some(avail_desc) = queue.pop(&self.mem) {
-            if inflate && valid_inflate_desc(&avail_desc) {
-                let num_addrs = avail_desc.len / 4;
-                for i in 0..num_addrs as usize {
-                    let addr = match avail_desc.addr.checked_add((i * 4) as u64) {
-                        Some(a) => a,
-                        None => break,
-                    };
-                    let guest_input: u32 = match self.mem.read_obj_from_addr(addr) {
-                        Ok(a) => a,
-                        Err(_) => continue,
+            let index = avail_desc.index;
+
+            if inflate {
+                let mut reader = Reader::new(&self.mem, avail_desc);
+                let data_length = reader.available_bytes();
+
+                if data_length % 4 != 0 {
+                    error!("invalid inflate buffer size: {}", data_length);
+                    continue;
+                }
+
+                let num_addrs = data_length / 4;
+                for _ in 0..num_addrs as usize {
+                    let guest_input = match reader.read_obj::<Le32>() {
+                        Ok(a) => a.to_native(),
+                        Err(err) => {
+                            error!("error while reading unused pages: {}", err);
+                            break;
+                        }
                     };
                     let guest_address =
-                        GuestAddress((guest_input as u64) << VIRTIO_BALLOON_PFN_SHIFT);
+                        GuestAddress((u64::from(guest_input)) << VIRTIO_BALLOON_PFN_SHIFT);
 
                     if self
                         .mem
@@ -120,7 +125,7 @@ impl Worker {
                 }
             }
 
-            queue.add_used(&self.mem, avail_desc.index, 0);
+            queue.add_used(&self.mem, index, 0);
             needs_interrupt = true;
         }
 
diff --git a/devices/src/virtio/rng.rs b/devices/src/virtio/rng.rs
index 1ad683a..4e6448f 100644
--- a/devices/src/virtio/rng.rs
+++ b/devices/src/virtio/rng.rs
@@ -11,9 +11,9 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::thread;
 
-use sys_util::{error, EventFd, GuestMemory, PollContext, PollToken};
+use sys_util::{error, warn, EventFd, GuestMemory, PollContext, PollToken};
 
-use super::{Queue, VirtioDevice, INTERRUPT_STATUS_USED_RING, TYPE_RNG};
+use super::{Queue, VirtioDevice, Writer, INTERRUPT_STATUS_USED_RING, TYPE_RNG};
 
 const QUEUE_SIZE: u16 = 256;
 const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE];
@@ -50,21 +50,17 @@ impl Worker {
 
         let mut needs_interrupt = false;
         while let Some(avail_desc) = queue.pop(&self.mem) {
-            let mut len = 0;
-
-            // Drivers can only read from the random device.
-            if avail_desc.is_write_only() {
-                // Fill the read with data from the random device on the host.
-                if self
-                    .mem
-                    .read_to_memory(avail_desc.addr, &self.random_file, avail_desc.len as usize)
-                    .is_ok()
-                {
-                    len = avail_desc.len;
+            let index = avail_desc.index;
+            let mut writer = Writer::new(&self.mem, avail_desc);
+            // Fill the entire descriptor chain buffer with random bytes.
+            let written = match writer.write_from(&self.random_file, std::usize::MAX) {
+                Ok(n) => n,
+                Err(e) => {
+                    warn!("Failed to write random data to the guest: {}", e);
+                    0
                 }
-            }
-
-            queue.add_used(&self.mem, avail_desc.index, len);
+            };
+            queue.add_used(&self.mem, index, written as u32);
             needs_interrupt = true;
         }
 
diff --git a/devices/src/virtio/tpm.rs b/devices/src/virtio/tpm.rs
index e7088ed..a3cdf48 100644
--- a/devices/src/virtio/tpm.rs
+++ b/devices/src/virtio/tpm.rs
@@ -12,10 +12,13 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::thread;
 
-use sys_util::{error, EventFd, GuestMemory, GuestMemoryError, PollContext, PollToken};
+use sys_util::{error, EventFd, GuestMemory, PollContext, PollToken};
 use tpm2;
 
-use super::{DescriptorChain, Queue, VirtioDevice, INTERRUPT_STATUS_USED_RING, TYPE_TPM};
+use super::{
+    DescriptorChain, DescriptorError, Queue, Reader, VirtioDevice, Writer,
+    INTERRUPT_STATUS_USED_RING, TYPE_TPM,
+};
 
 // A single queue of size 2. The guest kernel driver will enqueue a single
 // descriptor chain containing one command buffer and one response buffer at a
@@ -43,41 +46,19 @@ struct Device {
     simulator: tpm2::Simulator,
 }
 
-// Checks that the input descriptor chain holds a read-only descriptor followed
-// by a write-only descriptor, as required of the guest's virtio tpm driver.
-//
-// Returns those descriptors as a tuple: `(read_desc, write_desc)`.
-fn two_input_descriptors(desc: DescriptorChain) -> Result<(DescriptorChain, DescriptorChain)> {
-    let read_desc = desc;
-    if !read_desc.is_read_only() {
-        return Err(Error::ExpectedReadOnly);
-    }
-
-    let write_desc = match read_desc.next_descriptor() {
-        Some(desc) => desc,
-        None => return Err(Error::ExpectedSecondBuffer),
-    };
-
-    if !write_desc.is_write_only() {
-        return Err(Error::ExpectedWriteOnly);
-    }
-
-    Ok((read_desc, write_desc))
-}
-
 impl Device {
     fn perform_work(&mut self, mem: &GuestMemory, desc: DescriptorChain) -> Result<u32> {
-        let (read_desc, write_desc) = two_input_descriptors(desc)?;
+        let mut reader = Reader::new(mem, desc.clone());
+        let mut writer = Writer::new(mem, desc);
 
-        if read_desc.len > TPM_BUFSIZE as u32 {
+        if reader.available_bytes() > TPM_BUFSIZE {
             return Err(Error::CommandTooLong {
-                size: read_desc.len as usize,
+                size: reader.available_bytes(),
             });
         }
 
-        let mut command = vec![0u8; read_desc.len as usize];
-        mem.read_exact_at_addr(&mut command, read_desc.addr)
-            .map_err(Error::Read)?;
+        let mut command = vec![0u8; reader.available_bytes() as usize];
+        reader.read_exact(&mut command).map_err(Error::Read)?;
 
         let response = self.simulator.execute_command(&command);
 
@@ -87,17 +68,16 @@ impl Device {
             });
         }
 
-        if response.len() > write_desc.len as usize {
+        if response.len() > writer.available_bytes() {
             return Err(Error::BufferTooSmall {
-                size: write_desc.len as usize,
+                size: writer.available_bytes(),
                 required: response.len(),
             });
         }
 
-        mem.write_all_at_addr(&response, write_desc.addr)
-            .map_err(Error::Write)?;
+        writer.write_all(&response).map_err(Error::Write)?;
 
-        Ok(response.len() as u32)
+        Ok(writer.bytes_written() as u32)
     }
 }
 
@@ -306,14 +286,11 @@ impl BitOrAssign for NeedsInterrupt {
 type Result<T> = std::result::Result<T, Error>;
 
 enum Error {
-    ExpectedReadOnly,
-    ExpectedSecondBuffer,
-    ExpectedWriteOnly,
     CommandTooLong { size: usize },
-    Read(GuestMemoryError),
+    Read(DescriptorError),
     ResponseTooLong { size: usize },
     BufferTooSmall { size: usize, required: usize },
-    Write(GuestMemoryError),
+    Write(DescriptorError),
 }
 
 impl Display for Error {
@@ -321,9 +298,6 @@ impl Display for Error {
         use self::Error::*;
 
         match self {
-            ExpectedReadOnly => write!(f, "vtpm expected first descriptor to be read-only"),
-            ExpectedSecondBuffer => write!(f, "vtpm expected a second descriptor"),
-            ExpectedWriteOnly => write!(f, "vtpm expected second descriptor to be write-only"),
             CommandTooLong { size } => write!(
                 f,
                 "vtpm command is too long: {} > {} bytes",