From 80c83a6995db1914cc7db012f6b8f3e223e45847 Mon Sep 17 00:00:00 2001
From: Ruiyu Ni <ruiyu.ni@intel.com>
Date: Tue, 17 Sep 2013 05:11:30 +0000
Subject: [PATCH] Fix 3 bugs in DiskIoDxe and PartitionDxe drivers introduced
 in DiskIo2 implementation. 1. DiskIo2 shouldn't signal the event when the *Ex
 interface returns failure status per the UEFI spec. 2. PartitionDxe should
 close DiskIo2 protocol when error happens in DriverBindingStart() otherwise
 Fat driver cannot open the DiskIo2 BY_DRIVER. 3. PartitionDxe should create
 event using TPL_NOTIFY instead of TPL_CALLBACK otherwise asynchronous FileIo
 may be blocked.

Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Feng Tian <feng.tian@intel.com>

git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@14680 6f19259b-4bc3-4df7-8a09-765794883524
---
 .../Universal/Disk/DiskIoDxe/DiskIo.c         | 29 +++++--------------
 .../Universal/Disk/PartitionDxe/Partition.c   |  6 ++--
 2 files changed, 10 insertions(+), 25 deletions(-)

diff --git a/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c b/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
index ba61300911..f53c3d91d0 100644
--- a/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
+++ b/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
@@ -403,8 +403,6 @@ DiskIo2OnReadWriteComplete (
   EFI_STATUS            TransactionStatus;
   DISK_IO_PRIVATE_DATA  *Instance;
 
-  gBS->CloseEvent (Event);
-
   Subtask           = (DISK_IO_SUBTASK *) Context;
   TransactionStatus = Subtask->BlockIo2Token.TransactionStatus;
   Task              = Subtask->Task;
@@ -414,26 +412,13 @@ DiskIo2OnReadWriteComplete (
   ASSERT (Instance->Signature == DISK_IO_PRIVATE_DATA_SIGNATURE);
   ASSERT (Task->Signature     == DISK_IO2_TASK_SIGNATURE);
 
-  if (Subtask->WorkingBuffer != NULL) {
-    if (!EFI_ERROR (TransactionStatus) && (Task->Token != NULL) && !Subtask->Write) {
-      CopyMem (Subtask->Buffer, Subtask->WorkingBuffer + Subtask->Offset, Subtask->Length);
-    }
-
-    //
-    // The WorkingBuffer of blocking subtask either points to SharedWorkingBuffer
-    // or will be used by non-blocking subtask which will be freed below.
-    //
-    if (!Subtask->Blocking) {
-      FreeAlignedPages (
-        Subtask->WorkingBuffer, 
-        Subtask->Length < Instance->BlockIo->Media->BlockSize
-        ? EFI_SIZE_TO_PAGES (Instance->BlockIo->Media->BlockSize)
-        : EFI_SIZE_TO_PAGES (Subtask->Length)
-        );
-    }
+  if ((Subtask->WorkingBuffer != NULL) && !EFI_ERROR (TransactionStatus) && 
+      (Task->Token != NULL) && !Subtask->Write
+     ) {
+    CopyMem (Subtask->Buffer, Subtask->WorkingBuffer + Subtask->Offset, Subtask->Length);
   }
-  RemoveEntryList (&Subtask->Link);
-  FreePool (Subtask);
+
+  DiskIoDestroySubtask (Instance, Subtask);
 
   if (EFI_ERROR (TransactionStatus) || IsListEmpty (&Task->Subtasks)) {
     if (Task->Token != NULL) {
@@ -961,7 +946,7 @@ DiskIo2ReadWriteDisk (
     RemoveEntryList (&Task->Link);
     EfiReleaseLock (&Instance->TaskQueueLock);
 
-    if (Task->Token != NULL) {
+    if (!EFI_ERROR (Status) && (Task->Token != NULL)) {
       //
       // Task->Token should be set to NULL by the DiskIo2OnReadWriteComplete
       // It it's not, that means the non-blocking request was downgraded to blocking request.
diff --git a/MdeModulePkg/Universal/Disk/PartitionDxe/Partition.c b/MdeModulePkg/Universal/Disk/PartitionDxe/Partition.c
index 5cc1397e7c..7fb2c0dfe0 100644
--- a/MdeModulePkg/Universal/Disk/PartitionDxe/Partition.c
+++ b/MdeModulePkg/Universal/Disk/PartitionDxe/Partition.c
@@ -352,11 +352,11 @@ PartitionDriverBindingStart (
           ControllerHandle
           );
     //
-    // Close Parent BlockIO2 if has.
+    // Close Parent DiskIo2 if has.
     //    
     gBS->CloseProtocol (
            ControllerHandle,
-           &gEfiBlockIo2ProtocolGuid,
+           &gEfiDiskIo2ProtocolGuid,
            This->DriverBindingHandle,
            ControllerHandle
            );
@@ -822,7 +822,7 @@ PartitionCreateAccessTask (
 
   Status = gBS->CreateEvent (
                   EVT_NOTIFY_SIGNAL,
-                  TPL_CALLBACK,
+                  TPL_NOTIFY,
                   PartitionOnAccessComplete,
                   Task,
                   &Task->DiskIo2Token.Event