Controller delay

Cargo.lock

@@ -4688,6 +4688,7 @@ name = "nvme_driver"
 version = "0.0.0"
 dependencies = [
  "anyhow",
+ "async-trait",
  "chipset_device",
  "disklayer_ram",
  "event-listener",
@@ -4699,6 +4700,7 @@ dependencies = [
  "mesh",
  "nvme",
  "nvme_spec",
+ "nvme_test",
  "pal_async",
  "parking_lot",
  "pci_core",
@@ -4736,6 +4738,40 @@ dependencies = [
  "zerocopy 0.8.24",
 ]
 
+[[package]]
+name = "nvme_test"
+version = "0.0.0"
+dependencies = [
+ "async-trait",
+ "chipset_device",
+ "device_emulators",
+ "disk_backend",
+ "event-listener",
+ "futures",
+ "futures-concurrency",
+ "guestmem",
+ "guid",
+ "inspect",
+ "mesh",
+ "nvme_common",
+ "nvme_resources",
+ "nvme_spec",
+ "pal_async",
+ "parking_lot",
+ "pci_core",
+ "pci_resources",
+ "scsi_buffers",
+ "task_control",
+ "thiserror 2.0.12",
+ "tracelimit",
+ "tracing",
+ "unicycle",
+ "user_driver",
+ "vm_resource",
+ "vmcore",
+ "zerocopy 0.8.24",
+]
+
 [[package]]
 name = "object"
 version = "0.36.7"

Cargo.toml

@@ -45,7 +45,7 @@ members = [
   "petri/petri-tool",
   "vm/loader/igvmfilegen",
   "vm/vmgs/vmgs_lib",
-  "vm/vmgs/vmgstool",
+  "vm/vmgs/vmgstool", 
 ]
 exclude = [
   "xsync",
@@ -248,6 +248,7 @@ nvme_common = { path = "vm/devices/storage/nvme_common" }
 nvme_driver = { path = "vm/devices/storage/disk_nvme/nvme_driver" }
 nvme_resources = { path = "vm/devices/storage/nvme_resources" }
 nvme_spec = { path = "vm/devices/storage/nvme_spec" }
+nvme_test = { path = "vm/devices/storage/nvme_test" }
 storage_string = { path = "vm/devices/storage/storage_string" }
 vmswitch = { path = "vm/devices/net/vmswitch" }
 pci_bus = { path = "vm/devices/pci/pci_bus" }

vm/devices/storage/disk_nvme/nvme_driver/Cargo.toml

@@ -36,6 +36,8 @@ pci_core.workspace = true
 scsi_buffers.workspace = true
 test_with_tracing.workspace = true
 user_driver_emulated_mock.workspace = true
+nvme_test.workspace = true
+async-trait.workspace = true
 
 guid.workspace = true
 

vm/devices/storage/disk_nvme/nvme_driver/src/tests.rs

@@ -10,13 +10,19 @@ use inspect::Inspect;
 use inspect::InspectMut;
 use nvme::NvmeControllerCaps;
 use nvme_spec::Cap;
+use nvme_spec::Command;
+use nvme_spec::Completion;
 use nvme_spec::nvm::DsmRange;
+use nvme_test::FaultConfiguration;
+use nvme_test::QueueFaultBehavior;
 use pal_async::DefaultDriver;
 use pal_async::async_test;
+use pal_async::timer::PolledTimer;
 use parking_lot::Mutex;
 use pci_core::msi::MsiInterruptSet;
 use scsi_buffers::OwnedRequestBuffers;
 use std::sync::Arc;
+use std::time::Duration;
 use test_with_tracing::test;
 use user_driver::DeviceBacking;
 use user_driver::DeviceRegisterIo;
@@ -26,9 +32,57 @@ use user_driver_emulated_mock::DeviceTestMemory;
 use user_driver_emulated_mock::EmulatedDevice;
 use user_driver_emulated_mock::Mapping;
 use vmcore::vm_task::SingleDriverBackend;
+use vmcore::vm_task::VmTaskDriver;
 use vmcore::vm_task::VmTaskDriverSource;
 use zerocopy::IntoBytes;
 
+struct AdminQueueFault {
+    pub driver: VmTaskDriver,
+}
+
+#[async_trait::async_trait]
+impl nvme_test::QueueFault for AdminQueueFault {
+    async fn fault_submission_queue(&self, mut command: Command) -> QueueFaultBehavior<Command> {
+        tracing::info!("Faulting submission queue using cid sequence number mismatch");
+        let opcode = nvme_spec::AdminOpcode(command.cdw0.opcode());
+        match opcode {
+            nvme_spec::AdminOpcode::CREATE_IO_COMPLETION_QUEUE => {
+                // Overwrite the previous cid to cause a panic.
+                command.cdw0.set_cid(0);
+                QueueFaultBehavior::Update(command)
+            }
+            _ => QueueFaultBehavior::Default,
+        }
+    }
+
+    async fn fault_completion_queue(
+        &self,
+        _completion: Completion,
+    ) -> QueueFaultBehavior<Completion> {
+        tracing::info!("Faulting completion queue using delay");
+        PolledTimer::new(&self.driver)
+            .sleep(Duration::from_millis(100))
+            .await;
+        QueueFaultBehavior::Default
+    }
+}
+
+#[async_test]
+#[should_panic(expected = "assertion `left == right` failed: cid sequence number mismatch:")]
+async fn test_nvme_command_fault(driver: DefaultDriver) {
+    let task_driver = VmTaskDriverSource::new(SingleDriverBackend::new(driver.clone())).simple();
+
+    test_nvme_fault_injection(
+        driver,
+        FaultConfiguration {
+            admin_fault: Some(Box::new(AdminQueueFault {
+                driver: task_driver,
+            })),
+        },
+    )
+    .await;
+}
+
 #[async_test]
 async fn test_nvme_driver_direct_dma(driver: DefaultDriver) {
     test_nvme_driver(driver, true).await;
@@ -309,6 +363,62 @@ async fn test_nvme_save_restore_inner(driver: DefaultDriver) {
     //     .unwrap();
 }
 
+async fn test_nvme_fault_injection(driver: DefaultDriver, fault_configuration: FaultConfiguration) {
+    const MSIX_COUNT: u16 = 2;
+    const IO_QUEUE_COUNT: u16 = 64;
+    const CPU_COUNT: u32 = 64;
+
+    // Arrange: Create 8MB of space. First 4MB for the device and second 4MB for the payload.
+    let pages = 1024; // 4MB
+    let device_test_memory = DeviceTestMemory::new(pages * 2, false, "test_nvme_driver");
+    let guest_mem = device_test_memory.guest_memory(); // Access to 0-8MB
+    let dma_client = device_test_memory.dma_client(); // Access 0-4MB
+    let payload_mem = device_test_memory.payload_mem(); // allow_dma is false, so this will follow the 'normal' test path (i.e. with bounce buffering behind the scenes)
+
+    // Arrange: Create the NVMe controller and driver.
+    let driver_source = VmTaskDriverSource::new(SingleDriverBackend::new(driver));
+    let mut msi_set = MsiInterruptSet::new();
+    let nvme = nvme_test::NvmeFaultController::new(
+        &driver_source,
+        guest_mem.clone(),
+        &mut msi_set,
+        &mut ExternallyManagedMmioIntercepts,
+        nvme_test::NvmeFaultControllerCaps {
+            msix_count: MSIX_COUNT,
+            max_io_queues: IO_QUEUE_COUNT,
+            subsystem_id: Guid::new_random(),
+        },
+        fault_configuration,
+    );
+
+    nvme.client() // 2MB namespace
+        .add_namespace(1, disklayer_ram::ram_disk(2 << 20, false).unwrap())
+        .await
+        .unwrap();
+    let device = NvmeTestEmulatedDevice::new(nvme, msi_set, dma_client.clone());
+    let driver = NvmeDriver::new(&driver_source, CPU_COUNT, device, false)
+        .await
+        .unwrap();
+    let namespace = driver.namespace(1).await.unwrap();
+
+    // Act: Write 1024 bytes of data to disk starting at LBA 1.
+    let buf_range = OwnedRequestBuffers::linear(0, 16384, true); // 32 blocks
+    payload_mem.write_at(0, &[0xcc; 4096]).unwrap();
+    namespace
+        .write(
+            0,
+            1,
+            2,
+            false,
+            &payload_mem,
+            buf_range.buffer(&payload_mem).range(),
+        )
+        .await
+        .unwrap();
+
+    driver.shutdown().await;
+}
+
 #[derive(Inspect)]
 pub struct NvmeTestEmulatedDevice<T: InspectMut, U: DmaClient> {
     device: EmulatedDevice<T, U>,

vm/devices/storage/nvme/Cargo.toml

@@ -26,12 +26,12 @@ guid.workspace = true
 inspect.workspace = true
 mesh.workspace = true
 pal_async.workspace = true
-task_control.workspace = true
 async-trait.workspace = true
 event-listener.workspace = true
 futures.workspace = true
 futures-concurrency.workspace = true
 parking_lot.workspace = true
+task_control.workspace = true
 thiserror.workspace = true
 tracelimit.workspace = true
 tracing.workspace = true

vm/devices/storage/nvme_spec/src/lib.rs

@@ -209,7 +209,7 @@ open_enum! {
 }
 
 #[repr(C)]
-#[derive(Debug, IntoBytes, Immutable, KnownLayout, FromBytes)]
+#[derive(Debug, Clone, IntoBytes, Immutable, KnownLayout, FromBytes)]
 pub struct Completion {
     pub dw0: u32,
     pub dw1: u32,

vm/devices/storage/nvme_test/Cargo.toml

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+[package]
+name = "nvme_test"
+edition.workspace = true
+rust-version.workspace = true
+
+[dependencies]
+disk_backend.workspace = true
+nvme_common.workspace = true
+nvme_resources.workspace = true
+nvme_spec.workspace = true
+scsi_buffers.workspace = true
+
+device_emulators.workspace = true
+pci_core.workspace = true
+pci_resources.workspace = true
+
+chipset_device.workspace = true
+guestmem.workspace = true
+vmcore.workspace = true
+vm_resource.workspace = true
+
+guid.workspace = true
+inspect.workspace = true
+mesh.workspace = true
+pal_async.workspace = true
+async-trait.workspace = true
+event-listener.workspace = true
+futures.workspace = true
+futures-concurrency.workspace = true
+parking_lot.workspace = true
+task_control.workspace = true
+thiserror.workspace = true
+tracelimit.workspace = true
+tracing.workspace = true
+unicycle.workspace = true
+zerocopy = { workspace = true, features = ["alloc"] }
+
+[dev-dependencies]
+user_driver.workspace = true
+
+[lints]
+workspace = true

vm/devices/storage/nvme_test/src/error.rs

@@ -0,0 +1,78 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+//! Error and result related types.
+
+use crate::spec;
+use std::error::Error;
+
+/// An NVMe error, consisting of a status code and optional error source.
+#[derive(Debug)]
+pub struct NvmeError {
+    status: spec::Status,
+    source: Option<Box<dyn Error + Send + Sync>>,
+}
+
+impl NvmeError {
+    pub fn new(status: spec::Status, source: impl Into<Box<dyn Error + Send + Sync>>) -> Self {
+        Self {
+            status,
+            source: Some(source.into()),
+        }
+    }
+}
+
+impl Error for NvmeError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        self.source.as_ref().map(|x| x.as_ref() as _)
+    }
+}
+
+impl std::fmt::Display for NvmeError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.status.status_code_type() {
+            spec::StatusCodeType::GENERIC => {
+                write!(f, "general error {:#x?}", self.status)
+            }
+            spec::StatusCodeType::COMMAND_SPECIFIC => {
+                write!(f, "command-specific error {:#x?}", self.status)
+            }
+            spec::StatusCodeType::MEDIA_ERROR => {
+                write!(f, "media error {:#x?}", self.status)
+            }
+            _ => write!(f, "{:#x?}", self.status),
+        }
+    }
+}
+
+impl From<spec::Status> for NvmeError {
+    fn from(status: spec::Status) -> Self {
+        NvmeError {
+            status,
+            source: None,
+        }
+    }
+}
+
+/// The result of an NVMe command.
+#[derive(Default)]
+pub struct CommandResult {
+    pub status: spec::Status,
+    pub dw: [u32; 2],
+}
+
+impl<T: Into<NvmeError>> From<T> for CommandResult {
+    fn from(status: T) -> Self {
+        CommandResult::new(status, [0; 2])
+    }
+}
+
+impl CommandResult {
+    pub fn new(status: impl Into<NvmeError>, dw: [u32; 2]) -> Self {
+        let status = status.into();
+        Self {
+            status: status.status,
+            dw,
+        }
+    }
+}