The patch enhances the Partition driver to return the media status (EFI_NO_MEDIA/EFI_MEDIA_CHANGED) even the other parameters are not valid for BlockIo2 interfaces.

Signed-off-by: niruiyu
Reviewed-by: erictian
Reviewed-by: qianouyang

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@12378 6f19259b-4bc3-4df7-8a09-765794883524

MdeModulePkg/Universal/Disk/PartitionDxe/Partition.c

@@ -712,6 +712,45 @@ PartitionFlushBlocks (
   return Private->ParentBlockIo->FlushBlocks (Private->ParentBlockIo);
 }
 
+/**
+  Probe the media status and return EFI_NO_MEDIA or EFI_MEDIA_CHANGED
+  for no media or media change case. Otherwise DefaultStatus is returned.
+
+  @param BlockIo2           Pointer to the BlockIo2 instance.
+  @param MediaId            Id of the media, changes every time the media is replaced.
+  @param DefaultStatus      The default status to return when it's not the no media
+                            or media change case.
+
+  @retval EFI_NO_MEDIA      There is no media.
+  @retval EFI_MEDIA_CHANGED The media was changed.
+  @retval others            The default status to return.
+**/
+EFI_STATUS
+ProbeMediaStatusEx (
+  IN EFI_BLOCK_IO2_PROTOCOL  *BlockIo2,
+  IN UINT32                  MediaId,
+  IN EFI_STATUS              DefaultStatus
+  )
+{
+  EFI_STATUS                 Status;
+
+  //
+  // Read from LBA 0 but passing NULL as buffer pointer to detect the media status.
+  //
+  Status = BlockIo2->ReadBlocksEx (
+                       BlockIo2,
+                       MediaId,
+                       0,
+                       NULL,
+                       0,
+                       NULL
+                       );
+  if ((Status == EFI_NO_MEDIA) || (Status == EFI_MEDIA_CHANGED)) {
+    return Status;
+  }
+  return DefaultStatus;
+}
+
 /**
   Reset the Block Device throught Block I/O2 protocol.
 
@@ -788,18 +827,19 @@ PartitionReadBlocksEx (
   UINT64                  Offset;
   UINT32                  UnderRun;
 
+  Private = PARTITION_DEVICE_FROM_BLOCK_IO2_THIS (This);
+
   if (Token == NULL) {
-    return EFI_INVALID_PARAMETER;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_INVALID_PARAMETER);
   }
 
-  Private = PARTITION_DEVICE_FROM_BLOCK_IO2_THIS (This);
   if (BufferSize % Private->BlockSize != 0) {
-    return EFI_BAD_BUFFER_SIZE;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_BAD_BUFFER_SIZE);
   }
 
   Offset = MultU64x32 (Lba, Private->BlockSize) + Private->Start;
   if (Offset + BufferSize > Private->End) {
-    return EFI_INVALID_PARAMETER;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_INVALID_PARAMETER);
   }
 
   //
@@ -809,7 +849,7 @@ PartitionReadBlocksEx (
   //
   Lba = DivU64x32Remainder (Offset, Private->BlockSize, &UnderRun);
   if (UnderRun != 0) {
-    return EFI_UNSUPPORTED;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_UNSUPPORTED);
   }
 
   //
@@ -817,7 +857,7 @@ PartitionReadBlocksEx (
   // device, in that case the Block I/O2 couldn't be called.
   //
   if (Private->BlockSize != Private->ParentBlockIo->Media->BlockSize) {
-    return EFI_UNSUPPORTED;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_UNSUPPORTED);
   }
 
   return Private->ParentBlockIo2->ReadBlocksEx (Private->ParentBlockIo2, MediaId, Lba, Token, BufferSize, Buffer);
@@ -869,18 +909,19 @@ PartitionWriteBlocksEx (
   UINT64                  Offset;
   UINT32                  UnderRun;
 
+  Private = PARTITION_DEVICE_FROM_BLOCK_IO2_THIS (This);
+
   if (Token == NULL) {
-    return EFI_INVALID_PARAMETER;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_INVALID_PARAMETER);
   }
 
-  Private = PARTITION_DEVICE_FROM_BLOCK_IO2_THIS (This);
   if (BufferSize % Private->BlockSize != 0) {
-    return EFI_BAD_BUFFER_SIZE;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_BAD_BUFFER_SIZE);
   }
 
   Offset = MultU64x32 (Lba, Private->BlockSize) + Private->Start;
   if (Offset + BufferSize > Private->End) {
-    return EFI_INVALID_PARAMETER;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_INVALID_PARAMETER);
   }
 
   //
@@ -890,7 +931,7 @@ PartitionWriteBlocksEx (
   //
   Lba = DivU64x32Remainder (Offset, Private->BlockSize, &UnderRun);
   if (UnderRun != 0) {
-    return EFI_UNSUPPORTED;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_UNSUPPORTED);
   }
 
   //
@@ -898,7 +939,7 @@ PartitionWriteBlocksEx (
   // in that case it couldn't call parent Block I/O2. 
   //
   if (Private->BlockSize != Private->ParentBlockIo->Media->BlockSize) {
-    return EFI_UNSUPPORTED;
+    return ProbeMediaStatusEx (Private->ParentBlockIo2, MediaId, EFI_UNSUPPORTED);
   }
 
   return Private->ParentBlockIo2->WriteBlocksEx (Private->ParentBlockIo2, MediaId, Lba, Token, BufferSize, Buffer);