MdeModulePkg/Usb: Read a large number of blocks

[mirror_edk2.git] / MdeModulePkg / Bus / Usb / UsbBotPei / PeiAtapi.c

diff --git a/MdeModulePkg/Bus/Usb/UsbBotPei/PeiAtapi.c b/MdeModulePkg/Bus/Usb/UsbBotPei/PeiAtapi.c
index 422ac5fec996429a259f1c84025e58919c0ab6b2..5111e4579e25d4421afb43411faa8f0d06227926 100644 (file)
--- a/MdeModulePkg/Bus/Usb/UsbBotPei/PeiAtapi.c
+++ b/MdeModulePkg/Bus/Usb/UsbBotPei/PeiAtapi.c
@@ -2,6 +2,7 @@
 Pei USB ATAPI command implementations.\r
 \r
 Copyright (c) 1999 - 2018, Intel Corporation. All rights reserved.<BR>\r
+Copyright (C) 2022 Advanced Micro Devices, Inc. All rights reserved.<BR>\r
 \r
 SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
@@ -382,14 +383,14 @@ PeiUsbRead10 (
   ATAPI_PACKET_COMMAND  Packet;\r
   ATAPI_READ10_CMD      *Read10Packet;\r
   UINT16                MaxBlock;\r
-  UINT16                BlocksRemaining;\r
-  UINT16                SectorCount;\r
+  UINT32                BlocksRemaining;\r
+  UINT32                SectorCount;\r
   UINT32                Lba32;\r
   UINT32                BlockSize;\r
   UINT32                ByteCount;\r
   VOID                  *PtrBuffer;\r
   EFI_STATUS            Status;\r
-  UINT16                TimeOut;\r
+  UINT32                TimeOut;\r
 \r
   //\r
   // prepare command packet for the Inquiry Packet Command.\r
@@ -401,16 +402,13 @@ PeiUsbRead10 (
 \r
   BlockSize = (UINT32)PeiBotDevice->Media.BlockSize;\r
 \r
-  MaxBlock        = (UINT16)(65535 / BlockSize);\r
-  BlocksRemaining = (UINT16)NumberOfBlocks;\r
+  MaxBlock = (UINT16)(MAX_UINT16 / BlockSize);\r
+  ASSERT (NumberOfBlocks < MAX_UINT32);\r
+  BlocksRemaining = (UINT32)NumberOfBlocks;\r
 \r
   Status = EFI_SUCCESS;\r
   while (BlocksRemaining > 0) {\r
-    if (BlocksRemaining <= MaxBlock) {\r
-      SectorCount = BlocksRemaining;\r
-    } else {\r
-      SectorCount = MaxBlock;\r
-    }\r
+    SectorCount = MIN (BlocksRemaining, MaxBlock);\r
 \r
     //\r
     // fill the Packet data structure\r
@@ -435,7 +433,7 @@ PeiUsbRead10 (
 \r
     ByteCount = SectorCount * BlockSize;\r
 \r
-    TimeOut = (UINT16)(SectorCount * 2000);\r
+    TimeOut = SectorCount * 2000;\r
 \r
     //\r
     // send command packet\r
@@ -448,16 +446,17 @@ PeiUsbRead10 (
                (VOID *)PtrBuffer,\r
                ByteCount,\r
                EfiUsbDataIn,\r
-               TimeOut\r
+               (UINT16)MIN (TimeOut, MAX_UINT16)\r
                );\r
 \r
     if (Status != EFI_SUCCESS) {\r
       return Status;\r
     }\r
 \r
+    ASSERT (Lba32 <= (MAX_UINT32-SectorCount));\r
     Lba32          += SectorCount;\r
     PtrBuffer       = (UINT8 *)PtrBuffer + SectorCount * BlockSize;\r
-    BlocksRemaining = (UINT16)(BlocksRemaining - SectorCount);\r
+    BlocksRemaining = BlocksRemaining - SectorCount;\r
   }\r
 \r
   return Status;\r