Produce DiskInfo for usb mass storage device

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

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

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/ComponentName.c

@@ -1,7 +1,7 @@
 /** @file
   UEFI Component Name(2) protocol implementation for USB Mass Storage Driver.
 
-Copyright (c) 2004 - 2007, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2004 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
 
-#include "UsbMassImpl.h"
+#include "UsbMass.h"
 
 //
 // EFI Component Name Protocol

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMass.h

@@ -2,7 +2,7 @@
   Definition of USB Mass Storage Class and its value, USB Mass Transport Protocol, 
   and other common definitions.
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -22,7 +22,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Protocol/BlockIo.h>
 #include <Protocol/UsbIo.h>
 #include <Protocol/DevicePath.h>
-
+#include <Protocol/DiskInfo.h>
 #include <Library/BaseLib.h>
 #include <Library/DebugLib.h>
 #include <Library/BaseMemoryLib.h>
@@ -32,6 +32,15 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/MemoryAllocationLib.h>
 #include <Library/DevicePathLib.h>
 
+typedef struct _USB_MASS_TRANSPORT USB_MASS_TRANSPORT;
+typedef struct _USB_MASS_DEVICE    USB_MASS_DEVICE;
+
+#include "UsbMassBot.h"
+#include "UsbMassCbi.h"
+#include "UsbMassBoot.h"
+#include "UsbMassDiskInfo.h"
+#include "UsbMassImpl.h"
+
 #define USB_IS_IN_ENDPOINT(EndPointAddr)      (((EndPointAddr) & BIT7) == BIT7)
 #define USB_IS_OUT_ENDPOINT(EndPointAddr)     (((EndPointAddr) & BIT7) == 0)
 #define USB_IS_BULK_ENDPOINT(Attribute)       (((Attribute) & (BIT0 | BIT1)) == USB_ENDPOINT_BULK)
@@ -177,16 +186,16 @@ EFI_STATUS
 /// structure so that the CBI protocol can be easily removed when
 /// it is no longer necessary.
 ///
-typedef struct {
+struct _USB_MASS_TRANSPORT {
   UINT8                   Protocol;
   USB_MASS_INIT_TRANSPORT Init;        ///< Initialize the mass storage transport protocol
   USB_MASS_EXEC_COMMAND   ExecCommand; ///< Transport command to the device then get result
   USB_MASS_RESET          Reset;       ///< Reset the device
   USB_MASS_GET_MAX_LUN    GetMaxLun;   ///< Get max lun, only for bot
   USB_MASS_CLEAN_UP       CleanUp;     ///< Clean up the resources.
-} USB_MASS_TRANSPORT;
+};
 
-typedef struct {
+struct _USB_MASS_DEVICE {
   UINT32                    Signature;
   EFI_HANDLE                Controller;
   EFI_USB_IO_PROTOCOL       *UsbIo;
@@ -197,7 +206,9 @@ typedef struct {
   UINT8                     Lun;          ///< Logical Unit Number
   UINT8                     Pdt;          ///< Peripheral Device Type
   USB_MASS_TRANSPORT        *Transport;   ///< USB mass storage transport protocol
-  VOID                      *Context;     
+  VOID                      *Context;
-} USB_MASS_DEVICE;
+  EFI_DISK_INFO_PROTOCOL    DiskInfo;
+  USB_BOOT_INQUIRY_DATA     InquiryData;
+};
 
 #endif

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c

@@ -2,7 +2,7 @@
   Implementation of the command set of USB Mass Storage Specification
   for Bootability, Revision 1.0.
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
 
-#include "UsbMassImpl.h"
+#include "UsbMass.h"
 
 /**
   Execute REQUEST SENSE Command to retrieve sense data from device.
@@ -314,25 +314,24 @@ UsbBootInquiry (
   )
 {
   USB_BOOT_INQUIRY_CMD        InquiryCmd;
-  USB_BOOT_INQUIRY_DATA       InquiryData;
   EFI_BLOCK_IO_MEDIA          *Media;
   EFI_STATUS                  Status;
 
   Media = &(UsbMass->BlockIoMedia);
 
   ZeroMem (&InquiryCmd, sizeof (USB_BOOT_INQUIRY_CMD));
-  ZeroMem (&InquiryData, sizeof (USB_BOOT_INQUIRY_DATA));
+  ZeroMem (&UsbMass->InquiryData, sizeof (USB_BOOT_INQUIRY_DATA));
 
   InquiryCmd.OpCode   = USB_BOOT_INQUIRY_OPCODE;
   InquiryCmd.Lun      = (UINT8) (USB_BOOT_LUN (UsbMass->Lun));
-  InquiryCmd.AllocLen = (UINT8) sizeof (InquiryData);
+  InquiryCmd.AllocLen = (UINT8) sizeof (USB_BOOT_INQUIRY_DATA);
 
   Status = UsbBootExecCmdWithRetry (
              UsbMass,
              &InquiryCmd,
              (UINT8) sizeof (USB_BOOT_INQUIRY_CMD),
              EfiUsbDataIn,
-             &InquiryData,
+             &UsbMass->InquiryData,
              sizeof (USB_BOOT_INQUIRY_DATA),
              USB_BOOT_GENERAL_CMD_TIMEOUT
              );
@@ -344,8 +343,8 @@ UsbBootInquiry (
   // Get information from PDT (Peripheral Device Type) field and Removable Medium Bit
   // from the inquiry data.
   //
-  UsbMass->Pdt          = (UINT8) (USB_BOOT_PDT (InquiryData.Pdt));
+  UsbMass->Pdt          = (UINT8) (USB_BOOT_PDT (UsbMass->InquiryData.Pdt));
-  Media->RemovableMedia = (BOOLEAN) (USB_BOOT_REMOVABLE (InquiryData.Removable));
+  Media->RemovableMedia = (BOOLEAN) (USB_BOOT_REMOVABLE (UsbMass->InquiryData.Removable));
   //
   // Set block size to the default value of 512 Bytes, in case no media is present at first time.
   //

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.h

@@ -2,7 +2,7 @@
   Definition of the command set of USB Mass Storage Specification
   for Bootability, Revision 1.0.
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -16,8 +16,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #ifndef _EFI_USB_MASS_BOOT_H_
 #define _EFI_USB_MASS_BOOT_H_
 
-#include "UsbMass.h"
-
 //
 // The opcodes of various USB boot commands:
 // INQUIRY/REQUEST_SENSE are "No Timeout Commands" as specified

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBot.c

@@ -2,7 +2,7 @@
   Implementation of the USB mass storage Bulk-Only Transport protocol,
   according to USB Mass Storage Class Bulk-Only Transport, Revision 1.0.
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -13,8 +13,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
 
-#include "UsbMassBoot.h"
+#include "UsbMass.h"
-#include "UsbMassBot.h"
 
 //
 // Definition of USB BOT Transport Protocol

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBot.h

@@ -3,7 +3,7 @@
   based on the "Universal Serial Bus Mass Storage Class Bulk-Only
   Transport" Revision 1.0, September 31, 1999.
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -17,8 +17,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #ifndef _EFI_USBMASS_BOT_H_
 #define _EFI_USBMASS_BOT_H_
 
-#include "UsbMass.h"
-
 extern USB_MASS_TRANSPORT mUsbBotTransport;
 
 //

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassCbi.c

@@ -4,7 +4,7 @@
   Notice: it is being obsoleted by the standard body in favor of the BOT
   (Bulk-Only Transport).
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -15,8 +15,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
 
-#include "UsbMassBoot.h"
+#include "UsbMass.h"
-#include "UsbMassCbi.h"
 
 //
 // Definition of USB CBI0 Transport Protocol

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassCbi.h

@@ -2,7 +2,7 @@
   Defination for the USB mass storage Control/Bulk/Interrupt (CBI) transport,
   according to USB Mass Storage Class Control/Bulk/Interrupt (CBI) Transport, Revision 1.1.
 
-Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -16,8 +16,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #ifndef _EFI_USBMASS_CBI_H_
 #define _EFI_USBMASS_CBI_H_
 
-#include "UsbMass.h"
-
 extern USB_MASS_TRANSPORT mUsbCbi0Transport;
 extern USB_MASS_TRANSPORT mUsbCbi1Transport;
 

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassDiskinfo.c

@@ -0,0 +1,162 @@
+/** @file
+  This file is used to implement the EFI_DISK_INFO_PROTOCOL interface.
+
+Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+This program and the accompanying materials
+are licensed and made available under the terms and conditions of the BSD License
+which accompanies this distribution.  The full text of the license may be found at
+http://opensource.org/licenses/bsd-license.php
+
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+
+#include "UsbMass.h"
+
+EFI_DISK_INFO_PROTOCOL gUsbDiskInfoProtocolTemplate = {
+  EFI_DISK_INFO_USB_INTERFACE_GUID,
+  UsbDiskInfoInquiry,
+  UsbDiskInfoIdentify,
+  UsbDiskInfoSenseData,
+  UsbDiskInfoWhichIde
+};
+
+/**
+  Initialize the installation of DiskInfo protocol.
+
+  This function prepares for the installation of DiskInfo protocol on the child handle.
+  By default, it installs DiskInfo protocol with USB interface GUID. 
+
+  @param[in]  UsbMass  The pointer of USB_MASS_DEVICE.
+
+**/
+VOID
+InitializeDiskInfo (
+  IN  USB_MASS_DEVICE   *UsbMass
+  )
+{
+  CopyMem (&UsbMass->DiskInfo, &gUsbDiskInfoProtocolTemplate, sizeof (gUsbDiskInfoProtocolTemplate));
+}
+
+
+/**
+  Provides inquiry information for the controller type.
+  
+  This function is used to get inquiry data.  Data format
+  of Identify data is defined by the Interface GUID.
+
+  @param[in]      This              Pointer to the EFI_DISK_INFO_PROTOCOL instance.
+  @param[in, out] InquiryData       Pointer to a buffer for the inquiry data.
+  @param[in, out] InquiryDataSize   Pointer to the value for the inquiry data size.
+
+  @retval EFI_SUCCESS            The command was accepted without any errors.
+  @retval EFI_NOT_FOUND          Device does not support this data class 
+  @retval EFI_DEVICE_ERROR       Error reading InquiryData from device 
+  @retval EFI_BUFFER_TOO_SMALL   InquiryDataSize not big enough 
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoInquiry (
+  IN     EFI_DISK_INFO_PROTOCOL   *This,
+  IN OUT VOID                     *InquiryData,
+  IN OUT UINT32                   *InquiryDataSize
+  )
+{
+  EFI_STATUS        Status;
+  USB_MASS_DEVICE   *UsbMass;
+
+  UsbMass  = USB_MASS_DEVICE_FROM_DISK_INFO (This);
+
+  Status = EFI_BUFFER_TOO_SMALL;
+  if (*InquiryDataSize >= sizeof (UsbMass->InquiryData)) {
+    Status = EFI_SUCCESS;
+    CopyMem (InquiryData, &UsbMass->InquiryData, sizeof (UsbMass->InquiryData));
+  }
+  *InquiryDataSize = sizeof (UsbMass->InquiryData);
+  return Status;
+}
+
+
+/**
+  Provides identify information for the controller type.
+
+  This function is used to get identify data.  Data format
+  of Identify data is defined by the Interface GUID.
+
+  @param[in]      This              Pointer to the EFI_DISK_INFO_PROTOCOL 
+                                    instance.
+  @param[in, out] IdentifyData      Pointer to a buffer for the identify data.
+  @param[in, out] IdentifyDataSize  Pointer to the value for the identify data
+                                    size.
+
+  @retval EFI_SUCCESS            The command was accepted without any errors.
+  @retval EFI_NOT_FOUND          Device does not support this data class 
+  @retval EFI_DEVICE_ERROR       Error reading IdentifyData from device 
+  @retval EFI_BUFFER_TOO_SMALL   IdentifyDataSize not big enough 
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoIdentify (
+  IN     EFI_DISK_INFO_PROTOCOL   *This,
+  IN OUT VOID                     *IdentifyData,
+  IN OUT UINT32                   *IdentifyDataSize
+  )
+{
+  return EFI_NOT_FOUND;
+}
+
+/**
+  Provides sense data information for the controller type.
+  
+  This function is used to get sense data. 
+  Data format of Sense data is defined by the Interface GUID.
+
+  @param[in]      This              Pointer to the EFI_DISK_INFO_PROTOCOL instance.
+  @param[in, out] SenseData         Pointer to the SenseData.
+  @param[in, out] SenseDataSize     Size of SenseData in bytes.
+  @param[out]     SenseDataNumber   Pointer to the value for the sense data size.
+
+  @retval EFI_SUCCESS            The command was accepted without any errors.
+  @retval EFI_NOT_FOUND          Device does not support this data class.
+  @retval EFI_DEVICE_ERROR       Error reading SenseData from device.
+  @retval EFI_BUFFER_TOO_SMALL   SenseDataSize not big enough.
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoSenseData (
+  IN     EFI_DISK_INFO_PROTOCOL   *This,
+  IN OUT VOID                     *SenseData,
+  IN OUT UINT32                   *SenseDataSize,
+  OUT    UINT8                    *SenseDataNumber
+  )
+{
+  return EFI_NOT_FOUND;
+}
+
+
+/**
+  This function is used to get controller information.
+
+  @param[in]  This         Pointer to the EFI_DISK_INFO_PROTOCOL instance. 
+  @param[out] IdeChannel   Pointer to the Ide Channel number.  Primary or secondary.
+  @param[out] IdeDevice    Pointer to the Ide Device number.  Master or slave.
+
+  @retval EFI_SUCCESS       IdeChannel and IdeDevice are valid.
+  @retval EFI_UNSUPPORTED   This is not an IDE device.
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoWhichIde (
+  IN  EFI_DISK_INFO_PROTOCOL   *This,
+  OUT UINT32                   *IdeChannel,
+  OUT UINT32                   *IdeDevice
+  )
+{
+  return EFI_UNSUPPORTED;
+}
+

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassDiskinfo.h

@@ -0,0 +1,125 @@
+/** @file
+  Header file for EFI_DISK_INFO_PROTOCOL interface.
+
+Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+This program and the accompanying materials
+are licensed and made available under the terms and conditions of the BSD License
+which accompanies this distribution.  The full text of the license may be found at
+http://opensource.org/licenses/bsd-license.php
+
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+
+/**
+  Initialize the installation of DiskInfo protocol.
+
+  This function prepares for the installation of DiskInfo protocol on the child handle.
+  By default, it installs DiskInfo protocol with USB interface GUID. 
+
+  @param  UsbMass  The pointer of USB_MASS_DEVICE.
+
+**/
+VOID
+InitializeDiskInfo (
+  IN  USB_MASS_DEVICE   *UsbMass
+  );
+
+
+/**
+  Provides inquiry information for the controller type.
+  
+  This function is used to get inquiry data.  Data format
+  of Identify data is defined by the Interface GUID.
+
+  @param[in]      This              Pointer to the EFI_DISK_INFO_PROTOCOL instance.
+  @param[in, out] InquiryData       Pointer to a buffer for the inquiry data.
+  @param[in, out] InquiryDataSize   Pointer to the value for the inquiry data size.
+
+  @retval EFI_SUCCESS            The command was accepted without any errors.
+  @retval EFI_NOT_FOUND          Device does not support this data class 
+  @retval EFI_DEVICE_ERROR       Error reading InquiryData from device 
+  @retval EFI_BUFFER_TOO_SMALL   InquiryDataSize not big enough 
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoInquiry (
+  IN     EFI_DISK_INFO_PROTOCOL   *This,
+  IN OUT VOID                     *InquiryData,
+  IN OUT UINT32                   *InquiryDataSize
+  );
+
+/**
+  Provides identify information for the controller type.
+
+  This function is used to get identify data.  Data format
+  of Identify data is defined by the Interface GUID.
+
+  @param[in]      This              Pointer to the EFI_DISK_INFO_PROTOCOL 
+                                    instance.
+  @param[in, out] IdentifyData      Pointer to a buffer for the identify data.
+  @param[in, out] IdentifyDataSize  Pointer to the value for the identify data
+                                    size.
+
+  @retval EFI_SUCCESS            The command was accepted without any errors.
+  @retval EFI_NOT_FOUND          Device does not support this data class 
+  @retval EFI_DEVICE_ERROR       Error reading IdentifyData from device 
+  @retval EFI_BUFFER_TOO_SMALL   IdentifyDataSize not big enough 
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoIdentify (
+  IN     EFI_DISK_INFO_PROTOCOL   *This,
+  IN OUT VOID                     *IdentifyData,
+  IN OUT UINT32                   *IdentifyDataSize
+  );
+
+/**
+  Provides sense data information for the controller type.
+  
+  This function is used to get sense data. 
+  Data format of Sense data is defined by the Interface GUID.
+
+  @param[in]      This              Pointer to the EFI_DISK_INFO_PROTOCOL instance.
+  @param[in, out] SenseData         Pointer to the SenseData.
+  @param[in, out] SenseDataSize     Size of SenseData in bytes.
+  @param[out]     SenseDataNumber   Pointer to the value for the sense data size.
+
+  @retval EFI_SUCCESS            The command was accepted without any errors.
+  @retval EFI_NOT_FOUND          Device does not support this data class.
+  @retval EFI_DEVICE_ERROR       Error reading SenseData from device.
+  @retval EFI_BUFFER_TOO_SMALL   SenseDataSize not big enough.
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoSenseData (
+  IN     EFI_DISK_INFO_PROTOCOL   *This,
+  IN OUT VOID                     *SenseData,
+  IN OUT UINT32                   *SenseDataSize,
+  OUT    UINT8                    *SenseDataNumber
+  );
+
+
+/**
+  This function is used to get controller information.
+
+  @param[in]  This         Pointer to the EFI_DISK_INFO_PROTOCOL instance. 
+  @param[out] IdeChannel   Pointer to the Ide Channel number.  Primary or secondary.
+  @param[out] IdeDevice    Pointer to the Ide Device number.  Master or slave.
+
+  @retval EFI_SUCCESS       IdeChannel and IdeDevice are valid.
+  @retval EFI_UNSUPPORTED   This is not an IDE device.
+
+**/
+EFI_STATUS
+EFIAPI
+UsbDiskInfoWhichIde (
+  IN  EFI_DISK_INFO_PROTOCOL   *This,
+  OUT UINT32                   *IdeChannel,
+  OUT UINT32                   *IdeDevice
+  );
+

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassImpl.c

@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
 
-#include "UsbMassImpl.h"
+#include "UsbMass.h"
 
 #define USB_MASS_TRANSPORT_COUNT    3
 //
@@ -554,6 +554,8 @@ UsbMassInitMultiLun (
       goto ON_ERROR;
     }
 
+    InitializeDiskInfo (UsbMass);
+
     //
     // Create a new handle for each LUN, and install Block I/O Protocol and Device Path Protocol.
     //
@@ -563,6 +565,8 @@ UsbMassInitMultiLun (
                     UsbMass->DevicePath,
                     &gEfiBlockIoProtocolGuid,
                     &UsbMass->BlockIo,
+                    &gEfiDiskInfoProtocolGuid,
+                    &UsbMass->DiskInfo,
                     NULL
                     );
     
@@ -591,6 +595,8 @@ UsbMassInitMultiLun (
              UsbMass->DevicePath,
              &gEfiBlockIoProtocolGuid,
              &UsbMass->BlockIo,
+             &gEfiDiskInfoProtocolGuid,
+             &UsbMass->DiskInfo,
              NULL
              );
       goto ON_ERROR;
@@ -702,11 +708,15 @@ UsbMassInitNonLun (
     goto ON_ERROR;
   }
     
-  Status = gBS->InstallProtocolInterface (
+  InitializeDiskInfo (UsbMass);
+
+  Status = gBS->InstallMultipleProtocolInterfaces (
                   &Controller,
                   &gEfiBlockIoProtocolGuid,
-                  EFI_NATIVE_INTERFACE,
+                  &UsbMass->BlockIo,
-                  &UsbMass->BlockIo
+                  &gEfiDiskInfoProtocolGuid,
+                  &UsbMass->DiskInfo,
+                  NULL
                   );
   if (EFI_ERROR (Status)) {
     goto ON_ERROR;
@@ -1007,10 +1017,13 @@ USBMassDriverBindingStop (
     // Uninstall Block I/O protocol from the device handle,
     // then call the transport protocol to stop itself.
     //
-    Status = gBS->UninstallProtocolInterface (
+    Status = gBS->UninstallMultipleProtocolInterfaces (
                     Controller,
                     &gEfiBlockIoProtocolGuid,
-                    &UsbMass->BlockIo
+                    &UsbMass->BlockIo,
+                    &gEfiDiskInfoProtocolGuid,
+                    &UsbMass->DiskInfo,
+                    NULL
                     );
     if (EFI_ERROR (Status)) {
       return Status;
@@ -1068,6 +1081,8 @@ USBMassDriverBindingStop (
                     UsbMass->DevicePath,
                     &gEfiBlockIoProtocolGuid,
                     &UsbMass->BlockIo,
+                    &gEfiDiskInfoProtocolGuid,
+                    &UsbMass->DiskInfo,
                     NULL
                     );
     

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassImpl.h

@@ -2,7 +2,7 @@
   Definitions of functions for Driver Binding Protocol and Block I/O Protocol,
   and other internal definitions.
 
-Copyright (c) 2007 - 2008, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -16,16 +16,15 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #ifndef _EFI_USBMASS_IMPL_H_
 #define _EFI_USBMASS_IMPL_H_
 
-#include "UsbMass.h"
-#include "UsbMassBot.h"
-#include "UsbMassCbi.h"
-#include "UsbMassBoot.h"
-
 #define  USB_MASS_SIGNATURE    SIGNATURE_32 ('U', 's', 'b', 'M')
 
 #define USB_MASS_DEVICE_FROM_BLOCK_IO(a) \
         CR (a, USB_MASS_DEVICE, BlockIo, USB_MASS_SIGNATURE)
 
+#define USB_MASS_DEVICE_FROM_DISK_INFO(a) \
+        CR (a, USB_MASS_DEVICE, DiskInfo, USB_MASS_SIGNATURE)
+
+
 extern EFI_COMPONENT_NAME_PROTOCOL   gUsbMassStorageComponentName;
 extern EFI_COMPONENT_NAME2_PROTOCOL  gUsbMassStorageComponentName2;
 

MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassStorageDxe.inf

@@ -14,7 +14,7 @@
 # 3. USB Mass Storage Class Bulk-Only Transport, Revision 1.0.
 # 4. UEFI Specification, v2.1
 #
-# Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+# Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>
 #
 #  This program and the accompanying materials
 #  are licensed and made available under the terms and conditions of the BSD License
@@ -56,6 +56,8 @@
   UsbMassCbi.h
   UsbMass.h
   UsbMassCbi.c
+  UsbMassDiskinfo.h
+  UsbMassDiskinfo.c
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -75,4 +77,4 @@
   gEfiUsbIoProtocolGuid                         ## TO_START
   gEfiDevicePathProtocolGuid                    ## TO_START
   gEfiBlockIoProtocolGuid                       ## BY_START
-
+  gEfiDiskInfoProtocolGuid                      ## BY_START