From: Pierre Schweitzer <pierre@reactos.org>
Date: Thu, 2 Mar 2017 20:03:55 +0000 (+0000)
Subject: [DISK]
X-Git-Tag: ReactOS-0.4.4-CLT2017~64
X-Git-Url: https://git.reactos.org/?p=reactos.git;a=commitdiff_plain;h=021ab6e00083933939add1bed72cbe25333e1e9a

[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 267c3cb3231..ffdd3b59e02 100644
--- 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;