[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

svn path=/trunk/; revision=74027

diff --git a/reactos/drivers/storage/class/disk/disk.c b/reactos/drivers/storage/class/disk/disk.c
index 267c3cb..ffdd3b5 100644 (file)
--- a/reactos/drivers/storage/class/disk/disk.c
+++ b/reactos/drivers/storage/class/disk/disk.c
@@ -2054,16 +2054,27 @@ Return Value:
         PDISK_DATA        physicalDiskData;
         BOOLEAN           removable = FALSE;
         BOOLEAN           listInitialized = FALSE;
+        ULONG             copyLength;
 
-        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))) {
+        if (irpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_DISK_GET_DRIVE_GEOMETRY) {
+            if (irpStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(DISK_GEOMETRY)) {
+                status = STATUS_BUFFER_TOO_SMALL;
+                break;
+            }
 
-            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 @@ Return Value:
 
             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;