From 37b3fa64e828185ad781f648a12ac45bbfd21bbe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Capelle?= <capelle.mikael@gmail.com>
Date: Tue, 16 Jul 2024 17:16:09 +0200
Subject: [PATCH] Try to limit buffer size to USHORT in NtQueryObject.

---
 src/usvfs_dll/hooks/ntdll.cpp | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/usvfs_dll/hooks/ntdll.cpp b/src/usvfs_dll/hooks/ntdll.cpp
index 4286a7b..b79f0a6 100644
--- a/src/usvfs_dll/hooks/ntdll.cpp
+++ b/src/usvfs_dll/hooks/ntdll.cpp
@@ -11,6 +11,8 @@
 #include <addrtools.h>
 #include <unicodestring.h>
 #include <queue>
+#include <limits>
+#include <algorithm>
 
 namespace ulog = usvfs::log;
 namespace ush  = usvfs::shared;
@@ -1096,10 +1098,22 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryObject(
       buffer =
           std::wstring(deviceName) + L'\\' + std::wstring(buffer.data() + 7, buffer.size() - 7);
 
-      // the name is put in the buffer AFTER the struct, so the required size if 
-      // sizeof(OBJECT_NAME_INFORMATION) + the number of bytes for the name + 2 bytes for a wide null character
-      const auto requiredLength = sizeof(OBJECT_NAME_INFORMATION) + buffer.size() * 2 + 2;
-      if (ObjectInformationLength < requiredLength) {
+      // the name is put in the object AFTER the struct, so the remaining length is the size of
+      // the object minus the size of the struct
+      //
+      // the buffer length is a USHORT so if we have more than the maximum possible value
+      // for USHORT, we need to truncate
+      //
+      const USHORT unicodeBufferLength =
+          static_cast<USHORT>(std::min(
+              static_cast<unsigned long long>(ObjectInformationLength - sizeof(OBJECT_NAME_INFORMATION)),
+              static_cast<unsigned long long>(std::numeric_limits<USHORT>::max())
+          ));
+
+      
+      // the number of bytes for the name + 2 bytes for a wide null character
+      const auto requiredLength = buffer.size() * 2 + 2;
+      if (unicodeBufferLength < requiredLength) {
         // if the status was info length mismatch, we keep it, we are just going to update
         // *ReturnLength
         if (res != STATUS_INFO_LENGTH_MISMATCH) {
@@ -1107,12 +1121,10 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryObject(
         }
 
         if (ReturnLength) {
-          *ReturnLength = requiredLength;
+          *ReturnLength = sizeof(OBJECT_NAME_INFORMATION) + requiredLength;
         }
       } else {
         // put the unicode buffer at the end of the object
-        const auto unicodeBufferLength =
-            ObjectInformationLength - sizeof(OBJECT_NAME_INFORMATION);
         LPWSTR unicodeBuffer = reinterpret_cast<LPWSTR>(
           static_cast<LPSTR>(ObjectInformation) + sizeof(OBJECT_NAME_INFORMATION));
 
@@ -1125,7 +1137,7 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryObject(
         // update the actual unicode string
         info->Name.Buffer        = unicodeBuffer;
         info->Name.Length        = static_cast<USHORT>(buffer.size() * 2);
-        info->Name.MaximumLength = static_cast<USHORT>(unicodeBufferLength);
+        info->Name.MaximumLength = unicodeBufferLength;
 
         res                      = STATUS_SUCCESS;
       }