MdeModulePkg PartitionDxe: Add Re-entry handling logic for BindingStop

There are scenario when the BindingStop service of PartitionDxe driver be
re-entered.

An example will be ejecting a DVD from a SATA DVDROM and then run
"reconnect -r" under shell. In this specific case, part of the calling
stack will be:

PartitionDriverBindingStop() (PartitionDxe) ->
Stop first child handle (PartitionDxe) ->
ScsiDiskFlushBlocksEx() (ScsiDiskDxe) ->
A media change is detected (ScsiDiskDxe) ->
Reinstall of BlockIO(2) protocols (ScsiDiskDxe) ->
Entering PartitionDriverBindingStop() again (PartitionDxe) ->
Potential risk of referencing already stopped child handle (PartitionDxe)
...

The current code has potential issue of referencing of already stopped
child handle. This commit adds re-entry handling logic to resolve such
issue.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Hao Wu <hao.a.wu@intel.com>
Reviewed-by: Feng Tian <feng.tian@intel.com>

MdeModulePkg/Universal/Disk/PartitionDxe/Partition.c

@@ -407,6 +407,16 @@ PartitionDriverBindingStop (
   Private = NULL;
 
   if (NumberOfChildren == 0) {
+    //
+    // In the case of re-entry of the PartitionDriverBindingStop, the
+    // NumberOfChildren may not reflect the actual number of children on the
+    // bus driver. Hence, additional check is needed here.
+    //
+    if (HasChildren (ControllerHandle)) {
+      DEBUG((EFI_D_ERROR, "PartitionDriverBindingStop: Still has child.\n"));
+      return EFI_DEVICE_ERROR;
+    }
+
     //
     // Close the bus driver
     //
@@ -459,8 +469,25 @@ PartitionDriverBindingStop (
 
 
     Private = PARTITION_DEVICE_FROM_BLOCK_IO_THIS (BlockIo);
+    if (Private->InStop) {
+      //
+      // If the child handle is going to be stopped again during the re-entry
+      // of DriverBindingStop, just do nothing.
+      //
+      break;
+    }
+    Private->InStop = TRUE;
 
-    Status = gBS->CloseProtocol (
+    BlockIo->FlushBlocks (BlockIo);
+
+    if (BlockIo2 != NULL) {
+      Status = BlockIo2->FlushBlocksEx (BlockIo2, NULL);
+      DEBUG((EFI_D_ERROR, "PartitionDriverBindingStop: FlushBlocksEx returned with %r\n", Status));
+    } else {
+      Status = EFI_SUCCESS;
+    }
+
+    gBS->CloseProtocol (
            ControllerHandle,
            &gEfiDiskIoProtocolGuid,
            This->DriverBindingHandle,
@@ -472,8 +499,13 @@ PartitionDriverBindingStop (
     // Remove the BlockIo2 Protocol if has.
     //
     if (BlockIo2 != NULL) {
-      BlockIo->FlushBlocks (BlockIo);
+      //
-      BlockIo2->FlushBlocksEx (BlockIo2, NULL);
+      // Some device drivers might re-install the BlockIO(2) protocols for a
+      // media change condition. Therefore, if the FlushBlocksEx returned with
+      // EFI_MEDIA_CHANGED, just let the BindingStop fail to avoid potential
+      // reference of already stopped child handle.
+      //
+      if (Status != EFI_MEDIA_CHANGED) {
         Status = gBS->UninstallMultipleProtocolInterfaces (
                          ChildHandleBuffer[Index],
                          &gEfiDevicePathProtocolGuid,
@@ -486,8 +518,8 @@ PartitionDriverBindingStop (
                          NULL,
                          NULL
                          );
+      }
     } else {
-      BlockIo->FlushBlocks (BlockIo);
       Status = gBS->UninstallMultipleProtocolInterfaces (
                        ChildHandleBuffer[Index],
                        &gEfiDevicePathProtocolGuid,
@@ -501,6 +533,7 @@ PartitionDriverBindingStop (
     }
 
     if (EFI_ERROR (Status)) {
+      Private->InStop = FALSE;
       gBS->OpenProtocol (
              ControllerHandle,
              &gEfiDiskIoProtocolGuid,
@@ -516,6 +549,9 @@ PartitionDriverBindingStop (
 
     if (EFI_ERROR (Status)) {
       AllChildrenStopped = FALSE;
+      if (Status == EFI_MEDIA_CHANGED) {
+        break;
+      }
     }
   }
 
@@ -1263,3 +1299,41 @@ InitializePartition (
   return Status;
 }
 
+
+/**
+  Test to see if there is any child on ControllerHandle.
+
+  @param[in]  ControllerHandle    Handle of device to test.
+
+  @retval TRUE                    There are children on the ControllerHandle.
+  @retval FALSE                   No child is on the ControllerHandle.
+
+**/
+BOOLEAN
+HasChildren (
+  IN EFI_HANDLE           ControllerHandle
+  )
+{
+  EFI_OPEN_PROTOCOL_INFORMATION_ENTRY  *OpenInfoBuffer;
+  UINTN                                EntryCount;
+  EFI_STATUS                           Status;
+  UINTN                                Index;
+
+  Status = gBS->OpenProtocolInformation (
+                  ControllerHandle,
+                  &gEfiDiskIoProtocolGuid,
+                  &OpenInfoBuffer,
+                  &EntryCount
+                  );
+  ASSERT_EFI_ERROR (Status);
+
+  for (Index = 0; Index < EntryCount; Index++) {
+    if ((OpenInfoBuffer[Index].Attributes & EFI_OPEN_PROTOCOL_BY_CHILD_CONTROLLER) != 0) {
+      break;
+    }
+  }
+  FreePool (OpenInfoBuffer);
+
+  return (BOOLEAN) (Index < EntryCount);
+}
+

MdeModulePkg/Universal/Disk/PartitionDxe/Partition.h

@@ -61,6 +61,7 @@ typedef struct {
   UINT64                    Start;
   UINT64                    End;
   UINT32                    BlockSize;
+  BOOLEAN                   InStop;
 
   EFI_GUID                  *EspGuid;
 
@@ -345,6 +346,20 @@ PartitionInstallChildHandle (
   IN  BOOLEAN                      InstallEspGuid
   );
 
+/**
+  Test to see if there is any child on ControllerHandle.
+
+  @param[in]  ControllerHandle    Handle of device to test.
+
+  @retval TRUE                    There are children on the ControllerHandle.
+  @retval FALSE                   No child is on the ControllerHandle.
+
+**/
+BOOLEAN
+HasChildren (
+  IN EFI_HANDLE           ControllerHandle
+  );
+
 /**
   Install child handles if the Handle supports GPT partition structure.