Some fixes for check on buffer length.

ModOrganizer2 · Jun 23, 2024 · 7267310 · 7267310
1 parent 8e1985e
commit 7267310
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/src/usvfs_dll/hooks/ntdll.cpp b/src/usvfs_dll/hooks/ntdll.cpp
@@ -1068,7 +1068,7 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryObject(
                         ObjectInformationLength, ReturnLength);
   POST_REALCALL
 
-  if ((res == STATUS_SUCCESS || res == STATUS_INFO_LENGTH_MISMATCH)
+  if ((res == STATUS_SUCCESS || res == STATUS_BUFFER_OVERFLOW)
     && (ObjectInformationClass == ObjectNameInformation)) {
     const auto trackerInfo = ntdllHandleTracker.lookup(Handle);
     const auto redir       = applyReroute(READ_CONTEXT(), callContext, trackerInfo);
@@ -1094,7 +1094,7 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryObject(
 
       // TODO: check this...
       if (ObjectInformationLength < buffer.size() * 2 + sizeof(OBJECT_NAME_INFORMATION)) {
-        res = STATUS_INFO_LENGTH_MISMATCH;
+        res = STATUS_BUFFER_OVERFLOW;
 
         if (ReturnLength) {
           *ReturnLength = buffer.size() * 2 + sizeof(OBJECT_NAME_INFORMATION);
@@ -1116,6 +1116,8 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryObject(
         info->Name.Buffer        = unicodeBuffer;
         info->Name.Length        = buffer.size() * 2;
         info->Name.MaximumLength = unicodeBufferLength;
+
+        res                      = STATUS_SUCCESS;
       }
     }
 
@@ -1151,7 +1153,8 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryInformationFile(
                                  FileInformationClass);
   POST_REALCALL
 
-  if ((res == STATUS_SUCCESS || res == STATUS_INFO_LENGTH_MISMATCH) && (
+  if ((res == STATUS_SUCCESS || res == STATUS_BUFFER_OVERFLOW) &&
+      (
     FileInformationClass == FileNameInformation 
     || FileInformationClass == FileAllInformation 
     || FileInformationClass == FileNormalizedNameInformation)) {
@@ -1174,15 +1177,16 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryInformationFile(
 
     if (redir.redirected)
     {
-      if (maxNameSize < trackerInfo.size()) {
-        res = STATUS_INFO_LENGTH_MISMATCH;
+      if (maxNameSize < trackerInfo.size() - 6) {
+        res = STATUS_BUFFER_OVERFLOW;
       } else {
         LPCWSTR filenameFixed = static_cast<LPCWSTR>(trackerInfo);
         if (info->FileName[0] == L'\\') {
           // strip the \??\X: prefix (X being the drive name)
           filenameFixed = filenameFixed + 6;
         }
         SetInfoFilename(FileInformation, FileInformationClass, filenameFixed);
+        res = STATUS_SUCCESS;
       }
     }
 

diff --git a/test/tvfs_test/main.cpp b/test/tvfs_test/main.cpp
@@ -424,6 +424,20 @@ TEST_F(USVFSTest, NtQueryObjectVirtualFile)
     ASSERT_EQ(0, wcscmp(fileNameInfo->FileName, L"\\np.exe"));
   }
 
+  // buffer of size should be too small for the original path (\Windows\notepad.exe)
+  // but not for \np.exe
+  {
+    char buffer[sizeof(ULONG) + 8 * 2];
+    IO_STATUS_BLOCK status;
+    const auto res = usvfs::hook_NtQueryInformationFile(
+        hdl, &status, buffer, sizeof(buffer), FileNormalizedNameInformation);
+    ASSERT_EQ(STATUS_SUCCESS, status.Status);
+
+    FILE_NAME_INFORMATION* fileNameInfo =
+        reinterpret_cast<FILE_NAME_INFORMATION*>(buffer);
+    ASSERT_EQ(0, wcscmp(fileNameInfo->FileName, L"\\np.exe"));
+  }
+
   {
     char buffer[2048];
     const auto res = usvfs::hook_NtQueryObject(hdl, ObjectNameInformation, buffer,