[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;