MdeModulePkg/UsbMass: Fix integer overflow when BlockSize is 1

UsbBootReadWriteBlocks() and UsbBootReadWriteBlocks16() use a UINT16
local variable to hold the value of
USB_BOOT_MAX_CARRY_SIZE (=0x10000) / BlockSize.
When BlockSize is 1, the UINT16 local variable is set to 0x10000
but the high-16 bits are truncated resulting the final value be 0.

It causes the while-loop in the two functions accesses 0 block in
each loop, resulting the loop never ends.

The patch fixes the two functions to make sure no integer overflow
happens.

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Steven Shi <steven.shi@intel.com>
Reviewed-by: Star Zeng <star.zeng@intel.com>

diff --git a/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c b/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c
index 07a376440c65d4f9e90317aa06eb722aa09560c0..0b35cbacf08c929dc9258698c9c21e9405187f56 100644 (file)
--- a/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c
+++ b/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c
@@ -815,14 +815,14 @@ UsbBootReadWriteBlocks (
 {\r
   USB_BOOT_READ_WRITE_10_CMD Cmd;\r
   EFI_STATUS                 Status;\r
-  UINT16                     Count;\r
-  UINT16                     CountMax;\r
+  UINT32                     Count;\r
+  UINT32                     CountMax;\r
   UINT32                     BlockSize;\r
   UINT32                     ByteSize;\r
   UINT32                     Timeout;\r
 \r
   BlockSize = UsbMass->BlockIoMedia.BlockSize;\r
-  CountMax = (UINT16)(USB_BOOT_MAX_CARRY_SIZE / BlockSize);\r
+  CountMax  = USB_BOOT_MAX_CARRY_SIZE / BlockSize;\r
   Status    = EFI_SUCCESS;\r
 \r
   while (TotalBlock > 0) {\r
@@ -831,8 +831,9 @@ UsbBootReadWriteBlocks (
     // on the device. We must split the total block because the READ10\r
     // command only has 16 bit transfer length (in the unit of block).\r
     //\r
-    Count     = (UINT16)((TotalBlock < CountMax) ? TotalBlock : CountMax);\r
-    ByteSize  = (UINT32)Count * BlockSize;\r
+    Count    = (UINT32)MIN (TotalBlock, CountMax);\r
+    Count    = MIN (MAX_UINT16, Count);\r
+    ByteSize = Count * BlockSize;\r
 \r
     //\r
     // USB command's upper limit timeout is 5s. [USB2.0-9.2.6.1]\r
@@ -847,7 +848,7 @@ UsbBootReadWriteBlocks (
     Cmd.OpCode  = Write ? USB_BOOT_WRITE10_OPCODE : USB_BOOT_READ10_OPCODE;\r
     Cmd.Lun     = (UINT8) (USB_BOOT_LUN (UsbMass->Lun));\r
     WriteUnaligned32 ((UINT32 *) Cmd.Lba, SwapBytes32 (Lba));\r
-    WriteUnaligned16 ((UINT16 *) Cmd.TransferLen, SwapBytes16 (Count));\r
+    WriteUnaligned16 ((UINT16 *) Cmd.TransferLen, SwapBytes16 ((UINT16)Count));\r
 \r
     Status = UsbBootExecCmdWithRetry (\r
                UsbMass,\r
@@ -867,7 +868,7 @@ UsbBootReadWriteBlocks (
       Lba, Count\r
       ));\r
     Lba        += Count;\r
-    Buffer     += Count * BlockSize;\r
+    Buffer     += ByteSize;\r
     TotalBlock -= Count;\r
   }\r
 \r
@@ -897,22 +898,22 @@ UsbBootReadWriteBlocks16 (
 {\r
   UINT8                     Cmd[16];\r
   EFI_STATUS                Status;\r
-  UINT16                    Count;\r
-  UINT16                    CountMax;\r
+  UINT32                    Count;\r
+  UINT32                    CountMax;\r
   UINT32                    BlockSize;\r
   UINT32                    ByteSize;\r
   UINT32                    Timeout;\r
 \r
   BlockSize = UsbMass->BlockIoMedia.BlockSize;\r
-  CountMax = (UINT16)(USB_BOOT_MAX_CARRY_SIZE / BlockSize);\r
+  CountMax  = USB_BOOT_MAX_CARRY_SIZE / BlockSize;\r
   Status    = EFI_SUCCESS;\r
 \r
   while (TotalBlock > 0) {\r
     //\r
     // Split the total blocks into smaller pieces.\r
     //\r
-    Count     = (UINT16)((TotalBlock < CountMax) ? TotalBlock : CountMax);\r
-    ByteSize  = (UINT32)Count * BlockSize;\r
+    Count    = (UINT32)MIN (TotalBlock, CountMax);\r
+    ByteSize = Count * BlockSize;\r
 \r
     //\r
     // USB command's upper limit timeout is 5s. [USB2.0-9.2.6.1]\r
@@ -947,7 +948,7 @@ UsbBootReadWriteBlocks16 (
       Lba, Count\r
       ));\r
     Lba        += Count;\r
-    Buffer     += Count * BlockSize;\r
+    Buffer     += ByteSize;\r
     TotalBlock -= Count;\r
   }\r
 \r