MdeModulePkg/CdExpressPei: Add RecoveryBlockIo2Ppi support

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Feng Tian <feng.tian@intel.com>
Reviewed-by: Star Zeng <star.zeng@intel.com>
git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@17481 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/MdeModulePkg/Universal/Disk/CdExpressPei/CdExpressPei.inf b/MdeModulePkg/Universal/Disk/CdExpressPei/CdExpressPei.inf
index 4a44a5a9e3ccd94a3bb2b30eeb86bfb86699316e..e9d45e5b86b0251286b74f64a510d7d476f2c924 100644 (file)
--- a/MdeModulePkg/Universal/Disk/CdExpressPei/CdExpressPei.inf
+++ b/MdeModulePkg/Universal/Disk/CdExpressPei/CdExpressPei.inf
@@ -5,7 +5,7 @@
 # finds whether there is Recovery data in the device. If it finds recovery\r
 # data, it will install Device Recovery Module PPI.\r
 #\r
-# Copyright (c) 2006 - 2014, Intel Corporation. All rights reserved.<BR>\r
+# Copyright (c) 2006 - 2015, Intel Corporation. All rights reserved.<BR>\r
 #\r
 # This program and the accompanying materials\r
 # are licensed and made available under the terms and conditions\r
@@ -60,6 +60,9 @@
   ## NOTIFY\r
   ## CONSUMES\r
   gEfiPeiVirtualBlockIoPpiGuid\r
+  ## NOTIFY\r
+  ## CONSUMES\r
+  gEfiPeiVirtualBlockIo2PpiGuid\r
   gEfiPeiDeviceRecoveryModulePpiGuid            ## PRODUCES\r
 \r
 [FeaturePcd]\r

diff --git a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c
index dccb4605428ee37ac17c47ce7fa8693b258c387d..2f8d1abed09823f2fc932f971fb2c573d99e40b0 100644 (file)
--- a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c
+++ b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c
@@ -1,7 +1,7 @@
 /** @file\r
   Source file for CD recovery PEIM\r
 \r
-Copyright (c) 2006 - 2014, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2006 - 2015, Intel Corporation. All rights reserved.<BR>\r
 \r
 This program and the accompanying materials\r
 are licensed and made available under the terms and conditions\r
@@ -60,7 +60,8 @@ CdExpressPeimEntry (
   }\r
 \r
   PrivateData->CapsuleCount = 0;\r
-  Status                    = UpdateBlocksAndVolumes (PrivateData);\r
+  Status = UpdateBlocksAndVolumes (PrivateData, TRUE);\r
+  Status = UpdateBlocksAndVolumes (PrivateData, FALSE);\r
 \r
   //\r
   // Installs Ppi\r
@@ -69,9 +70,9 @@ CdExpressPeimEntry (
   PrivateData->DeviceRecoveryPpi.GetRecoveryCapsuleInfo     = GetRecoveryCapsuleInfo;\r
   PrivateData->DeviceRecoveryPpi.LoadRecoveryCapsule        = LoadRecoveryCapsule;\r
 \r
-  PrivateData->PpiDescriptor.Flags                          = (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST);\r
-  PrivateData->PpiDescriptor.Guid = &gEfiPeiDeviceRecoveryModulePpiGuid;\r
-  PrivateData->PpiDescriptor.Ppi = &PrivateData->DeviceRecoveryPpi;\r
+  PrivateData->PpiDescriptor.Flags = (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST);\r
+  PrivateData->PpiDescriptor.Guid  = &gEfiPeiDeviceRecoveryModulePpiGuid;\r
+  PrivateData->PpiDescriptor.Ppi   = &PrivateData->DeviceRecoveryPpi;\r
 \r
   Status = PeiServicesInstallPpi (&PrivateData->PpiDescriptor);\r
   if (EFI_ERROR (Status)) {\r
@@ -87,11 +88,19 @@ CdExpressPeimEntry (
   //\r
   PrivateData->NotifyDescriptor.Flags =\r
     (\r
-      EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK |\r
-      EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST\r
+      EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK\r
     );\r
   PrivateData->NotifyDescriptor.Guid    = &gEfiPeiVirtualBlockIoPpiGuid;\r
   PrivateData->NotifyDescriptor.Notify  = BlockIoNotifyEntry;\r
+\r
+  PrivateData->NotifyDescriptor2.Flags =\r
+    (\r
+      EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK |\r
+      EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST\r
+    );\r
+  PrivateData->NotifyDescriptor2.Guid    = &gEfiPeiVirtualBlockIo2PpiGuid;\r
+  PrivateData->NotifyDescriptor2.Notify  = BlockIoNotifyEntry;\r
+\r
   return PeiServicesNotifyPpi (&PrivateData->NotifyDescriptor);\r
 \r
 }\r
@@ -117,7 +126,11 @@ BlockIoNotifyEntry (
   IN VOID                       *Ppi\r
   )\r
 {\r
-  UpdateBlocksAndVolumes (mPrivateData);\r
+  if (CompareGuid (NotifyDescriptor->Guid, &gEfiPeiVirtualBlockIo2PpiGuid)) {\r
+    UpdateBlocksAndVolumes (mPrivateData, TRUE);\r
+  } else {\r
+    UpdateBlocksAndVolumes (mPrivateData, FALSE);\r
+  }\r
 \r
   return EFI_SUCCESS;\r
 }\r
@@ -126,22 +139,26 @@ BlockIoNotifyEntry (
   Finds out all the current Block IO PPIs in the system and add them into private data.\r
 \r
   @param PrivateData                    The private data structure that contains recovery module information.\r
+  @param BlockIo2                       Boolean to show whether using BlockIo2 or BlockIo.\r
 \r
   @retval EFI_SUCCESS                   The blocks and volumes are updated successfully.\r
 \r
 **/\r
 EFI_STATUS\r
 UpdateBlocksAndVolumes (\r
-  IN OUT PEI_CD_EXPRESS_PRIVATE_DATA     *PrivateData\r
+  IN OUT PEI_CD_EXPRESS_PRIVATE_DATA     *PrivateData,\r
+  IN     BOOLEAN                         BlockIo2\r
   )\r
 {\r
   EFI_STATUS                      Status;\r
   EFI_PEI_PPI_DESCRIPTOR          *TempPpiDescriptor;\r
   UINTN                           BlockIoPpiInstance;\r
   EFI_PEI_RECOVERY_BLOCK_IO_PPI   *BlockIoPpi;\r
+  EFI_PEI_RECOVERY_BLOCK_IO2_PPI  *BlockIo2Ppi;\r
   UINTN                           NumberBlockDevices;\r
   UINTN                           IndexBlockDevice;\r
   EFI_PEI_BLOCK_IO_MEDIA          Media;\r
+  EFI_PEI_BLOCK_IO2_MEDIA         Media2;\r
   EFI_PEI_SERVICES                **PeiServices;\r
 \r
   IndexBlockDevice = 0;\r
@@ -150,12 +167,21 @@ UpdateBlocksAndVolumes (
   // Assuming all device Block Io Peims are dispatched already\r
   //\r
   for (BlockIoPpiInstance = 0; BlockIoPpiInstance < PEI_CD_EXPRESS_MAX_BLOCK_IO_PPI; BlockIoPpiInstance++) {\r
-    Status = PeiServicesLocatePpi (\r
-                              &gEfiPeiVirtualBlockIoPpiGuid,\r
-                              BlockIoPpiInstance,\r
-                              &TempPpiDescriptor,\r
-                              (VOID **) &BlockIoPpi\r
-                              );\r
+    if (BlockIo2) {\r
+      Status = PeiServicesLocatePpi (\r
+                                &gEfiPeiVirtualBlockIo2PpiGuid,\r
+                                BlockIoPpiInstance,\r
+                                &TempPpiDescriptor,\r
+                                (VOID **) &BlockIo2Ppi\r
+                                );\r
+    } else {\r
+      Status = PeiServicesLocatePpi (\r
+                                &gEfiPeiVirtualBlockIoPpiGuid,\r
+                                BlockIoPpiInstance,\r
+                                &TempPpiDescriptor,\r
+                                (VOID **) &BlockIoPpi\r
+                                );\r
+    }\r
     if (EFI_ERROR (Status)) {\r
       //\r
       // Done with all Block Io Ppis\r
@@ -164,11 +190,19 @@ UpdateBlocksAndVolumes (
     }\r
 \r
     PeiServices = (EFI_PEI_SERVICES  **) GetPeiServicesTablePointer ();\r
-    Status = BlockIoPpi->GetNumberOfBlockDevices (\r
-                          PeiServices,\r
-                          BlockIoPpi,\r
-                          &NumberBlockDevices\r
-                          );\r
+    if (BlockIo2) {\r
+      Status = BlockIo2Ppi->GetNumberOfBlockDevices (\r
+                              PeiServices,\r
+                              BlockIo2Ppi,\r
+                              &NumberBlockDevices\r
+                              );\r
+    } else {\r
+      Status = BlockIoPpi->GetNumberOfBlockDevices (\r
+                            PeiServices,\r
+                            BlockIoPpi,\r
+                            &NumberBlockDevices\r
+                            );\r
+    }\r
     if (EFI_ERROR (Status) || (NumberBlockDevices == 0)) {\r
       continue;\r
     }\r
@@ -176,28 +210,51 @@ UpdateBlocksAndVolumes (
     // Just retrieve the first block, should emulate all blocks.\r
     //\r
     for (IndexBlockDevice = 1; IndexBlockDevice <= NumberBlockDevices && PrivateData->CapsuleCount < PEI_CD_EXPRESS_MAX_CAPSULE_NUMBER; IndexBlockDevice ++) {\r
-      Status = BlockIoPpi->GetBlockDeviceMediaInfo (\r
-                            PeiServices,\r
-                            BlockIoPpi,\r
-                            IndexBlockDevice,\r
-                            &Media\r
-                            );\r
-      if (EFI_ERROR (Status) ||\r
-          !Media.MediaPresent ||\r
-           ((Media.DeviceType != IdeCDROM) && (Media.DeviceType != UsbMassStorage)) ||\r
-          (Media.BlockSize != PEI_CD_BLOCK_SIZE)\r
-          ) {\r
-        continue;\r
+      if (BlockIo2) {\r
+        Status = BlockIo2Ppi->GetBlockDeviceMediaInfo (\r
+                                PeiServices,\r
+                                BlockIo2Ppi,\r
+                                IndexBlockDevice,\r
+                                &Media2\r
+                                );\r
+        if (EFI_ERROR (Status) ||\r
+            !Media2.MediaPresent ||\r
+             ((Media2.InterfaceType != MSG_ATAPI_DP) && (Media2.InterfaceType != MSG_USB_DP)) ||\r
+            (Media2.BlockSize != PEI_CD_BLOCK_SIZE)\r
+            ) {\r
+          continue;\r
+        }\r
+        DEBUG ((EFI_D_INFO, "PeiCdExpress InterfaceType is %d\n", Media2.InterfaceType));\r
+        DEBUG ((EFI_D_INFO, "PeiCdExpress MediaPresent is %d\n", Media2.MediaPresent));\r
+        DEBUG ((EFI_D_INFO, "PeiCdExpress BlockSize is  0x%x\n", Media2.BlockSize));\r
+      } else {\r
+        Status = BlockIoPpi->GetBlockDeviceMediaInfo (\r
+                              PeiServices,\r
+                              BlockIoPpi,\r
+                              IndexBlockDevice,\r
+                              &Media\r
+                              );\r
+        if (EFI_ERROR (Status) ||\r
+            !Media.MediaPresent ||\r
+             ((Media.DeviceType != IdeCDROM) && (Media.DeviceType != UsbMassStorage)) ||\r
+            (Media.BlockSize != PEI_CD_BLOCK_SIZE)\r
+            ) {\r
+          continue;\r
+        }\r
+        DEBUG ((EFI_D_INFO, "PeiCdExpress DeviceType is %d\n", Media.DeviceType));\r
+        DEBUG ((EFI_D_INFO, "PeiCdExpress MediaPresent is %d\n", Media.MediaPresent));\r
+        DEBUG ((EFI_D_INFO, "PeiCdExpress BlockSize is  0x%x\n", Media.BlockSize));\r
       }\r
 \r
-      DEBUG ((EFI_D_INFO, "PeiCdExpress DeviceType is   %d\n", Media.DeviceType));\r
-      DEBUG ((EFI_D_INFO, "PeiCdExpress MediaPresent is %d\n", Media.MediaPresent));\r
-      DEBUG ((EFI_D_INFO, "PeiCdExpress BlockSize is  0x%x\n", Media.BlockSize));\r
       DEBUG ((EFI_D_INFO, "PeiCdExpress Status is %d\n", Status));\r
 \r
       DEBUG ((EFI_D_INFO, "IndexBlockDevice is %d\n", IndexBlockDevice));\r
-      PrivateData->CapsuleData[PrivateData->CapsuleCount].IndexBlock = IndexBlockDevice;  \r
-      PrivateData->CapsuleData[PrivateData->CapsuleCount].BlockIo    = BlockIoPpi;\r
+      PrivateData->CapsuleData[PrivateData->CapsuleCount].IndexBlock = IndexBlockDevice;\r
+      if (BlockIo2) {\r
+        PrivateData->CapsuleData[PrivateData->CapsuleCount].BlockIo2 = BlockIo2Ppi;\r
+      } else {\r
+        PrivateData->CapsuleData[PrivateData->CapsuleCount].BlockIo  = BlockIoPpi;\r
+      }\r
       Status = FindRecoveryCapsules (PrivateData);\r
       DEBUG ((EFI_D_INFO, "Status is %d\n", Status));\r
 \r
@@ -231,6 +288,7 @@ FindRecoveryCapsules (
   EFI_STATUS                      Status;\r
   UINTN                           Lba;\r
   EFI_PEI_RECOVERY_BLOCK_IO_PPI   *BlockIoPpi;\r
+  EFI_PEI_RECOVERY_BLOCK_IO2_PPI  *BlockIo2Ppi;\r
   UINTN                           BufferSize;\r
   UINT8                           *Buffer;\r
   UINT8                           Type;\r
@@ -251,6 +309,7 @@ FindRecoveryCapsules (
   //\r
   IndexBlockDevice = PrivateData->CapsuleData[PrivateData->CapsuleCount].IndexBlock;\r
   BlockIoPpi       = PrivateData->CapsuleData[PrivateData->CapsuleCount].BlockIo;\r
+  BlockIo2Ppi      = PrivateData->CapsuleData[PrivateData->CapsuleCount].BlockIo2;\r
 \r
   VolumeSpaceSize = 0;\r
   StartOfVolume   = TRUE;\r
@@ -258,14 +317,25 @@ FindRecoveryCapsules (
 \r
   while (TRUE) {\r
     SetMem (Buffer, BufferSize, 0);\r
-    Status = BlockIoPpi->ReadBlocks (\r
-                          (EFI_PEI_SERVICES **) GetPeiServicesTablePointer (),\r
-                          BlockIoPpi,\r
-                          IndexBlockDevice,\r
-                          Lba,\r
-                          BufferSize,\r
-                          Buffer\r
-                          );\r
+    if (BlockIo2Ppi != NULL) {\r
+      Status = BlockIo2Ppi->ReadBlocks (\r
+                            (EFI_PEI_SERVICES **) GetPeiServicesTablePointer (),\r
+                            BlockIo2Ppi,\r
+                            IndexBlockDevice,\r
+                            Lba,\r
+                            BufferSize,\r
+                            Buffer\r
+                            );\r
+    } else {\r
+      Status = BlockIoPpi->ReadBlocks (\r
+                            (EFI_PEI_SERVICES **) GetPeiServicesTablePointer (),\r
+                            BlockIoPpi,\r
+                            IndexBlockDevice,\r
+                            Lba,\r
+                            BufferSize,\r
+                            Buffer\r
+                            );\r
+    }\r
     if (EFI_ERROR (Status)) {\r
       return Status;\r
     }\r
@@ -302,7 +372,7 @@ FindRecoveryCapsules (
     RoorDirRecord   = (PEI_CD_EXPRESS_DIR_FILE_RECORD *) (Buffer + PEI_CD_EXPRESS_ROOT_DIR_RECORD_OFFSET);\r
     RootDirLBA      = RoorDirRecord->LocationOfExtent[0];\r
 \r
-    Status          = RetrieveCapsuleFileFromRoot (PrivateData, BlockIoPpi, IndexBlockDevice, RootDirLBA);\r
+    Status          = RetrieveCapsuleFileFromRoot (PrivateData, BlockIoPpi, BlockIo2Ppi, IndexBlockDevice, RootDirLBA);\r
     if (!EFI_ERROR (Status)) {\r
       //\r
       // Just look for the first primary descriptor\r
@@ -321,6 +391,7 @@ FindRecoveryCapsules (
 \r
   @param PrivateData                    The private data structure that contains recovery module information.\r
   @param BlockIoPpi                     The Block IO PPI used to access the volume.\r
+  @param BlockIo2Ppi                    The Block IO 2 PPI used to access the volume.\r
   @param IndexBlockDevice               The index of current block device.\r
   @param Lba                            The starting logic block address to retrieve capsule.\r
 \r
@@ -334,6 +405,7 @@ EFIAPI
 RetrieveCapsuleFileFromRoot (\r
   IN OUT PEI_CD_EXPRESS_PRIVATE_DATA        *PrivateData,\r
   IN EFI_PEI_RECOVERY_BLOCK_IO_PPI          *BlockIoPpi,\r
+  IN EFI_PEI_RECOVERY_BLOCK_IO2_PPI         *BlockIo2Ppi,\r
   IN UINTN                                  IndexBlockDevice,\r
   IN UINT32                                 Lba\r
   )\r
@@ -349,14 +421,25 @@ RetrieveCapsuleFileFromRoot (
 \r
   SetMem (Buffer, BufferSize, 0);\r
 \r
-  Status = BlockIoPpi->ReadBlocks (\r
-                        (EFI_PEI_SERVICES **) GetPeiServicesTablePointer (),\r
-                        BlockIoPpi,\r
-                        IndexBlockDevice,\r
-                        Lba,\r
-                        BufferSize,\r
-                        Buffer\r
-                        );\r
+  if (BlockIo2Ppi != NULL) {\r
+    Status = BlockIo2Ppi->ReadBlocks (\r
+                          (EFI_PEI_SERVICES **) GetPeiServicesTablePointer (),\r
+                          BlockIo2Ppi,\r
+                          IndexBlockDevice,\r
+                          Lba,\r
+                          BufferSize,\r
+                          Buffer\r
+                          );\r
+  } else {\r
+    Status = BlockIoPpi->ReadBlocks (\r
+                          (EFI_PEI_SERVICES **) GetPeiServicesTablePointer (),\r
+                          BlockIoPpi,\r
+                          IndexBlockDevice,\r
+                          Lba,\r
+                          BufferSize,\r
+                          Buffer\r
+                          );\r
+  }\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
@@ -440,7 +523,8 @@ GetNumberRecoveryCapsules (
   PEI_CD_EXPRESS_PRIVATE_DATA *PrivateData;\r
 \r
   PrivateData = PEI_CD_EXPRESS_PRIVATE_DATA_FROM_THIS (This);\r
-  UpdateBlocksAndVolumes (PrivateData);\r
+  UpdateBlocksAndVolumes (PrivateData, TRUE);\r
+  UpdateBlocksAndVolumes (PrivateData, FALSE);\r
   *NumberRecoveryCapsules = PrivateData->CapsuleCount;\r
 \r
   if (*NumberRecoveryCapsules == 0) {\r
@@ -546,6 +630,7 @@ LoadRecoveryCapsule (
   EFI_STATUS                      Status;\r
   PEI_CD_EXPRESS_PRIVATE_DATA     *PrivateData;\r
   EFI_PEI_RECOVERY_BLOCK_IO_PPI   *BlockIoPpi;\r
+  EFI_PEI_RECOVERY_BLOCK_IO2_PPI  *BlockIo2Ppi;\r
   UINTN                           NumberRecoveryCapsules;\r
 \r
   Status = GetNumberRecoveryCapsules (PeiServices, This, &NumberRecoveryCapsules);\r
@@ -564,15 +649,27 @@ LoadRecoveryCapsule (
 \r
   PrivateData = PEI_CD_EXPRESS_PRIVATE_DATA_FROM_THIS (This);\r
   BlockIoPpi  = PrivateData->CapsuleData[CapsuleInstance - 1].BlockIo;\r
+  BlockIo2Ppi = PrivateData->CapsuleData[CapsuleInstance - 1].BlockIo2;\r
 \r
-  Status = BlockIoPpi->ReadBlocks (\r
-                        PeiServices,\r
-                        BlockIoPpi,\r
-                        PrivateData->CapsuleData[CapsuleInstance - 1].IndexBlock,\r
-                        PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleStartLBA,\r
-                        PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleSize,\r
-                        Buffer\r
-                        );\r
+  if (BlockIo2Ppi != NULL) {\r
+    Status = BlockIo2Ppi->ReadBlocks (\r
+                          PeiServices,\r
+                          BlockIo2Ppi,\r
+                          PrivateData->CapsuleData[CapsuleInstance - 1].IndexBlock,\r
+                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleStartLBA,\r
+                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleSize,\r
+                          Buffer\r
+                          );\r
+  } else {\r
+    Status = BlockIoPpi->ReadBlocks (\r
+                          PeiServices,\r
+                          BlockIoPpi,\r
+                          PrivateData->CapsuleData[CapsuleInstance - 1].IndexBlock,\r
+                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleStartLBA,\r
+                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleSize,\r
+                          Buffer\r
+                          );\r
+  }\r
   return Status;\r
 }\r
 \r

diff --git a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h
index b452747e704c6a37e702737590f3cc6b0824a228..5ca26722a742c8aa7bcc9e877af02e69a1b4039e 100644 (file)
--- a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h
+++ b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h
@@ -1,7 +1,7 @@
 /** @file\r
   Header file for CD recovery PEIM\r
 \r
-Copyright (c) 2006 - 2014, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2006 - 2015, Intel Corporation. All rights reserved.<BR>\r
 \r
 This program and the accompanying materials\r
 are licensed and made available under the terms and conditions\r
@@ -21,6 +21,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <PiPei.h>\r
 \r
 #include <Ppi/BlockIo.h>\r
+#include <Ppi/BlockIo2.h>\r
 #include <Guid/RecoveryDevice.h>\r
 #include <Ppi/DeviceRecoveryModule.h>\r
 \r
@@ -67,6 +68,7 @@ typedef struct {
   UINTN                           CapsuleSize;\r
   UINTN                           IndexBlock;\r
   EFI_PEI_RECOVERY_BLOCK_IO_PPI   *BlockIo;\r
+  EFI_PEI_RECOVERY_BLOCK_IO2_PPI  *BlockIo2;\r
 } PEI_CD_EXPRESS_CAPSULE_DATA;\r
 \r
 #define PEI_CD_EXPRESS_PRIVATE_DATA_SIGNATURE SIGNATURE_32 ('p', 'c', 'd', 'e')\r
@@ -77,6 +79,7 @@ typedef struct {
   EFI_PEI_DEVICE_RECOVERY_MODULE_PPI    DeviceRecoveryPpi;\r
   EFI_PEI_PPI_DESCRIPTOR                PpiDescriptor;\r
   EFI_PEI_NOTIFY_DESCRIPTOR             NotifyDescriptor;\r
+  EFI_PEI_NOTIFY_DESCRIPTOR             NotifyDescriptor2;\r
 \r
   UINT8                                 *BlockBuffer;\r
   UINTN                                 CapsuleCount;\r
@@ -130,13 +133,15 @@ BlockIoNotifyEntry (
   Finds out all the current Block IO PPIs in the system and add them into private data.\r
 \r
   @param PrivateData                    The private data structure that contains recovery module information.\r
+  @param BlockIo2                       Boolean to show whether using BlockIo2 or BlockIo.\r
 \r
   @retval EFI_SUCCESS                   The blocks and volumes are updated successfully.\r
 \r
 **/\r
 EFI_STATUS\r
 UpdateBlocksAndVolumes (\r
-  IN OUT PEI_CD_EXPRESS_PRIVATE_DATA     *PrivateData\r
+  IN OUT PEI_CD_EXPRESS_PRIVATE_DATA     *PrivateData,\r
+  IN     BOOLEAN                         BlockIo2\r
   );\r
 \r
 /**\r
@@ -253,6 +258,7 @@ FindRecoveryCapsules (
 \r
   @param PrivateData                    The private data structure that contains recovery module information.\r
   @param BlockIoPpi                     The Block IO PPI used to access the volume.\r
+  @param BlockIo2Ppi                    The Block IO 2 PPI used to access the volume.\r
   @param IndexBlockDevice               The index of current block device.\r
   @param Lba                            The starting logic block address to retrieve capsule.\r
 \r
@@ -266,6 +272,7 @@ EFIAPI
 RetrieveCapsuleFileFromRoot (\r
   IN OUT PEI_CD_EXPRESS_PRIVATE_DATA        *PrivateData,\r
   IN EFI_PEI_RECOVERY_BLOCK_IO_PPI          *BlockIoPpi,\r
+  IN EFI_PEI_RECOVERY_BLOCK_IO2_PPI         *BlockIo2Ppi,\r
   IN UINTN                                  IndexBlockDevice,\r
   IN UINT32                                 Lba\r
   );\r