Fix numerous nvme issues - see each commit

MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpress.c

@@ -1154,6 +1154,14 @@ NvmExpressDriverBindingStart (
                   );
 
   if (EFI_ERROR (Status) && (Status != EFI_ALREADY_STARTED)) {
+    DEBUG ((DEBUG_ERROR, "%a: failed to open PCI I/O protocol (%r)\n", __func__, Status));
+    // need to free the device path protocol if it was opened successfully
+    gBS->CloseProtocol (
+           Controller,
+           &gEfiDevicePathProtocolGuid,
+           This->DriverBindingHandle,
+           Controller
+           );
     return Status;
   }
 
@@ -1164,7 +1172,7 @@ NvmExpressDriverBindingStart (
     Private = AllocateZeroPool (sizeof (NVME_CONTROLLER_PRIVATE_DATA));
 
     if (Private == NULL) {
-      DEBUG ((DEBUG_ERROR, "NvmExpressDriverBindingStart: allocating pool for Nvme Private Data failed!\n"));
+      DEBUG ((DEBUG_ERROR, "%a: allocating pool for Nvme Private Data failed!\n", __func__));
       Status = EFI_OUT_OF_RESOURCES;
       goto Exit;
     }
@@ -1180,7 +1188,8 @@ NvmExpressDriverBindingStart (
                       );
 
     if (EFI_ERROR (Status)) {
-      return Status;
+      DEBUG ((DEBUG_ERROR, "%a: failed to get PCI attributes (%r)\n", __func__, Status));
+      goto Exit;
     }
 
     //

MdeModulePkg/Bus/Pci/NvmExpressDxe/NvmExpressMediaSanitize.c

@@ -352,27 +352,31 @@ NvmExpressMediaClear (
   }
 
   //
-  // If an invalid buffer or buffer size is sent, the Media Clear operation
-  // cannot be performed as it requires a native WRITE command. The overwrite
-  // buffer must have granularity of a namespace block size.
+  // If an invalid buffer is sent, the Media Clear operation
+  // cannot be performed as it requires a native WRITE command.
   //
   if (SectorOwBuffer == NULL) {
     return EFI_INVALID_PARAMETER;
   }
 
+  //
+  // Since the parameters do not include OwBuffer Size, the size of the overwrite buffer is assumed to be equal
+  // to the block size of the media.
+  //
+
   Status = EFI_SUCCESS;
 
   //
   // Per NIST 800-88r1, one or more pass of writes may be alteratively used.
   //
   for (TotalPassCount = 0; TotalPassCount < PassCount; TotalPassCount++) {
-    for (SectorOffset = 0; SectorOffset < Media->LastBlock; SectorOffset++ ) {
+    for (SectorOffset = 0; SectorOffset <= Media->LastBlock; SectorOffset++ ) {
       Status = Device->BlockIo.WriteBlocks (
                                  &Device->BlockIo,
                                  MediaId,
-                                 SectorOffset,  // Sector/LBA offset (increment each pass)
-                                 1,             // Write one sector per write
-                                 SectorOwBuffer // overwrite buffer
+                                 SectorOffset,                   // Sector/LBA offset (increment each pass)
+                                 (UINTN)Device->Media.BlockSize, // Write one block/sector at a time with overwrite buffer
+                                 SectorOwBuffer                  // overwrite buffer
                                  );
     }
 

MdeModulePkg/Bus/Pci/NvmExpressDxe/UnitTest/MediaSanitizeUnitTest.c

@@ -489,6 +489,52 @@ NvmeBlockIoWriteBlocksEx (
   return Status;
 }
 
+//
+// Global tracking variables for MediaClear WriteBlocks coverage tests
+//
+STATIC EFI_LBA  mMaxLbaWritten;
+STATIC UINT64   mWriteCallCount;
+STATIC UINTN    mLastBufferSize;
+
+/**
+  Tracking WriteBlocks mock that records LBA, BufferSize, and call count
+  but always returns EFI_SUCCESS. Used to verify MediaClear loop coverage.
+
+  @param[in]     This        BlockIo Protocol.
+  @param[in]     MediaId     Id of the media.
+  @param[in]     Lba         Logical Block Address.
+  @param[in]     BufferSize  Size of Buffer.
+  @param[in]     Buffer      Actual buffer to use to write.
+
+ **/
+EFI_STATUS
+EFIAPI
+NvmeBlockIoWriteBlocksTracking (
+  IN  EFI_BLOCK_IO_PROTOCOL  *This,
+  IN  UINT32                 MediaId,
+  IN  EFI_LBA                Lba,
+  IN  UINTN                  BufferSize,
+  IN  VOID                   *Buffer
+  )
+{
+  if ((This == NULL) || (Buffer == NULL)) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  if (MediaId != This->Media->MediaId) {
+    return EFI_MEDIA_CHANGED;
+  }
+
+  mWriteCallCount++;
+  mLastBufferSize = BufferSize;
+
+  if (Lba > mMaxLbaWritten) {
+    mMaxLbaWritten = Lba;
+  }
+
+  return EFI_SUCCESS;
+}
+
 /**
   MediaSanitizePurgeUnitTest to initialize a Private Namespace instance.
 
@@ -670,6 +716,129 @@ MediaSanitizePurgeUnitTest (
   return UNIT_TEST_PASSED;
 }
 
+/**
+  MediaClearBufferSizeUnitTest to verify correct BufferSize is passed to WriteBlocks.
+
+  This test detects Finding 5: NvmExpressMediaClear passes BufferSize=1 to
+  WriteBlocks instead of Device->Media.BlockSize. The existing NvmeBlockIoWriteBlocks
+  mock validates (BufferSize % BlockSize) == 0, so BufferSize=1 with BlockSize=512
+  returns EFI_BAD_BUFFER_SIZE.
+
+  @param[in]  Context  Unit test case context
+ **/
+UNIT_TEST_STATUS
+EFIAPI
+MediaClearBufferSizeUnitTest (
+  IN UNIT_TEST_CONTEXT  Context
+  )
+{
+  NVME_DEVICE_PRIVATE_DATA  *NvmeDevice;
+  UNIT_TEST_STATUS          UnitTestStatus;
+  EFI_STATUS                Status;
+  UINT8                     OverwriteBuffer[512];
+
+  NvmeDevice     = NULL;
+  UnitTestStatus = NvmeCreateDeviceInstance (&NvmeDevice);
+
+  UT_ASSERT_STATUS_EQUAL (UnitTestStatus, UNIT_TEST_PASSED);
+  UT_ASSERT_NOT_NULL (NvmeDevice);
+
+  //
+  // Use small LastBlock to keep test fast.
+  // The existing NvmeBlockIoWriteBlocks mock validates BufferSize alignment.
+  //
+  NvmeDevice->Media.LastBlock = 3;
+
+  SetMem (OverwriteBuffer, sizeof (OverwriteBuffer), 0xAA);
+
+  //
+  // Call MediaClear. If BufferSize=1 is passed to WriteBlocks instead of
+  // BlockSize (512), the mock returns EFI_BAD_BUFFER_SIZE.
+  //
+  Status = NvmExpressMediaClear (
+             &NvmeDevice->MediaSanitize,
+             NvmeDevice->Media.MediaId,
+             1,
+             OverwriteBuffer
+             );
+
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  UnitTestStatus = NvmeDestroyDeviceInstance (&NvmeDevice);
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  MediaClearLastBlockUnitTest to verify all blocks including LastBlock are written.
+
+  This test detects Finding 6: off-by-one where loop uses SectorOffset < LastBlock
+  instead of SectorOffset <= LastBlock, skipping the last sector.
+
+  Uses a tracking WriteBlocks mock to record the maximum LBA written and total
+  call count.
+
+  @param[in]  Context  Unit test case context
+ **/
+UNIT_TEST_STATUS
+EFIAPI
+MediaClearLastBlockUnitTest (
+  IN UNIT_TEST_CONTEXT  Context
+  )
+{
+  NVME_DEVICE_PRIVATE_DATA  *NvmeDevice;
+  UNIT_TEST_STATUS          UnitTestStatus;
+  EFI_STATUS                Status;
+  UINT8                     OverwriteBuffer[512];
+  EFI_LBA                   ExpectedLastBlock;
+
+  NvmeDevice     = NULL;
+  UnitTestStatus = NvmeCreateDeviceInstance (&NvmeDevice);
+
+  UT_ASSERT_STATUS_EQUAL (UnitTestStatus, UNIT_TEST_PASSED);
+  UT_ASSERT_NOT_NULL (NvmeDevice);
+
+  //
+  // Use small LastBlock and swap in the tracking WriteBlocks mock
+  //
+  ExpectedLastBlock               = 7;
+  NvmeDevice->Media.LastBlock     = ExpectedLastBlock;
+  NvmeDevice->BlockIo.WriteBlocks = NvmeBlockIoWriteBlocksTracking;
+
+  //
+  // Reset tracking state
+  //
+  mMaxLbaWritten  = 0;
+  mWriteCallCount = 0;
+  mLastBufferSize = 0;
+
+  SetMem (OverwriteBuffer, sizeof (OverwriteBuffer), 0xBB);
+
+  Status = NvmExpressMediaClear (
+             &NvmeDevice->MediaSanitize,
+             NvmeDevice->Media.MediaId,
+             1,
+             OverwriteBuffer
+             );
+
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  //
+  // Finding 6: If loop uses < instead of <=, mMaxLbaWritten will be
+  // (ExpectedLastBlock - 1) instead of ExpectedLastBlock.
+  //
+  UT_ASSERT_EQUAL (mMaxLbaWritten, ExpectedLastBlock);
+
+  //
+  // Verify correct number of WriteBlocks calls: (LastBlock + 1) sectors * 1 pass
+  //
+  UT_ASSERT_EQUAL (mWriteCallCount, (UINT64)(ExpectedLastBlock + 1));
+
+  UnitTestStatus = NvmeDestroyDeviceInstance (&NvmeDevice);
+
+  return UNIT_TEST_PASSED;
+}
+
 /**
   NvmeSanitizeUnitTest to Test calls to NvmExpressSanitize.
 
@@ -1091,6 +1260,32 @@ MediaSanitizeUnitTestEntry (
     NULL                                 // (Optional) UNIT_TEST_CONTEXT
     );
 
+  //
+  // Add test case for MediaClear BufferSize validation (Finding 5)
+  //
+  AddTestCase (
+    MediaSanitizeProtocolTestSuite,                     // Test Suite Handle
+    "MediaClear BufferSize Passed to WriteBlocks Test", // Test Description
+    "MediaClear",                                       // Test Class
+    MediaClearBufferSizeUnitTest,                       // UNIT_TEST_FUNCTION()
+    NULL,                                               // (Optional) UNIT_TEST_PREREQUISITE()
+    NULL,                                               // (Optional) UNIT_TEST_CLEANUP()
+    NULL                                                // (Optional) UNIT_TEST_CONTEXT
+    );
+
+  //
+  // Add test case for MediaClear last block coverage (Finding 6)
+  //
+  AddTestCase (
+    MediaSanitizeProtocolTestSuite,                    // Test Suite Handle
+    "MediaClear All Blocks Including LastBlock Test",  // Test Description
+    "MediaClear",                                      // Test Class
+    MediaClearLastBlockUnitTest,                       // UNIT_TEST_FUNCTION()
+    NULL,                                              // (Optional) UNIT_TEST_PREREQUISITE()
+    NULL,                                              // (Optional) UNIT_TEST_CLEANUP()
+    NULL                                               // (Optional) UNIT_TEST_CONTEXT
+    );
+
   //
   // Execute the tests.
   //