]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThruHci.c
MdeModulePkg: Add bRefClkFreq card attribute programming support
[mirror_edk2.git] / MdeModulePkg / Bus / Ufs / UfsPassThruDxe / UfsPassThruHci.c
index 822554cebbe29baaf6dd7f0596ce6f8761ff26f3..4a9fa01e7d8177f816283c23b3b1ad09011eaf7d 100644 (file)
@@ -2,14 +2,9 @@
   UfsPassThruDxe driver is used to produce EFI_EXT_SCSI_PASS_THRU protocol interface\r
   for upper layer application to execute UFS-supported SCSI cmds.\r
 \r
-  Copyright (c) 2014 - 2017, Intel Corporation. All rights reserved.<BR>\r
-  This program and the accompanying materials\r
-  are licensed and made available under the terms and conditions of the BSD License\r
-  which accompanies this distribution.  The full text of the license may be found at\r
-  http://opensource.org/licenses/bsd-license.php.\r
-\r
-  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+  Copyright (c) 2014 - 2022, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) Microsoft Corporation.<BR>\r
+  SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
@@ -30,9 +25,9 @@
 **/\r
 EFI_STATUS\r
 UfsMmioRead32 (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN     UINTN                        Offset,\r
-     OUT UINT32                       *Value\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     UINTN                       Offset,\r
+  OUT UINT32                         *Value\r
   )\r
 {\r
   EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
@@ -60,9 +55,9 @@ UfsMmioRead32 (
 **/\r
 EFI_STATUS\r
 UfsMmioWrite32 (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN  UINTN                        Offset,\r
-  IN  UINT32                       Value\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN  UINTN                       Offset,\r
+  IN  UINT32                      Value\r
   )\r
 {\r
   EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
@@ -92,17 +87,17 @@ UfsMmioWrite32 (
 **/\r
 EFI_STATUS\r
 UfsWaitMemSet (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN  UINTN                        Offset,\r
-  IN  UINT32                       MaskValue,\r
-  IN  UINT32                       TestValue,\r
-  IN  UINT64                       Timeout\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN  UINTN                       Offset,\r
+  IN  UINT32                      MaskValue,\r
+  IN  UINT32                      TestValue,\r
+  IN  UINT64                      Timeout\r
   )\r
 {\r
-  UINT32     Value;\r
-  UINT64     Delay;\r
-  BOOLEAN    InfiniteWait;\r
-  EFI_STATUS Status;\r
+  UINT32      Value;\r
+  UINT64      Delay;\r
+  BOOLEAN     InfiniteWait;\r
+  EFI_STATUS  Status;\r
 \r
   if (Timeout == 0) {\r
     InfiniteWait = TRUE;\r
@@ -133,7 +128,6 @@ UfsWaitMemSet (
     MicroSecondDelay (1);\r
 \r
     Delay--;\r
-\r
   } while (InfiniteWait || (Delay > 0));\r
 \r
   return EFI_TIMEOUT;\r
@@ -148,8 +142,8 @@ UfsWaitMemSet (
 **/\r
 VOID\r
 DumpUicCmdExecResult (\r
-  IN  UINT8     UicOpcode,\r
-  IN  UINT8     Result\r
+  IN  UINT8  UicOpcode,\r
+  IN  UINT8  Result\r
   )\r
 {\r
   if (UicOpcode <= UfsUicDmePeerSet) {\r
@@ -179,14 +173,14 @@ DumpUicCmdExecResult (
         break;\r
       case 0x08:\r
         DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - PEER_COMMUNICATION_FAILURE\n"));\r
-        break; \r
+        break;\r
       case 0x09:\r
         DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - BUSY\n"));\r
         break;\r
       case 0x0A:\r
         DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - DME_FAILURE\n"));\r
-        break;        \r
-      default :\r
+        break;\r
+      default:\r
         ASSERT (FALSE);\r
         break;\r
     }\r
@@ -196,8 +190,8 @@ DumpUicCmdExecResult (
         break;\r
       case 0x01:\r
         DEBUG ((DEBUG_VERBOSE, "UIC control command fails - FAILURE\n"));\r
-        break;     \r
-      default :\r
+        break;\r
+      default:\r
         ASSERT (FALSE);\r
         break;\r
     }\r
@@ -212,7 +206,7 @@ DumpUicCmdExecResult (
 **/\r
 VOID\r
 DumpQueryResponseResult (\r
-  IN  UINT8     Result\r
+  IN  UINT8  Result\r
   )\r
 {\r
   switch (Result) {\r
@@ -242,11 +236,11 @@ DumpQueryResponseResult (
       break;\r
     case 0xFE:\r
       DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Opcode\n"));\r
-      break; \r
+      break;\r
     case 0xFF:\r
       DEBUG ((DEBUG_VERBOSE, "Query Response with General Failure\n"));\r
       break;\r
-    default :\r
+    default:\r
       ASSERT (FALSE);\r
       break;\r
   }\r
@@ -262,18 +256,18 @@ DumpQueryResponseResult (
 **/\r
 VOID\r
 SwapLittleEndianToBigEndian (\r
-  IN OUT UINT8         *Buffer,\r
-  IN     UINT32        BufferSize\r
+  IN OUT UINT8   *Buffer,\r
+  IN     UINT32  BufferSize\r
   )\r
 {\r
-  UINT32 Index;\r
-  UINT8  Temp;\r
-  UINT32 SwapCount;\r
+  UINT32  Index;\r
+  UINT8   Temp;\r
+  UINT32  SwapCount;\r
 \r
   SwapCount = BufferSize / 2;\r
   for (Index = 0; Index < SwapCount; Index++) {\r
-    Temp = Buffer[Index];\r
-    Buffer[Index] = Buffer[BufferSize - 1 - Index];\r
+    Temp                           = Buffer[Index];\r
+    Buffer[Index]                  = Buffer[BufferSize - 1 - Index];\r
     Buffer[BufferSize - 1 - Index] = Temp;\r
   }\r
 }\r
@@ -292,32 +286,32 @@ SwapLittleEndianToBigEndian (
 **/\r
 VOID\r
 UfsFillTsfOfQueryReqUpiu (\r
-  IN OUT UTP_UPIU_TSF        *TsfBase,\r
-  IN     UINT8               Opcode,\r
-  IN     UINT8               DescId    OPTIONAL,\r
-  IN     UINT8               Index     OPTIONAL,\r
-  IN     UINT8               Selector  OPTIONAL,\r
-  IN     UINT16              Length    OPTIONAL,\r
-  IN     UINT32              Value     OPTIONAL\r
+  IN OUT UTP_UPIU_TSF  *TsfBase,\r
+  IN     UINT8         Opcode,\r
+  IN     UINT8         DescId    OPTIONAL,\r
+  IN     UINT8         Index     OPTIONAL,\r
+  IN     UINT8         Selector  OPTIONAL,\r
+  IN     UINT16        Length    OPTIONAL,\r
+  IN     UINT32        Value     OPTIONAL\r
   )\r
 {\r
   ASSERT (TsfBase != NULL);\r
   ASSERT (Opcode <= UtpQueryFuncOpcodeTogFlag);\r
 \r
-  TsfBase->Opcode   = Opcode;\r
+  TsfBase->Opcode = Opcode;\r
   if (Opcode != UtpQueryFuncOpcodeNop) {\r
     TsfBase->DescId   = DescId;\r
     TsfBase->Index    = Index;\r
     TsfBase->Selector = Selector;\r
 \r
     if ((Opcode == UtpQueryFuncOpcodeRdDesc) || (Opcode == UtpQueryFuncOpcodeWrDesc)) {\r
-      SwapLittleEndianToBigEndian ((UINT8*)&Length, sizeof (Length));\r
+      SwapLittleEndianToBigEndian ((UINT8 *)&Length, sizeof (Length));\r
       TsfBase->Length = Length;\r
     }\r
-  \r
+\r
     if (Opcode == UtpQueryFuncOpcodeWrAttr) {\r
-      SwapLittleEndianToBigEndian ((UINT8*)&Value, sizeof (Value));\r
-      TsfBase->Value  = Value;\r
+      SwapLittleEndianToBigEndian ((UINT8 *)&Value, sizeof (Value));\r
+      TsfBase->Value = Value;\r
     }\r
   }\r
 }\r
@@ -338,16 +332,16 @@ UfsFillTsfOfQueryReqUpiu (
 **/\r
 EFI_STATUS\r
 UfsInitCommandUpiu (\r
-  IN OUT UTP_COMMAND_UPIU              *Command,\r
-  IN     UINT8                         Lun,\r
-  IN     UINT8                         TaskTag,\r
-  IN     UINT8                         *Cdb,\r
-  IN     UINT8                         CdbLength,\r
-  IN     UFS_DATA_DIRECTION            DataDirection,\r
-  IN     UINT32                        ExpDataTranLen\r
+  IN OUT UTP_COMMAND_UPIU    *Command,\r
+  IN     UINT8               Lun,\r
+  IN     UINT8               TaskTag,\r
+  IN     UINT8               *Cdb,\r
+  IN     UINT8               CdbLength,\r
+  IN     UFS_DATA_DIRECTION  DataDirection,\r
+  IN     UINT32              ExpDataTranLen\r
   )\r
 {\r
-  UINT8                   Flags;\r
+  UINT8  Flags;\r
 \r
   ASSERT ((Command != NULL) && (Cdb != NULL));\r
 \r
@@ -370,7 +364,7 @@ UfsInitCommandUpiu (
   Command->Lun       = Lun;\r
   Command->TaskTag   = TaskTag;\r
   Command->CmdSet    = 0x00;\r
-  SwapLittleEndianToBigEndian ((UINT8*)&ExpDataTranLen, sizeof (ExpDataTranLen));\r
+  SwapLittleEndianToBigEndian ((UINT8 *)&ExpDataTranLen, sizeof (ExpDataTranLen));\r
   Command->ExpDataTranLen = ExpDataTranLen;\r
 \r
   CopyMem (Command->Cdb, Cdb, CdbLength);\r
@@ -390,27 +384,23 @@ UfsInitCommandUpiu (
 **/\r
 EFI_STATUS\r
 UfsInitUtpPrdt (\r
-  IN  UTP_TR_PRD                       *Prdt,\r
-  IN  VOID                             *Buffer,\r
-  IN  UINT32                           BufferSize\r
+  IN  UTP_TR_PRD  *Prdt,\r
+  IN  VOID        *Buffer,\r
+  IN  UINT32      BufferSize\r
   )\r
 {\r
-  UINT32     PrdtIndex;\r
-  UINT32     RemainingLen;\r
-  UINT8      *Remaining;\r
-  UINTN      PrdtNumber;\r
+  UINT32  PrdtIndex;\r
+  UINT32  RemainingLen;\r
+  UINT8   *Remaining;\r
+  UINTN   PrdtNumber;\r
 \r
-  if ((BufferSize & (BIT0 | BIT1)) != 0) {\r
-    BufferSize &= ~(BIT0 | BIT1);\r
-    DEBUG ((DEBUG_WARN, "UfsInitUtpPrdt: The BufferSize [%d] is not dword-aligned!\n", BufferSize));\r
-  }\r
+  ASSERT (((UINTN)Buffer & (BIT0 | BIT1)) == 0);\r
+  ASSERT ((BufferSize & (BIT1 | BIT0)) == 0);\r
 \r
   if (BufferSize == 0) {\r
     return EFI_SUCCESS;\r
   }\r
 \r
-  ASSERT (((UINTN)Buffer & (BIT0 | BIT1)) == 0);\r
-\r
   RemainingLen = BufferSize;\r
   Remaining    = Buffer;\r
   PrdtNumber   = (UINTN)DivU64x32 ((UINT64)BufferSize + UFS_MAX_DATA_LEN_PER_PRD - 1, UFS_MAX_DATA_LEN_PER_PRD);\r
@@ -424,8 +414,8 @@ UfsInitUtpPrdt (
 \r
     Prdt[PrdtIndex].DbAddr  = (UINT32)RShiftU64 ((UINT64)(UINTN)Remaining, 2);\r
     Prdt[PrdtIndex].DbAddrU = (UINT32)RShiftU64 ((UINT64)(UINTN)Remaining, 32);\r
-    RemainingLen -= UFS_MAX_DATA_LEN_PER_PRD;\r
-    Remaining    += UFS_MAX_DATA_LEN_PER_PRD;\r
+    RemainingLen           -= UFS_MAX_DATA_LEN_PER_PRD;\r
+    Remaining              += UFS_MAX_DATA_LEN_PER_PRD;\r
   }\r
 \r
   return EFI_SUCCESS;\r
@@ -448,14 +438,14 @@ UfsInitUtpPrdt (
 **/\r
 EFI_STATUS\r
 UfsInitQueryRequestUpiu (\r
-  IN OUT UTP_QUERY_REQ_UPIU            *QueryReq,\r
-  IN     UINT8                         TaskTag,\r
-  IN     UINT8                         Opcode,\r
-  IN     UINT8                         DescId,\r
-  IN     UINT8                         Index,\r
-  IN     UINT8                         Selector,\r
-  IN     UINTN                         DataSize   OPTIONAL,\r
-  IN     UINT8                         *Data      OPTIONAL\r
+  IN OUT UTP_QUERY_REQ_UPIU  *QueryReq,\r
+  IN     UINT8               TaskTag,\r
+  IN     UINT8               Opcode,\r
+  IN     UINT8               DescId,\r
+  IN     UINT8               Index,\r
+  IN     UINT8               Selector,\r
+  IN     UINTN               DataSize   OPTIONAL,\r
+  IN     UINT8               *Data      OPTIONAL\r
   )\r
 {\r
   ASSERT (QueryReq != NULL);\r
@@ -469,7 +459,7 @@ UfsInitQueryRequestUpiu (
   }\r
 \r
   if (Opcode == UtpQueryFuncOpcodeWrAttr) {\r
-    UfsFillTsfOfQueryReqUpiu (&QueryReq->Tsf, Opcode, DescId, Index, Selector, 0, *(UINT32*)Data);\r
+    UfsFillTsfOfQueryReqUpiu (&QueryReq->Tsf, Opcode, DescId, Index, Selector, 0, *(UINT32 *)Data);\r
   } else if ((Opcode == UtpQueryFuncOpcodeRdDesc) || (Opcode == UtpQueryFuncOpcodeWrDesc)) {\r
     UfsFillTsfOfQueryReqUpiu (&QueryReq->Tsf, Opcode, DescId, Index, Selector, (UINT16)DataSize, 0);\r
   } else {\r
@@ -478,6 +468,9 @@ UfsInitQueryRequestUpiu (
 \r
   if (Opcode == UtpQueryFuncOpcodeWrDesc) {\r
     CopyMem (QueryReq + 1, Data, DataSize);\r
+\r
+    SwapLittleEndianToBigEndian ((UINT8 *)&DataSize, sizeof (UINT16));\r
+    QueryReq->DataSegLen = (UINT16)DataSize;\r
   }\r
 \r
   return EFI_SUCCESS;\r
@@ -504,17 +497,17 @@ UfsCreateScsiCommandDesc (
   IN     UINT8                                       Lun,\r
   IN     EFI_EXT_SCSI_PASS_THRU_SCSI_REQUEST_PACKET  *Packet,\r
   IN     UTP_TRD                                     *Trd,\r
-     OUT VOID                                        **CmdDescHost,\r
-     OUT VOID                                        **CmdDescMapping\r
+  OUT VOID                                           **CmdDescHost,\r
+  OUT VOID                                           **CmdDescMapping\r
   )\r
 {\r
-  UINTN                             TotalLen;\r
-  UINTN                             PrdtNumber;\r
-  UTP_COMMAND_UPIU                  *CommandUpiu;\r
-  EFI_PHYSICAL_ADDRESS              CmdDescPhyAddr;\r
-  EFI_STATUS                        Status;\r
-  UINT32                            DataLen;\r
-  UFS_DATA_DIRECTION                DataDirection;\r
+  UINTN                 TotalLen;\r
+  UINTN                 PrdtNumber;\r
+  UTP_COMMAND_UPIU      *CommandUpiu;\r
+  EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;\r
+  EFI_STATUS            Status;\r
+  UINT32                DataLen;\r
+  UFS_DATA_DIRECTION    DataDirection;\r
 \r
   ASSERT ((Private != NULL) && (Packet != NULL) && (Trd != NULL));\r
 \r
@@ -532,14 +525,14 @@ UfsCreateScsiCommandDesc (
 \r
   PrdtNumber = (UINTN)DivU64x32 ((UINT64)DataLen + UFS_MAX_DATA_LEN_PER_PRD - 1, UFS_MAX_DATA_LEN_PER_PRD);\r
 \r
-  TotalLen   = ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)) + PrdtNumber * sizeof (UTP_TR_PRD);\r
+  TotalLen = ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)) + PrdtNumber * sizeof (UTP_TR_PRD);\r
 \r
   Status = UfsAllocateAlignCommonBuffer (Private, TotalLen, CmdDescHost, &CmdDescPhyAddr, CmdDescMapping);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
 \r
-  CommandUpiu = (UTP_COMMAND_UPIU*)*CmdDescHost;\r
+  CommandUpiu = (UTP_COMMAND_UPIU *)*CmdDescHost;\r
 \r
   UfsInitCommandUpiu (CommandUpiu, Lun, Private->TaskTag++, Packet->Cdb, Packet->CdbLength, DataDirection, DataLen);\r
 \r
@@ -581,18 +574,18 @@ UfsCreateDMCommandDesc (
   IN     UFS_PASS_THRU_PRIVATE_DATA            *Private,\r
   IN     UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet,\r
   IN     UTP_TRD                               *Trd,\r
-     OUT VOID                                  **CmdDescHost,\r
-     OUT VOID                                  **CmdDescMapping\r
+  OUT VOID                                     **CmdDescHost,\r
+  OUT VOID                                     **CmdDescMapping\r
   )\r
 {\r
-  UINTN                         TotalLen;\r
-  UTP_QUERY_REQ_UPIU            *QueryReqUpiu;\r
-  UINT8                         Opcode;\r
-  UINT32                        DataSize;\r
-  UINT8                         *Data;\r
-  UINT8                         DataDirection;\r
-  EFI_PHYSICAL_ADDRESS          CmdDescPhyAddr;\r
-  EFI_STATUS                    Status;\r
+  UINTN                 TotalLen;\r
+  UTP_QUERY_REQ_UPIU    *QueryReqUpiu;\r
+  UINT8                 Opcode;\r
+  UINT32                DataSize;\r
+  UINT8                 *Data;\r
+  UINT8                 DataDirection;\r
+  EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;\r
+  EFI_STATUS            Status;\r
 \r
   ASSERT ((Private != NULL) && (Packet != NULL) && (Trd != NULL));\r
 \r
@@ -602,32 +595,14 @@ UfsCreateDMCommandDesc (
   }\r
 \r
   DataDirection = Packet->DataDirection;\r
-  if (DataDirection == UfsDataIn) {\r
-    DataSize = Packet->InTransferLength;\r
-    Data     = Packet->InDataBuffer;\r
-  } else if (DataDirection == UfsDataOut) {\r
-    DataSize = Packet->OutTransferLength;\r
-    Data     = Packet->OutDataBuffer;\r
-  } else {\r
-    DataSize = 0;\r
-    Data     = NULL;\r
-  }\r
-\r
-  if (((Opcode != UtpQueryFuncOpcodeSetFlag) && (Opcode != UtpQueryFuncOpcodeClrFlag) && (Opcode != UtpQueryFuncOpcodeTogFlag))\r
-    && ((DataSize == 0) || (Data == NULL))) {\r
-    return EFI_INVALID_PARAMETER;\r
-  }\r
-\r
-  if (((Opcode == UtpQueryFuncOpcodeSetFlag) || (Opcode == UtpQueryFuncOpcodeClrFlag) || (Opcode == UtpQueryFuncOpcodeTogFlag))\r
-    && ((DataSize != 0) || (Data != NULL))) {\r
-    return EFI_INVALID_PARAMETER;\r
-  }\r
-\r
-  if ((Opcode == UtpQueryFuncOpcodeWrAttr) && (DataSize != sizeof (UINT32))) {\r
-    return EFI_INVALID_PARAMETER;\r
-  }\r
+  DataSize      = Packet->TransferLength;\r
+  Data          = Packet->DataBuffer;\r
 \r
   if ((Opcode == UtpQueryFuncOpcodeWrDesc) || (Opcode == UtpQueryFuncOpcodeRdDesc)) {\r
+    if ((DataSize == 0) || (Data == NULL)) {\r
+      return EFI_INVALID_PARAMETER;\r
+    }\r
+\r
     TotalLen = ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)) + ROUNDUP8 (DataSize);\r
   } else {\r
     TotalLen = ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU));\r
@@ -641,7 +616,7 @@ UfsCreateDMCommandDesc (
   //\r
   // Initialize UTP QUERY REQUEST UPIU\r
   //\r
-  QueryReqUpiu = (UTP_QUERY_REQ_UPIU*)*CmdDescHost;\r
+  QueryReqUpiu = (UTP_QUERY_REQ_UPIU *)*CmdDescHost;\r
   ASSERT (QueryReqUpiu != NULL);\r
   UfsInitQueryRequestUpiu (\r
     QueryReqUpiu,\r
@@ -665,11 +640,11 @@ UfsCreateDMCommandDesc (
   Trd->UcdBa  = (UINT32)RShiftU64 ((UINT64)CmdDescPhyAddr, 7);\r
   Trd->UcdBaU = (UINT32)RShiftU64 ((UINT64)CmdDescPhyAddr, 32);\r
   if (Opcode == UtpQueryFuncOpcodeWrDesc) {\r
-    Trd->RuL  = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)), sizeof (UINT32));\r
-    Trd->RuO  = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (DataSize), sizeof (UINT32));\r
+    Trd->RuL = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)), sizeof (UINT32));\r
+    Trd->RuO = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (DataSize), sizeof (UINT32));\r
   } else {\r
-    Trd->RuL  = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)) + ROUNDUP8 (DataSize), sizeof (UINT32));\r
-    Trd->RuO  = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)), sizeof (UINT32));\r
+    Trd->RuL = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)) + ROUNDUP8 (DataSize), sizeof (UINT32));\r
+    Trd->RuO = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)), sizeof (UINT32));\r
   }\r
 \r
   return EFI_SUCCESS;\r
@@ -690,16 +665,16 @@ UfsCreateDMCommandDesc (
 **/\r
 EFI_STATUS\r
 UfsCreateNopCommandDesc (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA        *Private,\r
-  IN     UTP_TRD                           *Trd,\r
-     OUT VOID                              **CmdDescHost,\r
-     OUT VOID                              **CmdDescMapping\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     UTP_TRD                     *Trd,\r
+  OUT VOID                           **CmdDescHost,\r
+  OUT VOID                           **CmdDescMapping\r
   )\r
 {\r
-  UINTN                    TotalLen;\r
-  UTP_NOP_OUT_UPIU         *NopOutUpiu;\r
-  EFI_STATUS               Status;\r
-  EFI_PHYSICAL_ADDRESS     CmdDescPhyAddr;\r
+  UINTN                 TotalLen;\r
+  UTP_NOP_OUT_UPIU      *NopOutUpiu;\r
+  EFI_STATUS            Status;\r
+  EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;\r
 \r
   ASSERT ((Private != NULL) && (Trd != NULL));\r
 \r
@@ -709,7 +684,7 @@ UfsCreateNopCommandDesc (
     return Status;\r
   }\r
 \r
-  NopOutUpiu = (UTP_NOP_OUT_UPIU*)*CmdDescHost;\r
+  NopOutUpiu = (UTP_NOP_OUT_UPIU *)*CmdDescHost;\r
   ASSERT (NopOutUpiu != NULL);\r
   NopOutUpiu->TaskTag = Private->TaskTag++;\r
 \r
@@ -741,23 +716,23 @@ UfsCreateNopCommandDesc (
 **/\r
 EFI_STATUS\r
 UfsFindAvailableSlotInTrl (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-     OUT UINT8                        *Slot\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  OUT UINT8                          *Slot\r
   )\r
 {\r
-  UINT8            Nutrs;\r
-  UINT8            Index;\r
-  UINT32           Data;\r
-  EFI_STATUS       Status;\r
+  UINT8       Nutrs;\r
+  UINT8       Index;\r
+  UINT32      Data;\r
+  EFI_STATUS  Status;\r
 \r
   ASSERT ((Private != NULL) && (Slot != NULL));\r
 \r
-  Status  = UfsMmioRead32 (Private, UFS_HC_UTRLDBR_OFFSET, &Data);\r
+  Status = UfsMmioRead32 (Private, UFS_HC_UTRLDBR_OFFSET, &Data);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
 \r
-  Nutrs   = (UINT8)((Private->Capabilities & UFS_HC_CAP_NUTRS) + 1);\r
+  Nutrs = (UINT8)((Private->UfsHcInfo.Capabilities & UFS_HC_CAP_NUTRS) + 1);\r
 \r
   for (Index = 0; Index < Nutrs; Index++) {\r
     if ((Data & (BIT0 << Index)) == 0) {\r
@@ -769,32 +744,6 @@ UfsFindAvailableSlotInTrl (
   return EFI_NOT_READY;\r
 }\r
 \r
-/**\r
-  Find out available slot in task management transfer list of a UFS device.\r
-\r
-  @param[in]  Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
-  @param[out] Slot          The available slot.\r
-\r
-  @retval EFI_SUCCESS       The available slot was found successfully.\r
-\r
-**/\r
-EFI_STATUS\r
-UfsFindAvailableSlotInTmrl (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-     OUT UINT8                        *Slot\r
-  )\r
-{\r
-  ASSERT ((Private != NULL) && (Slot != NULL));\r
-\r
-  //\r
-  // The simplest algo to always use slot 0.\r
-  // TODO: enhance it to support async transfer with multiple slot.\r
-  //\r
-  *Slot = 0;\r
-\r
-  return EFI_SUCCESS;\r
-}\r
-\r
 /**\r
   Start specified slot in transfer list of a UFS device.\r
 \r
@@ -804,12 +753,12 @@ UfsFindAvailableSlotInTmrl (
 **/\r
 EFI_STATUS\r
 UfsStartExecCmd (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN  UINT8                        Slot\r
-  ) \r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN  UINT8                       Slot\r
+  )\r
 {\r
-  UINT32        Data;\r
-  EFI_STATUS    Status;\r
+  UINT32      Data;\r
+  EFI_STATUS  Status;\r
 \r
   Status = UfsMmioRead32 (Private, UFS_HC_UTRLRSR_OFFSET, &Data);\r
   if (EFI_ERROR (Status)) {\r
@@ -840,12 +789,12 @@ UfsStartExecCmd (
 **/\r
 EFI_STATUS\r
 UfsStopExecCmd (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN  UINT8                        Slot\r
-  ) \r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN  UINT8                       Slot\r
+  )\r
 {\r
-  UINT32        Data;\r
-  EFI_STATUS    Status;\r
+  UINT32      Data;\r
+  EFI_STATUS  Status;\r
 \r
   Status = UfsMmioRead32 (Private, UFS_HC_UTRLDBR_OFFSET, &Data);\r
   if (EFI_ERROR (Status)) {\r
@@ -868,60 +817,97 @@ UfsStopExecCmd (
 }\r
 \r
 /**\r
-  Read or write specified device descriptor of a UFS device.\r
+  Extracts return data from query response upiu.\r
 \r
-  @param[in]      Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
-  @param[in]      Read          The boolean variable to show r/w direction.\r
-  @param[in]      DescId        The ID of device descriptor.\r
-  @param[in]      Index         The Index of device descriptor.\r
-  @param[in]      Selector      The Selector of device descriptor.\r
-  @param[in, out] Descriptor    The buffer of device descriptor to be read or written.\r
-  @param[in]      DescSize      The size of device descriptor buffer.\r
+  @param[in] Packet     Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.\r
+  @param[in] QueryResp  Pointer to the query response.\r
 \r
-  @retval EFI_SUCCESS           The device descriptor was read/written successfully.\r
-  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the device descriptor.\r
-  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the device descriptor.\r
+  @retval EFI_INVALID_PARAMETER Packet or QueryResp are empty or opcode is invalid.\r
+  @retval EFI_DEVICE_ERROR      Data returned from device is invalid.\r
+  @retval EFI_SUCCESS           Data extracted.\r
 \r
 **/\r
 EFI_STATUS\r
-UfsRwDeviceDesc (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN     BOOLEAN                      Read,\r
-  IN     UINT8                        DescId,\r
-  IN     UINT8                        Index,\r
-  IN     UINT8                        Selector,\r
-  IN OUT VOID                         *Descriptor,\r
-  IN     UINT32                       DescSize\r
+UfsGetReturnDataFromQueryResponse (\r
+  IN UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet,\r
+  IN UTP_QUERY_RESP_UPIU                   *QueryResp\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
-  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;\r
-  UINT8                                Slot;\r
-  UTP_TRD                              *Trd;\r
-  UTP_QUERY_RESP_UPIU                  *QueryResp;\r
-  UINT32                               CmdDescSize;\r
-  UINT16                               ReturnDataSize;\r
-  VOID                                 *CmdDescHost;\r
-  VOID                                 *CmdDescMapping;\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;\r
+  UINT16  ReturnDataSize;\r
+  UINT32  ReturnData;\r
 \r
-  ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));\r
+  if ((Packet == NULL) || (QueryResp == NULL)) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
 \r
-  if (Read) {\r
-    Packet.DataDirection     = UfsDataIn;\r
-    Packet.InDataBuffer      = Descriptor;\r
-    Packet.InTransferLength  = DescSize;\r
-    Packet.Opcode            = UtpQueryFuncOpcodeRdDesc;\r
-  } else {\r
-    Packet.DataDirection     = UfsDataOut;\r
-    Packet.OutDataBuffer     = Descriptor;\r
-    Packet.OutTransferLength = DescSize;\r
-    Packet.Opcode            = UtpQueryFuncOpcodeWrDesc;\r
+  switch (Packet->Opcode) {\r
+    case UtpQueryFuncOpcodeRdDesc:\r
+      ReturnDataSize = QueryResp->Tsf.Length;\r
+      SwapLittleEndianToBigEndian ((UINT8 *)&ReturnDataSize, sizeof (UINT16));\r
+      //\r
+      // Make sure the hardware device does not return more data than expected.\r
+      //\r
+      if (ReturnDataSize > Packet->TransferLength) {\r
+        return EFI_DEVICE_ERROR;\r
+      }\r
+\r
+      CopyMem (Packet->DataBuffer, (QueryResp + 1), ReturnDataSize);\r
+      Packet->TransferLength = ReturnDataSize;\r
+      break;\r
+    case UtpQueryFuncOpcodeWrDesc:\r
+      ReturnDataSize = QueryResp->Tsf.Length;\r
+      SwapLittleEndianToBigEndian ((UINT8 *)&ReturnDataSize, sizeof (UINT16));\r
+      Packet->TransferLength = ReturnDataSize;\r
+      break;\r
+    case UtpQueryFuncOpcodeRdFlag:\r
+    case UtpQueryFuncOpcodeSetFlag:\r
+    case UtpQueryFuncOpcodeClrFlag:\r
+    case UtpQueryFuncOpcodeTogFlag:\r
+      //\r
+      // The 'FLAG VALUE' field is at byte offset 3 of QueryResp->Tsf.Value\r
+      //\r
+      *((UINT8 *)(Packet->DataBuffer)) = *((UINT8 *)&(QueryResp->Tsf.Value) + 3);\r
+      break;\r
+    case UtpQueryFuncOpcodeRdAttr:\r
+    case UtpQueryFuncOpcodeWrAttr:\r
+      ReturnData = QueryResp->Tsf.Value;\r
+      SwapLittleEndianToBigEndian ((UINT8 *)&ReturnData, sizeof (UINT32));\r
+      CopyMem (Packet->DataBuffer, &ReturnData, sizeof (UINT32));\r
+      break;\r
+    default:\r
+      return EFI_INVALID_PARAMETER;\r
   }\r
-  Packet.DescId              = DescId;\r
-  Packet.Index               = Index;\r
-  Packet.Selector            = Selector;\r
-  Packet.Timeout             = UFS_TIMEOUT;\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Creates Transfer Request descriptor and sends Query Request to the device.\r
+\r
+  @param[in] Private Pointer to the UFS_PASS_THRU_PRIVATE_DATA.\r
+  @param[in] Packet  Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.\r
+\r
+  @retval EFI_SUCCESS           The device descriptor was read/written successfully.\r
+  @retval EFI_INVALID_PARAMETER The DescId, Index and Selector fields in Packet are invalid\r
+                                combination to point to a type of UFS device descriptor.\r
+  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the device descriptor.\r
+  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the device descriptor.\r
+\r
+**/\r
+EFI_STATUS\r
+UfsSendDmRequestRetry (\r
+  IN UFS_PASS_THRU_PRIVATE_DATA            *Private,\r
+  IN UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet\r
+  )\r
+{\r
+  UINT8                               Slot;\r
+  UTP_TRD                             *Trd;\r
+  VOID                                *CmdDescHost;\r
+  VOID                                *CmdDescMapping;\r
+  UINT32                              CmdDescSize;\r
+  EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
+  UTP_QUERY_RESP_UPIU                 *QueryResp;\r
+  EFI_STATUS                          Status;\r
 \r
   //\r
   // Find out which slot of transfer request list is available.\r
@@ -930,21 +916,19 @@ UfsRwDeviceDesc (
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
-  \r
-  Trd = ((UTP_TRD*)Private->UtpTrlBase) + Slot;\r
+\r
+  Trd = ((UTP_TRD *)Private->UtpTrlBase) + Slot;\r
   //\r
   // Fill transfer request descriptor to this slot.\r
   //\r
-  Status = UfsCreateDMCommandDesc (Private, &Packet, Trd, &CmdDescHost, &CmdDescMapping);\r
+  Status = UfsCreateDMCommandDesc (Private, Packet, Trd, &CmdDescHost, &CmdDescMapping);\r
   if (EFI_ERROR (Status)) {\r
+    DEBUG ((DEBUG_ERROR, "Failed to create DM command descriptor\n"));\r
     return Status;\r
   }\r
 \r
-  //\r
-  // Check the transfer request result.\r
-  //\r
-  UfsHc       = Private->UfsHostController;\r
-  QueryResp   = (UTP_QUERY_RESP_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));\r
+  UfsHc     = Private->UfsHostController;\r
+  QueryResp = (UTP_QUERY_RESP_UPIU *)((UINT8 *)CmdDescHost + Trd->RuO * sizeof (UINT32));\r
   ASSERT (QueryResp != NULL);\r
   CmdDescSize = Trd->RuO * sizeof (UINT32) + Trd->RuL * sizeof (UINT32);\r
 \r
@@ -955,30 +939,32 @@ UfsRwDeviceDesc (
 \r
   //\r
   // Wait for the completion of the transfer request.\r
-  //  \r
-  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, Packet.Timeout);\r
+  //\r
+  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0, 0, Packet->Timeout);\r
   if (EFI_ERROR (Status)) {\r
     goto Exit;\r
   }\r
 \r
-  if (QueryResp->QueryResp != 0) {\r
+  if ((Trd->Ocs != 0) || (QueryResp->QueryResp != UfsUtpQueryResponseSuccess)) {\r
+    DEBUG ((DEBUG_ERROR, "Failed to send query request, OCS = %X, QueryResp = %X\n", Trd->Ocs, QueryResp->QueryResp));\r
     DumpQueryResponseResult (QueryResp->QueryResp);\r
-    Status = EFI_DEVICE_ERROR;\r
-    goto Exit;\r
-  }\r
-\r
-  if (Trd->Ocs == 0) {\r
-    ReturnDataSize = QueryResp->Tsf.Length;\r
-    SwapLittleEndianToBigEndian ((UINT8*)&ReturnDataSize, sizeof (UINT16));\r
 \r
-    if (Read) {\r
-      CopyMem (Packet.InDataBuffer, (QueryResp + 1), ReturnDataSize);\r
-      Packet.InTransferLength = ReturnDataSize;\r
+    if ((QueryResp->QueryResp == UfsUtpQueryResponseInvalidSelector) ||\r
+        (QueryResp->QueryResp == UfsUtpQueryResponseInvalidIndex) ||\r
+        (QueryResp->QueryResp == UfsUtpQueryResponseInvalidIdn))\r
+    {\r
+      Status = EFI_INVALID_PARAMETER;\r
     } else {\r
-      Packet.OutTransferLength = ReturnDataSize;\r
+      Status = EFI_DEVICE_ERROR;\r
     }\r
-  } else {\r
-    Status = EFI_DEVICE_ERROR;\r
+\r
+    goto Exit;\r
+  }\r
+\r
+  Status = UfsGetReturnDataFromQueryResponse (Packet, QueryResp);\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG ((DEBUG_ERROR, "Failed to get return data from query response\n"));\r
+    goto Exit;\r
   }\r
 \r
 Exit:\r
@@ -989,6 +975,7 @@ Exit:
   if (CmdDescMapping != NULL) {\r
     UfsHc->Unmap (UfsHc, CmdDescMapping);\r
   }\r
+\r
   if (CmdDescHost != NULL) {\r
     UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (CmdDescSize), CmdDescHost);\r
   }\r
@@ -997,121 +984,151 @@ Exit:
 }\r
 \r
 /**\r
-  Read or write specified attribute of a UFS device.\r
+  Sends Query Request to the device. Query is sent until device responds correctly or counter runs out.\r
+\r
+  @param[in] Private Pointer to the UFS_PASS_THRU_PRIVATE_DATA.\r
+  @param[in] Packet  Pointer to the UFS_DEVICE_MANAGEMENT_PACKET.\r
+\r
+  @retval EFI_SUCCESS           The device responded correctly to the Query request.\r
+  @retval EFI_INVALID_PARAMETER The DescId, Index and Selector fields in Packet are invalid\r
+                                combination to point to a type of UFS device descriptor.\r
+  @retval EFI_DEVICE_ERROR      A device error occurred while waiting for the response from the device.\r
+  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of the operation.\r
+\r
+**/\r
+EFI_STATUS\r
+UfsSendDmRequest (\r
+  IN UFS_PASS_THRU_PRIVATE_DATA            *Private,\r
+  IN UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet\r
+  )\r
+{\r
+  EFI_STATUS  Status;\r
+  UINT8       Retry;\r
+\r
+  Status = EFI_SUCCESS;\r
+\r
+  for (Retry = 0; Retry < 5; Retry++) {\r
+    Status = UfsSendDmRequestRetry (Private, Packet);\r
+    if (!EFI_ERROR (Status)) {\r
+      return EFI_SUCCESS;\r
+    }\r
+  }\r
+\r
+  DEBUG ((DEBUG_ERROR, "Failed to get response from the device after %d retries\n", Retry));\r
+  return Status;\r
+}\r
+\r
+/**\r
+  Read or write specified device descriptor of a UFS device.\r
 \r
   @param[in]      Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
   @param[in]      Read          The boolean variable to show r/w direction.\r
-  @param[in]      AttrId        The ID of Attribute.\r
-  @param[in]      Index         The Index of Attribute.\r
-  @param[in]      Selector      The Selector of Attribute.\r
-  @param[in, out] Attributes    The value of Attribute to be read or written.\r
+  @param[in]      DescId        The ID of device descriptor.\r
+  @param[in]      Index         The Index of device descriptor.\r
+  @param[in]      Selector      The Selector of device descriptor.\r
+  @param[in, out] Descriptor    The buffer of device descriptor to be read or written.\r
+  @param[in, out] DescSize      The size of device descriptor buffer. On input, the size, in bytes,\r
+                                of the data buffer specified by Descriptor. On output, the number\r
+                                of bytes that were actually transferred.\r
 \r
-  @retval EFI_SUCCESS           The Attribute was read/written successfully.\r
-  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the Attribute.\r
-  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the Attribute.\r
+  @retval EFI_SUCCESS           The device descriptor was read/written successfully.\r
+  @retval EFI_INVALID_PARAMETER DescId, Index and Selector are invalid combination to point to a\r
+                                type of UFS device descriptor.\r
+  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the device descriptor.\r
+  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the device descriptor.\r
 \r
 **/\r
 EFI_STATUS\r
-UfsRwAttributes (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN     BOOLEAN                      Read,\r
-  IN     UINT8                        AttrId,\r
-  IN     UINT8                        Index,\r
-  IN     UINT8                        Selector,\r
-  IN OUT UINT32                       *Attributes\r
+UfsRwDeviceDesc (\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     BOOLEAN                     Read,\r
+  IN     UINT8                       DescId,\r
+  IN     UINT8                       Index,\r
+  IN     UINT8                       Selector,\r
+  IN OUT VOID                        *Descriptor,\r
+  IN OUT UINT32                      *DescSize\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
-  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;\r
-  UINT8                                Slot;\r
-  UTP_TRD                              *Trd;\r
-  UTP_QUERY_RESP_UPIU                  *QueryResp;\r
-  UINT32                               CmdDescSize;\r
-  UINT32                               ReturnData;\r
-  VOID                                 *CmdDescHost;\r
-  VOID                                 *CmdDescMapping;\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;\r
+  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  Packet;\r
+  EFI_STATUS                            Status;\r
+\r
+  if (DescSize == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
 \r
   ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));\r
 \r
   if (Read) {\r
-    Packet.DataDirection     = UfsDataIn;\r
-    Packet.Opcode            = UtpQueryFuncOpcodeRdAttr;\r
+    Packet.DataDirection = UfsDataIn;\r
+    Packet.Opcode        = UtpQueryFuncOpcodeRdDesc;\r
   } else {\r
-    Packet.DataDirection     = UfsDataOut;\r
-    Packet.Opcode            = UtpQueryFuncOpcodeWrAttr;\r
+    Packet.DataDirection = UfsDataOut;\r
+    Packet.Opcode        = UtpQueryFuncOpcodeWrDesc;\r
   }\r
-  Packet.DescId              = AttrId;\r
-  Packet.Index               = Index;\r
-  Packet.Selector            = Selector;\r
-  Packet.Timeout             = UFS_TIMEOUT;\r
 \r
-  //\r
-  // Find out which slot of transfer request list is available.\r
-  //\r
-  Status = UfsFindAvailableSlotInTrl (Private, &Slot);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-  \r
-  Trd = ((UTP_TRD*)Private->UtpTrlBase) + Slot;\r
-  //\r
-  // Fill transfer request descriptor to this slot.\r
-  //\r
-  Status = UfsCreateDMCommandDesc (Private, &Packet, Trd, &CmdDescHost, &CmdDescMapping);\r
+  Packet.DataBuffer     = Descriptor;\r
+  Packet.TransferLength = *DescSize;\r
+  Packet.DescId         = DescId;\r
+  Packet.Index          = Index;\r
+  Packet.Selector       = Selector;\r
+  Packet.Timeout        = UFS_TIMEOUT;\r
+\r
+  Status = UfsSendDmRequest (Private, &Packet);\r
   if (EFI_ERROR (Status)) {\r
-    return Status;\r
+    *DescSize = 0;\r
+  } else {\r
+    *DescSize = Packet.TransferLength;\r
   }\r
 \r
-  //\r
-  // Check the transfer request result.\r
-  //\r
-  UfsHc       = Private->UfsHostController;\r
-  QueryResp   = (UTP_QUERY_RESP_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));\r
-  ASSERT (QueryResp != NULL);\r
-  CmdDescSize = Trd->RuO * sizeof (UINT32) + Trd->RuL * sizeof (UINT32);\r
-\r
-  //\r
-  // Start to execute the transfer request.\r
-  //\r
-  UfsStartExecCmd (Private, Slot);\r
+  return Status;\r
+}\r
 \r
-  //\r
-  // Wait for the completion of the transfer request.\r
-  //  \r
-  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, Packet.Timeout);\r
-  if (EFI_ERROR (Status)) {\r
-    goto Exit;\r
-  }\r
+/**\r
+  Read or write specified attribute of a UFS device.\r
 \r
-  if (QueryResp->QueryResp != 0) {\r
-    DumpQueryResponseResult (QueryResp->QueryResp);\r
-    Status = EFI_DEVICE_ERROR;\r
-    goto Exit;\r
-  }\r
+  @param[in]      Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
+  @param[in]      Read          The boolean variable to show r/w direction.\r
+  @param[in]      AttrId        The ID of Attribute.\r
+  @param[in]      Index         The Index of Attribute.\r
+  @param[in]      Selector      The Selector of Attribute.\r
+  @param[in, out] Attributes    The value of Attribute to be read or written.\r
 \r
-  if (Trd->Ocs == 0) {\r
-    ReturnData = QueryResp->Tsf.Value;\r
-    SwapLittleEndianToBigEndian ((UINT8*)&ReturnData, sizeof (UINT32));\r
-    *Attributes = ReturnData;\r
-  } else {\r
-    Status = EFI_DEVICE_ERROR;\r
-  }\r
+  @retval EFI_SUCCESS           The Attribute was read/written successfully.\r
+  @retval EFI_INVALID_PARAMETER AttrId, Index and Selector are invalid combination to point to a\r
+                                type of UFS device descriptor.\r
+  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the Attribute.\r
+  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the Attribute.\r
 \r
-Exit:\r
-  UfsHc->Flush (UfsHc);\r
+**/\r
+EFI_STATUS\r
+UfsRwAttributes (\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     BOOLEAN                     Read,\r
+  IN     UINT8                       AttrId,\r
+  IN     UINT8                       Index,\r
+  IN     UINT8                       Selector,\r
+  IN OUT UINT32                      *Attributes\r
+  )\r
+{\r
+  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  Packet;\r
 \r
-  UfsStopExecCmd (Private, Slot);\r
+  ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));\r
 \r
-  if (CmdDescMapping != NULL) {\r
-    UfsHc->Unmap (UfsHc, CmdDescMapping);\r
+  if (Read) {\r
+    Packet.DataDirection = UfsDataIn;\r
+    Packet.Opcode        = UtpQueryFuncOpcodeRdAttr;\r
+  } else {\r
+    Packet.DataDirection = UfsDataOut;\r
+    Packet.Opcode        = UtpQueryFuncOpcodeWrAttr;\r
   }\r
 \r
-  if (CmdDescHost != NULL) {\r
-    UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (CmdDescSize), CmdDescHost);\r
-  }\r
+  Packet.DataBuffer = Attributes;\r
+  Packet.DescId     = AttrId;\r
+  Packet.Index      = Index;\r
+  Packet.Selector   = Selector;\r
+  Packet.Timeout    = UFS_TIMEOUT;\r
 \r
-  return Status;\r
+  return UfsSendDmRequest (Private, &Packet);\r
 }\r
 \r
 /**\r
@@ -1123,27 +1140,20 @@ Exit:
   @param[in, out] Value         The value to set or clear flag.\r
 \r
   @retval EFI_SUCCESS           The flag was read/written successfully.\r
+  @retval EFI_INVALID_PARAMETER FlagId is an invalid UFS flag ID.\r
   @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the flag.\r
   @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the flag.\r
 \r
 **/\r
 EFI_STATUS\r
 UfsRwFlags (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN     BOOLEAN                      Read,\r
-  IN     UINT8                        FlagId,\r
-  IN OUT UINT8                        *Value\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     BOOLEAN                     Read,\r
+  IN     UINT8                       FlagId,\r
+  IN OUT UINT8                       *Value\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
-  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;\r
-  UINT8                                Slot;\r
-  UTP_TRD                              *Trd;\r
-  UTP_QUERY_RESP_UPIU                  *QueryResp;\r
-  UINT32                               CmdDescSize;\r
-  VOID                                 *CmdDescHost;\r
-  VOID                                 *CmdDescMapping;\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;\r
+  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  Packet;\r
 \r
   if (Value == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
@@ -1153,86 +1163,26 @@ UfsRwFlags (
 \r
   if (Read) {\r
     ASSERT (Value != NULL);\r
-    Packet.DataDirection     = UfsDataIn;\r
-    Packet.Opcode            = UtpQueryFuncOpcodeRdFlag;\r
+    Packet.DataDirection = UfsDataIn;\r
+    Packet.Opcode        = UtpQueryFuncOpcodeRdFlag;\r
   } else {\r
-    Packet.DataDirection     = UfsDataOut;\r
+    Packet.DataDirection = UfsDataOut;\r
     if (*Value == 1) {\r
-      Packet.Opcode          = UtpQueryFuncOpcodeSetFlag;\r
+      Packet.Opcode = UtpQueryFuncOpcodeSetFlag;\r
     } else if (*Value == 0) {\r
-      Packet.Opcode          = UtpQueryFuncOpcodeClrFlag;\r
+      Packet.Opcode = UtpQueryFuncOpcodeClrFlag;\r
     } else {\r
       return EFI_INVALID_PARAMETER;\r
     }\r
   }\r
-  Packet.DescId              = FlagId;\r
-  Packet.Index               = 0;\r
-  Packet.Selector            = 0;\r
-  Packet.Timeout             = UFS_TIMEOUT;\r
-\r
-  //\r
-  // Find out which slot of transfer request list is available.\r
-  //\r
-  Status = UfsFindAvailableSlotInTrl (Private, &Slot);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-\r
-  //\r
-  // Fill transfer request descriptor to this slot.\r
-  //\r
-  Trd    = ((UTP_TRD*)Private->UtpTrlBase) + Slot;\r
-  Status = UfsCreateDMCommandDesc (Private, &Packet, Trd, &CmdDescHost, &CmdDescMapping);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-\r
-  //\r
-  // Check the transfer request result.\r
-  //\r
-  UfsHc       = Private->UfsHostController;\r
-  QueryResp   = (UTP_QUERY_RESP_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));\r
-  ASSERT (QueryResp != NULL);\r
-  CmdDescSize = Trd->RuO * sizeof (UINT32) + Trd->RuL * sizeof (UINT32);\r
-\r
-  //\r
-  // Start to execute the transfer request.\r
-  //\r
-  UfsStartExecCmd (Private, Slot);\r
-\r
-  //\r
-  // Wait for the completion of the transfer request.\r
-  //  \r
-  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, Packet.Timeout);\r
-  if (EFI_ERROR (Status)) {\r
-    goto Exit;\r
-  }\r
 \r
-  if (QueryResp->QueryResp != 0) {\r
-    DumpQueryResponseResult (QueryResp->QueryResp);\r
-    Status = EFI_DEVICE_ERROR;\r
-    goto Exit;\r
-  }\r
-\r
-  if (Trd->Ocs == 0) {\r
-    *Value = (UINT8)QueryResp->Tsf.Value;\r
-  } else {\r
-    Status = EFI_DEVICE_ERROR;\r
-  }\r
-\r
-Exit:\r
-  UfsHc->Flush (UfsHc);\r
-\r
-  UfsStopExecCmd (Private, Slot);\r
-\r
-  if (CmdDescMapping != NULL) {\r
-    UfsHc->Unmap (UfsHc, CmdDescMapping);\r
-  }\r
-  if (CmdDescHost != NULL) {\r
-    UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (CmdDescSize), CmdDescHost);\r
-  }\r
+  Packet.DataBuffer = Value;\r
+  Packet.DescId     = FlagId;\r
+  Packet.Index      = 0;\r
+  Packet.Selector   = 0;\r
+  Packet.Timeout    = UFS_TIMEOUT;\r
 \r
-  return Status;\r
+  return UfsSendDmRequest (Private, &Packet);\r
 }\r
 \r
 /**\r
@@ -1248,12 +1198,12 @@ Exit:
 **/\r
 EFI_STATUS\r
 UfsSetFlag (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN  UINT8                        FlagId\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN  UINT8                       FlagId\r
   )\r
 {\r
-  EFI_STATUS             Status;\r
-  UINT8                  Value;\r
+  EFI_STATUS  Status;\r
+  UINT8       Value;\r
 \r
   Value  = 1;\r
   Status = UfsRwFlags (Private, FALSE, FlagId, &Value);\r
@@ -1261,32 +1211,6 @@ UfsSetFlag (
   return Status;\r
 }\r
 \r
-/**\r
-  Clear specified flag to 0 on a UFS device.\r
-\r
-  @param[in]  Private           The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
-  @param[in]  FlagId            The ID of flag to be cleared.\r
-\r
-  @retval EFI_SUCCESS           The flag was cleared successfully.\r
-  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to clear the flag.\r
-  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of clearing the flag.\r
-\r
-**/\r
-EFI_STATUS\r
-UfsClearFlag (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN  UINT8                        FlagId\r
-  )\r
-{\r
-  EFI_STATUS             Status;\r
-  UINT8                  Value;\r
-\r
-  Value  = 0;\r
-  Status = UfsRwFlags (Private, FALSE, FlagId, &Value);\r
-\r
-  return Status;\r
-}\r
-\r
 /**\r
   Read specified flag from a UFS device.\r
 \r
@@ -1301,12 +1225,12 @@ UfsClearFlag (
 **/\r
 EFI_STATUS\r
 UfsReadFlag (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,\r
-  IN     UINT8                        FlagId,\r
-     OUT UINT8                        *Value\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     UINT8                       FlagId,\r
+  OUT UINT8                          *Value\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
+  EFI_STATUS  Status;\r
 \r
   Status = UfsRwFlags (Private, TRUE, FlagId, Value);\r
 \r
@@ -1328,17 +1252,17 @@ UfsReadFlag (
 **/\r
 EFI_STATUS\r
 UfsExecNopCmds (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA       *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
-  UINT8                                Slot;\r
-  UTP_TRD                              *Trd;\r
-  UTP_NOP_IN_UPIU                      *NopInUpiu;\r
-  UINT32                               CmdDescSize;\r
-  VOID                                 *CmdDescHost;\r
-  VOID                                 *CmdDescMapping;\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;\r
+  EFI_STATUS                          Status;\r
+  UINT8                               Slot;\r
+  UTP_TRD                             *Trd;\r
+  UTP_NOP_IN_UPIU                     *NopInUpiu;\r
+  UINT32                              CmdDescSize;\r
+  VOID                                *CmdDescHost;\r
+  VOID                                *CmdDescMapping;\r
+  EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
 \r
   //\r
   // Find out which slot of transfer request list is available.\r
@@ -1348,7 +1272,7 @@ UfsExecNopCmds (
     return Status;\r
   }\r
 \r
-  Trd    = ((UTP_TRD*)Private->UtpTrlBase) + Slot;\r
+  Trd    = ((UTP_TRD *)Private->UtpTrlBase) + Slot;\r
   Status = UfsCreateNopCommandDesc (Private, Trd, &CmdDescHost, &CmdDescMapping);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
@@ -1357,8 +1281,8 @@ UfsExecNopCmds (
   //\r
   // Check the transfer request result.\r
   //\r
-  UfsHc       = Private->UfsHostController;\r
-  NopInUpiu   = (UTP_NOP_IN_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));\r
+  UfsHc     = Private->UfsHostController;\r
+  NopInUpiu = (UTP_NOP_IN_UPIU *)((UINT8 *)CmdDescHost + Trd->RuO * sizeof (UINT32));\r
   ASSERT (NopInUpiu != NULL);\r
   CmdDescSize = Trd->RuO * sizeof (UINT32) + Trd->RuL * sizeof (UINT32);\r
 \r
@@ -1369,7 +1293,7 @@ UfsExecNopCmds (
 \r
   //\r
   // Wait for the completion of the transfer request.\r
-  //  \r
+  //\r
   Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, UFS_TIMEOUT);\r
   if (EFI_ERROR (Status)) {\r
     goto Exit;\r
@@ -1389,6 +1313,7 @@ Exit:
   if (CmdDescMapping != NULL) {\r
     UfsHc->Unmap (UfsHc, CmdDescMapping);\r
   }\r
+\r
   if (CmdDescHost != NULL) {\r
     UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (CmdDescSize), CmdDescHost);\r
   }\r
@@ -1396,6 +1321,149 @@ Exit:
   return Status;\r
 }\r
 \r
+/**\r
+  Cleanup data buffers after data transfer. This function\r
+  also takes care to copy all data to user memory pool for\r
+  unaligned data transfers.\r
+\r
+  @param[in] Private   Pointer to the UFS_PASS_THRU_PRIVATE_DATA\r
+  @param[in] TransReq  Pointer to the transfer request\r
+**/\r
+VOID\r
+UfsReconcileDataTransferBuffer (\r
+  IN UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN UFS_PASS_THRU_TRANS_REQ     *TransReq\r
+  )\r
+{\r
+  if (TransReq->DataBufMapping != NULL) {\r
+    Private->UfsHostController->Unmap (\r
+                                  Private->UfsHostController,\r
+                                  TransReq->DataBufMapping\r
+                                  );\r
+  }\r
+\r
+  //\r
+  // Check if unaligned transfer was performed. If it was and we read\r
+  // data from device copy memory to user data buffers before cleanup.\r
+  // The assumption is if auxiliary aligned data buffer is not NULL then\r
+  // unaligned transfer has been performed.\r
+  //\r
+  if (TransReq->AlignedDataBuf != NULL) {\r
+    if (TransReq->Packet->DataDirection == EFI_EXT_SCSI_DATA_DIRECTION_READ) {\r
+      CopyMem (TransReq->Packet->InDataBuffer, TransReq->AlignedDataBuf, TransReq->Packet->InTransferLength);\r
+    }\r
+\r
+    //\r
+    // Wipe out the transfer buffer in case it contains sensitive data.\r
+    //\r
+    ZeroMem (TransReq->AlignedDataBuf, TransReq->AlignedDataBufSize);\r
+    FreeAlignedPages (TransReq->AlignedDataBuf, EFI_SIZE_TO_PAGES (TransReq->AlignedDataBufSize));\r
+    TransReq->AlignedDataBuf = NULL;\r
+  }\r
+}\r
+\r
+/**\r
+  Prepare data buffer for transfer.\r
+\r
+  @param[in]      Private   Pointer to the UFS_PASS_THRU_PRIVATE_DATA\r
+  @param[in, out] TransReq  Pointer to the transfer request\r
+\r
+  @retval EFI_DEVICE_ERROR  Failed to prepare buffer for transfer\r
+  @retval EFI_SUCCESS       Buffer ready for transfer\r
+**/\r
+EFI_STATUS\r
+UfsPrepareDataTransferBuffer (\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN OUT UFS_PASS_THRU_TRANS_REQ     *TransReq\r
+  )\r
+{\r
+  EFI_STATUS                           Status;\r
+  VOID                                 *DataBuf;\r
+  UINT32                               DataLen;\r
+  UINTN                                MapLength;\r
+  EFI_PHYSICAL_ADDRESS                 DataBufPhyAddr;\r
+  EDKII_UFS_HOST_CONTROLLER_OPERATION  Flag;\r
+  UTP_TR_PRD                           *PrdtBase;\r
+\r
+  DataBufPhyAddr = 0;\r
+  DataBuf        = NULL;\r
+\r
+  //\r
+  // For unaligned data transfers we allocate auxiliary DWORD aligned memory pool.\r
+  // When command is finished auxiliary memory pool is copied into actual user memory.\r
+  // This is requiered to assure data transfer safety(DWORD alignment required by UFS spec.)\r
+  //\r
+  if (TransReq->Packet->DataDirection == EFI_EXT_SCSI_DATA_DIRECTION_READ) {\r
+    if (((UINTN)TransReq->Packet->InDataBuffer % 4 != 0) || (TransReq->Packet->InTransferLength % 4 != 0)) {\r
+      DataLen = TransReq->Packet->InTransferLength + (4 - (TransReq->Packet->InTransferLength % 4));\r
+      DataBuf = AllocateAlignedPages (EFI_SIZE_TO_PAGES (DataLen), 4);\r
+      if (DataBuf == NULL) {\r
+        return EFI_DEVICE_ERROR;\r
+      }\r
+\r
+      ZeroMem (DataBuf, DataLen);\r
+      TransReq->AlignedDataBuf     = DataBuf;\r
+      TransReq->AlignedDataBufSize = DataLen;\r
+    } else {\r
+      DataLen = TransReq->Packet->InTransferLength;\r
+      DataBuf = TransReq->Packet->InDataBuffer;\r
+    }\r
+\r
+    Flag = EdkiiUfsHcOperationBusMasterWrite;\r
+  } else {\r
+    if (((UINTN)TransReq->Packet->OutDataBuffer % 4 != 0) || (TransReq->Packet->OutTransferLength % 4 != 0)) {\r
+      DataLen = TransReq->Packet->OutTransferLength + (4 - (TransReq->Packet->OutTransferLength % 4));\r
+      DataBuf = AllocateAlignedPages (EFI_SIZE_TO_PAGES (DataLen), 4);\r
+      if (DataBuf == NULL) {\r
+        return EFI_DEVICE_ERROR;\r
+      }\r
+\r
+      CopyMem (DataBuf, TransReq->Packet->OutDataBuffer, TransReq->Packet->OutTransferLength);\r
+      TransReq->AlignedDataBuf     = DataBuf;\r
+      TransReq->AlignedDataBufSize = DataLen;\r
+    } else {\r
+      DataLen = TransReq->Packet->OutTransferLength;\r
+      DataBuf = TransReq->Packet->OutDataBuffer;\r
+    }\r
+\r
+    Flag = EdkiiUfsHcOperationBusMasterRead;\r
+  }\r
+\r
+  if (DataLen != 0) {\r
+    MapLength = DataLen;\r
+    Status    = Private->UfsHostController->Map (\r
+                                              Private->UfsHostController,\r
+                                              Flag,\r
+                                              DataBuf,\r
+                                              &MapLength,\r
+                                              &DataBufPhyAddr,\r
+                                              &TransReq->DataBufMapping\r
+                                              );\r
+\r
+    if (EFI_ERROR (Status) || (DataLen != MapLength)) {\r
+      if (TransReq->AlignedDataBuf != NULL) {\r
+        //\r
+        // Wipe out the transfer buffer in case it contains sensitive data.\r
+        //\r
+        ZeroMem (TransReq->AlignedDataBuf, TransReq->AlignedDataBufSize);\r
+        FreeAlignedPages (TransReq->AlignedDataBuf, EFI_SIZE_TO_PAGES (TransReq->AlignedDataBufSize));\r
+        TransReq->AlignedDataBuf = NULL;\r
+      }\r
+\r
+      return EFI_DEVICE_ERROR;\r
+    }\r
+  }\r
+\r
+  //\r
+  // Fill PRDT table of Command UPIU for executed SCSI cmd.\r
+  //\r
+  PrdtBase = (UTP_TR_PRD *)((UINT8 *)TransReq->CmdDescHost + ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)));\r
+  ASSERT (PrdtBase != NULL);\r
+  UfsInitUtpPrdt (PrdtBase, (VOID *)(UINTN)DataBufPhyAddr, DataLen);\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
 /**\r
   Sends a UFS-supported SCSI Request Packet to a UFS device that is attached to the UFS host controller.\r
 \r
@@ -1428,29 +1496,24 @@ UfsExecScsiCmds (
   IN     EFI_EVENT                                   Event    OPTIONAL\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
-  UTP_RESPONSE_UPIU                    *Response;\r
-  UINT16                               SenseDataLen;\r
-  UINT32                               ResTranCount;\r
-  VOID                                 *DataBuf;\r
-  EFI_PHYSICAL_ADDRESS                 DataBufPhyAddr;\r
-  UINT32                               DataLen;\r
-  UINTN                                MapLength;\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;\r
-  EDKII_UFS_HOST_CONTROLLER_OPERATION  Flag;\r
-  UTP_TR_PRD                           *PrdtBase;\r
-  EFI_TPL                              OldTpl;\r
-  UFS_PASS_THRU_TRANS_REQ              *TransReq;\r
+  EFI_STATUS                          Status;\r
+  UTP_RESPONSE_UPIU                   *Response;\r
+  UINT16                              SenseDataLen;\r
+  UINT32                              ResTranCount;\r
+  EFI_TPL                             OldTpl;\r
+  UFS_PASS_THRU_TRANS_REQ             *TransReq;\r
+  EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
 \r
-  TransReq       = AllocateZeroPool (sizeof (UFS_PASS_THRU_TRANS_REQ));\r
+  TransReq = AllocateZeroPool (sizeof (UFS_PASS_THRU_TRANS_REQ));\r
   if (TransReq == NULL) {\r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
   TransReq->Signature     = UFS_PASS_THRU_TRANS_REQ_SIG;\r
   TransReq->TimeoutRemain = Packet->Timeout;\r
-  DataBufPhyAddr = 0;\r
-  UfsHc          = Private->UfsHostController;\r
+  TransReq->Packet        = Packet;\r
+\r
+  UfsHc = Private->UfsHostController;\r
   //\r
   // Find out which slot of transfer request list is available.\r
   //\r
@@ -1459,7 +1522,7 @@ UfsExecScsiCmds (
     return Status;\r
   }\r
 \r
-  TransReq->Trd = ((UTP_TRD*)Private->UtpTrlBase) + TransReq->Slot;\r
+  TransReq->Trd = ((UTP_TRD *)Private->UtpTrlBase) + TransReq->Slot;\r
 \r
   //\r
   // Fill transfer request descriptor to this slot.\r
@@ -1478,44 +1541,16 @@ UfsExecScsiCmds (
 \r
   TransReq->CmdDescSize = TransReq->Trd->PrdtO * sizeof (UINT32) + TransReq->Trd->PrdtL * sizeof (UTP_TR_PRD);\r
 \r
-  if (Packet->DataDirection == EFI_EXT_SCSI_DATA_DIRECTION_READ) {\r
-    DataBuf       = Packet->InDataBuffer;\r
-    DataLen       = Packet->InTransferLength;\r
-    Flag          = EdkiiUfsHcOperationBusMasterWrite;\r
-  } else {\r
-    DataBuf       = Packet->OutDataBuffer;\r
-    DataLen       = Packet->OutTransferLength;\r
-    Flag          = EdkiiUfsHcOperationBusMasterRead;\r
-  }\r
-\r
-  if (DataLen != 0) {\r
-    MapLength = DataLen;\r
-    Status    = UfsHc->Map (\r
-                         UfsHc,\r
-                         Flag,\r
-                         DataBuf,\r
-                         &MapLength,\r
-                         &DataBufPhyAddr,\r
-                         &TransReq->DataBufMapping\r
-                         );\r
-\r
-    if (EFI_ERROR (Status) || (DataLen != MapLength)) {\r
-      goto Exit1;\r
-    }\r
+  Status = UfsPrepareDataTransferBuffer (Private, TransReq);\r
+  if (EFI_ERROR (Status)) {\r
+    goto Exit1;\r
   }\r
-  //\r
-  // Fill PRDT table of Command UPIU for executed SCSI cmd.\r
-  //\r
-  PrdtBase = (UTP_TR_PRD*)((UINT8*)TransReq->CmdDescHost + ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)));\r
-  ASSERT (PrdtBase != NULL);\r
-  UfsInitUtpPrdt (PrdtBase, (VOID*)(UINTN)DataBufPhyAddr, DataLen);\r
 \r
   //\r
   // Insert the async SCSI cmd to the Async I/O list\r
   //\r
   if (Event != NULL) {\r
-    OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
-    TransReq->Packet      = Packet;\r
+    OldTpl                = gBS->RaiseTPL (TPL_NOTIFY);\r
     TransReq->CallerEvent = Event;\r
     InsertTailList (&Private->Queue, &TransReq->TransferList);\r
     gBS->RestoreTPL (OldTpl);\r
@@ -1535,7 +1570,7 @@ UfsExecScsiCmds (
 \r
   //\r
   // Wait for the completion of the transfer request.\r
-  // \r
+  //\r
   Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << TransReq->Slot, 0, Packet->Timeout);\r
   if (EFI_ERROR (Status)) {\r
     goto Exit;\r
@@ -1544,14 +1579,21 @@ UfsExecScsiCmds (
   //\r
   // Get sense data if exists\r
   //\r
-  Response     = (UTP_RESPONSE_UPIU*)((UINT8*)TransReq->CmdDescHost + TransReq->Trd->RuO * sizeof (UINT32));\r
+  Response = (UTP_RESPONSE_UPIU *)((UINT8 *)TransReq->CmdDescHost + TransReq->Trd->RuO * sizeof (UINT32));\r
   ASSERT (Response != NULL);\r
   SenseDataLen = Response->SenseDataLen;\r
-  SwapLittleEndianToBigEndian ((UINT8*)&SenseDataLen, sizeof (UINT16));\r
-  \r
+  SwapLittleEndianToBigEndian ((UINT8 *)&SenseDataLen, sizeof (UINT16));\r
+\r
   if ((Packet->SenseDataLength != 0) && (Packet->SenseData != NULL)) {\r
-    CopyMem (Packet->SenseData, Response->SenseData, SenseDataLen);\r
-    Packet->SenseDataLength = (UINT8)SenseDataLen;\r
+    //\r
+    // Make sure the hardware device does not return more data than expected.\r
+    //\r
+    if (SenseDataLen <= Packet->SenseDataLength) {\r
+      CopyMem (Packet->SenseData, Response->SenseData, SenseDataLen);\r
+      Packet->SenseDataLength = (UINT8)SenseDataLen;\r
+    } else {\r
+      Packet->SenseDataLength = 0;\r
+    }\r
   }\r
 \r
   //\r
@@ -1568,13 +1610,13 @@ UfsExecScsiCmds (
     if (Packet->DataDirection == EFI_EXT_SCSI_DATA_DIRECTION_READ) {\r
       if ((Response->Flags & BIT5) == BIT5) {\r
         ResTranCount = Response->ResTranCount;\r
-        SwapLittleEndianToBigEndian ((UINT8*)&ResTranCount, sizeof (UINT32));\r
+        SwapLittleEndianToBigEndian ((UINT8 *)&ResTranCount, sizeof (UINT32));\r
         Packet->InTransferLength -= ResTranCount;\r
       }\r
     } else {\r
       if ((Response->Flags & BIT5) == BIT5) {\r
         ResTranCount = Response->ResTranCount;\r
-        SwapLittleEndianToBigEndian ((UINT8*)&ResTranCount, sizeof (UINT32));\r
+        SwapLittleEndianToBigEndian ((UINT8 *)&ResTranCount, sizeof (UINT32));\r
         Packet->OutTransferLength -= ResTranCount;\r
       }\r
     }\r
@@ -1587,45 +1629,38 @@ Exit:
 \r
   UfsStopExecCmd (Private, TransReq->Slot);\r
 \r
-  if (TransReq->DataBufMapping != NULL) {\r
-    UfsHc->Unmap (UfsHc, TransReq->DataBufMapping);\r
-  }\r
+  UfsReconcileDataTransferBuffer (Private, TransReq);\r
 \r
 Exit1:\r
   if (TransReq->CmdDescMapping != NULL) {\r
     UfsHc->Unmap (UfsHc, TransReq->CmdDescMapping);\r
   }\r
+\r
   if (TransReq->CmdDescHost != NULL) {\r
     UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (TransReq->CmdDescSize), TransReq->CmdDescHost);\r
   }\r
+\r
   if (TransReq != NULL) {\r
     FreePool (TransReq);\r
   }\r
+\r
   return Status;\r
 }\r
 \r
-\r
 /**\r
-  Sent UIC DME_LINKSTARTUP command to start the link startup procedure.\r
+  Send UIC command.\r
 \r
-  @param[in] Private          The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
-  @param[in] UicOpcode        The opcode of the UIC command.\r
-  @param[in] Arg1             The value for 1st argument of the UIC command.\r
-  @param[in] Arg2             The value for 2nd argument of the UIC command.\r
-  @param[in] Arg3             The value for 3rd argument of the UIC command.\r
+  @param[in]      Private     The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.\r
+  @param[in, out] UicCommand  UIC command descriptor. On exit contains UIC command results.\r
 \r
   @return EFI_SUCCESS      Successfully execute this UIC command and detect attached UFS device.\r
   @return EFI_DEVICE_ERROR Fail to execute this UIC command and detect attached UFS device.\r
-  @return EFI_NOT_FOUND    The presence of the UFS device isn't detected.\r
 \r
 **/\r
 EFI_STATUS\r
 UfsExecUicCommands (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA    *Private,\r
-  IN  UINT8                         UicOpcode,\r
-  IN  UINT32                        Arg1,\r
-  IN  UINT32                        Arg2,\r
-  IN  UINT32                        Arg3\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN OUT EDKII_UIC_COMMAND        *UicCommand\r
   )\r
 {\r
   EFI_STATUS  Status;\r
@@ -1651,17 +1686,17 @@ UfsExecUicCommands (
   // only after all the UIC command argument registers (UICCMDARG1, UICCMDARG2 and UICCMDARG3)\r
   // are set.\r
   //\r
-  Status = UfsMmioWrite32 (Private, UFS_HC_UCMD_ARG1_OFFSET, Arg1);\r
+  Status = UfsMmioWrite32 (Private, UFS_HC_UCMD_ARG1_OFFSET, UicCommand->Arg1);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
 \r
-  Status = UfsMmioWrite32 (Private, UFS_HC_UCMD_ARG2_OFFSET, Arg2);\r
+  Status = UfsMmioWrite32 (Private, UFS_HC_UCMD_ARG2_OFFSET, UicCommand->Arg2);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
 \r
-  Status = UfsMmioWrite32 (Private, UFS_HC_UCMD_ARG3_OFFSET, Arg3);\r
+  Status = UfsMmioWrite32 (Private, UFS_HC_UCMD_ARG3_OFFSET, UicCommand->Arg3);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
@@ -1674,51 +1709,39 @@ UfsExecUicCommands (
     return Status;\r
   }\r
 \r
-  Status = UfsMmioWrite32 (Private, UFS_HC_UIC_CMD_OFFSET, (UINT32)UicOpcode);\r
+  Status = UfsMmioWrite32 (Private, UFS_HC_UIC_CMD_OFFSET, UicCommand->Opcode);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
 \r
   //\r
   // UFS 2.0 spec section 5.3.1 Offset:0x20 IS.Bit10 UIC Command Completion Status (UCCS)\r
-  // This bit is set to '1' by the host controller upon completion of a UIC command. \r
+  // This bit is set to '1' by the host controller upon completion of a UIC command.\r
   //\r
-  Status  = UfsWaitMemSet (Private, UFS_HC_IS_OFFSET, UFS_HC_IS_UCCS, UFS_HC_IS_UCCS, UFS_TIMEOUT);\r
+  Status = UfsWaitMemSet (Private, UFS_HC_IS_OFFSET, UFS_HC_IS_UCCS, UFS_HC_IS_UCCS, UFS_TIMEOUT);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
 \r
-  if (UicOpcode != UfsUicDmeReset) {\r
-    Status = UfsMmioRead32 (Private, UFS_HC_UCMD_ARG2_OFFSET, &Data);\r
+  if (UicCommand->Opcode != UfsUicDmeReset) {\r
+    Status = UfsMmioRead32 (Private, UFS_HC_UCMD_ARG2_OFFSET, &UicCommand->Arg2);\r
     if (EFI_ERROR (Status)) {\r
       return Status;\r
     }\r
-    if ((Data & 0xFF) != 0) {\r
-      DEBUG_CODE_BEGIN();\r
-        DumpUicCmdExecResult (UicOpcode, (UINT8)(Data & 0xFF));\r
-      DEBUG_CODE_END();\r
-      return EFI_DEVICE_ERROR;\r
-    }\r
-  }\r
-\r
-  //\r
-  // Check value of HCS.DP and make sure that there is a device attached to the Link.\r
-  //\r
-  Status = UfsMmioRead32 (Private, UFS_HC_STATUS_OFFSET, &Data);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
 \r
-  if ((Data & UFS_HC_HCS_DP) == 0) {\r
-    Status  = UfsWaitMemSet (Private, UFS_HC_IS_OFFSET, UFS_HC_IS_ULSS, UFS_HC_IS_ULSS, UFS_TIMEOUT);\r
+    Status = UfsMmioRead32 (Private, UFS_HC_UCMD_ARG3_OFFSET, &UicCommand->Arg3);\r
     if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+\r
+    if ((UicCommand->Arg2 & 0xFF) != 0) {\r
+      DEBUG_CODE_BEGIN ();\r
+      DumpUicCmdExecResult ((UINT8)UicCommand->Opcode, (UINT8)(UicCommand->Arg2 & 0xFF));\r
+      DEBUG_CODE_END ();\r
       return EFI_DEVICE_ERROR;\r
     }\r
-    return EFI_NOT_FOUND;\r
   }\r
 \r
-  DEBUG ((DEBUG_INFO, "UfsPassThruDxe: found a attached UFS device\n"));\r
-\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -1738,19 +1761,19 @@ UfsExecUicCommands (
 **/\r
 EFI_STATUS\r
 UfsAllocateAlignCommonBuffer (\r
-  IN     UFS_PASS_THRU_PRIVATE_DATA    *Private,\r
-  IN     UINTN                         Size,\r
-     OUT VOID                          **CmdDescHost,\r
-     OUT EFI_PHYSICAL_ADDRESS          *CmdDescPhyAddr,\r
-     OUT VOID                          **CmdDescMapping\r
+  IN     UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN     UINTN                       Size,\r
+  OUT VOID                           **CmdDescHost,\r
+  OUT EFI_PHYSICAL_ADDRESS           *CmdDescPhyAddr,\r
+  OUT VOID                           **CmdDescMapping\r
   )\r
 {\r
-  EFI_STATUS                           Status;\r
-  UINTN                                Bytes;\r
-  BOOLEAN                              Is32BitAddr;\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;\r
+  EFI_STATUS                          Status;\r
+  UINTN                               Bytes;\r
+  BOOLEAN                             Is32BitAddr;\r
+  EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
 \r
-  if ((Private->Capabilities & UFS_HC_CAP_64ADDR) == UFS_HC_CAP_64ADDR) {\r
+  if ((Private->UfsHcInfo.Capabilities & UFS_HC_CAP_64ADDR) == UFS_HC_CAP_64ADDR) {\r
     Is32BitAddr = FALSE;\r
   } else {\r
     Is32BitAddr = TRUE;\r
@@ -1825,11 +1848,19 @@ UfsAllocateAlignCommonBuffer (
 **/\r
 EFI_STATUS\r
 UfsEnableHostController (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA     *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  EFI_STATUS             Status;\r
-  UINT32                 Data;\r
+  EFI_STATUS  Status;\r
+  UINT32      Data;\r
+\r
+  if ((mUfsHcPlatform != NULL) && (mUfsHcPlatform->Callback != NULL)) {\r
+    Status = mUfsHcPlatform->Callback (Private->Handle, EdkiiUfsHcPreHce, &Private->UfsHcDriverInterface);\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((DEBUG_ERROR, "Failure from platform driver during EdkiiUfsHcPreHce, Status = %r\n", Status));\r
+      return Status;\r
+    }\r
+  }\r
 \r
   //\r
   // UFS 2.0 spec section 7.1.1 - Host Controller Initialization\r
@@ -1849,6 +1880,7 @@ UfsEnableHostController (
     if (EFI_ERROR (Status)) {\r
       return Status;\r
     }\r
+\r
     //\r
     // Wait until HCE is read as '0' before continuing.\r
     //\r
@@ -1874,6 +1906,14 @@ UfsEnableHostController (
     return EFI_DEVICE_ERROR;\r
   }\r
 \r
+  if ((mUfsHcPlatform != NULL) && (mUfsHcPlatform->Callback != NULL)) {\r
+    Status = mUfsHcPlatform->Callback (Private->Handle, EdkiiUfsHcPostHce, &Private->UfsHcDriverInterface);\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((DEBUG_ERROR, "Failure from platform driver during EdkiiUfsHcPostHce, Status = %r\n", Status));\r
+      return Status;\r
+    }\r
+  }\r
+\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -1889,34 +1929,52 @@ UfsEnableHostController (
 **/\r
 EFI_STATUS\r
 UfsDeviceDetection (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA     *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  UINTN                  Retry;\r
-  EFI_STATUS             Status;\r
+  UINTN              Retry;\r
+  EFI_STATUS         Status;\r
+  UINT32             Data;\r
+  EDKII_UIC_COMMAND  LinkStartupCommand;\r
+\r
+  if ((mUfsHcPlatform != NULL) && (mUfsHcPlatform->Callback != NULL)) {\r
+    Status = mUfsHcPlatform->Callback (Private->Handle, EdkiiUfsHcPreLinkStartup, &Private->UfsHcDriverInterface);\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((DEBUG_ERROR, "Failure from platform driver during EdkiiUfsHcPreLinkStartup, Status = %r\n", Status));\r
+      return Status;\r
+    }\r
+  }\r
 \r
   //\r
   // Start UFS device detection.\r
   // Try up to 3 times for establishing data link with device.\r
   //\r
   for (Retry = 0; Retry < 3; Retry++) {\r
-    Status = UfsExecUicCommands (Private, UfsUicDmeLinkStartup, 0, 0, 0);\r
-    if (!EFI_ERROR (Status)) {\r
-      break;\r
+    LinkStartupCommand.Opcode = UfsUicDmeLinkStartup;\r
+    LinkStartupCommand.Arg1   = 0;\r
+    LinkStartupCommand.Arg2   = 0;\r
+    LinkStartupCommand.Arg3   = 0;\r
+    Status                    = UfsExecUicCommands (Private, &LinkStartupCommand);\r
+    if (EFI_ERROR (Status)) {\r
+      return EFI_DEVICE_ERROR;\r
     }\r
 \r
-    if (Status == EFI_NOT_FOUND) {\r
-      continue;\r
+    Status = UfsMmioRead32 (Private, UFS_HC_STATUS_OFFSET, &Data);\r
+    if (EFI_ERROR (Status)) {\r
+      return EFI_DEVICE_ERROR;\r
     }\r
 \r
-    return EFI_DEVICE_ERROR;\r
-  }\r
-\r
-  if (Retry == 3) {\r
-    return EFI_NOT_FOUND;\r
+    if ((Data & UFS_HC_HCS_DP) == 0) {\r
+      Status = UfsWaitMemSet (Private, UFS_HC_IS_OFFSET, UFS_HC_IS_ULSS, UFS_HC_IS_ULSS, UFS_TIMEOUT);\r
+      if (EFI_ERROR (Status)) {\r
+        return EFI_DEVICE_ERROR;\r
+      }\r
+    } else {\r
+      return EFI_SUCCESS;\r
+    }\r
   }\r
 \r
-  return EFI_SUCCESS;\r
+  return EFI_NOT_FOUND;\r
 }\r
 \r
 /**\r
@@ -1930,16 +1988,15 @@ UfsDeviceDetection (
 **/\r
 EFI_STATUS\r
 UfsInitTaskManagementRequestList (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA     *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  UINT32                 Data;\r
-  UINT8                  Nutmrs;\r
-  VOID                   *CmdDescHost;\r
-  EFI_PHYSICAL_ADDRESS   CmdDescPhyAddr;\r
-  VOID                   *CmdDescMapping;\r
-  EFI_STATUS             Status;\r
-  \r
+  UINT8                 Nutmrs;\r
+  VOID                  *CmdDescHost;\r
+  EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;\r
+  VOID                  *CmdDescMapping;\r
+  EFI_STATUS            Status;\r
+\r
   //\r
   // Initial h/w and s/w context for future operations.\r
   //\r
@@ -1947,17 +2004,10 @@ UfsInitTaskManagementRequestList (
   CmdDescMapping = NULL;\r
   CmdDescPhyAddr = 0;\r
 \r
-  Status = UfsMmioRead32 (Private, UFS_HC_CAP_OFFSET, &Data);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-\r
-  Private->Capabilities = Data;\r
-\r
   //\r
   // Allocate and initialize UTP Task Management Request List.\r
   //\r
-  Nutmrs = (UINT8) (RShiftU64 ((Private->Capabilities & UFS_HC_CAP_NUTMRS), 16) + 1);\r
+  Nutmrs = (UINT8)(RShiftU64 ((Private->UfsHcInfo.Capabilities & UFS_HC_CAP_NUTMRS), 16) + 1);\r
   Status = UfsAllocateAlignCommonBuffer (Private, Nutmrs * sizeof (UTP_TMRD), &CmdDescHost, &CmdDescPhyAddr, &CmdDescMapping);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
@@ -1976,6 +2026,7 @@ UfsInitTaskManagementRequestList (
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
+\r
   Private->UtpTmrlBase = CmdDescHost;\r
   Private->Nutmrs      = Nutmrs;\r
   Private->TmrlMapping = CmdDescMapping;\r
@@ -2003,15 +2054,14 @@ UfsInitTaskManagementRequestList (
 **/\r
 EFI_STATUS\r
 UfsInitTransferRequestList (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA     *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  UINT32                 Data;\r
-  UINT8                  Nutrs;\r
-  VOID                   *CmdDescHost;\r
-  EFI_PHYSICAL_ADDRESS   CmdDescPhyAddr;\r
-  VOID                   *CmdDescMapping;  \r
-  EFI_STATUS             Status;\r
+  UINT8                 Nutrs;\r
+  VOID                  *CmdDescHost;\r
+  EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;\r
+  VOID                  *CmdDescMapping;\r
+  EFI_STATUS            Status;\r
 \r
   //\r
   // Initial h/w and s/w context for future operations.\r
@@ -2020,17 +2070,10 @@ UfsInitTransferRequestList (
   CmdDescMapping = NULL;\r
   CmdDescPhyAddr = 0;\r
 \r
-  Status = UfsMmioRead32 (Private, UFS_HC_CAP_OFFSET, &Data);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-\r
-  Private->Capabilities = Data;\r
-\r
   //\r
   // Allocate and initialize UTP Transfer Request List.\r
   //\r
-  Nutrs  = (UINT8)((Private->Capabilities & UFS_HC_CAP_NUTRS) + 1);\r
+  Nutrs  = (UINT8)((Private->UfsHcInfo.Capabilities & UFS_HC_CAP_NUTRS) + 1);\r
   Status = UfsAllocateAlignCommonBuffer (Private, Nutrs * sizeof (UTP_TRD), &CmdDescHost, &CmdDescPhyAddr, &CmdDescMapping);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
@@ -2051,7 +2094,7 @@ UfsInitTransferRequestList (
   }\r
 \r
   Private->UtpTrlBase = CmdDescHost;\r
-  Private->Nutrs      = Nutrs;  \r
+  Private->Nutrs      = Nutrs;\r
   Private->TrlMapping = CmdDescMapping;\r
 \r
   //\r
@@ -2077,10 +2120,10 @@ UfsInitTransferRequestList (
 **/\r
 EFI_STATUS\r
 UfsControllerInit (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA     *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  EFI_STATUS             Status;\r
+  EFI_STATUS  Status;\r
 \r
   Status = UfsEnableHostController (Private);\r
   if (EFI_ERROR (Status)) {\r
@@ -2121,11 +2164,11 @@ UfsControllerInit (
 **/\r
 EFI_STATUS\r
 UfsControllerStop (\r
-  IN  UFS_PASS_THRU_PRIVATE_DATA     *Private\r
+  IN  UFS_PASS_THRU_PRIVATE_DATA  *Private\r
   )\r
 {\r
-  EFI_STATUS             Status;\r
-  UINT32                 Data;\r
+  EFI_STATUS  Status;\r
+  UINT32      Data;\r
 \r
   //\r
   // Enable the UTP Task Management Request List by setting the UTP Task Management\r
@@ -2152,6 +2195,7 @@ UfsControllerStop (
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
+\r
   ASSERT ((Data & UFS_HC_HCE_EN) == UFS_HC_HCE_EN);\r
 \r
   Status = UfsMmioWrite32 (Private, UFS_HC_ENABLE_OFFSET, 0);\r
@@ -2172,7 +2216,6 @@ UfsControllerStop (
   return EFI_SUCCESS;\r
 }\r
 \r
-\r
 /**\r
   Internal helper function which will signal the caller event and clean up\r
   resources.\r
@@ -2186,17 +2229,17 @@ UfsControllerStop (
 VOID\r
 EFIAPI\r
 SignalCallerEvent (\r
-  IN UFS_PASS_THRU_PRIVATE_DATA      *Private,\r
-  IN UFS_PASS_THRU_TRANS_REQ         *TransReq\r
+  IN UFS_PASS_THRU_PRIVATE_DATA  *Private,\r
+  IN UFS_PASS_THRU_TRANS_REQ     *TransReq\r
   )\r
 {\r
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL *UfsHc;\r
-  EFI_EVENT                          CallerEvent;\r
+  EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;\r
+  EFI_EVENT                           CallerEvent;\r
 \r
   ASSERT ((Private != NULL) && (TransReq != NULL));\r
 \r
-  UfsHc        = Private->UfsHostController;\r
-  CallerEvent  = TransReq->CallerEvent;\r
+  UfsHc       = Private->UfsHostController;\r
+  CallerEvent = TransReq->CallerEvent;\r
 \r
   RemoveEntryList (&TransReq->TransferList);\r
 \r
@@ -2204,13 +2247,12 @@ SignalCallerEvent (
 \r
   UfsStopExecCmd (Private, TransReq->Slot);\r
 \r
-  if (TransReq->DataBufMapping != NULL) {\r
-    UfsHc->Unmap (UfsHc, TransReq->DataBufMapping);\r
-  }\r
+  UfsReconcileDataTransferBuffer (Private, TransReq);\r
 \r
   if (TransReq->CmdDescMapping != NULL) {\r
     UfsHc->Unmap (UfsHc, TransReq->CmdDescMapping);\r
   }\r
+\r
   if (TransReq->CmdDescHost != NULL) {\r
     UfsHc->FreeBuffer (\r
              UfsHc,\r
@@ -2235,36 +2277,37 @@ SignalCallerEvent (
 VOID\r
 EFIAPI\r
 ProcessAsyncTaskList (\r
-  IN EFI_EVENT          Event,\r
-  IN VOID               *Context\r
+  IN EFI_EVENT  Event,\r
+  IN VOID       *Context\r
   )\r
 {\r
-  UFS_PASS_THRU_PRIVATE_DATA                    *Private;\r
-  LIST_ENTRY                                    *Entry;\r
-  LIST_ENTRY                                    *NextEntry;\r
-  UFS_PASS_THRU_TRANS_REQ                       *TransReq;\r
-  EFI_EXT_SCSI_PASS_THRU_SCSI_REQUEST_PACKET    *Packet;\r
-  UTP_RESPONSE_UPIU                             *Response;\r
-  UINT16                                        SenseDataLen;\r
-  UINT32                                        ResTranCount;\r
-  UINT32                                        SlotsMap;\r
-  UINT32                                        Value;\r
-  EFI_STATUS                                    Status;\r
-\r
-  Private   = (UFS_PASS_THRU_PRIVATE_DATA*) Context;\r
-  SlotsMap  = 0;\r
+  UFS_PASS_THRU_PRIVATE_DATA                  *Private;\r
+  LIST_ENTRY                                  *Entry;\r
+  LIST_ENTRY                                  *NextEntry;\r
+  UFS_PASS_THRU_TRANS_REQ                     *TransReq;\r
+  EFI_EXT_SCSI_PASS_THRU_SCSI_REQUEST_PACKET  *Packet;\r
+  UTP_RESPONSE_UPIU                           *Response;\r
+  UINT16                                      SenseDataLen;\r
+  UINT32                                      ResTranCount;\r
+  UINT32                                      SlotsMap;\r
+  UINT32                                      Value;\r
+  EFI_STATUS                                  Status;\r
+\r
+  Private  = (UFS_PASS_THRU_PRIVATE_DATA *)Context;\r
+  SlotsMap = 0;\r
 \r
   //\r
   // Check the entries in the async I/O queue are done or not.\r
   //\r
-  if (!IsListEmpty(&Private->Queue)) {\r
-    EFI_LIST_FOR_EACH_SAFE (Entry, NextEntry, &Private->Queue) {\r
-      TransReq  = UFS_PASS_THRU_TRANS_REQ_FROM_THIS (Entry);\r
-      Packet    = TransReq->Packet;\r
+  if (!IsListEmpty (&Private->Queue)) {\r
+    BASE_LIST_FOR_EACH_SAFE (Entry, NextEntry, &Private->Queue) {\r
+      TransReq = UFS_PASS_THRU_TRANS_REQ_FROM_THIS (Entry);\r
+      Packet   = TransReq->Packet;\r
 \r
       if ((SlotsMap & (BIT0 << TransReq->Slot)) != 0) {\r
         return;\r
       }\r
+\r
       SlotsMap |= BIT0 << TransReq->Slot;\r
 \r
       Status = UfsMmioRead32 (Private, UFS_HC_UTRLDBR_OFFSET, &Value);\r
@@ -2301,14 +2344,21 @@ ProcessAsyncTaskList (
         //\r
         // Get sense data if exists\r
         //\r
-        Response = (UTP_RESPONSE_UPIU*)((UINT8*)TransReq->CmdDescHost + TransReq->Trd->RuO * sizeof (UINT32));\r
+        Response = (UTP_RESPONSE_UPIU *)((UINT8 *)TransReq->CmdDescHost + TransReq->Trd->RuO * sizeof (UINT32));\r
         ASSERT (Response != NULL);\r
         SenseDataLen = Response->SenseDataLen;\r
-        SwapLittleEndianToBigEndian ((UINT8*)&SenseDataLen, sizeof (UINT16));\r
+        SwapLittleEndianToBigEndian ((UINT8 *)&SenseDataLen, sizeof (UINT16));\r
 \r
         if ((Packet->SenseDataLength != 0) && (Packet->SenseData != NULL)) {\r
-          CopyMem (Packet->SenseData, Response->SenseData, SenseDataLen);\r
-          Packet->SenseDataLength = (UINT8)SenseDataLen;\r
+          //\r
+          // Make sure the hardware device does not return more data than expected.\r
+          //\r
+          if (SenseDataLen <= Packet->SenseDataLength) {\r
+            CopyMem (Packet->SenseData, Response->SenseData, SenseDataLen);\r
+            Packet->SenseDataLength = (UINT8)SenseDataLen;\r
+          } else {\r
+            Packet->SenseDataLength = 0;\r
+          }\r
         }\r
 \r
         //\r
@@ -2325,13 +2375,13 @@ ProcessAsyncTaskList (
           if (Packet->DataDirection == EFI_EXT_SCSI_DATA_DIRECTION_READ) {\r
             if ((Response->Flags & BIT5) == BIT5) {\r
               ResTranCount = Response->ResTranCount;\r
-              SwapLittleEndianToBigEndian ((UINT8*)&ResTranCount, sizeof (UINT32));\r
+              SwapLittleEndianToBigEndian ((UINT8 *)&ResTranCount, sizeof (UINT32));\r
               Packet->InTransferLength -= ResTranCount;\r
             }\r
           } else {\r
             if ((Response->Flags & BIT5) == BIT5) {\r
               ResTranCount = Response->ResTranCount;\r
-              SwapLittleEndianToBigEndian ((UINT8*)&ResTranCount, sizeof (UINT32));\r
+              SwapLittleEndianToBigEndian ((UINT8 *)&ResTranCount, sizeof (UINT32));\r
               Packet->OutTransferLength -= ResTranCount;\r
             }\r
           }\r
@@ -2348,3 +2398,71 @@ ProcessAsyncTaskList (
   }\r
 }\r
 \r
+/**\r
+  Execute UIC command.\r
+\r
+  @param[in]      This        Pointer to driver interface produced by the UFS controller.\r
+  @param[in, out] UicCommand  Descriptor of the command that will be executed.\r
+\r
+  @retval EFI_SUCCESS            Command executed successfully.\r
+  @retval EFI_INVALID_PARAMETER  This or UicCommand is NULL.\r
+  @retval Others                 Command failed to execute.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+UfsHcDriverInterfaceExecUicCommand (\r
+  IN     EDKII_UFS_HC_DRIVER_INTERFACE  *This,\r
+  IN OUT EDKII_UIC_COMMAND              *UicCommand\r
+  )\r
+{\r
+  UFS_PASS_THRU_PRIVATE_DATA  *Private;\r
+\r
+  if ((This == NULL) || (UicCommand == NULL)) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  Private = UFS_PASS_THRU_PRIVATE_DATA_FROM_DRIVER_INTF (This);\r
+\r
+  return UfsExecUicCommands (Private, UicCommand);\r
+}\r
+\r
+/**\r
+  Initializes UfsHcInfo field in private data.\r
+\r
+  @param[in] Private  Pointer to host controller private data.\r
+\r
+  @retval EFI_SUCCESS  UfsHcInfo initialized successfully.\r
+  @retval Others       Failed to initalize UfsHcInfo.\r
+**/\r
+EFI_STATUS\r
+GetUfsHcInfo (\r
+  IN UFS_PASS_THRU_PRIVATE_DATA  *Private\r
+  )\r
+{\r
+  UINT32      Data;\r
+  EFI_STATUS  Status;\r
+\r
+  Status = UfsMmioRead32 (Private, UFS_HC_VER_OFFSET, &Data);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  Private->UfsHcInfo.Version = Data;\r
+\r
+  Status = UfsMmioRead32 (Private, UFS_HC_CAP_OFFSET, &Data);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  Private->UfsHcInfo.Capabilities = Data;\r
+\r
+  if ((mUfsHcPlatform != NULL) && (mUfsHcPlatform->OverrideHcInfo != NULL)) {\r
+    Status = mUfsHcPlatform->OverrideHcInfo (Private->Handle, &Private->UfsHcInfo);\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((DEBUG_ERROR, "Failure from platform on OverrideHcInfo, Status = %r\n", Status));\r
+      return Status;\r
+    }\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r