From 20efea8fa47f1e6527bcdeb874aecc581d9e1845 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20S=C5=82abo=C5=84?= <asaillen@protonmail.com>
Date: Fri, 16 Feb 2024 16:48:33 +0100
Subject: [PATCH] [USBSTOR] Fix PdoHandleQueryInstanceId and increase serial
 number descriptor size to MAXIMUM_USB_STRING_LENGTH (#6413)

Serial number on some USB devices might exceed the number of 100 characters
(e.g. 120 characters on "SanDisk Ultra 3.2Gen1" pendrive) and cause buffer
overflow, resulting in usbstor.sys crash.

- Use pool allocation for instance ID generation.
  Fixes stack overflow on USB storage devices with large serial number.
- Print the LUN number as a hexadecimal, not as a character.
- Verify the serial number descriptor before using it.
- Increase the max descriptor size for serial number to
  MAXIMUM_USB_STRING_LENGTH. This fixes serial number string truncation.

Based on suggestions by disean and ThFabba.

CORE-17625
---
 drivers/usb/usbstor/descriptor.c |  2 +-
 drivers/usb/usbstor/pdo.c        | 53 +++++++++++++++++++++++---------
 drivers/usb/usbstor/usbstor.h    |  1 +
 3 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/drivers/usb/usbstor/descriptor.c b/drivers/usb/usbstor/descriptor.c
index 885cc7f6b72..788da8de4a2 100644
--- a/drivers/usb/usbstor/descriptor.c
+++ b/drivers/usb/usbstor/descriptor.c
@@ -118,7 +118,7 @@ USBSTOR_GetDescriptors(
      if (DeviceExtension->DeviceDescriptor->iSerialNumber)
      {
          // get serial number
-         Status = USBSTOR_GetDescriptor(DeviceExtension->LowerDeviceObject, USB_STRING_DESCRIPTOR_TYPE, 100 * sizeof(WCHAR), DeviceExtension->DeviceDescriptor->iSerialNumber, 0x0409, (PVOID*)&DeviceExtension->SerialNumber);
+         Status = USBSTOR_GetDescriptor(DeviceExtension->LowerDeviceObject, USB_STRING_DESCRIPTOR_TYPE, MAXIMUM_USB_STRING_LENGTH, DeviceExtension->DeviceDescriptor->iSerialNumber, 0x0409, (PVOID*)&DeviceExtension->SerialNumber);
          if (!NT_SUCCESS(Status))
          {
              FreeItem(DeviceExtension->DeviceDescriptor);
diff --git a/drivers/usb/usbstor/pdo.c b/drivers/usb/usbstor/pdo.c
index 28039993340..cec5c793795 100644
--- a/drivers/usb/usbstor/pdo.c
+++ b/drivers/usb/usbstor/pdo.c
@@ -437,37 +437,60 @@ USBSTOR_PdoHandleQueryInstanceId(
 {
     PPDO_DEVICE_EXTENSION PDODeviceExtension;
     PFDO_DEVICE_EXTENSION FDODeviceExtension;
-    WCHAR Buffer[100];
-    ULONG Length;
+    PUSB_STRING_DESCRIPTOR Descriptor;
+    ULONG CharCount;
     LPWSTR InstanceId;
+    NTSTATUS Status;
 
-    PDODeviceExtension = (PPDO_DEVICE_EXTENSION)DeviceObject->DeviceExtension;
-    FDODeviceExtension = (PFDO_DEVICE_EXTENSION)PDODeviceExtension->LowerDeviceObject->DeviceExtension;
+    PDODeviceExtension = DeviceObject->DeviceExtension;
+    FDODeviceExtension = PDODeviceExtension->LowerDeviceObject->DeviceExtension;
 
-    // format instance id
-    if (FDODeviceExtension->SerialNumber)
+    Descriptor = FDODeviceExtension->SerialNumber;
+    if (Descriptor && (Descriptor->bLength >= sizeof(USB_COMMON_DESCRIPTOR) + sizeof(WCHAR)))
     {
-        // using serial number from device
-        swprintf(Buffer, L"%s&%c", FDODeviceExtension->SerialNumber->bString, PDODeviceExtension->LUN);
+        /* Format the serial number descriptor only if supported by the device */
+        CharCount = (Descriptor->bLength - sizeof(USB_COMMON_DESCRIPTOR)) / sizeof(WCHAR) +
+                    (sizeof("&") - 1) +
+                    (sizeof("F") - 1) + // LUN: 1 char (MAX_LUN)
+                    sizeof(ANSI_NULL);
     }
     else
     {
-        // use instance count and LUN
-        swprintf(Buffer, L"%04lu&%c", FDODeviceExtension->InstanceCount, PDODeviceExtension->LUN);
+        /* Use the instance count and LUN as a fallback */
+        CharCount = (sizeof("99999999") - 1) + // Instance Count: 8 chars
+                    (sizeof("&") - 1) +
+                    (sizeof("F") - 1) + // LUN: 1 char (MAX_LUN)
+                    sizeof(ANSI_NULL);
     }
 
-    Length = wcslen(Buffer) + 1;
-
-    InstanceId = ExAllocatePoolWithTag(PagedPool, Length * sizeof(WCHAR), USB_STOR_TAG);
+    InstanceId = ExAllocatePoolUninitialized(PagedPool, CharCount * sizeof(WCHAR), USB_STOR_TAG);
     if (!InstanceId)
     {
         Irp->IoStatus.Information = 0;
         return STATUS_INSUFFICIENT_RESOURCES;
     }
 
-    wcscpy(InstanceId, Buffer);
+    if (Descriptor && (Descriptor->bLength >= sizeof(USB_COMMON_DESCRIPTOR) + sizeof(WCHAR)))
+    {	
+        Status = RtlStringCchPrintfW(InstanceId,
+                                     CharCount,
+                                     L"%s&%x",
+                                     Descriptor->bString,
+                                     PDODeviceExtension->LUN);
+    }
+    else
+    {
+        Status = RtlStringCchPrintfW(InstanceId,
+                                     CharCount,
+                                     L"%04lu&%x",
+                                     FDODeviceExtension->InstanceCount,
+                                     PDODeviceExtension->LUN);
+    }
 
-    DPRINT("USBSTOR_PdoHandleQueryInstanceId %S\n", InstanceId);
+    /* This should not happen */
+    ASSERT(NT_SUCCESS(Status));
+
+    DPRINT("USBSTOR_PdoHandleQueryInstanceId '%S'\n", InstanceId);
 
     Irp->IoStatus.Information = (ULONG_PTR)InstanceId;
     return STATUS_SUCCESS;
diff --git a/drivers/usb/usbstor/usbstor.h b/drivers/usb/usbstor/usbstor.h
index 8a65e2bd521..11bb2d2c41c 100644
--- a/drivers/usb/usbstor/usbstor.h
+++ b/drivers/usb/usbstor/usbstor.h
@@ -2,6 +2,7 @@
 #define _USBSTOR_H_
 
 #include <wdm.h>
+#include <ntstrsafe.h>
 #include <usbdi.h>
 #include <usbbusif.h>
 #include <usbdlib.h>