devices: PCI: add capabilities callback

virtio devices should be able to specify capabilities

BUG=chromium:936567
TEST=boot vm

Change-Id: I049f9967eb59a7904528fff5aea844e30c636e28
Reviewed-on: https://chromium-review.googlesource.com/1493012
Commit-Ready: Daniel Prilik <prilik@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Stephen Barber <smbarber@chromium.org>

4 files changed, 43 insertions, 21 deletions

diff --git a/arch/src/lib.rs b/arch/src/lib.rs
index aa6c8a3..38bc2c7 100644
--- a/arch/src/lib.rs
+++ b/arch/src/lib.rs
@@ -177,6 +177,7 @@ pub fn generate_pci_root(
         let device_ranges = device
             .allocate_device_bars(resources)
             .map_err(DeviceRegistrationError::AllocateDeviceAddrs)?;
+        device.register_device_capabilities();
         for (event, addr, datamatch) in device.ioeventfds() {
             let io_addr = IoeventAddress::Mmio(addr);
             vm.register_ioevent(&event, io_addr, datamatch)
diff --git a/devices/src/pci/pci_device.rs b/devices/src/pci/pci_device.rs
index d5bc741..bd811b5 100644
--- a/devices/src/pci/pci_device.rs
+++ b/devices/src/pci/pci_device.rs
@@ -76,6 +76,11 @@ pub trait PciDevice: Send {
         Ok(Vec::new())
     }
 
+    /// Register any capabilties specified by the device.
+    fn register_device_capabilities(&mut self) -> Result<()> {
+        Ok(())
+    }
+
     /// Gets a list of ioeventfds that should be registered with the running VM. The list is
     /// returned as a Vec of (eventfd, addr, datamatch) tuples.
     fn ioeventfds(&self) -> Vec<(&EventFd, u64, Datamatch)> {
@@ -160,6 +165,9 @@ impl<T: PciDevice + ?Sized> PciDevice for Box<T> {
     fn allocate_device_bars(&mut self, resources: &mut SystemAllocator) -> Result<Vec<(u64, u64)>> {
         (**self).allocate_device_bars(resources)
     }
+    fn register_device_capabilities(&mut self) -> Result<()> {
+        (**self).register_device_capabilities()
+    }
     fn ioeventfds(&self) -> Vec<(&EventFd, u64, Datamatch)> {
         (**self).ioeventfds()
     }
diff --git a/devices/src/virtio/virtio_device.rs b/devices/src/virtio/virtio_device.rs
index 64018bd..d6fe9b5 100644
--- a/devices/src/virtio/virtio_device.rs
+++ b/devices/src/virtio/virtio_device.rs
@@ -7,7 +7,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 
 use super::*;
-use pci::PciBarConfiguration;
+use pci::{PciBarConfiguration, PciCapability};
 use sys_util::{EventFd, GuestMemory};
 
 /// Trait for virtio devices to be driven by a virtio transport.
@@ -79,4 +79,9 @@ pub trait VirtioDevice: Send {
     fn get_device_bars(&self) -> Option<Vec<PciBarConfiguration>> {
         None
     }
+
+    /// Optionally returns additional capabiltiies needed by device
+    fn get_device_caps(&self) -> Option<Vec<Box<PciCapability>>> {
+        None
+    }
 }
diff --git a/devices/src/virtio/virtio_pci_device.rs b/devices/src/virtio/virtio_pci_device.rs
index bfff317..35f227d 100644
--- a/devices/src/virtio/virtio_pci_device.rs
+++ b/devices/src/virtio/virtio_pci_device.rs
@@ -19,7 +19,7 @@ use sys_util::{EventFd, GuestMemory, Result};
 
 use self::virtio_pci_common_config::VirtioPciCommonConfig;
 
-enum PciCapabilityType {
+pub enum PciCapabilityType {
     CommonConfig = 1,
     NotifyConfig = 2,
     IsrConfig = 3,
@@ -69,7 +69,7 @@ impl VirtioPciCap {
 #[allow(dead_code)]
 #[repr(packed)]
 #[derive(Clone, Copy)]
-struct VirtioPciNotifyCap {
+pub struct VirtioPciNotifyCap {
     cap: VirtioPciCap,
     notify_off_multiplier: Le32,
 }
@@ -231,7 +231,7 @@ impl VirtioPciDevice {
         }
     }
 
-    fn add_pci_capabilities(
+    fn add_settings_pci_capabilities(
         &mut self,
         settings_bar: u8,
     ) -> std::result::Result<(), PciDeviceError> {
@@ -338,7 +338,7 @@ impl PciDevice for VirtioPciDevice {
         ranges.push((settings_config_addr, CAPABILITY_BAR_SIZE));
 
         // Once the BARs are allocated, the capabilities can be added to the PCI configuration.
-        self.add_pci_capabilities(settings_bar)?;
+        self.add_settings_pci_capabilities(settings_bar)?;
 
         Ok(ranges)
     }
@@ -348,26 +348,34 @@ impl PciDevice for VirtioPciDevice {
         resources: &mut SystemAllocator,
     ) -> std::result::Result<Vec<(u64, u64)>, PciDeviceError> {
         let mut ranges = Vec::new();
-        let configs = self.device.get_device_bars();
-        match configs {
-            Some(configs) => {
-                for mut config in configs {
-                    let device_addr = resources
-                        .allocate_device_addresses(config.get_size())
-                        .ok_or(PciDeviceError::IoAllocationFailed(config.get_size()))?;
-                    config.set_address(device_addr);
-                    let _device_bar = self
-                        .config_regs
-                        .add_pci_bar(&config)
-                        .map_err(|e| PciDeviceError::IoRegistrationFailed(device_addr, e))?;
-                    ranges.push((device_addr, config.get_size()));
-                }
+        if let Some(configs) = self.device.get_device_bars() {
+            for mut config in configs {
+                let device_addr = resources
+                    .allocate_device_addresses(config.get_size())
+                    .ok_or(PciDeviceError::IoAllocationFailed(config.get_size()))?;
+                config.set_address(device_addr);
+                let _device_bar = self
+                    .config_regs
+                    .add_pci_bar(&config)
+                    .map_err(|e| PciDeviceError::IoRegistrationFailed(device_addr, e))?;
+                ranges.push((device_addr, config.get_size()));
             }
-            None => (),
-        };
+        }
         Ok(ranges)
     }
 
+    fn register_device_capabilities(&mut self) -> std::result::Result<(), PciDeviceError> {
+        if let Some(caps) = self.device.get_device_caps() {
+            for cap in caps {
+                self.config_regs
+                    .add_capability(&*cap)
+                    .map_err(|e| PciDeviceError::CapabilitiesSetup(e))?;
+            }
+        }
+
+        Ok(())
+    }
+
     fn ioeventfds(&self) -> Vec<(&EventFd, u64, Datamatch)> {
         let bar0 = self.config_regs.get_bar_addr(self.settings_bar as usize) as u64;
         let notify_base = bar0 + NOTIFICATION_BAR_OFFSET;