[pschweitzer] 74027: [DISK] Fix size checking for IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, and only copy what's possible to fit in the buffer. Also, make the code a bit more readable. CORE-12858

Author: pschweitzer
Date: Thu Mar  2 20:03:55 2017
New Revision: 74027

URL: http://svn.reactos.org/svn/reactos?rev=74027&view=rev
Log:
[DISK]
Fix size checking for IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, and only copy what's possible to fit in the buffer.
Also, make the code a bit more readable.

CORE-12858

Modified:
    trunk/reactos/drivers/storage/class/disk/disk.c

Modified: trunk/reactos/drivers/storage/class/disk/disk.c
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/drivers/storage/class/disk/...
==============================================================================
--- trunk/reactos/drivers/storage/class/disk/disk.c	[iso-8859-1] (original)
+++ trunk/reactos/drivers/storage/class/disk/disk.c	[iso-8859-1] Thu Mar  2 20:03:55 2017
@@ -2054,16 +2054,27 @@
         PDISK_DATA        physicalDiskData;
         BOOLEAN           removable = FALSE;
         BOOLEAN           listInitialized = FALSE;
-
-        if ((irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY &&
-             irpStack->Parameters.DeviceIoControl.OutputBufferLength <
-            sizeof(DISK_GEOMETRY)) ||
-             (irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY_EX &&
-             irpStack->Parameters.DeviceIoControl.OutputBufferLength <
-            sizeof(DISK_GEOMETRY_EX))) {
-
-            status = STATUS_BUFFER_TOO_SMALL;
-            break;
+        ULONG             copyLength;
+
+        if (irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY) {
+            if (irpStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(DISK_GEOMETRY)) {
+                status = STATUS_BUFFER_TOO_SMALL;
+                break;
+            }
+
+            copyLength = sizeof(DISK_GEOMETRY);
+        } else {
+            ASSERT(irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY_EX);
+            if (irpStack->Parameters.DeviceIoControl.OutputBufferLength < FIELD_OFFSET(DISK_GEOMETRY_EX, Data)) {
+                status = STATUS_BUFFER_TOO_SMALL;
+                break;
+            }
+
+            if (irpStack->Parameters.DeviceIoControl.OutputBufferLength >= sizeof(DISK_GEOMETRY_EX)) {
+                copyLength = sizeof(DISK_GEOMETRY_EX);
+            } else {
+                copyLength = FIELD_OFFSET(DISK_GEOMETRY_EX, Data);
+            }
         }

status = STATUS_SUCCESS;
@@ -2121,15 +2132,10 @@

RtlMoveMemory(Irp->AssociatedIrp.SystemBuffer,
                           deviceExtension->DiskGeometry,
-                          (irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY) ?
-                          sizeof(DISK_GEOMETRY) :
-                          sizeof(DISK_GEOMETRY_EX));
+                          copyLength);

status = STATUS_SUCCESS;
-            Irp->IoStatus.Information =
-               (irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY) ?
-               sizeof(DISK_GEOMETRY) :
-               sizeof(DISK_GEOMETRY_EX);
+            Irp->IoStatus.Information = copyLength;
         }

break;