From 72673107217f361d2556d727357a2aafb2f6ba96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Capelle?= <capelle.mikael@gmail.com>
Date: Sun, 23 Jun 2024 14:15:01 +0200
Subject: [PATCH] Some fixes for check on buffer length.

---
 src/usvfs_dll/hooks/ntdll.cpp | 14 +++++++++-----
 test/tvfs_test/main.cpp       | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/usvfs_dll/hooks/ntdll.cpp b/src/usvfs_dll/hooks/ntdll.cpp
index fac0578..1637c5e 100644
--- 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,8 +1177,8 @@ 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'\\') {
@@ -1183,6 +1186,7 @@ DLLEXPORT NTSTATUS WINAPI usvfs::hook_NtQueryInformationFile(
           filenameFixed = filenameFixed + 6;
         }
         SetInfoFilename(FileInformation, FileInformationClass, filenameFixed);
+        res = STATUS_SUCCESS;
       }
     }
 
diff --git a/test/tvfs_test/main.cpp b/test/tvfs_test/main.cpp
index 4294a53..fc76773 100644
--- 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,