MdeModulePkg/UsbMass: Fix USB key write failure

[mirror_edk2.git] / MdeModulePkg / Bus / Usb / UsbMassStorageDxe / UsbMassBoot.c

diff --git a/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c b/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c
index 0b35cbacf08c929dc9258698c9c21e9405187f56..69e2252134064b1a80574a0da73d80aa86664e87 100644 (file)
--- a/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c
+++ b/MdeModulePkg/Bus/Usb/UsbMassStorageDxe/UsbMassBoot.c
@@ -736,6 +736,13 @@ UsbBootDetectMedia (
     return Status;\r
   }\r
 \r
+  //\r
+  // Simply reject device whose block size is unacceptable small (==0) or large (>64K).\r
+  //\r
+  if ((Media->BlockSize == 0) || (Media->BlockSize > USB_BOOT_MAX_CARRY_SIZE)) {\r
+    return EFI_DEVICE_ERROR;\r
+  }\r
+\r
   //\r
   // Detect whether it is necessary to reinstall the Block I/O Protocol.\r
   //\r
@@ -854,7 +861,7 @@ UsbBootReadWriteBlocks (
                UsbMass,\r
                &Cmd,\r
                (UINT8) sizeof (USB_BOOT_READ_WRITE_10_CMD),\r
-               EfiUsbDataIn,\r
+               Write ? EfiUsbDataOut : EfiUsbDataIn,\r
                Buffer,\r
                ByteSize,\r
                Timeout\r
@@ -934,7 +941,7 @@ UsbBootReadWriteBlocks16 (
                UsbMass,\r
                Cmd,\r
                (UINT8) sizeof (Cmd),\r
-               EfiUsbDataIn,\r
+               Write ? EfiUsbDataOut : EfiUsbDataIn,\r
                Buffer,\r
                ByteSize,\r
                Timeout\r