fix(snapshot): handle errors in PCI cfg on snap restore

src/vmm/src/devices/virtio/transport/pci/device.rs

@@ -32,7 +32,9 @@ use crate::devices::virtio::transport::pci::common_config::{
 use crate::devices::virtio::transport::pci::device_status::*;
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
 use crate::logger::{debug, error, warn};
-use crate::pci::configuration::{PciCapability, PciConfiguration, PciConfigurationState};
+use crate::pci::configuration::{
+    PciCapability, PciConfiguration, PciConfigurationError, PciConfigurationState,
+};
 use crate::pci::msix::{MsixCap, MsixConfig, MsixConfigState};
 use crate::pci::{
     BarReprogrammingParams, DeviceRelocationError, PciCapabilityId, PciClassCode, PciDevice,
@@ -239,6 +241,8 @@ pub enum VirtioPciDeviceError {
     CreateVirtioPciDevice(#[from] DeviceRelocationError),
     /// Error creating MSI configuration: {0}
     Msi(#[from] InterruptError),
+    /// Invalid PCI configuration state: {0}
+    PciConfiguration(#[from] PciConfigurationError),
 }
 
 pub struct VirtioPciDevice {
@@ -412,11 +416,13 @@ impl VirtioPciDevice {
         let pci_config = PciConfiguration::type0_from_state(
             state.pci_configuration_state,
             Some(msix_config.clone()),
-        );
+        )?;
         let virtio_common_config = VirtioPciCommonConfig::new(state.pci_dev_state);
         let cap_pci_cfg_info = VirtioPciCfgCapInfo {
             offset: state.cap_pci_cfg_offset,
-            cap: *VirtioPciCfgCap::from_slice(&state.cap_pci_cfg).unwrap(),
+            cap: *VirtioPciCfgCap::from_slice(&state.cap_pci_cfg).ok_or(
+                PciConfigurationError::InvalidCapPciCfgLength(state.cap_pci_cfg.len()),
+            )?,
         };
 
         let interrupt = Arc::new(VirtioInterruptMsix::new(

src/vmm/src/pci/configuration.rs

@@ -15,6 +15,19 @@ use super::msix::MsixConfig;
 use crate::logger::{info, warn};
 use crate::pci::{PciCapabilityId, PciClassCode};
 
+/// Errors from restoring PCI configuration state.
+#[derive(Debug, thiserror::Error, displaydoc::Display)]
+pub enum PciConfigurationError {
+    /// Invalid registers length: {0}
+    InvalidRegistersLength(usize),
+    /// Invalid writable_bits length: {0}
+    InvalidWritableBitsLength(usize),
+    /// Invalid bars length: {0}
+    InvalidBarsLength(usize),
+    /// Invalid cap_pci_cfg length: {0}
+    InvalidCapPciCfgLength(usize),
+}
+
 // The number of 32bit registers in the config space, 4096 bytes.
 const NUM_CONFIGURATION_REGISTERS: usize = 1024;
 
@@ -129,15 +142,33 @@ impl PciConfiguration {
     pub fn type0_from_state(
         state: PciConfigurationState,
         msix_config: Option<Arc<Mutex<MsixConfig>>>,
-    ) -> Self {
-        PciConfiguration {
-            registers: state.registers.try_into().unwrap(),
-            writable_bits: state.writable_bits.try_into().unwrap(),
-            bars: state.bars.try_into().unwrap(),
+    ) -> Result<Self, PciConfigurationError> {
+        let reg_len = state.registers.len();
+        let registers = state
+            .registers
+            .try_into()
+            .map_err(|_| PciConfigurationError::InvalidRegistersLength(reg_len))?;
+
+        let wb_len = state.writable_bits.len();
+        let writable_bits = state
+            .writable_bits
+            .try_into()
+            .map_err(|_| PciConfigurationError::InvalidWritableBitsLength(wb_len))?;
+
+        let bar_len = state.bars.len();
+        let bars = state
+            .bars
+            .try_into()
+            .map_err(|_| PciConfigurationError::InvalidBarsLength(bar_len))?;
+
+        Ok(PciConfiguration {
+            registers,
+            writable_bits,
+            bars,
             last_capability: state.last_capability,
             msix_cap_reg_idx: state.msix_cap_reg_idx,
             msix_config,
-        }
+        })
     }
 
     /// Create PCI configuration space state