devices/vsock: fix vsock fuzz test

Signed-off-by: Jiaqi Gao <[email protected]>

src/devices/pci/src/config.rs

@@ -386,6 +386,7 @@ impl PciDevice {
         }
     }
 
+    #[cfg(not(feature = "fuzz"))]
     pub fn init(&mut self) -> Result<()> {
         let (vendor_id, device_id) =
             ConfigSpace::get_device_details(self.bus, self.device, self.func);
@@ -473,6 +474,77 @@ impl PciDevice {
         Ok(())
     }
 
+    #[cfg(feature = "fuzz")]
+    pub fn init(&mut self) -> Result<()> {
+        let (vendor_id, device_id) =
+            ConfigSpace::get_device_details(self.bus, self.device, self.func);
+        self.common_header.vendor_id = vendor_id;
+        self.common_header.device_id = device_id;
+        let command = self.read_u16(0x4);
+        let status = self.read_u16(0x6);
+
+        let mut current_bar_offset = 0x10;
+        let mut current_bar = 0;
+
+        //0x24 offset is last bar
+        while current_bar_offset <= 0x24 {
+            let bar = self.read_u32(current_bar_offset);
+
+            // lsb is 1 for I/O space bars
+            if bar & 1 == 1 {
+                self.bars[current_bar].bar_type = PciBarType::IoSpace;
+                self.bars[current_bar].address = u64::from(bar & 0xffff_fffc);
+            } else {
+                // bits 2-1 are the type 0 is 32-but, 2 is 64 bit
+                match bar >> 1 & 3 {
+                    0 => {
+                        let size = self.read_u32(current_bar_offset);
+
+                        let addr = if size > 0 {
+                            let addr = alloc_mmio32(size)?;
+                            self.set_bar_addr(current_bar_offset, addr);
+                            addr
+                        } else {
+                            bar
+                        };
+
+                        self.bars[current_bar].bar_type = PciBarType::MemorySpace32;
+                        self.bars[current_bar].address =
+                            u64::from(addr & PCI_MEM32_BASE_ADDRESS_MASK);
+                    }
+                    2 => {
+                        self.bars[current_bar].bar_type = PciBarType::MemorySpace64;
+
+                        let mut size = self.read_u64(current_bar_offset);
+                        let addr = if size > 0 {
+                            let addr = alloc_mmio64(size)?;
+                            self.set_bar_addr(current_bar_offset, addr as u32);
+                            self.set_bar_addr(current_bar_offset + 4, (addr >> 32) as u32);
+                            addr
+                        } else {
+                            bar as u64
+                        };
+
+                        self.bars[current_bar].address = addr & PCI_MEM64_BASE_ADDRESS_MASK;
+                        current_bar_offset += 4;
+                    }
+                    _ => return Err(PciError::InvalidBarType),
+                }
+            }
+
+            current_bar += 1;
+            current_bar_offset += 4;
+        }
+
+        // Enable the bits 0 (IO Space) and 1 (Memory Space) to activate the bar configuration
+        self.write_u16(
+            0x4,
+            (PciCommand::IO_SPACE | PciCommand::MEMORY_SPACE | PciCommand::BUS_MASTER).bits(),
+        );
+
+        Ok(())
+    }
+
     fn set_bar_addr(&self, offset: u8, addr: u32) {
         self.write_u32(offset, addr);
     }

src/devices/pci/src/mmio.rs

@@ -15,6 +15,11 @@ lazy_static! {
     static ref MMIO64: Mutex<u64> = Mutex::new(0);
 }
 
+#[cfg(feature = "fuzz")]
+lazy_static! {
+    static ref MMIO_OFFSET: Mutex<u64> = Mutex::new(0);
+}
+
 pub fn init_mmio(end_of_ram: u64) {
     *MMIO32.lock() = MMIO32_START;
 
@@ -29,8 +34,13 @@ pub fn init_mmio(end_of_ram: u64) {
 
 #[cfg(feature = "fuzz")]
 pub fn alloc_mmio32(size: u32) -> Result<u32> {
-    let addr = crate::get_fuzz_seed_address() + 0x10c;
-    Ok(addr as u32)
+    let cur = *MMIO_OFFSET.lock();
+    let addr = align_up(cur, size as u64).ok_or(PciError::InvalidParameter)?;
+
+    let addr = u32::try_from(addr).map_err(|_| PciError::InvalidParameter)?;
+
+    *MMIO_OFFSET.lock() = addr.checked_add(size).ok_or(PciError::InvalidParameter)? as u64;
+    Ok(addr)
 }
 
 #[cfg(not(feature = "fuzz"))]
@@ -48,7 +58,10 @@ pub fn alloc_mmio32(size: u32) -> Result<u32> {
 
 #[cfg(feature = "fuzz")]
 pub fn alloc_mmio64(size: u64) -> Result<u64> {
-    let addr = crate::get_fuzz_seed_address() + 0x10c;
+    let cur: u64 = *MMIO_OFFSET.lock();
+    let addr = align_up(cur, size).ok_or(PciError::InvalidParameter)?;
+
+    *MMIO_OFFSET.lock() = addr.checked_add(size).ok_or(PciError::InvalidParameter)?;
     Ok(addr)
 }
 

src/devices/virtio/fuzz/Cargo.toml

@@ -24,6 +24,7 @@ fuzz = ["afl"]
 [dependencies.virtio]
 path = ".."
 default-features = false
+features = ["fuzz"]
 
 # Prevent this from interfering with workspaces
 [workspace]

src/devices/virtio/fuzz/fuzz_targets/afl-virtio.rs

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: BSD-2-Clause-Patent
 
 mod fuzzlib;
-use fuzzlib::{get_fuzz_seed_address, init, COMMON_HEADER, fuzz_virtio};
+use fuzzlib::{fuzz_virtio, get_fuzz_seed_address, init, COMMON_HEADER};
 use pci::PciDevice;
 use virtio::{virtio_pci::VirtioPciTransport, virtqueue::VirtQueue, VirtioTransport};
 
@@ -26,7 +26,7 @@ fn main() {
     let ptr = (ptr.as_ptr() as u64 & PTR_ALIGN_VAR) + PTR_OFFSET as u64;
     let data = unsafe { core::slice::from_raw_parts_mut(ptr as *mut u8, DATA_LEN) };
 
-    let common_addr = ptr + 0x10c;
+    let common_addr = 0;
     let paddr = ptr + PAGE_SIZE as u64;
     init(paddr as usize, TD_PAYLOAD_DMA_SIZE);
     COMMON_HEADER.try_init_once(|| ptr).expect("init error");
@@ -38,15 +38,10 @@ fn main() {
         if let Some(arg) = args.next() {
             println!("{}", arg);
             let paths = std::path::Path::new(&arg);
-            
+
             if paths.is_file() {
                 let tmp = std::fs::read(paths).expect("read crash file fail");
                 data[..tmp.len()].clone_from_slice(&tmp);
-                unsafe {
-                    std::ptr::write_volatile((ptr + BARU64_1_OFFSET) as *mut u64, 0);
-                    std::ptr::write_volatile((ptr + BARU64_2_OFFSET) as *mut u64, 0);
-                    std::ptr::write_volatile((ptr + BARU64_3_OFFSET) as *mut u64, common_addr);
-                }
                 fuzz_virtio(paddr);
             } else if paths.is_dir() {
                 for path in std::fs::read_dir(paths).unwrap() {
@@ -56,11 +51,6 @@ fn main() {
                     }
                     let tmp = std::fs::read(&path).expect("read crash file fail");
                     data[..tmp.len()].clone_from_slice(&tmp);
-                    unsafe {
-                        std::ptr::write_volatile((ptr + BARU64_1_OFFSET) as *mut u64, 0);
-                        std::ptr::write_volatile((ptr + BARU64_2_OFFSET) as *mut u64, 0);
-                        std::ptr::write_volatile((ptr + BARU64_3_OFFSET) as *mut u64, common_addr);
-                    }
                     fuzz_virtio(paddr);
                 }
             } else {
@@ -71,11 +61,6 @@ fn main() {
     #[cfg(feature = "fuzz")]
     afl::fuzz!(|tmp: &[u8]| {
         data[..tmp.len()].clone_from_slice(&tmp);
-        unsafe {
-            std::ptr::write_volatile((ptr + BARU64_1_OFFSET) as *mut u64, 0);
-            std::ptr::write_volatile((ptr + BARU64_2_OFFSET) as *mut u64, 0);
-            std::ptr::write_volatile((ptr + BARU64_3_OFFSET) as *mut u64, common_addr);
-        }
         fuzz_virtio(paddr);
     });
 }

src/devices/virtio/fuzz/fuzz_targets/fuzzlib.rs

@@ -9,7 +9,7 @@ pub use dma_alloc::init;
 pub use pci::{get_fuzz_seed_address, PciDevice, COMMON_HEADER};
 pub use virtio::{
     virtio_pci::VirtioPciTransport,
-    virtqueue::{VirtqueueBuf, VirtQueue, VirtQueueLayout},
+    virtqueue::{VirtQueue, VirtQueueLayout, VirtqueueBuf},
     VirtioTransport,
 };
 

src/devices/vsock/fuzz/Cargo.toml

@@ -30,11 +30,11 @@ members = ["."]
 default = ["libfuzzer-sys"]
 fuzz = ["afl"]
 
-# [[bin]]
-# name = "afl_vsock"
-# path = "fuzz_targets/afl-vsock.rs"
-# test = false
-# doc = false
+[[bin]]
+name = "afl_vsock"
+path = "fuzz_targets/afl-vsock.rs"
+test = false
+doc = false
 
 # [[bin]]
 # name = "vsock"